fix Nan number bug in matrix multiplication (#17)

* done

* add tests

* add tests

Author: GiggleLiu

Project.toml

@@ -1,7 +1,7 @@
 name = "TropicalGEMM"
 uuid = "a4ad3063-64a7-4bad-8738-34ed09bc0236"
 authors = ["GiggleLiu <cacate0129@gmail.com> and contributors"]
-version = "0.1.6"
+version = "0.1.7"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

README.md

@@ -31,3 +31,9 @@ The benchmark and plotting scripts could be found in the benchmarks folder.
 
 ![Float64](benchmarks/benchmark-float64.png)
 ![Float32](benchmarks/benchmark-float32.png)
+
+
+## Warnings
+
+It is expected to have an ambiguity error when one uses both `TropicalGEMM` and `CUDA`.
+If you see these errors, please include `example/cudapatch.jl` in your project.

examples/cudapatch.jl

@@ -0,0 +1,37 @@
+# you will need this to use Tropical numbers in CUDA.
+using CUDA
+CUDA.allowscalar(false)
+using TropicalGEMM: XTranspose, NativeTypes, Tropical, TropicalTypes, CountingTropical
+using LinearAlgebra
+using Test
+
+const CTranspose{T} = Transpose{T, <:StridedCuVecOrMat{T}}
+for TT in [:(Tropical{<:NativeTypes}), :TropicalTypes]
+    for RT in [TT, :Real]
+        for (TA, CTA) in [(:CuMatrix, :CuMatrix), (:CTranspose, :(Transpose{<:Any, <:StridedCuVecOrMat}))]
+            for (TB, CTB) in [(:CuMatrix, :CuMatrix), (:CTranspose, :(Transpose{<:Any, <:StridedCuVecOrMat}))]
+                @eval function LinearAlgebra.mul!(o::CuMatrix{T}, a::$TA{T}, b::$TB{T}, α::$RT, β::$RT) where {T<:$TT}
+                    CUDA.CUBLAS.gemm_dispatch!(o, a, b, α, β)
+                end
+            end
+        end
+    end
+end
+
+@testset "cuda patch" begin
+    for T in [Tropical{Float64}, CountingTropical{Float64,Float64}]
+        a = T.(CUDA.randn(4, 4))
+        b = T.(CUDA.randn(4))
+        for A in [transpose(a), a, transpose(b)]
+            for B in [transpose(a), a, b]
+                if !(size(A) == (1,4) && size(B) == (4,))
+                    res0 = Array(A) * Array(B)
+                    res1 = A * B
+                    res2 = mul!(CUDA.zeros(T, size(res0)...), A, B, true, false)
+                    @test Array(res1) ≈ res0
+                    @test Array(res2) ≈ res0
+                end
+            end
+        end
+    end
+end

src/gemm.jl

@@ -149,6 +149,11 @@ for TA in [:AbstractMatrix, :XTranspose]
         @eval function LinearAlgebra.mul!(o::AbstractMatrix{T}, a::$TA{T}, b::$TB{T}, α::Number, β::Number) where {T<:Tropical{<:NativeTypes}}
             α = _convert_to_tropical(T, α)
             β = _convert_to_tropical(T, β)
+            if iszero(β)
+                @avx for j=1:size(o, 2), i=1:size(o, 1)
+                    o[i,j] = zero(T)
+                end
+            end
             Octavian.matmul!(o, a, b, α, β)
         end
     end

test/gemm.jl

@@ -1,5 +1,5 @@
 using TropicalGEMM, Octavian, LoopVectorization
-using TropicalNumbers
+using TropicalNumbers, LinearAlgebra
 using VectorizationBase: VecUnroll, Vec
 using TropicalGEMM: naive_mul!
 using Test
@@ -78,4 +78,9 @@ end
 @testset "fix julia-1.5" begin
     x=Tropical(Vec(1.0, 2.0))
     @test VecUnroll((x, x)) === Tropical(VecUnroll((Vec(1.0, 2.0), Vec(1.0, 2.0))))
+end
+
+@testset "fix nan bug" begin
+    res = LinearAlgebra.mul!(Tropical.(fill(NaN, 2, 2)), transpose(Tropical.(randn(2,2))), Tropical.(randn(2,2)), 1, 0)
+    @test !any(isnan, res)
 end