Re-enable precompilation

docs/src/man/precompilation.md

@@ -17,18 +17,19 @@ faster precompile times for fast TTFX for a wider range of inputs.
 
 ## 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:
 
 * `T` is either `Float64` or `ComplexF64`
 * `tensoradd!` is precompiled up to `N = 5`
 * `tensortrace!` is precompiled up to `4` free output indices and `2` pairs of traced indices
 * `tensorcontract!` is precompiled up to `3` free output indices on both inputs, and `2` contracted indices
 
-To enable precompilation with these default settings, you can *locally* change the `"precompile_workload"` key in the preferences.
+To disable precompilation altogether, for example during development or when you prefer to have small binaries,
+you can *locally* change the `"precompile_workload"` key in the preferences.
 
 ```julia
 using TensorOperations, Preferences
-set_preferences!(TensorOperations, "precompile_workload" => true; force=true)
+set_preferences!(TensorOperations, "precompile_workload" => false; force=true)
 ```
 
 ## Custom settings
@@ -43,12 +44,31 @@ set_preferences!(TensorOperations, "setting" => value; force=true)
 
 Here **setting** and **value** can take on the following:
 
-* `"precomple_eltypes"`: a `Vector{String}` that evaluate to the desired values of `T<:Number`
+* `"precompile_eltypes"`: a `Vector{String}` that evaluate to the desired values of `T<:Number`
 * `"precompile_add_ndims"`: an `Int` to specify the maximum `N` for `tensoradd!`
 * `"precompile_trace_ndims"`: a `Vector{Int}` of length 2 to specify the maximal number of free and traced indices for `tensortrace!`.
 * `"precompile_contract_ndims"`: a `Vector{Int}` of length 2 to specify the maximal number of free and contracted indices for `tensorcontract!`.
 
-!!! note "Backends"
+## Reuse in downstream packages
 
-    Currently, there is no support for precompiling methods that do not use the default backend. If this is a
-    feature you would find useful, feel free to contact us or open an issue.
+The default workload is factored into three reusable functions, one per operation family, which
+a downstream package can call from its own `PrecompileTools.@compile_workload` to precompile for a
+different backend, allocator, or array type:
+
+* `TensorOperations.precompile_tensoradd(T, N, backend, allocator)`
+* `TensorOperations.precompile_tensortrace(T, (N1, N2), backend, allocator)`
+* `TensorOperations.precompile_tensorcontract(T, (N1, N2, N3), backend, allocator)`
+
+Each takes a single scalar type `T`, a single index-rank specification, and (optionally) a
+`backend` and `allocator` selecting the implementation to precompile for. The tensors are built
+by `TensorOperations.precompile_maketensor(T, N)`, which returns a `Array{T,N}`. For example, to
+precompile the `tensoradd!` path for a custom backend and allocator:
+
+```julia
+using PrecompileTools, TensorOperations
+@compile_workload begin
+    for T in (Float64, ComplexF64), N in 0:3
+        TensorOperations.precompile_tensoradd(T, N, MyBackend(), MyAllocator())
+    end
+end
+```

src/implementation/abstractarray.jl

@@ -63,19 +63,23 @@ _stridedordiag(A::Diagonal) = A
 
 Check that `C` has `numind(pC)` indices and that `pC` constitutes a valid permutation.
 """
-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)
+    ndims(C) == numind(pC) && isperm(pC) ||
         throw(IndexError(lazy"invalid permutation of length $(ndims(C)): $pC"))
+    return nothing
 end
 
 """
     argcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::Index2Tuple)
 
 Check that `C` and `A` have `numind(pA)` indices and that `pA` constitutes a valid permutation.
 """
-function argcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::Index2Tuple)
+argcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::Index2Tuple) =
+    argcheck_tensoradd(C, A, linearize(pA))
+function argcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::IndexTuple)
     ndims(C) == ndims(A) || throw(IndexError("non-matching number of dimensions"))
-    argcheck_index2tuple(A, pA)
+    argcheck_indextuple(A, pA)
     return nothing
 end
 
@@ -85,14 +89,16 @@ end
 Check that the partial trace of `A` over indices `q` and with permutation of the remaining
 indices `p` is compatible with output `C`.
 """
+argcheck_tensortrace(C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple) =
+    argcheck_tensortrace(C, A, linearize(p), q)
 function argcheck_tensortrace(
-        C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple
+        C::AbstractArray, A::AbstractArray, p::IndexTuple, q::Index2Tuple
     )
     ndims(C) == numind(p) ||
         throw(IndexError(lazy"invalid selection of length $(ndims(C)): $p"))
     2 * numin(q) == 2 * numout(q) == ndims(A) - ndims(C) ||
         throw(IndexError("invalid number of trace dimensions"))
-    argcheck_index2tuple(A, ((p[1]..., q[1]...), (p[2]..., q[2]...)))
+    argcheck_indextuple(A, (p..., q[1]..., q[2]...))
     return nothing
 end
 
@@ -108,7 +114,15 @@ function argcheck_tensorcontract(
         B::AbstractArray, pB::Index2Tuple,
         pAB::Index2Tuple
     )
-    argcheck_index2tuple(C, pAB)
+    return argcheck_tensorcontract(C, A, pA, B, pB, linearize(pAB))
+end
+function argcheck_tensorcontract(
+        C::AbstractArray,
+        A::AbstractArray, pA::Index2Tuple,
+        B::AbstractArray, pB::Index2Tuple,
+        pAB::IndexTuple
+    )
+    argcheck_indextuple(C, pAB)
     argcheck_index2tuple(A, pA)
     argcheck_index2tuple(B, pB)
     numout(pA) + numin(pB) == ndims(C) ||
@@ -123,27 +137,28 @@ end
 
 Check that `C` and `A` have compatible sizes for the addition specified by `pA`.
 """
-function dimcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::Index2Tuple)
+dimcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::Index2Tuple) = dimcheck_tensoradd(C, A, linearize(pA))
+function dimcheck_tensoradd(C::AbstractArray, A::AbstractArray, pA::IndexTuple)
     szA, szC = size(A), size(C)
-    TupleTools.getindices(szA, linearize(pA)) == szC ||
+    TupleTools.getindices(szA, pA) == szC ||
         throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))
     return nothing
 end
 
 """
-    dimcheck_tensorcontract(C::AbstractArray, A::AbstractArray,
-                            p::Index2Tuple, q::Index2Tuple)
+    dimcheck_tensorcontract(C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple)
 
-Check that `C` and `A` have compatible sizes for the trace and addition specified by `p` and
-`q`.
+Check that `C` and `A` have compatible sizes for the trace and addition specified by `p` and `q`.
 """
+dimcheck_tensortrace(C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple) =
+    dimcheck_tensortrace(C, A, linearize(p), q)
 function dimcheck_tensortrace(
-        C::AbstractArray, A::AbstractArray, p::Index2Tuple, q::Index2Tuple
+        C::AbstractArray, A::AbstractArray, p::IndexTuple, q::Index2Tuple
     )
     szA, szC = size(A), size(C)
     TupleTools.getindices(szA, q[1]) == TupleTools.getindices(szA, q[2]) ||
         throw(DimensionMismatch("non-matching sizes in traced dimensions"))
-    TupleTools.getindices(szA, linearize(p)) == szC ||
+    TupleTools.getindices(szA, p) == szC ||
         throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))
     return nothing
 end
@@ -163,11 +178,19 @@ function dimcheck_tensorcontract(
         B::AbstractArray, pB::Index2Tuple,
         pAB::Index2Tuple
     )
+    return dimcheck_tensorcontract(C, A, pA, B, pB, linearize(pAB))
+end
+function dimcheck_tensorcontract(
+        C::AbstractArray,
+        A::AbstractArray, pA::Index2Tuple,
+        B::AbstractArray, pB::Index2Tuple,
+        pAB::IndexTuple
+    )
     szA, szB, szC = size(A), size(B), size(C)
     TupleTools.getindices(szA, pA[2]) == TupleTools.getindices(szB, pB[1]) ||
         throw(DimensionMismatch("non-matching sizes in contracted dimensions"))
     szAB = (TupleTools.getindices(szA, pA[1])..., TupleTools.getindices(szB, pB[2])...)
-    TupleTools.getindices(szAB, linearize(pAB)) == szC ||
+    TupleTools.getindices(szAB, pAB) == szC ||
         throw(DimensionMismatch("non-matching sizes in uncontracted dimensions"))
     return nothing
 end

src/implementation/blascontract.jl

@@ -6,14 +6,13 @@
 function blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
     rpA = reverse(pA)
     rpB = reverse(pB)
-    indCinoBA = let N₁ = numout(pA), N₂ = numin(pB)
-        map(n -> ifelse(n > N₁, n - N₁, n + N₂), linearize(pAB))
+    # 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)
+    rpAB = let N₁ = numout(pA), N₂ = numin(pB)
+        map(n -> ifelse(n > N₁, n - N₁, n + N₂), pAB)
     end
-    tpAB = trivialpermutation(pAB)
-    rpAB = (
-        TupleTools.getindices(indCinoBA, tpAB[1]),
-        TupleTools.getindices(indCinoBA, tpAB[2]),
-    )
     cp = allocator_checkpoint!(allocator)
     if contract_memcost(C, A, pA, B, pB, pAB) <= contract_memcost(C, B, rpB, A, rpA, rpAB)
         C = _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
@@ -28,15 +27,15 @@ function blas_contract!(
         C::StridedView{T, 2},
         A::StridedView{T, 2}, pA::Index2Tuple{1, 1},
         B::StridedView{T, 2}, pB::Index2Tuple{1, 1},
-        pAB::Index2Tuple{1, 1},
+        pAB::IndexTuple{2},
         α::Number, β::Number,
         backend, allocator
     ) where {T}
     A′ = pA == ((1,), (2,)) ? A : transpose(A)
     B′ = pB == ((1,), (2,)) ? B : transpose(B)
-    if pAB == ((1,), (2,))
+    if pAB == (1, 2)
         mul!(C, A′, B′, α, β)
-    elseif pAB == ((2,), (1,))
+    elseif pAB == (2, 1)
         mul!(C, transpose(B′), transpose(A′), α, β)
     end
     return C
@@ -49,18 +48,23 @@ function _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
     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)
         )
-        tensoradd!(C, C_, pAB, false, α, β, backend, allocator)
+        # `C` already exists, so only the permutation of `pAB` matters here and it can be
+        # handed over in the canonical `{N,0}` partition
+        tensoradd!(C, C_, (pAB, ()), false, α, β, backend, allocator)
         tensorfree!(C_.parent, allocator)
     end
     flagA || tensorfree!(A_.parent, allocator)
@@ -74,7 +78,7 @@ function _unsafe_blas_contract!(
         C::StridedView{T},
         A::StridedView{T}, pA,
         B::StridedView{T}, pB,
-        pAB, α, β
+        pAB::IndexTuple, α, β
     ) where {T <: BlasFloat}
     sizeA = size(A)
     sizeB = size(B)
@@ -84,7 +88,7 @@ function _unsafe_blas_contract!(
     osizeB = TupleTools.getindices(sizeB, pB[2])
 
     mul!(
-        sreshape(permutedims(C, linearize(pAB)), (prod(osizeA), prod(osizeB))),
+        sreshape(permutedims(C, pAB), (prod(osizeA), prod(osizeB))),
         sreshape(permutedims(A, linearize(pA)), (prod(osizeA), prod(csizeA))),
         sreshape(permutedims(B, linearize(pB)), (prod(csizeB), prod(osizeB))),
         α, β