Repository navigation
storagetype-adapted treetransformers - #556
Conversation
Codecov Report❌ Patch coverage is
... and 2 files with indirect coverage changes 🚀 New features to boost your workflow:
|
2c02176 to
bebe0d6
Compare
|
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 |
bebe0d6 to
5fcb5f4
Compare
56f134c to
33931a1
Compare
lkdvos
left a comment
lkdvos
left a comment
There was a problem hiding this comment.
I've rebased this on otp of the GPU kernels and attempted to simplify a bit.
| return idst | ||
| end | ||
| end | ||
| U = convert(Matrix{T}, U₀) |
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?
There was a problem hiding this comment.
The problem is that the GPU implementations don't keep a vector of U, and rather have a single vector of all coeffs, as vector-of-matrix datastructures don't play nicely with GPU implementations. See also my reply for the next comment, I am thinking of trying to get that working for both CPU and GPU implementations such that they all follow similar codepaths, and only require a dedicated kernel, but would do that in a follow-up.
| end | ||
|
|
||
| const GPUStridedSubblocks = StridedSubblocks{<:AnyGPUArray} | ||
| # we don't want/need the `StridedSubblocks` here because that information is already in the transformer |
There was a problem hiding this comment.
Would it make sense to do a similar thing for the "normal" (i.e. CPU) transformers (i.e. simply pass the flat data vector to add_transform_kernel! and pack the strided subblock information in the transformers)? This would make the workflow a bit more uniform. Currently it is quite complex and confusing to keep track of what is going where along the different codepaths, though I am fine with it for now.
There was a problem hiding this comment.
I absolutely agree, and had the same thought, although I'll defer this to a follow-up PR to avoid making too many changes at once here. I am somewhat convinced I can actually get both transformer types to align again by making some other changes as well, but to keep things reviewable I'll separate that out.
Jutho
left a comment
Jutho
left a comment
There was a problem hiding this comment.
Ok, generally looks great. I left a few questions to help my understanding.
|
I think one possible improvement is that currently on the host side, the argument |
`@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
16726e9 to
18bd84c
Compare
|
GPU failures are unrelated, will fix those in a follow-up. |
3 participants
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.