Support Float8 types through DLFP8Types.jl (#36)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: AntonOresten

Project.toml

@@ -1,7 +1,7 @@
 name = "cuTile"
 uuid = "0dea8319-8c4a-4662-a73d-20234d115b9a"
-authors = ["Tim Besard <tim.besard@gmail.com>"]
 version = "0.1.0"
+authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [deps]
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
@@ -12,13 +12,15 @@ IRStructurizer = "93e32bba-5bb8-402b-805d-ffb066edee93"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
 
 [sources]
 CompilerCaching = {url = "https://github.com/maleadt/CompilerCaching.jl", rev="main"}
 IRStructurizer = {url = "https://github.com/maleadt/IRStructurizer.jl", rev = "main"}
 
 [extensions]
 CUDAExt = "CUDA"
+DLFP8TypesExt = "DLFP8Types"
 
 [compat]
 julia = "1.11"

ext/DLFP8TypesExt.jl

@@ -0,0 +1,15 @@
+module DLFP8TypesExt
+
+import cuTile as ct
+
+using DLFP8Types: Float8_E4M3FN, Float8_E5M2
+
+function ct.julia_to_tile_dtype!(table::ct.TypeTable, ::Type{Float8_E4M3FN})
+    return ct.F8E4M3FN(table)
+end
+
+function ct.julia_to_tile_dtype!(table::ct.TypeTable, ::Type{Float8_E5M2})
+    return ct.F8E5M2(table)
+end
+
+end

src/bytecode/types.jl

@@ -127,6 +127,8 @@ BF16(table::TypeTable) = simple_type!(table, SimpleType.BF16)
 F32(table::TypeTable) = simple_type!(table, SimpleType.F32)
 TF32(table::TypeTable) = simple_type!(table, SimpleType.TF32)
 F64(table::TypeTable) = simple_type!(table, SimpleType.F64)
+F8E4M3FN(table::TypeTable) = simple_type!(table, SimpleType.F8E4M3FN)
+F8E5M2(table::TypeTable) = simple_type!(table, SimpleType.F8E5M2)
 Token(table::TypeTable) = simple_type!(table, SimpleType.Token)
 
 function tile_type!(table::TypeTable, dtype::TypeId, shape::AbstractVector{<:Integer})

src/compiler/codegen/values.jl

@@ -146,6 +146,8 @@ function constant_to_bytes(@nospecialize(value), @nospecialize(T::Type))
         return collect(reinterpret(UInt8, [Int32(value)]))
     elseif T === Int64 || T === UInt64
         return collect(reinterpret(UInt8, [Int64(value)]))
+    elseif T === Float16
+        return collect(reinterpret(UInt8, [Float16(value)]))
     elseif T === Float32
         return collect(reinterpret(UInt8, [Float32(value)]))
     elseif T === Float64

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 cuTile = "0dea8319-8c4a-4662-a73d-20234d115b9a"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FileCheck = "4e644321-382b-4b05-b0b6-5d23c3d944fb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/ext/DLFP8TypesExt.jl

@@ -0,0 +1,33 @@
+using DLFP8Types: Float8_E4M3FN, Float8_E5M2
+
+@testset "DLFP8Types extension" begin
+
+spec1d = ct.ArraySpec{1}(16, true)
+
+# Float32 -> Float8_E4M3FN
+@test @filecheck begin
+    @check_label "entry"
+    code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}, ct.TileArray{Float32,1,spec1d}}) do a, b
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        @check "ftof"
+        converted = convert(ct.Tile{Float8_E4M3FN}, tile)
+        ct.store(b, pid, ct.astype(converted, Float32))
+        return
+    end
+end
+
+# Float32 -> Float8_E5M2
+@test @filecheck begin
+    @check_label "entry"
+    code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}, ct.TileArray{Float32,1,spec1d}}) do a, b
+        pid = ct.bid(1)
+        tile = ct.load(a, pid, (16,))
+        @check "ftof"
+        converted = convert(ct.Tile{Float8_E5M2}, tile)
+        ct.store(b, pid, ct.astype(converted, Float32))
+        return
+    end
+end
+
+end

test/runtests.jl

@@ -41,8 +41,6 @@ end
 # Only include executing tests when CUDA is functional
 args = parse_args(ARGS)
 if filter_tests!(testsuite, args)
-    delete!(testsuite, "filecheck")
-
     cuda_functional = CUDA.functional()
     filter!(testsuite) do (test, _)
         if in(test, ["execution"]) || startswith(test, "examples/")