feat: add support for DLFP8Types (#1491)

* feat: add support for DLFP8Types

* Update Project.toml

Author: avik-pal

Project.toml

@@ -30,6 +30,8 @@ Sockets = "6462fe0b-24de-5631-8697-dd941f90decc"
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
+Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
@@ -43,13 +45,15 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
-[sources.ReactantCore]
-path = "lib/ReactantCore"
+[sources]
+ReactantCore = {path = "lib/ReactantCore"}
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantCUDAExt = ["CUDA", "GPUCompiler", "KernelAbstractions", "LLVM"]
+ReactantDLFP8TypesExt = "DLFP8Types"
+ReactantFloat8sExt = "Float8s"
 ReactantKernelAbstractionsExt = "KernelAbstractions"
 ReactantMPIExt = "MPI"
 ReactantNNlibExt = ["NNlib", "Statistics"]
@@ -67,10 +71,12 @@ Adapt = "4.1"
 ArrayInterface = "7.17.1"
 CEnum = "0.5"
 CUDA = "5.6"
+DLFP8Types = "0.1"
 Downloads = "1.6"
 EnumX = "1"
 Enzyme = "0.13.49"
 EnzymeCore = "0.8.11"
+Float8s = "0.1"
 Functors = "0.5"
 GPUArraysCore = "0.2"
 GPUCompiler = "1.3"

ext/ReactantDLFP8TypesExt.jl

@@ -0,0 +1,11 @@
+module ReactantDLFP8TypesExt
+
+using DLFP8Types: Float8_E4M3FN, Float8_E4M3FNUZ, Float8_E5M2, Float8_E5M2FNUZ
+using Reactant: Reactant
+
+Reactant.reactant_primitive(::Type{Float8_E4M3FN}) = Reactant.F8E4M3FN
+Reactant.reactant_primitive(::Type{Float8_E4M3FNUZ}) = Reactant.F8E4M3FNUZ
+Reactant.reactant_primitive(::Type{Float8_E5M2}) = Reactant.F8E5M2
+Reactant.reactant_primitive(::Type{Float8_E5M2FNUZ}) = Reactant.F8E5M2FNUZ
+
+end

ext/ReactantFloat8sExt.jl

@@ -0,0 +1,8 @@
+module ReactantFloat8sExt
+
+using Float8s: Float8_4
+using Reactant: Reactant
+
+Reactant.reactant_primitive(::Type{Float8_4}) = Reactant.F8E4M3FN
+
+end

test/Project.toml

@@ -2,6 +2,7 @@
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
@@ -34,6 +35,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Adapt = "4.1"
 ArrayInterface = "7.17.1"
 CUDA = "5.5"
+DLFP8Types = "0.1"
 Distributions = "0.25"
 Enzyme = "0.13.28"
 FFTW = "1.8"

test/custom_number_types.jl

@@ -1,38 +1,38 @@
-using Float8s, Reactant
+using Float8s, DLFP8Types, Reactant
 using Reactant: TracedRNumber
 
-Reactant.reactant_primitive(::Type{Float8_4}) = Reactant.F8E4M3FN
+@testset "Custom number types: $(T)" for T in [Float8_4, Float8_E4M3FN, Float8_E4M3FNUZ]
+    x = T[
+        -1.125 -0.21875 1.12
+        1.875 0.4375 1.0
+        0.5625 -1.0 0.937
+        -0.375 -0.34375 -0.6875
+        0.46875 0.75 -0.23437
+        -0.6875 -0.203125 0.375
+        0.875 -0.8125 2.5
+        -0.6875 -0.1171875 -1.625
+        0.75 0.9375 1.0
+        0.5 0.203125 1.75
+    ]
+    x_64 = Float64.(x)
+    x_ra = Reactant.to_rarray(x)
 
-x = Float8_4[
-    -1.125 -0.21875 1.12
-    1.875 0.4375 1.0
-    0.5625 -1.0 0.937
-    -0.375 -0.34375 -0.6875
-    0.46875 0.75 -0.23437
-    -0.6875 -0.203125 0.375
-    0.875 -0.8125 2.5
-    -0.6875 -0.1171875 -1.625
-    0.75 0.9375 1.0
-    0.5 0.203125 1.75
-]
-x_64 = Float64.(x)
-x_ra = Reactant.to_rarray(x)
+    @testset "Reductions" begin
+        sumall(x) = TracedRNumber{Float64}(sum(x))
 
-@testset "Reductions" begin
-    sumall(x) = TracedRNumber{Float64}(sum(x))
+        @test @jit(sumall(x_ra)) ≈ sum(x_64) atol = 1e-1 rtol = 1e-1
 
-    @test @jit(sumall(x_ra)) ≈ sum(x_64) atol = 1e-1 rtol = 1e-1
+        sum1(x) = TracedRNumber{Float64}.(sum(x; dims=1))
+        sum2(x) = TracedRNumber{Float64}.(sum(x; dims=2))
+        sum12(x) = TracedRNumber{Float64}.(sum(x; dims=(1, 2)))
 
-    sum1(x) = TracedRNumber{Float64}.(sum(x; dims=1))
-    sum2(x) = TracedRNumber{Float64}.(sum(x; dims=2))
-    sum12(x) = TracedRNumber{Float64}.(sum(x; dims=(1, 2)))
+        @test @jit(sum1(x_ra)) ≈ sum(x_64; dims=1) atol = 1e-1 rtol = 1e-1
+        @test @jit(sum2(x_ra)) ≈ sum(x_64; dims=2) atol = 1e-1 rtol = 1e-1
+        @test @jit(sum12(x_ra)) ≈ sum(x_64; dims=(1, 2)) atol = 1e-1 rtol = 1e-1
+    end
 
-    @test @jit(sum1(x_ra)) ≈ sum(x_64; dims=1) atol = 1e-1 rtol = 1e-1
-    @test @jit(sum2(x_ra)) ≈ sum(x_64; dims=2) atol = 1e-1 rtol = 1e-1
-    @test @jit(sum12(x_ra)) ≈ sum(x_64; dims=(1, 2)) atol = 1e-1 rtol = 1e-1
-end
-
-@testset "Broadcasting" begin
-    fn(x) = TracedRNumber{Float64}.(x .+ 1)
-    @test @jit(fn(x_ra)) ≈ (x_64 .+ 1) atol = 1e-1 rtol = 1e-1
+    @testset "Broadcasting" begin
+        fn(x) = TracedRNumber{Float64}.(x .+ 1)
+        @test @jit(fn(x_ra)) ≈ (x_64 .+ 1) atol = 1e-1 rtol = 1e-1
+    end
 end