Make precompilation opt-in (#213)

Author: lkdvos

docs/src/man/precompilation.md

@@ -17,27 +17,31 @@ 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:
+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:
 
 * `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
+To enable precompilation with these default settings, you can *locally* change the `"precompile_workload"` key in the preferences.
+
+```julia
+using TensorOperations, Preferences
+set_preferences!(TensorOperations, "precompile_workload" => true; force=true)
+```
 
-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.
+## Custom settings
 
-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.
+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.
 
 ```julia
 using TensorOperations, Preferences
-set_preferences!(TensorOperations, "precompile_workload" => false; force=true)
+set_preferences!(TensorOperations, "setting" => value; force=true)
 ```
 
-Alternatively, you can keep precompilation enabled, change the settings above through the same machinery, via:
+Here **setting** and **value** can take on the following:
 
 * `"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!`

src/TensorOperations.jl

@@ -75,6 +75,20 @@ const costcache = LRU{Any,Any}(; maxsize=10^5)
 using PackageExtensionCompat
 function __init__()
     @require_extensions
+
+    @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.
+        """ maxlog = 1
 end
 
 include("precompile.jl")

src/precompile.jl

@@ -1,5 +1,5 @@
 using PrecompileTools: PrecompileTools
-using Preferences: @load_preference
+using Preferences: @load_preference, load_preference
 
 # Validate preferences input
 # --------------------------
@@ -47,11 +47,24 @@ const PRECOMPILE_TRACE_NDIMS = validate_trace_ndims(@load_preference("precompile
 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__)
+    try
+        if load_preference(PrecompileTools, "precompile_workloads", true)
+            return load_preference(mod, "precompile_workload", false)
+        else
+            return false
+        end
+    catch
+        false
+    end
+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 PrecompileTools.workload_enabled(@__MODULE__)
+if workload_enabled()
     # tensoradd!
     # ----------
     for T in PRECOMPILE_ELTYPES