diff --git a/docs/src/man/precompilation.md b/docs/src/man/precompilation.md index a9ab9dbd..a11fcf20 100644 --- a/docs/src/man/precompilation.md +++ b/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 +``` diff --git a/src/implementation/abstractarray.jl b/src/implementation/abstractarray.jl index 017b16f3..77a1dcdf 100644 --- a/src/implementation/abstractarray.jl +++ b/src/implementation/abstractarray.jl @@ -63,9 +63,11 @@ _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 """ @@ -73,9 +75,11 @@ end 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 diff --git a/src/implementation/strided.jl b/src/implementation/strided.jl index c9669fd1..614108d9 100644 --- a/src/implementation/strided.jl +++ b/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 diff --git a/src/indices.jl b/src/indices.jl index 5e5bfeb7..8c8ecb1c 100644 --- a/src/indices.jl +++ b/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) diff --git a/src/interface.jl b/src/interface.jl index bdb5002f..b32ef117 100644 --- a/src/interface.jl +++ b/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]) diff --git a/src/precompile.jl b/src/precompile.jl index 829e27e8..9ce13274 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -1,5 +1,5 @@ -using PrecompileTools: PrecompileTools -using Preferences: @load_preference, load_preference +using PrecompileTools: PrecompileTools, @setup_workload, @compile_workload +using Preferences: @load_preference # Validate preferences input # -------------------------- @@ -49,92 +49,110 @@ const PRECOMPILE_CONTRACT_NDIMS = validate_contract_ndims( @load_preference("precompile_contract_ndims", [4, 2]) ) -# Copy from PrecompileTools.workload_enabled but default to false -function workload_enabled(mod::Module = @__MODULE__) - return try - if load_preference(PrecompileTools, "precompile_workloads", true) - return load_preference(mod, "precompile_workload", false) - else - return false - end - catch - false - end +# Precompilation workload +# ------------------------ +# The workload actually runs representative tensor operations so that PrecompileTools caches +# the specializations that get compiled. Each operation family is factored into a reusable +# `precompile_*` function that can also be called from a downstream package's own +# `@compile_workload` to precompile for a different backend, allocator, or array type. +# +# `@compile_workload` is enabled by default and honors the standard `precompile_workload` +# preference, which can be flipped to disable precompilation: +# +# using TensorOperations, Preferences +# set_preferences!(TensorOperations, "precompile_workload" => false; force=true) + +# Tensor constructor used by the precompile workloads: build a rank-`N` tensor of scalar type +# `T`. Downstream callers can add methods for other array types to reuse the `precompile_*` +# functions below. +precompile_maketensor(T, N) = zeros(T, ntuple(Returns(2), N)) + +""" + precompile_tensoradd(T, N, backend, allocator) + +Run [`tensoradd!`](@ref) and [`tensoralloc_add`](@ref) for scalar type `T` and output rank `N`, +using `backend` and `allocator`, so that their specializations are precompiled. +""" +function precompile_tensoradd( + T, N, backend = DefaultBackend(), allocator = DefaultAllocator() + ) + C = precompile_maketensor(T, N) + A = precompile_maketensor(T, N) + pA = (ntuple(identity, N), ()) + + 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) + + tensoralloc_add(T, A, pA, false, Val(true), allocator) + tensoralloc_add(T, A, pA, false, Val(false), allocator) + return nothing end -# Using explicit precompile statements here instead of @compile_workload: -# Actually running the precompilation through PrecompileTools leads to longer compile times -# Keeping the workload_enabled functionality to have the option of disabling precompilation -# in a compatible manner with the rest of the ecosystem -if workload_enabled() - # tensoradd! - # ---------- - for T in PRECOMPILE_ELTYPES - for N in 0:PRECOMPILE_ADD_NDIMS - C = Array{T, N} - A = Array{T, N} - pA = Index2Tuple{N, 0} - - precompile(tensoradd!, (C, A, pA, Bool, One, Zero)) - precompile(tensoradd!, (C, A, pA, Bool, T, Zero)) - precompile(tensoradd!, (C, A, pA, Bool, T, T)) - - precompile(tensoralloc_add, (T, A, pA, Bool, Val{true})) - precompile(tensoralloc_add, (T, A, pA, Bool, Val{false})) - end - end - - # tensortrace! - # ------------ - for T in PRECOMPILE_ELTYPES - for N1 in 0:PRECOMPILE_TRACE_NDIMS[1], N2 in 0:PRECOMPILE_TRACE_NDIMS[2] - C = Array{T, N1} - A = Array{T, N1 + 2N2} - p = Index2Tuple{N1, 0} - q = Index2Tuple{N2, N2} - - precompile(tensortrace!, (C, A, p, q, Bool, One, Zero)) - precompile(tensortrace!, (C, A, p, q, Bool, T, Zero)) - precompile(tensortrace!, (C, A, p, q, Bool, T, T)) +""" + precompile_tensortrace(T, (N1, N2), backend, allocator) + +Run [`tensortrace!`](@ref) for scalar type `T`, output rank `N1`, and `N2` traced index pairs, +using `backend` and `allocator`, so that their specializations are precompiled. +""" +function precompile_tensortrace( + T, (N1, N2), backend = DefaultBackend(), allocator = DefaultAllocator() + ) + C = precompile_maketensor(T, N1) + A = precompile_maketensor(T, N1 + 2N2) + p = (ntuple(identity, N1), ()) + q = (ntuple(i -> N1 + i, N2), ntuple(i -> N1 + N2 + i, N2)) + + tensortrace!(C, A, p, q, false, One(), Zero(), backend, allocator) + tensortrace!(C, A, p, q, false, one(T), Zero(), backend, allocator) + tensortrace!(C, A, p, q, false, one(T), zero(T), backend, allocator) + + # allocation re-uses tensoralloc_add + return nothing +end - # allocation re-uses tensoralloc_add - end - end +""" + precompile_tensorcontract(T, (N1, N2, N3), backend, allocator) + +Run [`tensorcontract!`](@ref) and [`tensoralloc_contract`](@ref) for scalar type `T`, with `N1` +and `N3` free output indices on the two inputs and `N2` contracted indices, using `backend` and +`allocator`, so that their specializations are precompiled. +""" +function precompile_tensorcontract( + T, (N1, N2, N3), backend = DefaultBackend(), allocator = DefaultAllocator() + ) + NA = N1 + N2 + NB = N2 + N3 + NC = N1 + N3 + C = precompile_maketensor(T, NC) + A = precompile_maketensor(T, NA) + B = precompile_maketensor(T, NB) + pA = (ntuple(identity, N1), ntuple(i -> N1 + i, N2)) + pB = (ntuple(identity, N2), ntuple(i -> N2 + i, N3)) + pAB = (ntuple(identity, NC), ()) + + tensorcontract!(C, A, pA, false, B, pB, false, pAB, One(), Zero(), backend, allocator) + tensorcontract!(C, A, pA, false, B, pB, false, pAB, one(T), Zero(), backend, allocator) + tensorcontract!(C, A, pA, false, B, pB, false, pAB, one(T), zero(T), backend, allocator) + + tensoralloc_contract(T, A, pA, false, B, pB, false, pAB, Val(true), allocator) + tensoralloc_contract(T, A, pA, false, B, pB, false, pAB, Val(false), allocator) + return nothing +end - # tensorcontract! - # --------------- - for T in PRECOMPILE_ELTYPES - for N1 in 0:PRECOMPILE_CONTRACT_NDIMS[1], N2 in 0:PRECOMPILE_CONTRACT_NDIMS[2], - N3 in 0:PRECOMPILE_CONTRACT_NDIMS[1] - - NA = N1 + N2 - NB = N2 + N3 - NC = N1 + N3 - C, A, B = Array{T, NC}, Array{T, NA}, Array{T, NB} - pA = Index2Tuple{N1, N2} - pB = Index2Tuple{N2, N3} - pAB = Index2Tuple{NC, 0} - - precompile(tensorcontract!, (C, A, pA, Bool, B, pB, Bool, pAB, One, Zero)) - precompile(tensorcontract!, (C, A, pA, Bool, B, pB, Bool, pAB, T, Zero)) - precompile(tensorcontract!, (C, A, pA, Bool, B, pB, Bool, pAB, T, T)) - - precompile(tensoralloc_contract, (T, A, pA, Bool, B, pB, Bool, pAB, Val{true})) - precompile(tensoralloc_contract, (T, A, pA, Bool, B, pB, Bool, pAB, Val{false})) +@setup_workload begin + @compile_workload begin + for T in PRECOMPILE_ELTYPES + for N in 0:PRECOMPILE_ADD_NDIMS + precompile_tensoradd(T, N) + end + for N1 in 0:PRECOMPILE_TRACE_NDIMS[1], N2 in 0:PRECOMPILE_TRACE_NDIMS[2] + precompile_tensortrace(T, (N1, N2)) + end + for N1 in 0:PRECOMPILE_CONTRACT_NDIMS[1], N2 in 0:PRECOMPILE_CONTRACT_NDIMS[2], + N3 in 0:PRECOMPILE_CONTRACT_NDIMS[1] + precompile_tensorcontract(T, (N1, N2, N3)) + end end end -else - @info """ - TensorOperations can optionally be instructed to precompile several functions, which can be used to reduce the time to first execution (TTFX). - This is disabled by default as this can take a while on some machines, and is only relevant for contraction-heavy workloads. - - To enable or disable precompilation, you can use the following script: - - ```julia - using TensorOperations, Preferences - set_preferences!(TensorOperations, "precompile_workload" => true; force=true) - ``` - - This will create a `LocalPreferences.toml` file next to your current `Project.toml` file to store this setting in a persistent way. - """ end