Add documentation

Author: lkdvos

docs/make.jl

@@ -11,7 +11,8 @@ makedocs(; modules=[TensorOperations],
                              "man/interface.md",
                              "man/backends.md",
                              "man/autodiff.md",
-                             "man/implementation.md"],
+                             "man/implementation.md",
+                             "man/precompilation.md"],
                 "Index" => "index/index.md"])
 
 # Documenter can also automatically deploy documentation to gh-pages.

docs/src/index.md

@@ -5,7 +5,7 @@
 ## Table of contents
 
 ```@contents
-Pages = ["index.md", "man/indexnotation.md", "man/functions.md", "man/interface.md", "man/backends.md", "man/autodiff.md", "man/implementation.md"]
+Pages = ["index.md", "man/indexnotation.md", "man/functions.md", "man/interface.md", "man/backends.md", "man/autodiff.md", "man/implementation.md", "man/precompilation.md"]
 Depth = 4
 ```
 

docs/src/man/precompilation.md

@@ -0,0 +1,46 @@
+# Precompilation
+
+Since version `v5.1.5`, TensorOperations.jl has some support for precompiling commonly called functions.
+The guiding philosophy is that often, tensor contractions are (part of) the bottlenecks of typical workflows,
+and as such we want to maximize performance. As a result, we are choosing to specialize many functions which
+may lead to a rather large time-to-first-execution (TTFX). In order to mitigate this, some of that work can
+be moved to precompile-time, avoiding the need to re-compile these specializations for every fresh Julia session.
+
+Nevertheless, TensorOperations is designed to work with a large variety of input types, and simply enumerating
+all of these tends to lead to prohibitively large precompilation times, as well as large system images.
+Therefore, there is some customization possible to tweak the desired level of precompilation, trading in
+faster precompile times for fast TTFX for a wider range of inputs.
+
+## Defaults
+
+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
+
+## Custom settings
+
+The default precompilation settings can be tweaked to allow for more or less expansive coverage. This is achieved
+through a combination of `PrecompileTools`- and `Preferences`-based functionality.
+
+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" => false; force=true)
+```
+
+Alternatively, you can keep precompilation enabled, change the settings above through the same machinery, via:
+
+* `"precomple_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"
+
+    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.

src/precompile.jl

@@ -29,8 +29,8 @@ function validate_trace_ndims(trace_ndims)
 end
 
 function validate_contract_ndims(contract_ndims)
-    contract_ndims isa Vector{Int} && length(contract_ndims) == 3 ||
-        throw(ArgumentError("`precompile_contract_ndims` should be a `Vector{Int}` of length 3, got `$contract_ndims`"))
+    contract_ndims isa Vector{Int} && length(contract_ndims) == 2 ||
+        throw(ArgumentError("`precompile_contract_ndims` should be a `Vector{Int}` of length 2, got `$contract_ndims`"))
     all(≥(0), contract_ndims) ||
         error("Invalid precompile_contract_ndims: `$contract_ndims`")
     return contract_ndims
@@ -45,7 +45,7 @@ const PRECOMPILE_ADD_NDIMS = validate_add_ndims(@load_preference("precompile_add
 const PRECOMPILE_TRACE_NDIMS = validate_trace_ndims(@load_preference("precompile_trace_ndims",
                                                                      [4, 2]))
 const PRECOMPILE_CONTRACT_NDIMS = validate_contract_ndims(@load_preference("precompile_contract_ndims",
-                                                                           [3, 2, 3]))
+                                                                           [4, 2]))
 
 # Using explicit precompile statements here instead of @compile_workload:
 # Actually running the precompilation through PrecompileTools leads to longer compile times
@@ -90,7 +90,7 @@ if PrecompileTools.workload_enabled(@__MODULE__)
     # ---------------
     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[3]
+            N3 in 0:PRECOMPILE_CONTRACT_NDIMS[1]
 
             NA = N1 + N2
             NB = N2 + N3