Confusion when using cuTENSOR

[lpawela]

I have some code using @cutensor. I have some problems with importing cuTENSOR. In my package I have

module SpinGlassTensors
using cuTENSOR
using CUDA, CUDA.CUSPARSE
using NNlib
using LinearAlgebra, MKL
using TensorOperations
# rest of module

When I run julia --project and run ] test everything works correctly. If I setup a new env with an empty Project.toml

mkdir new_env
touch new_env/Project.toml
cd new_env
julia --project

and run ] add <path_to_package> and after that precompile I get

Precompiling SpinGlassTensors
        Info Given SpinGlassTensors was explicitly requested, output will be shown live 
ERROR: LoadError: ArgumentError: cuTENSOR not loaded: add `using cuTENSOR` or `import cuTENSOR` before using `@cutensor`
Stacktrace:
 [1] var"@cutensor"(__source__::LineNumberNode, __module__::Module, ex::Expr)
   @ TensorOperations ~/.julia/packages/TensorOperations/LAzcX/src/indexnotation/tensormacros.jl:300
 [2] include(mod::Module, _path::String)
   @ Base ./Base.jl:495
 [3] include(x::String)
   @ SpinGlassTensors ~/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:2
 [4] top-level scope
   @ ~/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:27
 [5] include
   @ ./Base.jl:495 [inlined]
 [6] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2222
 [7] top-level scope
   @ stdin:3
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/mps/utils.jl:46
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/mps/utils.jl:43
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:2
in expression starting at stdin:3
  ✗ SpinGlassTensors
  0 dependencies successfully precompiled in 10 seconds. 81 already precompiled.

ERROR: The following 1 direct dependency failed to precompile:

SpinGlassTensors [7584fc6a-5a23-4eeb-8277-827aab0146ea]

Failed to precompile SpinGlassTensors [7584fc6a-5a23-4eeb-8277-827aab0146ea] to "/home/lpawela/.julia/compiled/v1.10/SpinGlassTensors/jl_de5whs".
ERROR: LoadError: ArgumentError: cuTENSOR not loaded: add `using cuTENSOR` or `import cuTENSOR` before using `@cutensor`
Stacktrace:
 [1] var"@cutensor"(__source__::LineNumberNode, __module__::Module, ex::Expr)
   @ TensorOperations ~/.julia/packages/TensorOperations/LAzcX/src/indexnotation/tensormacros.jl:300
 [2] include(mod::Module, _path::String)
   @ Base ./Base.jl:495
 [3] include(x::String)
   @ SpinGlassTensors ~/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:2
 [4] top-level scope
   @ ~/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:27
 [5] include
   @ ./Base.jl:495 [inlined]
 [6] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
   @ Base ./loading.jl:2222
 [7] top-level scope
   @ stdin:3
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/mps/utils.jl:46
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/mps/utils.jl:43
in expression starting at /home/lpawela/.julia/packages/SpinGlassTensors/O6rmo/src/SpinGlassTensors.jl:2
in expression starting at stdin:
Stacktrace:
  [1] pkgerror(msg::String)
    @ Pkg.Types ~/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/share/julia/stdlib/v1.10/Pkg/src/Types.jl:70
  [2] precompile(ctx::Pkg.Types.Context, pkgs::Vector{…}; internal_call::Bool, strict::Bool, warn_loaded::Bool, already_instantiated::Bool, timing::Bool, _from_loading::Bool, kwargs::@Kwargs{…})
    @ Pkg.API ~/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/share/julia/stdlib/v1.10/Pkg/src/API.jl:1659
  [3] precompile(pkgs::Vector{Pkg.Types.PackageSpec}; io::Base.TTY, kwargs::@Kwargs{_from_loading::Bool})
    @ Pkg.API ~/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/share/julia/stdlib/v1.10/Pkg/src/API.jl:159
  [4] precompile
    @ ~/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/share/julia/stdlib/v1.10/Pkg/src/API.jl:147 [inlined]
  [5] #precompile#114
    @ ~/.julia/juliaup/julia-1.10.2+0.x64.linux.gnu/share/julia/stdlib/v1.10/Pkg/src/API.jl:146 [inlined]
  [6] #invokelatest#2
    @ ./essentials.jl:894 [inlined]
  [7] invokelatest
    @ ./essentials.jl:889 [inlined]
  [8] _require(pkg::Base.PkgId, env::String)
    @ Base ./loading.jl:1963
  [9] __require_prelocked(uuidkey::Base.PkgId, env::String)
    @ Base ./loading.jl:1812
 [10] #invoke_in_world#3
    @ ./essentials.jl:926 [inlined]
 [11] invoke_in_world
    @ ./essentials.jl:923 [inlined]
 [12] _require_prelocked(uuidkey::Base.PkgId, env::String)
    @ Base ./loading.jl:1803
 [13] macro expansion
    @ ./loading.jl:1790 [inlined]
 [14] macro expansion
    @ ./lock.jl:267 [inlined]
 [15] __require(into::Module, mod::Symbol)
    @ Base ./loading.jl:1753
 [16] #invoke_in_world#3
    @ ./essentials.jl:926 [inlined]
 [17] invoke_in_world
    @ ./essentials.jl:923 [inlined]
 [18] require(into::Module, mod::Symbol)
    @ Base ./loading.jl:1746
Some type information was truncated. Use `show(err)` to see complete types.

Obviously, cuTENSOR is not visible to the macro. This causes issues with publishing the package in the Julia registry. Am I missing something?

[lkdvos]

The setup is done through a package extension for both cuTENSOR and CUDA, such that it only gets loaded if there has been an import or using statement for both of these at some point.

My best guess is to just add import cuTENSOR (and thus add this as a dependency to your package)

[lpawela]

Ok, I somehow didn't paste the using cuTENSOR line. I edited my original post. Now I am even more confused. The things I described above happen in Julia 1.9 on my local machine, on 1.10 it seems ok. Yet, during the automerge run in Julia Registry I get

  ✓ Strided
753
  ✓ NNlib → NNlibCUDAExt
754
  ✓ TensorOperations
755
  ✓ cuTENSOR
756
  ✓ TensorOperations → TensorOperationscuTENSORExt
757
  ✓ TensorOperations → TensorOperationsChainRulesCoreExt
758
  ✗ SpinGlassTensors
759
  82 dependencies successfully precompiled in 88 seconds
761
ERROR: The following 1 direct dependency failed to precompile:
762
SpinGlassTensors [7584fc6a-5a23-4eeb-8277-827aab0146ea]
763
Failed to precompile SpinGlassTensors [7584fc6a-5a23-4eeb-8277-827aab0146ea] to "/tmp/jl_pIyevc/compiled/v1.10/SpinGlassTensors/jl_MqPj3t".
764
ERROR: LoadError: ArgumentError: cuTENSOR not loaded: add `using cuTENSOR` or `import cuTENSOR` before using `@cutensor`
765
Stacktrace:
766
 [1] var"@cutensor"(__source__::LineNumberNode, __module__::Module, ex::Expr)
767
   @ TensorOperations /tmp/jl_pIyevc/packages/TensorOperations/LAzcX/src/indexnotation/tensormacros.jl:300
768
 [2] include(mod::Module, _path::String)
769
   @ Base ./Base.jl:495
770
 [3] include(x::String)
771
   @ SpinGlassTensors /tmp/jl_pIyevc/packages/SpinGlassTensors/1PbPU/src/SpinGlassTensors.jl:1
772
 [4] top-level scope
773
   @ /tmp/jl_pIyevc/packages/SpinGlassTensors/1PbPU/src/SpinGlassTensors.jl:26
774
 [5] include
775
   @ ./Base.jl:495 [inlined]
776
 [6] include_package_for_output(pkg::Base.PkgId, input::String, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String}, concrete_deps::Vector{Pair{Base.PkgId, UInt128}}, source::Nothing)
777
   @ Base ./loading.jl:2222
778
 [7] top-level scope
779
   @ stdin:3
780
in expression starting at /tmp/jl_pIyevc/packages/SpinGlassTensors/1PbPU/src/mps/utils.jl:46
781
in expression starting at /tmp/jl_pIyevc/packages/SpinGlassTensors/1PbPU/src/mps/utils.jl:43
782
in expression starting at /tmp/jl_pIyevc/packages/SpinGlassTensors/1PbPU/src/SpinGlassTensors.jl:1
783
in expression starting at stdin:
784
Stacktrace:
785
 [1] pkgerror(msg::String)
786
   @ Pkg.Types /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/Types.jl:70
787
 [2] precompile(ctx::Pkg.Types.Context, pkgs::Vector{Pkg.Types.PackageSpec}; internal_call::Bool, strict::Bool, warn_loaded::Bool, already_instantiated::Bool, timing::Bool, _from_loading::Bool, kwargs::@Kwargs{io::IOContext{Base.PipeEndpoint}})
788
   @ Pkg.API /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/API.jl:1659
789
 [3] precompile(pkgs::Vector{Pkg.Types.PackageSpec}; io::IOContext{Base.PipeEndpoint}, kwargs::@Kwargs{})
790
   @ Pkg.API /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/API.jl:159
791
 [4] precompile(pkgs::Vector{Pkg.Types.PackageSpec})
792
   @ Pkg.API /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/API.jl:148
793
 [5] precompile(; name::Nothing, uuid::Nothing, version::Nothing, url::Nothing, rev::Nothing, path::Nothing, mode::Pkg.Types.PackageMode, subdir::Nothing, kwargs::@Kwargs{})
794
   @ Pkg.API /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/API.jl:174
795
 [6] precompile()
796
   @ Pkg.API /opt/hostedtoolcache/julia/1.10.2/x64/share/julia/stdlib/v1.10/Pkg/src/API.jl:165
797
 [7] top-level scope
798
   @ none:17

[lkdvos]

I'm honestly not entirely sure what is going wrong here. The only thing I can think of is that the check I wrote is expanded before cuTENSOR is being loaded because it happens at compile time.

Could you try and replace the @cutensor ex... calls with @tensor backend=cuTENSOR allocator=cuTENSOR ex...? This should effectively bypass the check that cuTENSOR is loaded. If this works, I'll have to come up with a different way of ensuring the dependencies are loaded.

[lpawela]

Ok, this worked. Thank you!

[lkdvos]

This should be fixed as of version v5. In particular, it should no longer be necessary to explicitly specify the backend and allocator in the way I suggested before, and @cutensor should now just work even in the environment of a package.