Re-enable precompilation - #286
Conversation
Codecov Report❌ Patch coverage is
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
`standardize_scalartype` converts `β` to `scalartype(C)`, which turns `VectorInterface.Zero()` into a plain zero and thereby drops its strong-zero semantics. All kernels but one funnel `β` through `scale`/`scale!`/`Scaler`, which are strong-zero by construction; the `C::Diagonal, A::Diagonal, B::Diagonal` method instead used a raw broadcast, so `0 * C` propagated NaNs from uninitialized output memory. Resolve `conj` into the `StridedView`s at the top level, as the other methods in this file already do, and add `_diagdiagdiagcontract!` with an explicit `iszero(β)` branch following `stridedtensoradd!`. There is no reduction here, so the zero branch passes `nothing` for both `op` and `initop`, which assigns into `C` without ever reading it. This also replaces the broadcast with `Strided._mapreducedim!`, matching the other two diagonal kernels. Document the hazard on `standardize_scalartype` and add regression tests that poison the output with NaNs for all four diagonal `tensorcontract!` methods. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
| ## Defaults | ||
|
|
||
| By default, precompilation is disabled, but can be enabled for "tensors" of type `Array{T,N}`, where `T` and `N` range over the following values: | ||
| By default, precompilation is enabled for "tensors" of type `Array{T,N}`, where `T` and `N` range over the following values: |
There was a problem hiding this comment.
Is there a way to precompile for GPU arrays too?
There was a problem hiding this comment.
From what I understood yes and no, we can add a precompilation module to the package extensions to compile the "bookkeeping code", but from what I understood the actual GPU kernels themselves not yet (although it seemed like Tim was working on this). For cuTENSOR this is probably worth it anyways, as that library is compiled
kshyatt
left a comment
kshyatt
left a comment
There was a problem hiding this comment.
LGTM modulo my doc question (feel free to punt on this)
| function argcheck_index2tuple(C::AbstractArray, pC::Index2Tuple) | ||
| return ndims(C) == numind(pC) && isperm(linearize(pC)) || | ||
| argcheck_index2tuple(C::AbstractArray, pC::Index2Tuple) = argcheck_indextuple(C, linearize(pC)) | ||
| function argcheck_indextuple(C::AbstractArray, pC::IndexTuple) |
There was a problem hiding this comment.
I wonder if such functions need a @noinline, to avoid them being inlined in the Index2Tuple method body, and then yielding the same compilation cost again (as far as I understand how compilation works).
There was a problem hiding this comment.
Same for the other argchecks below
There was a problem hiding this comment.
I think it is probably hard to reason about this, since the alternative is that the Index2Tuple version gets inlined because it is small. From what I can measure this is anyways not really the dominant contribution, so I wanted to avoid doing too much compiler forcing
| # note: `pAB` is only ever consumed through `linearize`/`invperm` from here on, so it is | ||
| # canonicalized to an `IndexTuple` and the reversed permutation does not need to be | ||
| # repartitioned the way `pAB` is | ||
| pAB = linearize(pAB) |
There was a problem hiding this comment.
Would have been nice to maybe also denote this one with a prime, i.e.
| pAB = linearize(pAB) | |
| pAB′ = linearize(pAB) |
(Then also needs to be changed below)
|
|
||
| tensoradd!(C, A, pA, false, One(), Zero(), backend, allocator) | ||
| tensoradd!(C, A, pA, false, one(T), Zero(), backend, allocator) | ||
| tensoradd!(C, A, pA, false, one(T), zero(T), backend, allocator) |
There was a problem hiding this comment.
I wonder if doing conj=true also triggers additional compilation paths, since this will at some point make a StridedView where the conjugation flag is encoded as a type parameter, and these stridedviews go all the way down to blas_contract!.
There was a problem hiding this comment.
At least for the tensoradd case I'm sure that most of this is already compiled, the main reason being that we manually union-split
TensorOperations.jl/src/implementation/strided.jl
Lines 29 to 33 in 8fd0f8f
As a result, both branches are already precompiled because everything is type stable (I think, and while I was playing around with this I did indeed find the method count to be in line with that)
|
Sorry, left some comments after merging. |
3 participants
This closes #226.
The idea is that with the latest versions of Strided, it might be reasonable to just enable precompilation always, hopefully helping out with the CI runtimes in downstream packages like TensorKit and MPSKit.
I am not entirely sure if #208 is still happening, but I'd be willing to give it another try?
Additionally, I've slightly reorganized the
tensoradd/tensortrace/tensorcontractcode to reduce the total number of Strided kernels that have to be generated by canonicalizing the alpha and beta number types, and no longer specializing on unused backend or allocator arguments.I've also found that the
@tensor C[i; k] := A[i; j] * B[j; k]always yieldspAB = ((1,), (2,)), and therefore the fast path for matrix multiplication was actually never taken.