From 42e1ae11efb1fa2d9db1f3c9ad91a3f1fafcc92d Mon Sep 17 00:00:00 2001 From: lkdvos Date: Tue, 13 Oct 2026 13:00:28 -0400 Subject: [PATCH] Avoid boxing the BLAS-contractable inputs in `_blas_contract!` `makeblascontractable` returns either its input or a permuted copy, whose types differ when the copy is allocated by an allocator such as Bumper's. The resulting union is not an `isbits` type and was therefore heap allocated on every contraction, even when no copy was required. `_blas_contract!` now makes the inputs BLAS-contractable one at a time, passing them on to the next step in each branch rather than returning them, such that the union is never materialized. Co-Authored-By: Claude Opus 5.5 --- docs/src/changelog.md | 1 + src/implementation/blascontract.jl | 40 +++++++++++++++++++++++------- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/docs/src/changelog.md b/docs/src/changelog.md index b45884b6..399232cc 100644 --- a/docs/src/changelog.md +++ b/docs/src/changelog.md @@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `ncon` can be differentiated with Zygote, including networks with traces. - `ncon` accepts traces in networks whose index lists are given as row matrices, e.g. `[1 1 -1]`. +- Contractions using the `StridedBLAS` backend no longer heap allocate a view of an input tensor when an allocator such as Bumper's is used ([#273](https://github.com/QuantumKitHub/TensorOperations.jl/issues/273)). ### Performance diff --git a/src/implementation/blascontract.jl b/src/implementation/blascontract.jl index 289ae0c2..1521972c 100644 --- a/src/implementation/blascontract.jl +++ b/src/implementation/blascontract.jl @@ -43,23 +43,37 @@ end # implement necessary permutations function _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator) + # note: the inputs are made BLAS-contractable one at a time, passing them on to the next + # step rather than returning them, since their types depend on whether a copy was made, + # and returning such a union would require heap allocating the result + isblascontractable(A, pA) && eltype(A) == eltype(C) && + return _blas_contract_B!(C, A, pA, B, pB, pAB, α, β, backend, allocator) + A_ = _blascontractable_copy(A, pA, eltype(C), backend, allocator) + _blas_contract_B!(C, SV(A_), trivialpermutation(pA), B, pB, pAB, α, β, backend, allocator) + tensorfree!(A_, allocator) + return C +end +function _blas_contract_B!(C, A, pA, B, pB, pAB, α, β, backend, allocator) + isblascontractable(B, pB) && eltype(B) == eltype(C) && + return _blas_contract_C!(C, A, pA, B, pB, pAB, α, β, backend, allocator) + B_ = _blascontractable_copy(B, pB, eltype(C), backend, allocator) + _blas_contract_C!(C, A, pA, SV(B_), trivialpermutation(pB), pAB, α, β, backend, allocator) + tensorfree!(B_, allocator) + return C +end +function _blas_contract_C!(C, A, pA, B, pB, pAB, α, β, backend, allocator) TC = eltype(C) - - A_, pA, flagA = makeblascontractable(A, pA, TC, backend, allocator) - B_, pB, flagB = makeblascontractable(B, pB, TC, backend, allocator) - # the partition of `ipAB` is required by `isblasdestination` and `tensoralloc_add`, but # the actual contraction only needs the linearized permutation ipAB = oindABinC(pAB, pA, pB) ipAB′ = linearize(ipAB) flagC = isblasdestination(C, ipAB) if flagC - C_ = C - _unsafe_blas_contract!(C_, A_, pA, B_, pB, ipAB′, α, β) + _unsafe_blas_contract!(C, A, pA, B, pB, ipAB′, α, β) else C_ = SV(tensoralloc_add(TC, C, ipAB, false, Val(true), allocator)) _unsafe_blas_contract!( - C_, A_, pA, B_, pB, trivialpermutation(ipAB′), + C_, A, pA, B, pB, trivialpermutation(ipAB′), one(TC), zero(TC) ) # `C` already exists, so only the permutation of `pAB` matters here and it can be @@ -67,11 +81,19 @@ function _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator) tensoradd!(C, C_, (pAB, ()), false, α, β, backend, allocator) tensorfree!(C_.parent, allocator) end - flagA || tensorfree!(A_.parent, allocator) - flagB || tensorfree!(B_.parent, allocator) return C end +# Make a permuted copy of `A` with scalar type `TC`, that can be used directly in a BLAS +# contraction with the trivial partition. The copy should be freed afterwards. +# note: inlined such that the wrapper of the copy does not need to be allocated if it does not +# escape the caller, such as for an `Array` that only serves as the parent of a `StridedView` +@inline function _blascontractable_copy(A, pA, ::Type{TC}, backend, allocator) where {TC} + A_ = tensoralloc_add(TC, A, pA, false, Val(true), allocator) + tensoradd!(SV(A_), A, pA, false, One(), Zero(), backend, allocator) + return A_ +end + # perform the actual contraction, assuming it can be done as matrix multiplication by simply # reshaping without any further allocations function _unsafe_blas_contract!(