Don't materialize an AdjointTensorMap in TO.tensoradd! when conjA = true - #520
Don't materialize an AdjointTensorMap in TO.tensoradd! when conjA = true#520leburgel wants to merge 1 commit into
AdjointTensorMap in TO.tensoradd! when conjA = true#520Conversation
leburgel
changed the title
AdjointTensorMap in `AdjointTensorMap in TO.tensoradd! when conjA = true
Codecov Report❌ Patch coverage is
🚀 New features to boost your workflow:
|
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
Overall this looks great, it's definitely the right call to also specialize the adjointtensor, I've just been too lazy and putting it off, hoping that the actual runtime fraction is not too high. I somehow didn't realize that the structure mapping in the current state is actually not that complicated, so this does indeed open up many optimization possibilities.
On a larger scale comment, I'm not sure if this is really a route we want to take, but it does look like we could conceivably reduce the cache sizes by simply making the "treetransformers" only store which tokens go to which other tokens, and therefore decouple that from the actual subblockstructure dictionaries that can then fill out which data is where. This might be a bit larger of a refactor though, so I'll see what I can do next week and if this really makes sense
| tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler; | ||
| conjsrc::Bool = false |
There was a problem hiding this comment.
| tdst, tsrc, p, transformer, α, β, backend, allocator, scheduler; | |
| conjsrc::Bool = false | |
| tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator, scheduler |
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.
| A′ = adjoint(A) | ||
| pA′ = adjointtensorindices(A, _canonicalize(pA, C)) | ||
| permute!(C, A′, pA′, α, β, backend, allocator) | ||
| if C isa TensorMap && A isa TensorMap |
There was a problem hiding this comment.
It somehow feels like this specialization might make more sense on the permute!-level directly, where we could intercept TensorMap-AdjointTensorMap combinations directly? That would also immediately catch manual calls to permute!(tdst, tsrc', ...) in the same way.
To achieve this, I'm wondering if it makes sense to just have treebraider(Vdst, Vsrc, p, levels, conj) as the entrypoint, which then also fixes the caching keys being distinct? There's some subtleties left with this, mostly with respect to other combinations like permute!(tdst', tsrc, ...), which we would then probably want to map to the other case, but that might be reasonable?
There was a problem hiding this comment.
I think the first part of this comment is more similar to what is sketched in #519, as opposed to the tensoradd!-level fix that is sketched here.
|
|
||
| See also [`subblockstructure`](@ref). | ||
| """ | ||
| function adjoint_subblockstructure(W::HomSpace) |
There was a problem hiding this comment.
code-organization wise, would it be easier to make this subblockstructure(W::HomSpace, conjW::Bool=false)?
| newkeys = map(((f₁, f₂),) -> (f₂, f₁), collect(keys(structs))) | ||
| newvals = map(collect(values(structs))) do (sz, str, off) |
There was a problem hiding this comment.
it feels like it should be possible to avoid the collect calls here, but I'd have to dig through the Dictionaries interface(s) to see what is possible?
2 participants
A second option for a proper fix for #516, which avoids materializing an
AdjointTensorMapinTO.tensoradd!altogether, but rather passes theconjflag along to the array-leveltensoradd!Basically the same underlying idea as #519, but since I don't exactly know what I'm doing here I'm posting both, mostly just for inspiration.
Some benchmark timings based on the reproducer of #516: