Switch from Requires.jl to extensions (#72)

* switch from Requires.jl to extensions

* update test

* update test

Author: GiggleLiu

Project.toml

@@ -5,7 +5,6 @@ version = "1.3.6"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -19,15 +18,20 @@ Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SIMDTypes = "94e857df-77ce-4151-89e5-788b33177be4"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+GenericTensorNetworksCUDAExt = "CUDA"
+
 [compat]
 AbstractTrees = "0.3, 0.4"
-CUDA = "4"
+CUDA = "4, 5"
 DelimitedFiles = "1"
 DocStringExtensions = "0.8, 0.9"
 FFTW = "1.4"
@@ -37,15 +41,14 @@ Mods = "1.3"
 OMEinsum = "0.7"
 Polynomials = "4"
 Primes = "0.5"
-Requires = "1"
 SIMDTypes = "0.1"
 StatsBase = "0.33, 0.34"
 TropicalNumbers = "0.4, 0.5, 0.6"
-julia = "1"
+julia = "1.9"
 
 [extras]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Documenter"]
+test = ["Test", "Documenter", "CUDA"]

docs/make.jl

@@ -54,6 +54,7 @@ makedocs(;
         "References" => "ref.md",
     ],
     doctest=false,
+    warnonly = :missing_docs,
 )
 
 deploydocs(;

docs/src/ref.md

@@ -33,6 +33,8 @@ get_weights
 chweights
 nflavor
 fixedvertices
+
+extract_result
 ```
 
 #### Graph Problem Utilities
@@ -121,8 +123,6 @@ optimize_code
 getixsv
 getiyv
 contraction_complexity
-timespace_complexity
-timespacereadwrite_complexity
 estimate_memory
 @ein_str
 GreedyMethod

ext/GenericTensorNetworksCUDAExt.jl

@@ -0,0 +1,23 @@
+module GenericTensorNetworksCUDAExt
+
+using CUDA
+using GenericTensorNetworks
+import GenericTensorNetworks: onehotmask!, togpu
+
+# patch
+# Base.ndims(::Base.Broadcast.Broadcasted{CUDA.CuArrayStyle{0}}) = 0
+
+togpu(x::AbstractArray) = CuArray(x)
+
+function onehotmask!(A::CuArray{T}, X::CuArray{T}) where T
+    @assert length(A) == length(X)
+    mask = X .≈ inv.(A)
+    ci = argmax(mask)
+    mask .= false
+    CUDA.@allowscalar mask[ci] = true
+    # set some elements in X to zero to help back propagating.
+    X[(!).(mask)] .= Ref(zero(T))
+    return mask
+end
+
+end

src/GenericTensorNetworks.jl

@@ -79,9 +79,4 @@ include("deprecate.jl")
 include("multiprocessing.jl")
 include("visualize.jl")
 
-using Requires
-function __init__()
-    @require CUDA="052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
-end
-
 end

src/arithematics.jl

@@ -204,22 +204,22 @@ Examples
 ------------------------------
 ```jldoctest; setup=(using GenericTensorNetworks)
 julia> x = ExtendedTropical{3}(Tropical.([1.0, 2, 3]))
-ExtendedTropical{3, TropicalF64}(TropicalF64[1.0ₜ, 2.0ₜ, 3.0ₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[1.0ₜ, 2.0ₜ, 3.0ₜ])
 
 julia> y = ExtendedTropical{3}(Tropical.([-Inf, 2, 5]))
-ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, 2.0ₜ, 5.0ₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[-Infₜ, 2.0ₜ, 5.0ₜ])
 
 julia> x * y
-ExtendedTropical{3, TropicalF64}(TropicalF64[6.0ₜ, 7.0ₜ, 8.0ₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[6.0ₜ, 7.0ₜ, 8.0ₜ])
 
 julia> x + y
-ExtendedTropical{3, TropicalF64}(TropicalF64[2.0ₜ, 3.0ₜ, 5.0ₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[2.0ₜ, 3.0ₜ, 5.0ₜ])
 
 julia> one(x)
-ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, -Infₜ, 0.0ₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[-Infₜ, -Infₜ, 0.0ₜ])
 
 julia> zero(x)
-ExtendedTropical{3, TropicalF64}(TropicalF64[-Infₜ, -Infₜ, -Infₜ])
+ExtendedTropical{3, Tropical{Float64}}(Tropical{Float64}[-Infₜ, -Infₜ, -Infₜ])
 ```
 """
 struct ExtendedTropical{K,TO} <: Number

src/configurations.jl

@@ -27,8 +27,8 @@ function best_solutions(gp::GraphProblem; all=false, usecuda=false, invert=false
     xst = generate_tensors(_x(Tropical{T}; invert), gp)
     ymask = trues(fill(2, length(getiyv(gp.code)))...)
     if usecuda
-        xst = CuArray.(xst)
-        ymask = CuArray(ymask)
+        xst = togpu.(xst)
+        ymask = togpu(ymask)
     end
     if all
         # we use `Float64` as default because we want to support weighted graphs.

src/cuda.jl

@@ -114,7 +114,7 @@ julia> getixsv(gp.code)
  [8, 10]
 
 julia> gp.code(GenericTensorNetworks.generate_tensors(Tropical(1.0), gp)...)
-0-dimensional Array{TropicalF64, 0}:
+0-dimensional Array{Tropical{Float64}, 0}:
 4.0ₜ
 ```
 """
@@ -150,7 +150,7 @@ function contractx(gp::GraphProblem, x; usecuda=false)
     xs = generate_tensors(x, gp)
     @debug "contracting tensors ..."
     if usecuda
-        gp.code([CuArray(x) for x in xs]...)
+        gp.code([togpu(x) for x in xs]...)
     else
         gp.code(xs...)
     end

src/networks/networks.jl

@@ -14,3 +14,6 @@ end
 # a unified interface to optimize the contraction code
 _optimize_code(code, size_dict, optimizer::Nothing, simplifier) = code
 _optimize_code(code, size_dict, optimizer, simplifier) = optimize_code(code, size_dict, optimizer, simplifier)
+
+# upload tensors to GPU
+function togpu end