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
 
@@ -123,27 +129,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

src/implementation/strided.jl

@@ -18,12 +18,20 @@ function tensoradd!(
         α::Number, β::Number,
         backend::StridedBackend, allocator = DefaultAllocator()
     )
+    @nospecialize backend allocator
+
+    # standardize input types for compilation time
+    α′ = standardize_scalartype(C, α)
+    β′ = standardize_scalartype(C, β)
+    p = linearize(pA)
+
     # resolve conj flags and absorb into StridedView constructor to avoid type instabilities later on
     if conjA
-        stridedtensoradd!(SV(C), conj(SV(A)), pA, α, β, backend, allocator)
+        stridedtensoradd!(SV(C), conj(SV(A)), p, α′, β′)
     else
-        stridedtensoradd!(SV(C), SV(A), pA, α, β, backend, allocator)
+        stridedtensoradd!(SV(C), SV(A), p, α′, β′)
     end
+
     return C
 end
 
@@ -33,12 +41,20 @@ function tensortrace!(
         α::Number, β::Number,
         backend::StridedBackend, allocator = DefaultAllocator()
     )
+    @nospecialize backend allocator
+
+    # standardize input types for compilation time
+    α′ = standardize_scalartype(C, α)
+    β′ = standardize_scalartype(C, β)
+    p′ = linearize(p)
+
     # resolve conj flags and absorb into StridedView constructor to avoid type instabilities later on
     if conjA
-        stridedtensortrace!(SV(C), conj(SV(A)), p, q, α, β, backend, allocator)
+        stridedtensortrace!(SV(C), conj(SV(A)), p′, q, α′, β′)
     else
-        stridedtensortrace!(SV(C), SV(A), p, q, α, β, backend, allocator)
+        stridedtensortrace!(SV(C), SV(A), p′, q, α′, β′)
     end
+
     return C
 end
 
@@ -83,40 +99,34 @@ end
 (s::Scaler)(x, y) = scale(x * y, s.α)
 
 function stridedtensoradd!(
-        C::StridedView,
-        A::StridedView, pA::Index2Tuple,
-        α::Number, β::Number,
-        ::StridedBackend, allocator = DefaultAllocator()
+        C::StridedView, A::StridedView, pA::IndexTuple, α::Number, β::Number,
     )
     argcheck_tensoradd(C, A, pA)
     dimcheck_tensoradd(C, A, pA)
-    if !istrivialpermutation(pA) && Base.mightalias(C, A)
+    !istrivialpermutation(pA) && Base.mightalias(C, A) &&
         throw(ArgumentError("output tensor must not be aliased with input tensor"))
+    Ap = permutedims(A, pA)
+    if iszero(β)
+        Strided._mapreducedim!(Scaler(α), nothing, nothing, size(C), (C, Ap))
+    else
+        Strided._mapreducedim!(Scaler(α), Adder(), Scaler(β), size(C), (C, Ap))
     end
-
-    A′ = permutedims(A, linearize(pA))
-    Strided._mapreducedim!(Scaler(α), Adder(), Scaler(β), size(C), (C, A′))
     return C
 end
 
 function stridedtensortrace!(
-        C::StridedView,
-        A::StridedView, p::Index2Tuple, q::Index2Tuple,
-        α::Number, β::Number,
-        ::StridedBackend, allocator = DefaultAllocator()
+        C::StridedView, A::StridedView, p::IndexTuple, q::Index2Tuple, α::Number, β::Number,
     )
     argcheck_tensortrace(C, A, p, q)
     dimcheck_tensortrace(C, A, p, q)
-
     Base.mightalias(C, A) &&
         throw(ArgumentError("output tensor must not be aliased with input tensor"))
-
-    sizeA = i -> size(A, i)
-    strideA = i -> stride(A, i)
-    tracesize = sizeA.(q[1])
-    newstrides = (strideA.(linearize(p))..., (strideA.(q[1]) .+ strideA.(q[2]))...)
-    newsize = (size(C)..., tracesize...)
-
+    newsize = linearize(size(C), TupleTools.getindices(size(A), q[1]))
+    stA = strides(A)
+    newstrides = linearize(
+        TupleTools.getindices(stA, p),
+        TupleTools.getindices(stA, q[1]) .+ TupleTools.getindices(stA, q[2])
+    )
     A′ = SV(A.parent, newsize, newstrides, A.offset, A.op)
     Strided._mapreducedim!(Scaler(α), Adder(), Scaler(β), newsize, (C, A′))
     return C

src/indices.jl

@@ -18,9 +18,11 @@ and `N₂` right indices.
 const Index2Tuple{N₁, N₂} = Tuple{IndexTuple{N₁}, IndexTuple{N₂}}
 
 linearize(p::Index2Tuple) = (p[1]..., p[2]...)
+linearize(a::Tuple, b::Tuple) = (a..., b...)
 numout(p::Index2Tuple) = length(p[1])
 numin(p::Index2Tuple) = length(p[2])
 numind(p::Index2Tuple) = numout(p) + numin(p)
+numind(p::IndexTuple) = length(p)
 
 trivialpermutation(p::IndexTuple{N}) where {N} = ntuple(identity, Val(N))
 function trivialpermutation(p::Index2Tuple)

src/interface.jl

@@ -246,6 +246,17 @@ Obtain the type information of `C`, where `C` would be the output of
 """
 function tensorcontract_type end
 
+"""
+    standardize_scalartype(C, α::Number) -> α′
+
+Convert the scalar `α` into a standardized type suitable for combining with the output tensor `C`.
+This hook can be used to reduce the number of compiled specializations of some kernels,
+by ensuring that the scalars always arrive in a canonical type.
+"""
+standardize_scalartype(C, α::Number) = _standardize_scalartype(scalartype(C), α)
+_standardize_scalartype(::Type{T}, α::Number) where {T <: Number} = convert(T, α)
+_standardize_scalartype(::Type, α::Number) = α      # ensure polynomials and symbolic types are left alone
+
 """
     tensoralloc(ttype, structure, [istemp=false, allocator])