Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Bfloat16 support #446

Draft
wants to merge 5 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ version = "1.4.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"
Expand All @@ -26,11 +27,9 @@ 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]
BFloat16sExt = "BFloat16s"
SpecialFunctionsExt = "SpecialFunctions"

[compat]
Expand Down
14 changes: 0 additions & 14 deletions ext/BFloat16sExt.jl

This file was deleted.

4 changes: 3 additions & 1 deletion lib/mps/MPS.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,9 @@ using ObjectiveC, .Foundation

import GPUArrays

const MtlFloat = Union{Float32, Float16}
using BFloat16s: BFloat16

const MtlFloat = Union{Float32, Float16, BFloat16}

is_supported(dev::MTLDevice) = ccall(:MPSSupportsMTLDevice, Bool, (id{MTLDevice},), dev)

Expand Down
2 changes: 1 addition & 1 deletion lib/mps/matrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ Base.convert(::Type{MPSDataType}, x::Integer) = MPSDataType(x)

# Conversions for MPSDataTypes with Julia equivalents
const jl_mps_to_typ = Dict{MPSDataType, DataType}()
for type in [UInt8,UInt16,UInt32,UInt64,Int8,Int16,Int32,Int64,Float16,Float32,ComplexF16,ComplexF32,Bool]
for type in [:UInt8,:UInt16,:UInt32,:UInt64,:Int8,:Int16,:Int32,:Int64,:Float16,:BFloat16,:Float32,:ComplexF16,:ComplexF32,:Bool]
@eval Base.convert(::Type{MPSDataType}, ::Type{$type}) = $(Symbol(:MPSDataType, type))
@eval jl_mps_to_typ[$(Symbol(:MPSDataType, type))] = $type
end
Expand Down
1 change: 1 addition & 0 deletions src/Metal.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ using ExprTools: splitdef, combinedef
using Artifacts
using ObjectiveC, .CoreFoundation, .Foundation, .Dispatch, .OS
import KernelAbstractions
using BFloat16s

include("version.jl")

Expand Down
3 changes: 2 additions & 1 deletion src/compiler/compilation.jl
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@ function GPUCompiler.finish_ir!(@nospecialize(job::MetalCompilerJob),
# pointer type information for typed intrinsics
# (this is consumed by the LLVM IR downgrader)
for (jltyp, llvmtyp) in (Int32 => :i32, Int64 => :i64,
Float16 => :f16, Float32 => :f32),
Float16 => :f16, Float32 => :f32,
BFloat16 => :bf16),
(as, asname) in (AS.Device => "global", AS.ThreadGroup => "local")

# map of intrinsics to pointer operand indices and eltypes
Expand Down
4 changes: 2 additions & 2 deletions src/device/intrinsics/math.jl
Original file line number Diff line number Diff line change
Expand Up @@ -418,7 +418,7 @@ end
j = fma(1.442695f0, a, 12582912.0f0)
j = j - 12582912.0f0
i = unsafe_trunc(Int32, j)
f = fma(j, -6.93145752f-1, a) # log_2_hi
f = fma(j, -6.93145752f-1, a) # log_2_hi
f = fma(j, -1.42860677f-6, f) # log_2_lo

# approximate r = exp(f)-1 on interval [-log(2)/2, +log(2)/2]
Expand Down Expand Up @@ -460,4 +460,4 @@ end
end

return r
end
end
11 changes: 6 additions & 5 deletions src/device/intrinsics/simd.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ function convert_origin(origin::NTuple{2, Int64})
return (VecElement{Int64}(origin[1]-1), VecElement{Int64}(origin[2]-1))
end

for (jltype, suffix) in ((:Float16, "f16"), (:Float32, "f32"))
for (jltype, suffix) in ((:Float16, "f16"), (:Float32, "f32"), (:BFloat16, "bf16"))
for as in (AS.Device, AS.ThreadGroup)
@eval begin
@device_function simdgroup_load(
Expand Down Expand Up @@ -55,7 +55,7 @@ end
simdgroup_load(data::MtlDeviceArray{T}, matrix_origin=(1, 1))

Loads data from device or threadgroup memory into an 8x8 SIMD-group matrix
and returns it. `T` must be either `Float16` or `Float32`.
and returns it. `T` must be either `Float16`, `Float32`, or `BFloat16`.

# Arguments
- `matrix_origin::NTuple{2, Int64}=(1, 1)`: origin in the source memory to load from.
Expand All @@ -65,7 +65,7 @@ and returns it. `T` must be either `Float16` or `Float32`.
simdgroup_store(src, dest::MtlDeviceArray{T}, matrix_origin=(1, 1))

Stores data from an 8x8 SIMD-group matrix into device or threadgroup memory.
`T` must be either `Float16` or `Float32`.
`T` must be either `Float16`, `Float32`, `BFloat16`.

# Arguments
- `matrix_origin::NTuple{2, Int64}=(1, 1)`: origin in the destination memory to store to.
Expand All @@ -88,6 +88,7 @@ Returns `a * b + c`.

simd_shuffle_map = ((Float32, "f32"),
(Float16, "f16"),
(BFloat16, "bf16"),
(Int32, "s.i32"),
(UInt32, "u.i32"),
(Int16, "s.i16"),
Expand Down Expand Up @@ -118,7 +119,7 @@ The value for delta must be the same for all threads in the SIMD-group. This fun
doesn’t modify the upper delta lanes of data because it doesn’t wrap values around
the SIMD-group.

T must be one of the following: Float32, Float16, Int32, UInt32, Int16, UInt16, Int8, or UInt8
T must be one of the following: Float32, Float16, BFloat16, Int32, UInt32, Int16, UInt16, Int8, or UInt8
"""
simd_shuffle_down

Expand All @@ -131,6 +132,6 @@ lane ID minus delta.
The value of delta must be the same for all threads in a SIMD-group. This function doesn’t
modify the lower delta lanes of data because it doesn’t wrap values around the SIMD-group.

T must be one of the following: Float32, Float16, Int32, UInt32, Int16, UInt16, Int8, or UInt8
T must be one of the following: Float32, Float16, BFloat16, Int32, UInt32, Int16, UInt16, Int8, or UInt8
"""
simd_shuffle_up
15 changes: 9 additions & 6 deletions test/array.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
STORAGEMODES = [Metal.PrivateStorage, Metal.SharedStorage, Metal.ManagedStorage]

const FILL_TYPES = [Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64,
Float16, Float32]
Metal.metal_support() >= v"3.1" && push!(FILL_TYPES, BFloat16)

@testset "array" begin

let arr = MtlVector{Int}(undef, 1)
Expand Down Expand Up @@ -27,8 +31,7 @@ end
@test mtl(1:3) === 1:3


# Page 22 of https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
# Only bfloat missing
# Section 2.1 of https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf
supported_number_types = [Float16 => Float16,
Float32 => Float32,
Float64 => Float32,
Expand All @@ -41,6 +44,8 @@ end
UInt32 => UInt32,
UInt64 => UInt64,
UInt8 => UInt8]
Metal.metal_support() >= v"3.1" && push!(supported_number_types, BFloat16 => BFloat16)

# Test supported types and ensure only Float64 get converted to Float32
for (SrcType, TargType) in supported_number_types
@test mtl(SrcType[1]) isa MtlArray{TargType}
Expand Down Expand Up @@ -227,8 +232,7 @@ end

end

@testset "fill($T)" for T in [Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64,
Float16, Float32]
@testset "fill($T)" for T in FILL_TYPES
broken466a = T ∉ [Int8,UInt8]
broken466b = (Base.JLOptions().check_bounds != 1 || shader_validation)

Expand Down Expand Up @@ -267,8 +271,7 @@ end
end
end

@testset "fill!($T)" for T in [Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64,
Float16, Float32]
@testset "fill!($T)" for T in FILL_TYPES
broken466a = T ∉ [Int8,UInt8]
broken466b = (Base.JLOptions().check_bounds != 1 || shader_validation)

Expand Down
26 changes: 16 additions & 10 deletions test/device/intrinsics.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
using SpecialFunctions
using BFloat16s
using Metal: metal_support

@testset "arguments" begin
Expand Down Expand Up @@ -274,8 +275,9 @@ end
end

@testset "parametrically typed" begin
typs = [Int32, Int64, Float32]
@testset for typ in typs
types = [Int32, Int64, Float32]
metal_support() >= v"3.1" && push!(types, BFloat16)
@testset for typ in types
function kernel(d::MtlDeviceArray{T}, n) where {T}
t = thread_position_in_threadgroup_1d()
tr = n-t+1
Expand Down Expand Up @@ -308,8 +310,9 @@ end
############################################################################################

@testset "simd intrinsics" begin

@testset "shuffle($typ)" for typ in [Float32, Float16, Int32, UInt32, Int16, UInt16, Int8, UInt8]
types = [Float32, Float16, Int32, UInt32, Int16, UInt16, Int8, UInt8]
metal_support() >= v"3.1" && push!(types, BFloat16)
@testset "shuffle($typ)" for typ in types
function kernel(a::MtlDeviceVector{T}, b::MtlDeviceVector{T}) where T
idx = thread_position_in_grid_1d()
idx_in_simd = thread_index_in_simdgroup()
Expand Down Expand Up @@ -344,7 +347,9 @@ end
end

@testset "matrix functions" begin
@testset "load_store($typ)" for typ in [Float16, Float32]
simdgroup_types = [Float16, Float32]
metal_support() >= v"3.1" && push!(simdgroup_types, BFloat16)
@testset "load_store($typ)" for typ in simdgroup_types
function kernel(a::MtlDeviceArray{T}, b::MtlDeviceArray{T},
origin_a=(1, 1), origin_b=(1, 1)) where {T}
sg_a = simdgroup_load(a, origin_a)
Expand All @@ -367,7 +372,7 @@ end
end
end

@testset "load_store_tg($typ)" for typ in [Float16, Float32]
@testset "load_store_tg($typ)" for typ in simdgroup_types
function kernel(a::MtlDeviceArray{T}, b::MtlDeviceArray{T}) where {T}
pos = thread_position_in_threadgroup_2d()

Expand All @@ -391,7 +396,7 @@ end
@test Array(a) == Array(b)
end

@testset "mul($typ)" for typ in [Float16, Float32]
@testset "mul($typ)" for typ in simdgroup_types
function kernel(a::MtlDeviceArray{T}, b::MtlDeviceArray{T}, c::MtlDeviceArray{T}) where {T}
sg_a = simdgroup_load(a)
sg_b = simdgroup_load(b)
Expand All @@ -400,14 +405,15 @@ end
return
end

a = MtlArray(rand(typ, 8, 8))
b = MtlArray(rand(typ, 8, 8))
#Use `ones` for figuring out issues
a = MtlArray(ones(typ, 8, 8))
b = MtlArray(ones(typ, 8, 8))
c = MtlArray(zeros(typ, 8, 8))
@metal threads=(8, 8) kernel(a, b, c)
@test Array(a) * Array(b) ≈ Array(c)
end

@testset "mad($typ)" for typ in [Float16, Float32]
@testset "mad($typ)" for typ in simdgroup_types
function kernel(a::MtlDeviceArray{T}, b::MtlDeviceArray{T}, c::MtlDeviceArray{T},
d::MtlDeviceArray{T}) where {T}
sg_a = simdgroup_load(a)
Expand Down
17 changes: 11 additions & 6 deletions test/mps/linalg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -140,9 +140,11 @@ end
return perm, y
end
end
@testset "$ftype" for ftype in (Float16, Float32)
@testset "$ftype" ftypes = [Float16, Float32]

@testset "$ftype" for ftype in ftypes
# Normal operation
for (shp,k) in [((3,1), 2), ((20,30), 5)]
@testset "$shp, k=$k" for (shp,k) in [((3,1), 2), ((20,30), 5)]
cpu_a = rand(ftype, shp...)

#topk
Expand All @@ -163,11 +165,13 @@ end
@test Array(i) == cpu_i
@test Array(v) == cpu_v
end

# test too big `k`
shp = (20,30)
k = 17

cpu_a = rand(ftype, shp...)
cpu_i, cpu_v = cpu_topk(cpu_a, k)
@testset "$shp, k=$k" begin
cpu_a = rand(ftype, shp...)
cpu_i, cpu_v = cpu_topk(cpu_a, k)

a = MtlMatrix(cpu_a)
@test_throws "MPSMatrixFindTopK does not support values of k > 16" i, v = MPS.topk(a, k)
Expand All @@ -176,7 +180,8 @@ end
i = MtlMatrix{UInt32}(undef, (k, shp[2]))
v = MtlMatrix{ftype}(undef, (k, shp[2]))

@test_throws "MPSMatrixFindTopK does not support values of k > 16" i, v = MPS.topk!(a, i, v, k)
@test_throws "MPSMatrixFindTopK does not support values of k > 16" i, v = MPS.topk!(a, i, v, k)
end
end
end

Expand Down
21 changes: 20 additions & 1 deletion test/runtests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -73,14 +73,33 @@ for (rootpath, dirs, files) in walkdir(@__DIR__)
test_runners[file] = ()->include("$(@__DIR__)/$file.jl")
end
end

## GPUArrays testsuite
const gpuarr_eltypes = [Int16, Int32, Int64,
Complex{Int16}, Complex{Int32}, Complex{Int64},
Float16, Float32,
ComplexF16, ComplexF32]
const gpuarr_eltypes_nobf16 = copy(gpuarr_eltypes)

# don't test BFloat16 for unsupported operations
nobf16_tests = ["random", "reductions/reducedim!",
"reductions/mapreducedim!_large", "reductions/mapreduce",
"reductions/== isequal", "reductions/minimum maximum extrema",
"reductions/sum prod", "reductions/mapreducedim!", "reductions/reduce"]

# Add BFloat16 for tests that use it
Metal.metal_support() >= v"3.1" && push!(gpuarr_eltypes, BFloat16)

for name in keys(TestSuite.tests)
if Metal.DefaultStorageMode != Metal.PrivateStorage && name == "indexing scalar"
# GPUArrays' scalar indexing tests assume that indexing is not supported
continue
end

tmp_eltypes = name in nobf16_tests ? gpuarr_eltypes_nobf16 : gpuarr_eltypes

push!(tests, "gpuarrays$(Base.Filesystem.path_separator)$name")
test_runners["gpuarrays$(Base.Filesystem.path_separator)$name"] = ()->TestSuite.tests[name](MtlArray)
test_runners["gpuarrays$(Base.Filesystem.path_separator)$name"] = ()->TestSuite.tests[name](MtlArray;eltypes=tmp_eltypes)
end
unique!(tests)

Expand Down
10 changes: 2 additions & 8 deletions test/setup.jl
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
using Distributed, Test, Metal, Adapt, ObjectiveC, ObjectiveC.Foundation
using Distributed, Test, Metal, BFloat16s, Adapt, ObjectiveC, ObjectiveC.Foundation

Metal.functional() || error("Metal.jl is not functional on this system")

Expand All @@ -10,12 +10,6 @@ gpuarrays_root = dirname(dirname(gpuarrays))
include(joinpath(gpuarrays_root, "test", "testsuite.jl"))
testf(f, xs...; kwargs...) = TestSuite.compare(f, MtlArray, xs...; kwargs...)

const eltypes = [Int16, Int32, Int64,
Complex{Int16}, Complex{Int32}, Complex{Int64},
Float16, Float32,
ComplexF16, ComplexF32]
TestSuite.supported_eltypes(::Type{<:MtlArray}) = eltypes

const runtime_validation = get(ENV, "MTL_DEBUG_LAYER", "0") != "0"
const shader_validation = get(ENV, "MTL_SHADER_VALIDATION", "0") != "0"

Expand All @@ -32,7 +26,7 @@ function runtests(f, name)
# generate a temporary module to execute the tests in
mod_name = Symbol("Test", rand(1:100), "Main_", replace(name, '/' => '_'))
mod = @eval(Main, module $mod_name end)
@eval(mod, using Test, Random, Metal)
@eval(mod, using Test, Random, Metal, BFloat16s)

let id = myid()
wait(@spawnat 1 print_testworker_started(name, id))
Expand Down