From 4ecbd524a2abdc577607f35dfc816cb5c5459410 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 3 Nov 2026 14:24:59 +0100 Subject: [PATCH] SPIR-V: perform bfloat arithmetic on float, and round conversions once SPV_KHR_bfloat16 only covers storage and conversions: arithmetic, comparisons and math functions on `bfloat` need SPV_INTEL_bfloat16_arithmetic, which isn't a registered extension and which both back-ends require for them. So with only `supports_bfloat16` set, kernels using BFloat16 failed to compile. Widen these operations to `float` in `mcgen`, rounding the result back, unless the new `supports_bfloat16_arithmetic` flag is set. This has to happen there, after `prepare_execution!`, as its InstCombine narrows the widened operations back to `bfloat`. Julia's `DemoteFloat16` pass doesn't fit: it skips fast-math instructions and doesn't handle intrinsics. `fneg`, `fabs` and `copysign` operate on the bits instead, as conversions may change NaNs. `fma` is computed in `double` with round-to-odd, like Julia's codegen does for `fma_float`, so targets without double precision reject it, and `llvm.fmuladd` is unfused. `validate_ir` rejects the intrinsics that can't be widened, such as reductions. SPV_KHR_bfloat16 also defines conversions to `bfloat` as going through `float`, which rounds twice for doubles and for integers wider than its significand. Those are now rounded to `float` with round-to-odd first, as Julia 1.13's codegen does for integers. --- src/spirv.jl | 329 +++++++++++++++++++++++++++++++++- test/spirv/bfloat16.jl | 398 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 724 insertions(+), 3 deletions(-) create mode 100644 test/spirv/bfloat16.jl diff --git a/src/spirv.jl b/src/spirv.jl index b5156f47..f5d72818 100644 --- a/src/spirv.jl +++ b/src/spirv.jl @@ -61,7 +61,13 @@ Base.@kwdef struct SPIRVCompilerTarget <: AbstractCompilerTarget atomics::SPIRVAtomics = SPIRVAtomics() supports_fp16::Bool = true supports_fp64::Bool = true + # `bfloat` values, for storage and conversions (SPV_KHR_bfloat16, which the front-end + # enables in `extensions`) supports_bfloat16::Bool = false + # arithmetic, comparisons and math functions on `bfloat` values (needs + # `supports_bfloat16`, and SPV_INTEL_bfloat16_arithmetic in `extensions`). otherwise, + # these are performed on `float` (see `widen_bfloat16!`). + supports_bfloat16_arithmetic::Bool = false backend::Symbol = isavailable(SPIRV_LLVM_Backend_jll) ? :llvm : :khronos @@ -183,17 +189,22 @@ function validate_ir(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module) if !job.config.target.supports_fp64 append!(errors, check_ir_values(mod, LLVM.DoubleType())) end - if !job.config.target.supports_bfloat16 && isdefined(LLVM, :BFloatType) + bfloat16 = job.config.target.supports_bfloat16 + if !bfloat16 append!(errors, check_ir_values(mod, LLVM.BFloatType())) end - # atomics that `lower_atomics!` cannot lower + # atomics that `lower_atomics!` cannot lower, and `bfloat` operations that + # `widen_bfloat16!` cannot widen for f in mod.functions, bb in f.blocks, inst in bb.instructions reason = if is_atomic_memop(inst) action = spirv_atomic_action(job, inst) action isa String ? action : nothing elseif inst isa LLVM.FenceInst && spirv_syncscope(inst.syncscope.name) === nothing "fence with synchronization scope $(repr(inst.syncscope.name))" + elseif bfloat16 + action = spirv_bfloat16_action(job, inst) + action isa String ? action : nothing else nothing end @@ -304,6 +315,9 @@ end format=LLVM.CodeGenFileType.Assembly) target = job.config.target + # (after `prepare_execution!`, whose InstCombine would narrow the widened operations) + widen_bfloat16!(job, mod) + # The SPIRV Tools don't handle Julia's debug info, rejecting DW_LANG_Julia... strip_debuginfo!(mod) @@ -1011,7 +1025,7 @@ function lower_minimum_maximum!(mod::LLVM.Module) for (call, is_min, quiet) in calls typ = call.value_type eltyp = typ isa LLVM.VectorType ? typ.element_type : typ - bits = if eltyp isa LLVM.HalfType + bits = if eltyp isa LLVM.HalfType || eltyp isa LLVM.BFloatType 16 elseif eltyp isa LLVM.FloatType 32 @@ -1073,6 +1087,315 @@ function lower_minimum_maximum!(mod::LLVM.Module) return changed end +# Unless the target supports arithmetic on `bfloat` values, perform it on `float`, rounding +# the result back to `bfloat`: SPV_KHR_bfloat16 only covers storage and conversions, and the +# back-ends translate arithmetic, comparisons and math functions on `bfloat` to instructions +# that need SPV_INTEL_bfloat16_arithmetic, which isn't a registered extension. As `float` has +# more than twice the precision of `bfloat`, rounding the correctly-rounded `float` result of +# an addition, subtraction, multiplication, division or square root of `bfloat` values gives +# the correctly-rounded `bfloat` one (though OpenCL only requires the first three to be +# correctly rounded by default). `fma` is computed in `double` with round-to-odd, like Julia +# does, so it is not supported without double-precision support. `fneg`, `fabs` and +# `copysign` operate on the bits, as conversions may change NaNs. +# +# This also corrects the conversions to `bfloat` that would round twice: SPV_KHR_bfloat16 +# defines them as converting to `float` first, which is inexact for doubles and for integers +# wider than its significand. Those are rounded to `float` with round-to-odd instead, after +# which rounding to `bfloat` is correct. +# +# Julia's `DemoteFloat16` pass widens arithmetic on CPUs, but leaves fast-math instructions, +# intrinsics and conversions alone. This has to run after `prepare_execution!`, whose +# InstCombine narrows `fptrunc(op(fpext x, fpext y))` back to a `bfloat` operation. +function widen_bfloat16!(@nospecialize(job::CompilerJob{SPIRVCompilerTarget}), + mod::LLVM.Module) + changed = false + job.config.target.supports_bfloat16 || return changed + @tracepoint "widen bfloat16" begin + + worklist = Tuple{LLVM.Instruction,Symbol}[] + unsupported = IRError[] + for f in mod.functions, bb in f.blocks, inst in bb.instructions + action = spirv_bfloat16_action(job, inst) + if action isa Symbol + push!(worklist, (inst, action)) + elseif action isa String + push!(unsupported, (action, backtrace(inst), string(inst))) + end + end + # `validate_ir` rejects these too, but optimization after it can introduce them, and + # validation can be disabled. fail here rather than in the back-ends. + isempty(unsupported) || throw(InvalidIRError(job, unsupported)) + + @dispose builder=IRBuilder() begin + for (inst, action) in worklist + position!(builder, LLVM.before(inst)) + builder.debug_location = inst.debug_location + new = if action === :convert + convert_to_bfloat16!(builder, inst) + else + widen_bfloat16_operation!(builder, inst) + end + name = inst.name + replace_uses!(inst, new) + erase!(inst) + new.name = name + changed = true + end + end + + end + return changed +end + +is_bfloat16(@nospecialize(typ::LLVMType)) = + (typ isa LLVM.VectorType ? typ.element_type : typ) isa LLVM.BFloatType + +# `typ`, with `float` elements instead of `bfloat` ones +function widened_type(@nospecialize(typ::LLVMType)) + is_bfloat16(typ) || return typ + typ isa LLVM.VectorType ? LLVM.VectorType(LLVM.FloatType(), Int(typ.length)) : + LLVM.FloatType() +end + +# the intrinsics on `bfloat` values for which performing them on `float` and rounding the +# result to `bfloat` gives the same result, or one as accurate (for the approximate ones) +const SPIRV_BFLOAT16_WIDENED_INTRINSICS = Set([ + # exact + "llvm.minnum", "llvm.maxnum", "llvm.minimum", "llvm.maximum", "llvm.minimumnum", + "llvm.maximumnum", "llvm.floor", "llvm.ceil", "llvm.trunc", "llvm.rint", "llvm.nearbyint", + "llvm.round", "llvm.roundeven", "llvm.ldexp", "llvm.fptosi.sat", "llvm.fptoui.sat", + "llvm.lround", "llvm.llround", "llvm.lrint", "llvm.llrint", + # correctly rounded + "llvm.sqrt", + # approximate + "llvm.sin", "llvm.cos", "llvm.tan", "llvm.asin", "llvm.acos", "llvm.atan", "llvm.atan2", + "llvm.sinh", "llvm.cosh", "llvm.tanh", "llvm.exp", "llvm.exp2", "llvm.exp10", "llvm.log", + "llvm.log2", "llvm.log10", "llvm.pow", "llvm.powi", + # with their own expansion (see `widen_bfloat16_operation!`) + "llvm.fabs", "llvm.copysign", "llvm.fma", "llvm.fmuladd", "llvm.is.fpclass", +]) + +# How `widen_bfloat16!` handles `inst`: `:widen` to perform it on `float`, `:convert` to +# round its operand to `float` first, `nothing` to leave it alone, or the reason why it is +# not supported (which `validate_ir` reports). +function spirv_bfloat16_action(@nospecialize(job::CompilerJob{SPIRVCompilerTarget}), + inst::LLVM.Instruction) + if inst isa LLVM.FPTruncInst || inst isa LLVM.SIToFPInst || inst isa LLVM.UIToFPInst + is_bfloat16(inst.value_type) || return nothing + typ = first(inst.operands).value_type + typ isa LLVM.VectorType && (typ = typ.element_type) + exact = if typ isa LLVM.IntegerType + typ.width <= 24 + (inst isa LLVM.SIToFPInst) + else + typ isa LLVM.FloatType + end + return exact ? nothing : :convert + end + + job.config.target.supports_bfloat16_arithmetic && return nothing + if inst isa Union{LLVM.FNegInst, LLVM.FAddInst, LLVM.FSubInst, LLVM.FMulInst, + LLVM.FDivInst, LLVM.FRemInst, LLVM.FCmpInst} + return is_bfloat16(first(inst.operands).value_type) ? :widen : nothing + elseif inst isa LLVM.CallInst + is_bfloat16(inst.value_type) || + any(arg -> is_bfloat16(arg.value_type), inst.arguments) || return nothing + callee = inst.called_operand + # (other functions take `bfloat` values like any others) + callee isa LLVM.Function && LLVM.isintrinsic(callee) || return nothing + name = LLVM.name(LLVM.Intrinsic(callee)) + if name == "llvm.fma" && !job.config.target.supports_fp64 + return "call to $(callee.name) (the target supports neither bfloat arithmetic " * + "nor double precision, which correctly rounding it needs)" + end + name in SPIRV_BFLOAT16_WIDENED_INTRINSICS && return :widen + return "call to $(callee.name) on bfloat values " * + "(the target does not support bfloat arithmetic)" + end + return nothing +end + +# perform `inst`, an operation on `bfloat` values, on `float` values +function widen_bfloat16_operation!(builder::IRBuilder, inst::LLVM.Instruction) + widen(val) = is_bfloat16(val.value_type) ? + fpext!(builder, val, widened_type(val.value_type)) : val + typ = inst.value_type + flags = supports_fast_math(inst) ? NamedTuple(inst.fast_math) : nothing + function keep_flags(new) + if flags !== nothing && new isa LLVM.Instruction && supports_fast_math(new) + new.fast_math = flags + end + return new + end + + new = if inst isa LLVM.FCmpInst + lhs, rhs = inst.operands + fcmp!(builder, inst.predicate, widen(lhs), widen(rhs)) + elseif inst isa LLVM.FNegInst + return set_bfloat16_sign!(builder, only(inst.operands), :flip) + elseif !(inst isa LLVM.CallInst) + lhs, rhs = inst.operands + binop!(builder, inst.opcode, widen(lhs), widen(rhs)) + else + callee = inst.called_operand::LLVM.Function + intr = LLVM.Intrinsic(callee) + args = collect(LLVM.Value, inst.arguments) + if intr.name == "llvm.fma" + return fma_bfloat16!(builder, args...) + elseif intr.name == "llvm.fabs" + return set_bfloat16_sign!(builder, only(args), :clear) + elseif intr.name == "llvm.copysign" + return set_bfloat16_sign!(builder, args...) + elseif intr.name == "llvm.fmuladd" + # unfused, like the bfloat `fmul` and `fadd` it may be + x, y, z = args + prod = fptrunc!(builder, keep_flags(fmul!(builder, widen(x), widen(y))), typ) + return fptrunc!(builder, keep_flags(fadd!(builder, widen(prod), widen(z))), typ) + else + # (`fpext` may quiet a signaling NaN, which `llvm.is.fpclass` distinguishes) + exact = intr.name == "llvm.is.fpclass" + args = LLVM.Value[exact ? extend_bfloat16_bits!(builder, arg) : widen(arg) + for arg in args] + ft = LLVM.FunctionType(widened_type(typ), LLVMType[arg.value_type for arg in args]) + mod = builder.insert_block.parent.parent + call!(builder, ft, LLVM.Function(mod, intr, ft), args) + end + end + keep_flags(new) + return is_bfloat16(typ) ? fptrunc!(builder, new, typ) : new +end + +# `val`, `bfloat` values, with the sign bit flipped (`:flip`), cleared (`:clear`), or taken +# from `sign`, leaving the other bits alone +function set_bfloat16_sign!(builder::IRBuilder, val::LLVM.Value, sign) + typ = val.value_type + T_int = typ isa LLVM.VectorType ? LLVM.VectorType(LLVM.Int16Type(), Int(typ.length)) : + LLVM.Int16Type() + bits = bitcast!(builder, val, T_int) + sign_mask = int_constant(T_int, 0x8000) + bits = if sign === :flip + xor!(builder, bits, sign_mask) + else + bits = and!(builder, bits, int_constant(T_int, 0x7fff)) + sign === :clear ? bits : + or!(builder, bits, and!(builder, bitcast!(builder, sign, T_int), sign_mask)) + end + return bitcast!(builder, bits, typ) +end + +# `float` values with the bits of the `bfloat` values `val`, keeping signaling NaNs +function extend_bfloat16_bits!(builder::IRBuilder, val::LLVM.Value) + typ = val.value_type + is_bfloat16(typ) || return val + vec(T) = typ isa LLVM.VectorType ? LLVM.VectorType(T, Int(typ.length)) : T + bits = zext!(builder, bitcast!(builder, val, vec(LLVM.Int16Type())), vec(LLVM.Int32Type())) + bits = shl!(builder, bits, int_constant(vec(LLVM.Int32Type()), 16)) + return bitcast!(builder, bits, widened_type(typ)) +end + +# the integer `val`, as a constant of the integer or integer vector type `typ` +int_constant(typ::LLVM.IntegerType, val::Integer) = LLVM.ConstantInt(typ, val) +int_constant(typ::LLVM.VectorType, val::Integer) = + LLVM.ConstantVector(fill(LLVM.ConstantInt(typ.element_type, val), Int(typ.length))) + +# `fma` on `bfloat` values, computed in `double` (where the product is exact) and rounded +# to odd using an exact TwoSum, so that rounding to `bfloat` afterwards is correct. this is +# what Julia's codegen emits for `fma_float` since 1.13 (`emit_bfloat_fma`). +function fma_bfloat16!(builder::IRBuilder, x::LLVM.Value, y::LLVM.Value, z::LLVM.Value) + typ = x.value_type + vec(T) = typ isa LLVM.VectorType ? LLVM.VectorType(T, Int(typ.length)) : T + T_double = vec(LLVM.DoubleType()) + T_int = vec(LLVM.Int64Type()) + + a, b, c = (fpext!(builder, v, T_double) for v in (x, y, z)) + ab = fmul!(builder, a, b) + res = fadd!(builder, ab, c) + # the exact error of `ab + c` + bb = fsub!(builder, res, ab) + err = fadd!(builder, fsub!(builder, ab, fsub!(builder, res, bb)), fsub!(builder, c, bb)) + # round to odd: if inexact with an even last bit, step towards the exact result + u = bitcast!(builder, res, T_int) + zero = LLVM.null(T_int) + inexact = fcmp!(builder, LLVM.RealPredicate.ONE, err, LLVM.null(T_double)) + even = icmp!(builder, LLVM.IntPredicate.EQ, and!(builder, u, int_constant(T_int, 1)), zero) + away = icmp!(builder, LLVM.IntPredicate.EQ, + icmp!(builder, LLVM.IntPredicate.SLT, bitcast!(builder, err, T_int), zero), + icmp!(builder, LLVM.IntPredicate.SLT, u, zero)) + step = select!(builder, away, int_constant(T_int, 1), int_constant(T_int, -1)) + u = select!(builder, and!(builder, inexact, even), add!(builder, u, step), u) + return convert_to_bfloat16!(builder, bitcast!(builder, u, T_double), typ) +end + +# the conversion `inst` to `bfloat`, rounding its operand to `float` with round-to-odd first +function convert_to_bfloat16!(builder::IRBuilder, inst::LLVM.Instruction) + val = only(inst.operands) + signed = inst isa LLVM.SIToFPInst + return convert_to_bfloat16!(builder, val, inst.value_type, signed) +end +function convert_to_bfloat16!(builder::IRBuilder, val::LLVM.Value, typ::LLVMType, + signed::Bool=false) + # (vectors one element at a time, as the expansions need scalar constants) + if typ isa LLVM.VectorType + res = LLVM.UndefValue(typ) + for i in 0:Int(typ.length)-1 + idx = LLVM.ConstantInt(LLVM.Int32Type(), i) + elem = convert_to_bfloat16!(builder, extract_element!(builder, val, idx), + typ.element_type, signed) + res = insert_element!(builder, res, elem, idx) + end + return res + end + + f = if val.value_type isa LLVM.IntegerType + round_to_odd_float!(builder, val, signed) + else + round_to_odd_float!(builder, val) + end + return fptrunc!(builder, f, typ) +end + +# round the `double` value `x` to a `float` with round-to-odd: towards zero, setting the +# last bit if that was inexact. this keeps enough information for rounding the result to +# `bfloat` (which has less than half the precision) to be correct. +function round_to_odd_float!(builder::IRBuilder, x::LLVM.Value) + T_float = LLVM.FloatType() + T_int = LLVM.Int32Type() + f = fptrunc!(builder, x, T_float) + back = fpext!(builder, f, x.value_type) + mod = builder.insert_block.parent.parent + fabs = LLVM.Function(mod, LLVM.Intrinsic("llvm.fabs"), LLVMType[x.value_type]) + abs(val) = call!(builder, fabs.function_type, fabs, LLVM.Value[val]) + away = fcmp!(builder, LLVM.RealPredicate.OGT, abs(back), abs(x)) + inexact = fcmp!(builder, LLVM.RealPredicate.ONE, back, x) # (false for NaN) + bits = bitcast!(builder, f, T_int) + bits = select!(builder, away, sub!(builder, bits, LLVM.ConstantInt(T_int, 1)), bits) + bits = select!(builder, inexact, or!(builder, bits, LLVM.ConstantInt(T_int, 1)), bits) + return bitcast!(builder, bits, T_float) +end + +# convert the integer `x` to a `float`, rounding to odd: the bits below the 24 most +# significant ones are replaced by a sticky bit, so the conversion to `float` is exact. this +# is what Julia's codegen emits for conversions to `BFloat16` since 1.13 (`emit_inttobfloat`). +function round_to_odd_float!(builder::IRBuilder, x::LLVM.Value, signed::Bool) + typ = x.value_type::LLVM.IntegerType + zero = LLVM.ConstantInt(typ, 0) + one = LLVM.ConstantInt(typ, 1) + neg = signed ? icmp!(builder, LLVM.IntPredicate.SLT, x, zero) : nothing + # the magnitude, as an unsigned integer + m = signed ? select!(builder, neg, sub!(builder, zero, x), x) : x + mod = builder.insert_block.parent.parent + ctlz = LLVM.Function(mod, LLVM.Intrinsic("llvm.ctlz"), LLVMType[typ]) + lz = call!(builder, ctlz.function_type, ctlz, LLVM.Value[m, LLVM.ConstantInt(false)]) + shift = sub!(builder, LLVM.ConstantInt(typ, typ.width - 24), lz) + shift = select!(builder, icmp!(builder, LLVM.IntPredicate.SGT, shift, zero), shift, zero) + sticky = shl!(builder, one, shift) + lowmask = sub!(builder, sticky, one) + inexact = icmp!(builder, LLVM.IntPredicate.NE, and!(builder, m, lowmask), zero) + m = or!(builder, and!(builder, m, not!(builder, lowmask)), + select!(builder, inexact, sticky, zero)) + f = uitofp!(builder, m, LLVM.FloatType()) # exact + return signed ? select!(builder, neg, fneg!(builder, f), f) : f +end + # Align a global or local memory variable to 4 bytes, and pad it to a multiple of 4 bytes, so # that the 32-bit word containing any of its 8- or 16-bit values is part of it. Only internal # definitions can be padded: the size of others is up to the host or the runtime. Returns the diff --git a/test/spirv/bfloat16.jl b/test/spirv/bfloat16.jl new file mode 100644 index 00000000..3fb58b76 --- /dev/null +++ b/test/spirv/bfloat16.jl @@ -0,0 +1,398 @@ +# `bfloat` on targets that support it for storage and conversions only (SPV_KHR_bfloat16), +# where `widen_bfloat16!` performs arithmetic on `float` and rounds conversions once + +# a job for a SPIR-V target with `bfloat` storage, to process textual IR with. `validate` +# runs `spirv-val` on the output. +function bfloat16_job(; backend=:llvm, validate=true, kwargs...) + source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter()) + target = SPIRVCompilerTarget(; backend, supports_bfloat16=true, validate, + extensions="+SPV_KHR_bfloat16", kwargs...) + config = CompilerConfig(target, SPIRV.CompilerParams(); kernel=true) + CompilerJob(source, config) +end + +# a kernel storing the results of `body` (`%r1` to `%r$nresults`), which can use the +# `bfloat` values `%x`, `%y` and `%z`, the `double` `%d` and the integers `%i` and `%l` +function bfloat16_kernel(body; decls="", results=String[]) + stores = join([" store $T %r$i, ptr addrspace(1) %p" for (i, T) in enumerate(results)], + "\n") + """ + target datalayout = "$(GPUCompiler.llvm_datalayout(SPIRVCompilerTarget()))" + $decls + define spir_kernel void @kernel(ptr addrspace(1) %p, bfloat %x, bfloat %y, bfloat %z, + double %d, i32 %i, i64 %l) { + $body + $stores + ret void + } + """ +end + +# the floating-point operations on `bfloat` values in `mod`, which need +# SPV_INTEL_bfloat16_arithmetic +function bfloat16_arithmetic(mod::LLVM.Module) + isbf(T) = (T isa LLVM.VectorType ? T.element_type : T) isa LLVM.BFloatType + found = String[] + for f in mod.functions, bb in f.blocks, inst in bb.instructions + arith = inst isa Union{LLVM.FNegInst, LLVM.FAddInst, LLVM.FSubInst, LLVM.FMulInst, + LLVM.FDivInst, LLVM.FRemInst, LLVM.FCmpInst} || + (inst isa LLVM.CallInst && LLVM.isintrinsic(inst.called_operand)) + if arith && (isbf(inst.value_type) || any(op -> isbf(op.value_type), inst.operands)) + push!(found, string(inst)) + end + end + return found +end + +# widen `body` like `mcgen` does for the job, after lowering the minimum and maximum +# intrinsics like `finish_ir!` does. returns the IR, and the SPIR-V the job's back-end +# compiles it to (validated by `spirv-val`), unless `codegen` is false. +function widen_bfloat16(body; decls="", results=String[], codegen=true, kwargs...) + job = bfloat16_job(; kwargs...) + Context(; opaque_pointers=true) do ctx + mod = parse(LLVM.Module, bfloat16_kernel(body; decls, results)) + GPUCompiler.lower_minimum_maximum!(mod) + GPUCompiler.widen_bfloat16!(job, mod) + verify(mod) + ir = string(mod) + job.config.target.supports_bfloat16_arithmetic || @test bfloat16_arithmetic(mod) == [] + codegen || return ir, nothing + + mod.triple = GPUCompiler.llvm_triple(job.config.target) + # (before Julia 1.12, `mcgen` releases the type-inference lock its caller holds) + locked = VERSION < v"1.12.0-DEV.769" + locked && ccall(:jl_typeinf_lock_begin, Cvoid, ()) + asm = try + GPUCompiler.mcgen(job, mod) + finally + locked && ccall(:jl_typeinf_lock_end, Cvoid, ()) + end + ir, replace(asm, "\r\n" => "\n") + end +end + +for backend in (:khronos, :llvm) +@testset "$backend" begin + +@testset "arithmetic" begin + ir, asm = widen_bfloat16(""" + %r1 = fadd bfloat %x, %y + %a = fmul fast bfloat %r1, %z + %r2 = fneg bfloat %a + %r3 = fcmp olt bfloat %x, %y + %v = insertelement <2 x bfloat> poison, bfloat %x, i32 0 + %r4 = fdiv <2 x bfloat> %v, %v + """; results=["bfloat", "bfloat", "i1", "<2 x bfloat>"], backend) + + @test @filecheck begin + @check "[[X:%[0-9]+]] = fpext bfloat %x to float" + @check "[[Y:%[0-9]+]] = fpext bfloat %y to float" + @check "[[SUM:%[0-9]+]] = fadd float [[X]], [[Y]]" + @check "%r1 = fptrunc float [[SUM]] to bfloat" + @check "fmul fast float" + # (`fneg` on the bits, as conversions may change NaNs) + @check "xor i16 {{%[0-9]+}}, -32768" + @check "fcmp olt float" + @check "fdiv <2 x float>" + @check "fptrunc <2 x float> {{%[0-9]+}} to <2 x bfloat>" + ir + end + + @test @filecheck begin + @check "OpCapability BFloat16TypeKHR" + @check_not "BFloat16Arithmetic" + @check "OpFAdd %float" + @check "OpFConvert %bfloat16" + @check "OpBitwiseXor %ushort" + @check "OpFDiv %v2float" + asm + end +end + +@testset "intrinsics" begin + fns = ["fabs", "copysign", "minnum", "minimum", "maxnum", "floor", "sqrt", "sin", + "fma", "fmuladd"] + arity = Dict("copysign" => 2, "minnum" => 2, "minimum" => 2, "maxnum" => 2, "fma" => 3, + "fmuladd" => 3) + body = join(map(enumerate(fns)) do (i, fn) + args = join(["bfloat %" * "xyz"[j] for j in 1:get(arity, fn, 1)], ", ") + " %r$i = call bfloat @llvm.$fn.bf16($args)" + end, "\n") + decls = join(map(fns) do fn + "declare bfloat @llvm.$fn.bf16($(join(fill("bfloat", get(arity, fn, 1)), ", ")))" + end, "\n") + n = length(fns) + body *= """ + %r$(n+1) = call i1 @llvm.is.fpclass.bf16(bfloat %x, i32 1) + %r$(n+2) = call i32 @llvm.fptosi.sat.i32.bf16(bfloat %x) + %v = insertelement <2 x bfloat> poison, bfloat %x, i32 0 + %r$(n+3) = call <2 x bfloat> @llvm.fabs.v2bf16(<2 x bfloat> %v) + """ + decls *= """ + declare i1 @llvm.is.fpclass.bf16(bfloat, i32) + declare i32 @llvm.fptosi.sat.i32.bf16(bfloat) + declare <2 x bfloat> @llvm.fabs.v2bf16(<2 x bfloat>) + """ + results = [fill("bfloat", n); "i1"; "i32"; "<2 x bfloat>"] + ir, asm = widen_bfloat16(body; decls, results, backend) + + @test @filecheck begin + # (`fabs` and `copysign` on the bits, as conversions may change NaNs) + @check "%[[XB:[0-9]+]] = bitcast bfloat %x to i16" + @check "and i16 %[[XB]], 32767" + @check "%r1 = bitcast i16 {{%[0-9]+}} to bfloat" + @check "and i16 {{%[0-9]+}}, -32768" + @check "%r2 = bitcast i16 {{%[0-9]+}} to bfloat" + # (`lower_minimum_maximum!` fixes up the signs of zeros, comparing `bfloat` values) + @check "call nsz float @llvm.minnum.f32" + @check "fcmp oeq float" + @check "call float @llvm.minimum.f32" + @check "call float @llvm.floor.f32" + @check "call float @llvm.sqrt.f32" + @check "call float @llvm.sin.f32" + # `fma` in `double`, rounded to odd + @check "fmul double" + @check "fadd double" + @check "fptrunc double {{%[0-9]+}} to float" + # `fmuladd` unfused + @check "fmul float" + @check "fptrunc float {{%[0-9]+}} to bfloat" + @check "fadd float" + # the bits of the value, keeping signaling NaNs + @check "bitcast bfloat %x to i16" + @check "shl i32" + @check "call i1 @llvm.is.fpclass.f32" + @check "call i32 @llvm.fptosi.sat.i32.f32" + @check "and <2 x i16> {{%[0-9]+}}, {{splat .i16 32767.|}}" + ir + end + + @test @filecheck begin + @check_not "BFloat16Arithmetic" + @check "OpExtInst %float {{%[0-9]+}} fmin" + @check "OpExtInst %float {{%[0-9]+}} floor" + @check "OpExtInst %float {{%[0-9]+}} sqrt" + asm + end +end + +@testset "conversions" begin + # SPV_KHR_bfloat16 converts through `float`, so conversions from `double` and wide + # integers are rounded to odd first; the others are left alone + ir, asm = widen_bfloat16(""" + %r1 = fptrunc double %d to bfloat + %r2 = sitofp i32 %i to bfloat + %r3 = uitofp i64 %l to bfloat + %s = trunc i32 %i to i16 + %r4 = sitofp i16 %s to bfloat + %f = fpext bfloat %x to float + %r5 = fptrunc float %f to bfloat + %r6 = fpext bfloat %x to double + %r7 = fptosi bfloat %x to i64 + """; results=["bfloat", "bfloat", "bfloat", "bfloat", "bfloat", "double", "i64"], + backend) + + @test @filecheck begin + @check "fptrunc double %d to float" + @check "or i32" + @check "%r1 = fptrunc float {{%[0-9]+}} to bfloat" + @check "call i32 @llvm.ctlz.i32" + @check "uitofp i32 {{%[0-9]+}} to float" + @check "%r2 = fptrunc float {{%[0-9]+}} to bfloat" + @check "call i64 @llvm.ctlz.i64" + @check "uitofp i64 {{%[0-9]+}} to float" + @check "%r3 = fptrunc float {{%[0-9]+}} to bfloat" + @check "%r4 = sitofp i16 %s to bfloat" + @check "%r5 = fptrunc float %f to bfloat" + @check "%r6 = fpext bfloat %x to double" + @check "%r7 = fptosi bfloat %x to i64" + ir + end + + # (vectors one element at a time) + ir, _ = widen_bfloat16(""" + %v = insertelement <2 x double> poison, double %d, i32 0 + %r1 = fptrunc <2 x double> %v to <2 x bfloat> + """; results=["<2 x bfloat>"], backend) + @test count("fptrunc double", ir) == 2 +end + +@testset "native arithmetic" begin + # only the conversions that would round twice change + ir, _ = widen_bfloat16(""" + %r1 = fadd bfloat %x, %y + %r2 = fptrunc double %d to bfloat + """; results=["bfloat", "bfloat"], backend, supports_bfloat16_arithmetic=true, + # (SPIRV-Tools doesn't know SPV_INTEL_bfloat16_arithmetic) + codegen=false) + @test occursin("%r1 = fadd bfloat %x, %y", ir) + @test occursin("fptrunc double %d to float", ir) +end + +@testset "unsupported operations" begin + # e.g., ordered reductions, which would round differently + body = "%r1 = call bfloat @llvm.vector.reduce.fadd.v2bf16(bfloat %x, <2 x bfloat> zeroinitializer)" + decls = "declare bfloat @llvm.vector.reduce.fadd.v2bf16(bfloat, <2 x bfloat>)" + job = bfloat16_job(; backend) + Context(; opaque_pointers=true) do ctx + mod = parse(LLVM.Module, bfloat16_kernel(body; decls, results=["bfloat"])) + errors = GPUCompiler.validate_ir(job, mod) + @test startswith(only(errors)[1], "call to llvm.vector.reduce.fadd.v2bf16 on bfloat values") + + # `fma` without double precision, as it would round twice in `float` + fma = parse(LLVM.Module, bfloat16_kernel( + "%r1 = call bfloat @llvm.fma.bf16(bfloat %x, bfloat %y, bfloat %z)"; + decls="declare bfloat @llvm.fma.bf16(bfloat, bfloat, bfloat)", results=["bfloat"])) + errors = GPUCompiler.validate_ir(bfloat16_job(; backend, supports_fp64=false), fma) + @test startswith(only(errors)[1], "call to llvm.fma.bf16 (the target supports neither") + + # without validation, compilation fails before reaching the back-end + @test_throws GPUCompiler.InvalidIRError GPUCompiler.widen_bfloat16!(job, mod) + + # unless the target supports bfloat arithmetic + native_job = bfloat16_job(; backend, supports_bfloat16_arithmetic=true) + @test isempty(GPUCompiler.validate_ir(native_job, mod)) + end +end + +if isdefined(Core, :BFloat16) +@testset "kernels" begin + mod = @eval module $(gensym()) + const BF16 = Core.BFloat16 + function kernel(p::Ptr{BF16}, a::BF16, d::Float64) + x = unsafe_load(p) + unsafe_store!(p, Base.add_float(Base.mul_float(a, x), x), 1) + unsafe_store!(p, Base.lt_float(x, a) ? x : a, 2) + unsafe_store!(p, Base.abs_float(x), 3) + unsafe_store!(p, Base.fptrunc(BF16, d), 4) + return + end + reduce(p::Ptr{BF16}, x::NTuple{2,VecElement{BF16}}) = + (unsafe_store!(p, ccall("llvm.vector.reduce.fadd.v2bf16", llvmcall, BF16, + (BF16, NTuple{2,VecElement{BF16}}), reinterpret(BF16, 0x0000), x)); + return) + end + BF16 = Core.BFloat16 + + # InstCombine runs right before `mcgen`, so the widening happens in there + @test @filecheck begin + @check "OpCapability BFloat16TypeKHR" + @check_not "BFloat16Arithmetic" + @check "OpFMul %float" + @check "OpFAdd %float" + @check "OpFConvert %bfloat16" + SPIRV.code_native(mod.kernel, Tuple{Ptr{BF16}, BF16, Float64}; backend, kernel=true, + supports_bfloat16=true, extensions="+SPV_KHR_bfloat16") + end + + @test_throws_message(InvalidIRError, + SPIRV.code_execution(mod.reduce, Tuple{Ptr{BF16}, NTuple{2,VecElement{BF16}}}; + backend, supports_bfloat16=true, + extensions="+SPV_KHR_bfloat16")) do msg + occursin("call to llvm.vector.reduce.fadd.v2bf16 on bfloat values", msg) && + occursin(r"\[\d+\] reduce", msg) + end +end +end + +end +end + +# the conversions, evaluated on the host +if isdefined(Core, :BFloat16) && (Sys.ARCH === :x86_64 || Base.libllvm_version >= v"19") +@testset "rounding" begin + # compile `op` on `T` values, widened, to a host function returning the bits of the + # `bfloat` result + # (passing `bfloat` values as `i16`, as the host ABI may not support them) + function widened(op, Ts...) + args = join(["$(T == "bfloat" ? "i16" : T) %$(i-1)" for (i, T) in enumerate(Ts)], ", ") + casts = join(["%a$(i-1) = bitcast i16 %$(i-1) to bfloat" + for (i, T) in enumerate(Ts) if T == "bfloat"], "\n") + ir = Context(; opaque_pointers=true) do ctx + mod = parse(LLVM.Module, """ + declare bfloat @llvm.fma.bf16(bfloat, bfloat, bfloat) + define i16 @entry($args) { + $casts + $op + %bits = bitcast bfloat %r to i16 + ret i16 %bits + }""") + GPUCompiler.widen_bfloat16!(bfloat16_job(), mod) + replace(string(mod), r"^(;|source_filename|target|attributes).*$"m => "") + end + jlT = Dict("double" => Float64, "i32" => Int32, "i64" => Int64, "bfloat" => UInt16) + argtypes = Tuple{(jlT[T] for T in Ts)...} + args = [Symbol("a$i") for i in 1:length(Ts)] + f = @eval ($(args...),) -> Base.llvmcall(($ir, "entry"), UInt16, $argtypes, $(args...)) + bits(x) = x isa Core.BFloat16 ? reinterpret(UInt16, x) : x + (args...) -> Base.invokelatest(f, map(bits, args)...) + end + + # the bits of the `bfloat` nearest to `x` (ties to even), for values in the normal range + function nearest(x) + r = setprecision(BigFloat, 8) do + BigFloat(x) + end + (reinterpret(UInt32, Float32(r)) >> 16) % UInt16 + end + bf16(bits::UInt16) = reinterpret(Core.BFloat16, bits) + bf16(x::Float32) = bf16((reinterpret(UInt32, x) >> 16) % UInt16) # (exact values) + tobig(x::Core.BFloat16) = BigFloat(reinterpret(Float32, UInt32(reinterpret(UInt16, x)) << 16)) + + @testset "double" begin + f = widened("%r = fptrunc double %0 to bfloat", "double") + # values whose nearest `float` is halfway between two `bfloat`s + values = [1 + 2.0^-8 + 2.0^-30, 1 + 2.0^-8 - 2.0^-30, -(1 + 2.0^-8 + 2.0^-30), + 3 * 2.0^100 + 2.0^94 + 2.0^60, 1 + 2.0^-8, 1 + 3 * 2.0^-8] + for _ in 1:1000 + push!(values, randn() * 2.0^rand(-100:100)) + # near halfway between two `bfloat`s + h = (1 + rand(0:127) / 128 + 2.0^-8) * 2.0^rand(-100:100) + push!(values, h + h * 2.0^-rand(25:52) * rand((-1, 1))) + end + for x in values + @test f(x) == nearest(x) + end + # around the largest `bfloat`, and halfway to the next power of two + for x in (0x1.fe8p127, 0x1.ffp127 - 0x1p90, 0x1.ffp127, 0x1.ffp127 + 0x1p90) + @test f(x) == nearest(x) + end + @test f(1e39) == 0x7f80 + @test f(-1e39) == 0xff80 + @test f(Inf) == 0x7f80 + @test f(-0.0) == 0x8000 + @test f(NaN) & 0x7fff > 0x7f80 + end + + @testset "$T" for (T, irT) in ((Int32, "i32"), (Int64, "i64")) + fs = widened("%r = sitofp $irT %0 to bfloat", irT) + fu = widened("%r = uitofp $irT %0 to bfloat", irT) + U = unsigned(T) + values = T[typemin(T), typemax(T), 0, 1, -1, T(2)^30 + T(2)^22 + 1, + -(T(2)^30 + T(2)^22 + 1), rand(T, 10000)...] + for x in values + @test fs(x) == nearest(x) + @test fu(x) == nearest(reinterpret(U, x)) + end + end + + @testset "fma" begin + f = widened("%r = call bfloat @llvm.fma.bf16(bfloat %a0, bfloat %a1, bfloat %a2)", + "bfloat", "bfloat", "bfloat") + # a product halfway between two `bfloat`s, and a tiny addend + values = [(bf16(1.0625f0), bf16(1.0625f0), bf16(2f0^-60)), + (bf16(1.0625f0), bf16(1.0625f0), bf16(-2f0^-60))] + for _ in 1:1000 + # (normal values, of either sign) + a, b, c = (bf16(rand(UInt16) & 0xbfff | 0x3000) for _ in 1:3) + push!(values, (a, b, c)) + end + for (a, b, c) in values + exact = setprecision(BigFloat, 256) do + tobig(a) * tobig(b) + tobig(c) + end + @test f(a, b, c) == nearest(exact) + end + end +end +end