[WIP] cuTENSOR blocksparse implementation - #510
Draft
lkdvos wants to merge 4 commits into
Draft
Conversation
`TensorOperations.tensorcontract!` now delegates to `_tensorcontract!`, the designated hook for backends implementing a contraction by a route other than the default fusion-tree/BLAS path. Hooking there rather than at `contract!` keeps `conjA`/`conjB` as plain `Bool`s, since `_generic_tensorcontract!` folds them into `AdjointTensorMap`s whose data layout no longer matches the subblock structure of their own space. Also normalize the subblock strides in `degeneracystructure`, so that the don't-care strides of extent-1 axes come out monotone by construction. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Map a symmetric tensor contraction of CUDA-backed `TensorMap`s onto a single cuTENSOR block-sparse contraction of the underlying storage: the index tuples only determine mode labels, so there is no fusion tree transformation, no intermediate permutation and no temporary. Opt-in, since NVIDIA documents the block-sparse API as a public beta with no guarantee of stability. `blocksparse_compatible` gates the sector types this applies to. Fermionic sectors qualify: their raw block contraction differs from the categorical one by a per-block sign, which `blocksparse_contract_signs` derives in factorized form and the extension applies by scaling blocks. `plan_contract` hoists the data-independent part out of a hot loop; plans are also cached automatically. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`test/other/blocksparse.jl` pins down the host-side reasoning the backend relies on without needing a GPU: the sector ordering, the dual-space convention, the stride canonicalization, and the per-block sign correction, including a guard that each of its factors is load-bearing. `test/cuda/blocksparse.jl` compares the block-sparse path against the default one on raw flat data, sweeping output partitions, scalars, element types and conjugations, and checks that unsupported cases fall back silently but throw under `strict`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Add a manual page covering the contraction backends, block-sparse plan reuse and the block schedulers, and pull the corresponding docstrings into the library reference. Also refresh the stale `CuArrays.jl` note in the tensor manual, since GPU storage is now supported by loading CUDA.jl or AMDGPU.jl. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
kshyatt
self-requested a review
August 19, 2026 13:10
Member
Crab gang rise up!! |
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 an initial implementation of calling into cuTENSOR's blocksparse contraction, which should be able to speed up our
UniqueFusioncode. There is still a bit of annoyance for supporting conjugation flags, and non-trivial F or R symbol coefficients, where this has to be handled in pre/postprocessing, but it seems at least from my initial local benchmarks that this is still competitive and outperforms our naive implementation.Still requires quite a bit of cleanup, and there is one design question that I am still unsure about, which considers whether or not we should keep the "BlockSparseStructure" as something cuTENSOR specific moved into the extension, or as a generic thing backends can hook into.
On the one hand nothing about that is cuTENSOR specific, and there is a case to be made that we could even on CPU do this strategy which merges permutation and multiplication, (maybe someone has a RUST library somewhere that he is working on 😉) but unfortunately the annoying thing is that cuTENSOR is really particular about what integer types have to be where.
Anyways, I'll keep working on this but already wanted to open this up to let everyone know this is in progress.