Repository navigation
Conversation
Codecov Report❌ Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
lkdvos
force-pushed
the
ld-storagetype
branch
from
September 25, 2026 17:56
2c02176 to
bebe0d6
Compare
lkdvos
marked this pull request as ready for review
September 25, 2026 17:57
Member
Author
|
Benchmarks for the complex, non-abelian case with real recoupling coefficients: the
The small difference for the 4-leg permutes does not reproduce locally, where the branch is 2–6% faster for the same cases. These are dominated by the single-tree blocks (29 of 37 blocks), which go through the unchanged permuting Full results |
lkdvos
force-pushed
the
ld-storagetype
branch
from
September 26, 2026 01:58
bebe0d6 to
5fcb5f4
Compare
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
kshyatt
reviewed
Sep 28, 2026
lkdvos
force-pushed
the
ld-storagetype
branch
from
September 28, 2026 17:07
5fcb5f4 to
56f134c
Compare
`@cached` escaped its entire expansion, so that the cache styles, the global cache registry, the LRU constructor and the timer were resolved in the calling module, which only works within TensorKit. These are now referenced through `GlobalRef`s, and qualified function names such as `TensorKit.treebraider` are supported, such that package extensions can add cached methods to TensorKit functions. The global cache of such methods lives in the module that defines them, and is registered with its module name so that it is shown separately in `global_cache_info`. Also fixes the task-local cache key, which was spliced in as an identifier instead of as a symbol. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
`treebraider` and `treetransposer` now take the storagetype `A` of the destination tensor as their first argument. It is part of the cache key, and it is a dispatch point: a storage-specific method, with its own cache through `@cached`, can return a dedicated transformer type. The recoupling data is stored in the form the kernel needs for `A`: - `recoupling_scalartype` stores the coefficients in the precision of the storage. On CPU, real coefficients stay real for complex data. - Blocks of `GenericTreeTransformer` are stored as `RecouplingBlock`s, a concrete type holding either a host scalar (single tree) or a recoupling matrix in the storage of the destination. For GPU storage the matrices are thus converted once at construction, instead of adapted on every call. In the kernel, `α` is applied in the unpack step, such that the recoupling is a plain matrix product. For complex CPU data with real coefficients, the real and imaginary parts are recoupled in a single real `gemm` on a reinterpreted view of the buffer. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…orage Reinterpreting GPU storage as real yields native GPU arrays, such that the real and imaginary parts of complex data can be recoupled with real coefficients in a single `mul!`, which dispatches to the vendor BLAS. This adds a generic `_recouple!` method for complex `DenseVector` buffers with a real recoupling matrix, keeping the direct `BLAS.gemm!` call only for CPU storage, where views of reinterpreted arrays are not `StridedMatrix`. The CPU rule of `recoupling_scalartype`, keeping real coefficients real, is now used for all storage. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…coefficient promotion is now allowed
lkdvos
commented
Oct 5, 2026
lkdvos
left a comment
lkdvos
left a comment
Member
Author
There was a problem hiding this comment.
I've rebased this on otp of the GPU kernels and attempted to simplify a bit.
Jutho
reviewed
Oct 9, 2026
Comment on lines
+130
to
+134
| weights[i] = length(U₀) * prod(structure_dst[first(inds_dst)][1]) | ||
|
|
||
| @debug( | ||
| lazy"Created recoupling block for uncoupled: $(fs_src.uncoupled)", | ||
| sz = size(U), sparsity = count(!iszero, U) / length(U) | ||
| sz = size(U₀), sparsity = count(!iszero, U₀) / length(U₀) |
Member
There was a problem hiding this comment.
Is there a specific reason to make this U₀ ? Also U is defined at this point, no?
Jutho
reviewed
Oct 9, 2026
|
|
||
| # 3. Insert: scatter column j of buffer_dst into the destination, applying the | ||
| # actual index permutation p in the same tensoradd! call. | ||
| # actual index permutation p and the scaling α in the same tensoradd! call. |
Member
There was a problem hiding this comment.
Out of curiosity: I think this makes more sense, but was there a practical or performance reason for shifting where α is applied?
Jutho
reviewed
Oct 9, 2026
| return maximum(transformer.data; init = 0) do (U, _, inds_src) | ||
| return length(U) == 1 ? 0 : prod(structure_src[first(inds_src)][1]) * sum(size(U)) | ||
| length(U) == 1 && return 0 | ||
| return prod(structure_src[first(inds_src)][1]) * sum(size(U)) |
Member
There was a problem hiding this comment.
Any reason for this change? I don't think this really adds to readability.
Jutho
reviewed
Oct 9, 2026
Comment on lines
+175
to
+176
| FusionStyle(I) == UniqueFusion() && return UniqueTreeTransformer{T, N} | ||
| return GenericTreeTransformer{T, N} |
Jutho
reviewed
Oct 9, 2026
| return idst | ||
| end | ||
| end | ||
| U = convert(Matrix{T}, U₀) |
Member
There was a problem hiding this comment.
I would have thought that the storage type information would have been used to also immediately convert this to the proper type of array, i.e. host vs device/gpu?
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is preparatory work to hopefully simplify some of the logic in #533:
The goal is to make it possible and convenient to have a dispatched method for caching a different kind of treetransformer for device vs host code. This required 3 changes:
@cachedmacro needed a bit of hygiene to easily work from modules that aren't TensorKit (such as extensions)treebraidercalls etc now take the storagetype of thetdstas a first argument to facilitate dispatchGenericTreeTransformeris already slightly improved: it now stores the unitary recoupling coefficients in the correct storagetype to avoid data transfer.This should allow the GPU extension to simply overload these functions and return custom structs.