BFloat16 extension

Author: christiangnrd

Project.toml

@@ -6,6 +6,7 @@ version = "1.0.0"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
+BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
 CodecBzip2 = "523fee87-0ab8-5b00-afb7-3ecf72e48cfd"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
@@ -28,14 +29,17 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [weakdeps]
+BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
 SpecialFunctionsExt = "SpecialFunctions"
+BFloat16sExt = "BFloat16s"
 
 [compat]
 Adapt = "4"
 Artifacts = "1"
+BFloat16s = "0.5"
 CEnum = "0.4, 0.5"
 CodecBzip2 = "0.8"
 ExprTools = "0.1"

ext/BFloat16sExt.jl

@@ -0,0 +1,12 @@
+module BFloat16sExt
+
+using Metal: MPS.MPSDataType, MPS.MPSDataTypeBFloat16, MPS.jl_mps_to_typ, macos_version
+using BFloat16s
+
+# BFloat is only supported in MPS starting in MacOS 14
+if macos_version() >= v"14"
+    Base.convert(::Type{MPSDataType}, ::Type{BFloat16}) = MPSDataTypeBFloat16
+    jl_mps_to_typ[MPSDataTypeBFloat16] = BFloat16
+end
+
+end # module

lib/mps/matrix.jl

@@ -43,11 +43,7 @@ for type in [UInt8,UInt16,UInt32,UInt64,Int8,Int16,Int32,Int64,Float16,Float32,C
     @eval Base.convert(::Type{MPSDataType}, ::Type{$type}) = $(Symbol(:MPSDataType, type))
     @eval jl_mps_to_typ[$(Symbol(:MPSDataType, type))] = $type
 end
-# BFloat is only supported in MPS starting in MacOS 14
-if macos_version() >= v"14" && isdefined(Core, :BFloat16)
-    Base.convert(::Type{MPSDataType}, ::Type{Core.BFloat16}) = MPSDataTypeBFloat16
-    jl_mps_to_typ[MPSDataTypeBFloat16] = Core.BFloat16
-end
+
 Base.convert(::Type{DataType}, mpstyp::MPSDataType) = jl_mps_to_typ[mpstyp]
 
 

src/Metal.jl

@@ -67,4 +67,8 @@ include("MetalKernels.jl")
 import .MetalKernels: MetalBackend
 export MetalBackend
 
+@static if !isdefined(Base, :get_extension)
+    include("../ext/BFloat16sExt.jl")
+end
+
 end # module

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"