search

Jutho / TensorOperations.jl

Julia package for tensor contractions and related operations
https://jutho.github.io/TensorOperations.jl/stable/
Other
438 stars 55 forks source link

Confusion when using cuTENSOR #166

Open lpawela opened 3 months ago

lpawela commented 3 months ago

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 commented 3 months ago

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 commented 3 months ago

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 commented 3 months ago

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 commented 3 months ago

Ok, this worked. Thank you!