-
Notifications
You must be signed in to change notification settings - Fork 64
Don't materialize an AdjointTensorMap in TO.tensoradd! when conjA = true
#520
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -169,6 +169,30 @@ See also [`degeneracystructure`](@ref), [`blockstructure`](@ref). | |
| """ | ||
| subblockstructure(W::HomSpace) = Dictionary(fusiontrees(W), degeneracystructure(W).subblockstructure) | ||
|
|
||
| """ | ||
| adjoint_subblockstructure(W::HomSpace) -> Dictionary | ||
|
|
||
| Subblock structure of `W'`, expressed as strides into a buffer laid out for `W`. | ||
|
|
||
| This mirrors the relation used by `subblock(::AdjointTensorMap, ...)`: the adjoint shares its | ||
| parent's data, with the tree pair swapped and sizes and strides permuted by | ||
| `(domainind..., codomainind...)`. Permuting `adjoint(t)` can therefore read `t`'s own buffer | ||
| directly, instead of materialising an `AdjointTensorMap` and falling back to the uncached | ||
| tree-transformation path. | ||
|
|
||
| See also [`subblockstructure`](@ref). | ||
| """ | ||
| function adjoint_subblockstructure(W::HomSpace) | ||
| N₁, N₂ = numout(W), numin(W) | ||
| swap = (ntuple(i -> N₁ + i, N₂)..., ntuple(identity, N₁)...) | ||
| structs = subblockstructure(W) | ||
| newkeys = map(((f₁, f₂),) -> (f₂, f₁), collect(keys(structs))) | ||
| newvals = map(collect(values(structs))) do (sz, str, off) | ||
|
Comment on lines
+189
to
+190
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. it feels like it should be possible to avoid the |
||
| return (TupleTools.getindices(sz, swap), TupleTools.getindices(str, swap), off) | ||
| end | ||
| return Dictionary(newkeys, newvals) | ||
| end | ||
|
|
||
| """ | ||
| fusionblocks(W::HomSpace) | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change | ||||||
|---|---|---|---|---|---|---|---|---|
|
|
@@ -579,6 +579,29 @@ Base.@deprecate( | |||||||
| return tdst | ||||||||
| end | ||||||||
|
|
||||||||
| """ | ||||||||
| add_conj_transform!(tdst, tsrc, p, transformer, α, β, backend, allocator) -> tdst | ||||||||
|
|
||||||||
| Compute `tdst = β * tdst + α * permute(adjoint(tsrc), p)` while reading `tsrc`'s own data | ||||||||
| buffer, without materialising an `AdjointTensorMap`. | ||||||||
|
|
||||||||
| `transformer` must come from [`conj_treebraider`](@ref), whose source strides address the | ||||||||
| parent's storage; the conjugation itself is folded into the strided reads. | ||||||||
| """ | ||||||||
| @propagate_inbounds function add_conj_transform!( | ||||||||
| tdst::TensorMap, tsrc::TensorMap, p::Index2Tuple, transformer, | ||||||||
| α::Number, β::Number, backend, allocator | ||||||||
| ) | ||||||||
| @boundscheck spacecheck_transform(permute, tdst, tsrc', p) | ||||||||
| ntasks = use_threaded_transform(tdst, transformer) ? get_num_transformer_threads() : 1 | ||||||||
| scheduler = ntasks == 1 ? SerialScheduler() : DynamicScheduler(; ntasks, split = :roundrobin) | ||||||||
| add_transform_kernel!( | ||||||||
| tdst.data, tsrc.data, p, transformer, α, β, backend, allocator, scheduler; | ||||||||
| conjsrc = true | ||||||||
| ) | ||||||||
| return tdst | ||||||||
| end | ||||||||
|
|
||||||||
| function use_threaded_transform(t::TensorMap, transformer) | ||||||||
| return get_num_transformer_threads() > 1 && length(t.data) > Strided.MINTHREADLENGTH | ||||||||
| end | ||||||||
|
|
@@ -587,14 +610,15 @@ function use_threaded_transform(t::AbstractTensorMap, transformer) | |||||||
| end | ||||||||
|
|
||||||||
| function add_transform_kernel!( | ||||||||
| tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler | ||||||||
| tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler; | ||||||||
| conjsrc::Bool = false | ||||||||
|
Comment on lines
+613
to
+614
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
little bit of a nitpick, but I'd be inclined to just make this a mandatory (positional) argument instead, which is more in line with the remainder of these functions. Obviously doens't matter, just for consistency. |
||||||||
| ) | ||||||||
| I = sectortype(tdst) | ||||||||
| if FusionStyle(I) === UniqueFusion() | ||||||||
| tforeach(fusiontrees(tsrc); scheduler) do (f₁, f₂) | ||||||||
| (f₁′, f₂′), coeff = transformer((f₁, f₂)) | ||||||||
| @inbounds TO.tensoradd!( | ||||||||
| tdst[f₁′, f₂′], tsrc[f₁, f₂], p, false, α * coeff, β, backend, allocator | ||||||||
| tdst[f₁′, f₂′], tsrc[f₁, f₂], p, conjsrc, α * coeff, β, backend, allocator | ||||||||
| ) | ||||||||
| end | ||||||||
| return nothing | ||||||||
|
|
@@ -629,7 +653,7 @@ function add_transform_kernel!( | |||||||
| @inbounds for (i, (f₁, f₂)) in enumerate(fusiontrees(src)) | ||||||||
| TO.tensoradd!( | ||||||||
| sreshape(view(buffer_src, :, i), sz_src), tsrc[f₁, f₂], | ||||||||
| ptriv, false, One(), Zero(), backend, allocator | ||||||||
| ptriv, conjsrc, One(), Zero(), backend, allocator | ||||||||
| ) | ||||||||
| end | ||||||||
|
|
||||||||
|
|
@@ -657,19 +681,19 @@ end | |||||||
| # repeated specialization -- this only depends on `numind` and `eltype`. | ||||||||
| function add_transform_kernel!( | ||||||||
| data_dst::DenseVector, data_src::DenseVector, p, transformer::AbelianTreeTransformer, | ||||||||
| α, β, backend, allocator, scheduler | ||||||||
| α, β, backend, allocator, scheduler; conjsrc::Bool = false | ||||||||
| ) | ||||||||
| tforeach(transformer.data; scheduler) do (coeff, struct_dst, struct_src) | ||||||||
| TO.tensoradd!( | ||||||||
| StridedView(data_dst, struct_dst...), StridedView(data_src, struct_src...), | ||||||||
| p, false, α * coeff, β, backend, allocator | ||||||||
| p, conjsrc, α * coeff, β, backend, allocator | ||||||||
| ) | ||||||||
| end | ||||||||
| return nothing | ||||||||
| end | ||||||||
| function add_transform_kernel!( | ||||||||
| data_dst::DenseVector, data_src::DenseVector, p, transformer::GenericTreeTransformer, | ||||||||
| α, β, backend, allocator, scheduler | ||||||||
| α, β, backend, allocator, scheduler; conjsrc::Bool = false | ||||||||
| ) | ||||||||
| cp = TO.allocator_checkpoint!(allocator) | ||||||||
|
|
||||||||
|
|
@@ -686,7 +710,7 @@ function add_transform_kernel!( | |||||||
| TO.tensoradd!( | ||||||||
| StridedView(data_dst, sz_dst, only(structs_dst)...), | ||||||||
| StridedView(data_src, sz_src, only(structs_src)...), | ||||||||
| p, false, α * coeff, β, backend, allocator | ||||||||
| p, conjsrc, α * coeff, β, backend, allocator | ||||||||
| ) | ||||||||
| else # Multi-tree block: pack → recoupling matmul → unpack. | ||||||||
| rows, cols = size(U) | ||||||||
|
|
@@ -701,7 +725,7 @@ function add_transform_kernel!( | |||||||
| @inbounds for (i, struct_src_i) in enumerate(structs_src) | ||||||||
| TO.tensoradd!( | ||||||||
| sreshape(view(buffer_src, :, i), sz_src), StridedView(data_src, sz_src, struct_src_i...), | ||||||||
| ptriv, false, One(), Zero(), backend, allocator | ||||||||
| ptriv, conjsrc, One(), Zero(), backend, allocator | ||||||||
| ) | ||||||||
| end | ||||||||
|
|
||||||||
|
|
||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -52,9 +52,23 @@ function TO.tensoradd!( | |
| return C | ||
| end | ||
| if conjA | ||
| A′ = adjoint(A) | ||
| pA′ = adjointtensorindices(A, _canonicalize(pA, C)) | ||
| permute!(C, A′, pA′, α, β, backend, allocator) | ||
| if C isa TensorMap && A isa TensorMap | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It somehow feels like this specialization might make more sense on the To achieve this, I'm wondering if it makes sense to just have |
||
| # Both operands have a flat data buffer, so the adjoint never has to be | ||
| # materialised: `conj_treebraider` yields a memoized transformer whose source | ||
| # strides address `A`'s own storage. Materialising `adjoint(A)` here would push | ||
| # a perfectly ordinary TensorMap onto the uncached tree-transformation path. | ||
| n₁ = numin(A) # == numout(adjoint(A)) | ||
| levels = ntuple(identity, numind(A)) | ||
| levels′ = ( | ||
| TupleTools.getindices(levels, ntuple(identity, n₁)), | ||
| TupleTools.getindices(levels, n₁ .+ ntuple(identity, numout(A))), | ||
| ) | ||
| transformer = conj_treebraider(space(C), space(A), pA′, levels′) | ||
| @inbounds add_conj_transform!(C, A, pA′, transformer, α, β, backend, allocator) | ||
| else | ||
| permute!(C, adjoint(A), pA′, α, β, backend, allocator) | ||
| end | ||
| else | ||
| permute!(C, A, _canonicalize(pA, C), α, β, backend, allocator) | ||
| end | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
code-organization wise, would it be easier to make this
subblockstructure(W::HomSpace, conjW::Bool=false)?