Merge pull request #26 from TensorBFS/jg/fix-tropicalgemm

Add tropicalgemm and precompile

Author: mroavi

Project.toml

@@ -9,6 +9,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 TropicalGEMM = "a4ad3063-64a7-4bad-8738-34ed09bc0236"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
@@ -18,6 +19,7 @@ CUDA = "4"
 DocStringExtensions = "0.8.6, 0.9"
 OMEinsum = "0.7"
 Requires = "1"
+PrecompileTools = "1"
 TropicalGEMM = "0.1"
 TropicalNumbers = "0.5.4"
 julia = "1.3"

src/TensorInference.jl

@@ -3,6 +3,7 @@ module TensorInference
 using OMEinsum, LinearAlgebra
 using DocStringExtensions, TropicalNumbers
 using Artifacts
+# The Tropical GEMM support
 using TropicalGEMM
 
 # reexport OMEinsum functions
@@ -34,4 +35,13 @@ function __init__()
     @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
 end
 
+import PrecompileTools
+PrecompileTools.@setup_workload begin
+    # Putting some things in `@setup_workload` instead of `@compile_workload` can reduce the size of the
+    # precompile file and potentially make loading faster.
+    PrecompileTools.@compile_workload begin
+        include("../example/asia/asia.jl")
+    end
+end
+
 end # module

src/maxprob.jl

@@ -60,5 +60,5 @@ Returns an output array containing largest log-probabilities.
 function maximum_logp(tn::TensorNetworkModel; usecuda = false)::AbstractArray{<:Real}
     # generate tropical tensors with its elements being log(p).
     tensors = map(t -> Tropical.(log.(t)), adapt_tensors(tn; usecuda, rescale = false))
-    return map(content, tn.code(tensors...))
+    return broadcasted_content(tn.code(tensors...))
 end

src/mmap.jl

@@ -183,7 +183,7 @@ end
 
 function maximum_logp(mmap::MMAPModel; usecuda = false)::AbstractArray{<:Real}
     tensors = map(t -> OMEinsum.asarray(Tropical.(log.(t)), t), adapt_tensors(mmap; usecuda, rescale = false))
-    return map(content, mmap.code(tensors...))
+    return broadcasted_content(mmap.code(tensors...))
 end
 
 function log_probability(mmap::MMAPModel, config::Union{Dict, AbstractVector}; rescale = true, usecuda = false)::Real

src/utils.jl

@@ -192,3 +192,6 @@ function uai_problem_from_file(uai_filepath::String; uai_evid_filepath="", uai_m
     reference_marginals = isempty(uai_mar_filepath) ? Vector{eltype}[] : read_uai_mar_file(uai_mar_filepath)
     return UAIInstance(nvars, ncliques, cards, factors, obsvars, obsvals, reference_marginals)
 end
+
+# patch to get content by broadcasting into array, while keep array size unchanged.
+broadcasted_content(x) = asarray(content.(x), x)

test/cuda.jl

@@ -9,7 +9,7 @@ CUDA.allowscalar(false)
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40))
-    @info contraction_complexity(tn)
+    @debug contraction_complexity(tn)
     @time marginals2 = marginals(tn; usecuda = true)
     @test all(x -> x isa CuArray, marginals2)
     # for dangling vertices, the output size is 1.
@@ -26,10 +26,10 @@ end
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40))
-    @info contraction_complexity(tn)
+    @debug contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn; usecuda = true)
-    @test log_probability(tn, config) ≈ logp.n
+    @test log_probability(tn, config) ≈ logp
     culogp = maximum_logp(tn; usecuda = true)
     @test culogp isa CuArray
     @test Array(culogp)[] ≈ logp
@@ -54,10 +54,10 @@ end
     culogp = maximum_logp(tn2; usecuda = true)
     @test cup isa RescaledArray{T, N, <:CuArray} where {T, N}
     @test culogp isa CuArray
-    @test Array(cup)[] ≈ exp(Array(culogp)[].n)
+    @test Array(cup)[] ≈ exp(Array(culogp)[])
 
     # does not optimize over open vertices
     tn3 = MMAPModel(instance; marginalized = [2, 4, 6], optimizer)
     logp, config = most_probable_config(tn3; usecuda = true)
-    @test log_probability(tn3, config) ≈ logp.n
+    @test log_probability(tn3, config) ≈ logp
 end

test/inference.jl

@@ -63,8 +63,7 @@ end
                                      x -> map(first, x) # get the first capture of each element
 
                 for problem in problems
-                    @show problem
-
+                    @info "Testing: $problem"
                     @testset "$(problem)" begin
                         problem = read_uai_problem(problem)
 
@@ -74,8 +73,7 @@ end
                         if sc > 28
                             error("space complexity too large! got $(sc)")
                         end
-                        # @info(tn)
-                        @info contraction_complexity(tn)
+                        @debug contraction_complexity(tn)
                         marginals2 = marginals(tn)
                         # for dangling vertices, the output size is 1.
                         npass = 0

test/maxprob.jl

@@ -8,7 +8,7 @@ using TensorInference
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 3, niters = 2, βs = 1:0.1:80))
-    @info contraction_complexity(tn)
+    @debug contraction_complexity(tn)
     most_probable_config(tn)
     @time logp, config = most_probable_config(tn)
     @test log_probability(tn, config) ≈ logp

test/mmap.jl

@@ -15,17 +15,17 @@ end
     tn_ref = TensorNetworkModel(instance; optimizer)
     # does not marginalize any var
     mmap = MMAPModel(instance; marginalized = Int[], optimizer)
-    @info(mmap)
+    @debug(mmap)
     @test maximum_logp(tn_ref) ≈ maximum_logp(mmap)
 
     # marginalize all vars
     mmap2 = MMAPModel(instance; marginalized = collect(1:(instance.nvars)), optimizer)
-    @info(mmap2)
+    @debug(mmap2)
     @test Array(probability(tn_ref))[] ≈ exp(maximum_logp(mmap2)[])
 
     # does not optimize over open vertices
     mmap3 = MMAPModel(instance; marginalized = [2, 4, 6], optimizer)
-    @info(mmap3)
+    @debug(mmap3)
     logp, config = most_probable_config(mmap3)
     @test log_probability(mmap3, config) ≈ logp
 end