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
| fname = fcall.args[1] | ||
| # qualified names such as `TensorKit.treebraider` add methods to a function of another module | ||
| basename = Meta.isexpr(fname, :.) ? fname.args[end].value : fname | ||
| basename isa Symbol || error("cached macro can only be used on function definitions") |
Member
There was a problem hiding this comment.
any way we could embed the basename in here to give people more info?
kshyatt
reviewed
Sep 28, 2026
| dst::TransformSubblocks, src::TransformSubblocks, p, conjsrc::Bool, | ||
| blk::RecouplingBlock, buffer, α, β, backend, allocator | ||
| ) | ||
| U = blk.U |
Member
There was a problem hiding this comment.
is there no way we could dispatch on this? hugely ugly ternary here
kshyatt
reviewed
Sep 28, 2026
| buffer_dst = StridedView(buffer, (blocksize, rows), (1, blocksize), 0) | ||
| buffer_src = StridedView(buffer, (blocksize, cols), (1, blocksize), blocksize * rows) | ||
|
|
||
| # 1. Extract: copy each source block into column i of buffer_src as a flat vector, |
Member
There was a problem hiding this comment.
this might be a good spot for some ASCII art
kshyatt
reviewed
Sep 28, 2026
| mul!(_realview(rbuffer, buffer_dst), _realview(rbuffer, buffer_src), transpose(U)) | ||
| return nothing | ||
| end | ||
| # on the CPU, views of reinterpreted arrays are not `StridedMatrix`, so `mul!` would not use BLAS |
Member
There was a problem hiding this comment.
this comment seems kind of out of place to me. I don't know if it's needed, or maybe it should be rephrased? It's very Claudey
kshyatt
reviewed
Sep 28, 2026
| """ | ||
| struct RecouplingBlock{T, M <: AbstractMatrix{T}} | ||
| coeff::T | ||
| U::Union{Nothing, M} |
Member
There was a problem hiding this comment.
Why not make the type of U a type parameter?
kshyatt
reviewed
Sep 28, 2026
|
|
||
| # `StridedSubblocks` report their storage as the `StridedView` parent type, which is `Memory` there | ||
| @static if isdefined(Core, :Memory) | ||
| const CPUStorage{T} = Union{Array{T}, Memory{T}} |
Member
There was a problem hiding this comment.
This name seems weird to me. Maybe it should be DefaultStorage or something? It's defined in opposition to the GPU extensions which seems strange
kshyatt
reviewed
Sep 28, 2026
| recoupling_scalartype(A::Type{<:AbstractVector}, Tₛ::Type{<:Number}) -> Type{<:Number} | ||
|
|
||
| Scalar type used to store the recoupling coefficients with sector scalar type `Tₛ` in the | ||
| transformers for destination tensors with storagetype `A`. For storage with BLAS scalars, this is |
Member
There was a problem hiding this comment.
What about storage that isn't made of BLAS scalars?
kshyatt
reviewed
Sep 28, 2026
`@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>
lkdvos
force-pushed
the
ld-storagetype
branch
from
September 28, 2026 17:07
5fcb5f4 to
56f134c
Compare
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
2 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.