diff --git a/docs/src/man/element_types.md b/docs/src/man/element_types.md index 46680e81..7153989b 100644 --- a/docs/src/man/element_types.md +++ b/docs/src/man/element_types.md @@ -32,7 +32,24 @@ reductions or scans. Operations that need those reject a numeric float up front with an error, rather than letting it fail deeper down in `tileiras`. To compute with such values -element-wise, convert to an arithmetic float first. +element-wise, convert to an arithmetic float first. For example, `x .+ y` on two +`Float8_E4M3FN` tiles fails at kernel compile time with *"operations on a +restricted float element type are not supported"*; casting first is what makes +the intermediate precision explicit: + +```julia +f32(t) = convert(ct.Tile{Float32}, t) +sum = f32(x) .+ f32(y) +``` + +Comparisons and `ifelse` selection are the exceptions, and stay available: +comparisons upcast losslessly, having no result to round, and selection leaves +the values themselves alone. + +The rejection is per operation, not per function: applying a custom function +element-wise (`map`, or broadcasting a lambda) over a numeric-float tile is +rejected as well, even when every step inside it is a cast. Convert the tile +with `convert(ct.Tile{T}, tile)` or `T.(tile)` instead. This is why a `Float32` matmul that wants tensor cores converts its *operands* to `TFloat32` while leaving the accumulator `Float32`: the operands only ever diff --git a/ext/DLFP8TypesExt.jl b/ext/DLFP8TypesExt.jl index 841b44c3..c6fd517a 100644 --- a/ext/DLFP8TypesExt.jl +++ b/ext/DLFP8TypesExt.jl @@ -1,6 +1,7 @@ module DLFP8TypesExt import cuTile as ct +import DLFP8Types using DLFP8Types: Float8_E4M3FN, Float8_E5M2 @@ -42,4 +43,13 @@ for F8 in FP8Types end end +# FP8 is a storage / tensor-core operand format, not an arithmetic type: the +# Tile IR elementwise float ops only accept f16/bf16/f32/f64. Registered for +# every `FP8` subtype, not just the two with a Tile IR dtype, so that the +# broadcast/map gate and the tile-level guards recognize them all. The gate +# rejects the operations up front, so DLFP8Types' own scalar implementations — +# the Float32 round-trip arithmetic and the bit-level comparisons — are never +# consulted in a kernel; host-side FP8 stays untouched. +ct.is_restricted_float(::Type{<:DLFP8Types.FP8}) = true + end diff --git a/ext/MicrofloatsExt.jl b/ext/MicrofloatsExt.jl index e01ed539..8ce20d56 100644 --- a/ext/MicrofloatsExt.jl +++ b/ext/MicrofloatsExt.jl @@ -56,4 +56,13 @@ for MF in MicrofloatTypes end end +# Microfloats are storage / tensor-core operand formats, not arithmetic types: +# the Tile IR elementwise float ops only accept f16/bf16/f32/f64. Registered for +# every `Microfloat`, not just the four with a Tile IR dtype, so that the +# broadcast/map gate and the tile-level guards recognize them all. The gate +# rejects the operations up front, so Microfloats' own scalar implementations — +# arithmetic as a Float32 round-trip (src/ops.jl) — are never consulted in a +# kernel; host-side microfloat arithmetic stays untouched. +ct.is_restricted_float(::Type{<:Microfloats.Microfloat}) = true + end diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl index 225f122d..31f7d2bb 100644 --- a/src/compiler/intrinsics/core.jl +++ b/src/compiler/intrinsics/core.jl @@ -812,10 +812,12 @@ function emit_reduce!(ctx::CGCtx, args) for (k, tv) in enumerate(tile_tvs) etype = eltype(CC.widenconst(tv.jltype)) - # Restricted floats (TFloat32, future FP8/FP4) lack the arithmetic - # support that reduce body subprograms typically require, so reject - # them at the SCI boundary with a clear error. - is_restricted_float(etype) && + # Restricted floats (TFloat32, FP8/FP4 from the extensions) lack the + # arithmetic support that reduce body subprograms typically require, so + # reject them at the SCI boundary with a clear error. Resolved in the + # latest world for the same reason as `lookup_dtype!`: extensions + # register their types after the pipeline's world was frozen. + Base.invokelatest(is_restricted_float, etype)::Bool && throw(IRError("reduce: element type $etype is a restricted float and unsupported")) push!(elem_types, etype) dtype = lookup_dtype!(tt, etype) @@ -990,7 +992,8 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.scan), args) for (k, tv) in enumerate(tile_tvs) etype = eltype(CC.widenconst(tv.jltype)) - is_restricted_float(etype) && + # latest-world lookup, see `emit_reduce!` + Base.invokelatest(is_restricted_float, etype)::Bool && throw(IRError("scan: element type $etype is a restricted float and unsupported")) push!(elem_types, etype) dtype = lookup_dtype!(tt, etype) diff --git a/src/language/arithmetic.jl b/src/language/arithmetic.jl index 69582283..029981c8 100644 --- a/src/language/arithmetic.jl +++ b/src/language/arithmetic.jl @@ -90,13 +90,43 @@ conventions: the remainder takes the sign of the divisor. @inline divmod(x::Tile{T,S}, y::Tile{T,S}) where {T<:Integer, S} = (div.(x, y, RoundDown), mod.(x, y)) +""" + check_arithmetic(T) + +Reject arithmetic on a restricted float element type `T` (`TFloat32`, and the +FP8/FP4 types registered by the package extensions): the elementwise float +intrinsics only accept f16/bf16/f32/f64, so an unguarded operation produces an +opaque tileiras verifier failure. Fail early with an actionable error instead; +the collected diagnostic's stacktrace names the offending operator. The check +folds away for arithmetic floats. + +This guards the direct tile operators below; the broadcast and `map` paths are +gated in `_apply_broadcast` (language/broadcast.jl), which shares the message. +""" +function check_arithmetic(::Type{T}) where {T} + if is_restricted_float(T) + throw(ArgumentError(RESTRICTED_ARITHMETIC_MESSAGE)) + end + return nothing +end + +# The message must stay a compile-time constant (see `throw_constant` in +# transform/throws.jl): one assembled at run time from `T` degrades to +# "ArgumentError was thrown". The stacktrace names the offending operation. +const RESTRICTED_ARITHMETIC_MESSAGE = + "operations on a restricted float element type are not supported; " * + "perform an explicit cast instead, e.g. convert(Tile{Float32}, x)" + # direct operators (same shape required) -@inline Base.:(+)(a::Tile{T, S}, b::Tile{T, S}) where {T <: AbstractFloat, S} = Intrinsics.addf(a, b) +@inline Base.:(+)(a::Tile{T, S}, b::Tile{T, S}) where {T <: AbstractFloat, S} = + (check_arithmetic(T); Intrinsics.addf(a, b)) @inline Base.:(+)(a::Tile{T, S}, b::Tile{T, S}) where {T <: Integer, S} = Intrinsics.addi(a, b) -@inline Base.:(-)(a::Tile{T, S}, b::Tile{T, S}) where {T <: AbstractFloat, S} = Intrinsics.subf(a, b) +@inline Base.:(-)(a::Tile{T, S}, b::Tile{T, S}) where {T <: AbstractFloat, S} = + (check_arithmetic(T); Intrinsics.subf(a, b)) @inline Base.:(-)(a::Tile{T, S}, b::Tile{T, S}) where {T <: Integer, S} = Intrinsics.subi(a, b) -@inline Base.:(-)(a::Tile{T}) where {T <: AbstractFloat} = Intrinsics.negf(a) +@inline Base.:(-)(a::Tile{T}) where {T <: AbstractFloat} = + (check_arithmetic(T); Intrinsics.negf(a)) @inline Base.:(-)(a::Tile{T}) where {T <: Integer} = Intrinsics.negi(a) # All other tile arithmetic (*, -, /, ^, comparisons, ifelse, etc.) is handled @@ -117,6 +147,9 @@ end ## mixed arithmetic # direct operators (tile * scalar, tile / scalar) -@inline Base.:(*)(a::Tile{T}, b::Number) where {T <: AbstractFloat} = Intrinsics.mulf(a, broadcast_to(Tile(T(b)), size(a))) -@inline Base.:(*)(a::Number, b::Tile{T}) where {T <: AbstractFloat} = Intrinsics.mulf(broadcast_to(Tile(T(a)), size(b)), b) -@inline Base.:(/)(a::Tile{T}, b::Number) where {T <: AbstractFloat} = Intrinsics.divf(a, broadcast_to(Tile(T(b)), size(a))) +@inline Base.:(*)(a::Tile{T}, b::Number) where {T <: AbstractFloat} = + (check_arithmetic(T); Intrinsics.mulf(a, broadcast_to(Tile(T(b)), size(a)))) +@inline Base.:(*)(a::Number, b::Tile{T}) where {T <: AbstractFloat} = + (check_arithmetic(T); Intrinsics.mulf(broadcast_to(Tile(T(a)), size(b)), b)) +@inline Base.:(/)(a::Tile{T}, b::Number) where {T <: AbstractFloat} = + (check_arithmetic(T); Intrinsics.divf(a, broadcast_to(Tile(T(b)), size(a)))) diff --git a/src/language/broadcast.jl b/src/language/broadcast.jl index c737b8e4..7eb060e5 100644 --- a/src/language/broadcast.jl +++ b/src/language/broadcast.jl @@ -43,8 +43,9 @@ Base.Broadcast.broadcastable(t::Tile) = t # This handles all element-wise operations: scalar @overlay methods provide # the implementation for overlaid ops, while Julia's native scalar functions # (compiled to Core intrinsics) handle the rest. Mixed-type and type-changing -# operations (comparisons, ifelse) are supported by the mixed-type map methods -# in operations.jl. +# operations (comparisons, ifelse) need nothing extra — `f` decides the result +# element type. `Base.map` (operations.jl) enters the same path at +# `_apply_broadcast`, its tiles already sharing a shape. @inline function Base.copy(bc::Broadcasted{TileStyle}) args = _materialize_args(bc.args) promoted = _promote_to_tiles(args...) @@ -96,10 +97,69 @@ end (a, _broadcast_all(S, rest...)...) # Convert args to scalars, apply f, wrap result back into a Tile. +# +# Restricted floats (FP8/FP4/TFloat32) are gated here rather than by shadowing +# the upstream scalar methods one by one: kernels can only ever obtain a +# restricted-float scalar through this function (and `map`), so this is the one +# choke point that covers every present and future upstream method. Only +# conversion, selection and comparison are let through; everything else — +# arithmetic, math functions, user lambdas — is rejected before dispatch, so no +# upstream fallback implementation is ever consulted. @inline function _apply_broadcast(f, args...) - Intrinsics.from_scalar(f(map(_to_scalar, args)...), _result_shape(args...)) + if _restricted_args(args...) + _restricted_broadcast(f, args...) + else + _broadcast_scalars(f, args...) + end end +@inline _broadcast_scalars(f, args...) = + Intrinsics.from_scalar(f(map(_to_scalar, args)...), _result_shape(args...)) + +# Does any Tile argument have a restricted float element type? Tuple peeling +# rather than `any` with a closure: the argument tuple is heterogeneous (Tiles +# and Refs). Folds to `false` at inference time for arithmetic element types, +# keeping the common path branch-free. `is_restricted_float` is called directly, +# not through `invokelatest`: kernel inference sees the extension methods, and +# the fold depends on it. +@inline _restricted_args() = false +@inline _restricted_args(a::Tile, rest...) = + is_restricted_float(eltype(a)) || _restricted_args(rest...) +@inline _restricted_args(a, rest...) = _restricted_args(rest...) + +const ComparisonOps = Union{typeof(<), typeof(<=), typeof(>), typeof(>=), + typeof(==), typeof(!=), typeof(isless)} + +# Explicit element-type conversion (`Float32.(tile)`, `convert.(Float32, tile)`, +# and `convert(Tile{T}, tile)` via `map`) is the sanctioned escape hatch: pass it +# through to the constructor overlays, which lower it to a single `ftof`. +@inline _restricted_broadcast(f::Union{Type,typeof(convert)}, args...) = + _broadcast_scalars(f, args...) + +# `ifelse` selects between unmodified values and lowers via `Core.ifelse`, so no +# upstream restricted-float method is involved. +@inline _restricted_broadcast(f::typeof(ifelse), args...) = + _broadcast_scalars(f, args...) + +# Comparisons stay available (as they do in cuTile Python), but Tile IR has no +# native fp8/fp4 comparison: upcast the restricted operands to Float32 and +# re-apply. That is exact and injective for every restricted format (NaN → NaN, +# ±0 preserved), so the result matches the host's ordering. Going through the +# upcast rather than the upstream scalar `<` also keeps their implementation +# details (bit tricks, `isnan` guards) out of the kernel. +@inline _restricted_broadcast(f::ComparisonOps, args...) = + _apply_broadcast(f, map(_upcast_restricted, args)...) + +# Everything else — arithmetic, math functions, user lambdas — is rejected. +# Upstream implements scalar arithmetic on these formats as a Float32 round-trip, +# which would otherwise compile into a silent ftof/op/ftof with an extra rounding +# per operation, and no hint that a cast happened. +@inline _restricted_broadcast(f, args...) = throw(ArgumentError(RESTRICTED_ARITHMETIC_MESSAGE)) + +@inline _upcast_restricted(a::Tile{T}) where {T} = + is_restricted_float(T) ? convert(Tile{Float32}, a) : a +@inline _upcast_restricted(a) = a + # Reinterpret arguments as scalars for broadcast application: Tiles via # to_scalar, Refs via their contents. The Ref{Type{T}} method recovers the # Type from the type parameter, mirroring Base's `_broadcast_getindex`. diff --git a/src/language/operations.jl b/src/language/operations.jl index b9982a2e..881d804f 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1386,6 +1386,10 @@ The function `f` must be a zero-size callable (singleton or capture-free lambda) All tiles must have the same shape `S` — use broadcasting (`.+` etc.) or explicit `broadcast_to` for shape-mismatched operands. +Tiles with a restricted float element type only accept the same `f` as +broadcasting does (conversion, `ifelse`, comparisons); anything else, including +a lambda that merely casts, is rejected. + # Examples ```julia result = map(abs, tile) # Element-wise absolute value @@ -1393,9 +1397,8 @@ result = map(x -> x * x, tile) # Element-wise square result = map(+, a, b) # Element-wise addition (same shape required) ``` """ -@inline function Base.map(f, a::Tile{<:Any,S}, rest::Tile{<:Any,S}...) where {S} - Intrinsics.from_scalar(f(Intrinsics.to_scalar(a), map(Intrinsics.to_scalar, rest)...), S) -end +@inline Base.map(f, a::Tile{<:Any,S}, rest::Tile{<:Any,S}...) where {S} = + _apply_broadcast(f, a, rest...) """ mapreduce(identity, f, tile::Tile{T,S}; dims, init) -> Tile{T, reduced_shape} diff --git a/src/language/types.jl b/src/language/types.jl index 77540619..afecc10c 100644 --- a/src/language/types.jl +++ b/src/language/types.jl @@ -683,22 +683,22 @@ const ScalarInt = Union{Int8, UInt8, Int16, UInt16, Int32, UInt32, Int64, UInt64 """Scalar floating-point types supported by Tile IR (f16, bf16, tf32, f32, f64).""" const ScalarFloat = Union{Float16, BFloat16, Float32, Float64, TFloat32} -""" -Restricted floats — types whose op coverage is intentionally limited -(no general arithmetic, reductions, scans, …). Currently `TFloat32`; -future FP8/FP4 dtypes will join this union. Mirrors cuTile Python's -`NumericDTypeCategories.RestrictedFloat`. -""" -const RestrictedFloat = Union{TFloat32} - """ is_restricted_float(::Type) -> Bool -True if `T` is a restricted float. Used by `reduce` / `scan` (and other -arithmetic-requiring ops) to reject unsupported element types early -with a clear error rather than letting tileiras fail downstream. +True if `T` is a restricted float: a storage / tensor-core operand format whose +op coverage is intentionally limited (no general arithmetic, reductions or +scans). Used by the broadcast / `map` gate, by `reduce` / `scan`, and by the +tile arithmetic guards ([`check_arithmetic`](@ref)) to reject unsupported +element types early with a clear error rather than letting tileiras fail +downstream. + +`TFloat32` is built in; package extensions register their own types by adding +methods (e.g. `DLFP8TypesExt` for `Float8_E4M3FN`). Mirrors cuTile Python's +`NumericDTypeCategories.RestrictedFloat`. """ -@inline is_restricted_float(::Type{T}) where {T} = T <: RestrictedFloat +is_restricted_float(::Type) = false +is_restricted_float(::Type{TFloat32}) = true """Integer tile types.""" const IntTile{S} = Tile{T, S} where {T <: ScalarInt} diff --git a/test/codegen/operations.jl b/test/codegen/operations.jl index 4d2c413f..ffe468ab 100644 --- a/test/codegen/operations.jl +++ b/test/codegen/operations.jl @@ -2001,6 +2001,61 @@ end end end end + + # TFloat32 is a restricted float: a tensor-core operand format the + # elementwise float ops do not accept. The direct tile operators used to + # emit `addf`/`subf`/`negf`/`mulf`/`divf` on it anyway, which failed the + # tileiras verifier; they now reject it up front. + @testset "restricted float arithmetic" begin + spec_tf32 = ct.ArraySpec{1}(16, true) + AT = ct.TileArray{ct.TFloat32,1,spec_tf32} + + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) + ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) - ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + + @test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, -ct.load(a, pid, (16,))) + return + end, Tuple{AT, AT}) + + # tile × scalar and tile / scalar take the mixed-arithmetic path + @test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, ct.load(a, pid, (16,)) * 2.0f0) + return + end, Tuple{AT, AT}) + + @test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, ct.load(a, pid, (16,)) / 2.0f0) + return + end, Tuple{AT, AT}) + + # The broadcast path is gated before scalar dispatch, so it reports the + # same error as the direct operators (it used to fall through to Base's + # `no_op_err`, TFloat32 having no scalar `+` to begin with). + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) .+ ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + end end #========================================================================= diff --git a/test/extensions/DLFP8Types.jl b/test/extensions/DLFP8Types.jl index f18371e4..d787b057 100644 --- a/test/extensions/DLFP8Types.jl +++ b/test/extensions/DLFP8Types.jl @@ -60,6 +60,167 @@ end end +# FP8 is a restricted float: a storage / tensor-core operand format without +# general arithmetic. DLFP8Types' scalar fallbacks would otherwise let the +# broadcast path compile into a silent ftof/op/ftof round-trip, and the direct +# tile operators into an `addf` the tileiras verifier rejects. +@testset "restricted arithmetic" begin + +AT = ct.TileArray{Float8_E4M3FN,1,spec1d} + +# broadcast (upstream's Float32 round-trip fallback) +@test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) .+ ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + +# direct tile operator (previously an MLIR verifier failure) +@test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) + ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + +# unary math via broadcast +@test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, sqrt.(ct.load(a, pid, (16,)))) + return + end, Tuple{AT, AT}) + +# tile × scalar (the mixed-arithmetic guard) +@test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, ct.load(a, pid, (16,)) * 2.0f0) + return + end, Tuple{AT, AT}) + +# broadcast `muladd` expands to the scalar `x * y + z` +@test_throws "restricted float" code_tiled(devnull, + (a, b, c, d) -> begin + pid = ct.bid(1) + ct.store(d, pid, muladd.(ct.load(a, pid, (16,)), ct.load(b, pid, (16,)), + ct.load(c, pid, (16,)))) + return + end, Tuple{AT, AT, AT, AT}) + +# reduce and scan check `is_restricted_float` at codegen time, which must +# resolve in the latest world to see the extension's method (`invokelatest` +# in `emit_reduce!`/`emit_intrinsic!`; a frozen-world call would miss it and +# fall through to an opaque failure in the reduce body). +@test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ta = ct.load(a, pid, (16,)) + s = mapreduce(identity, max, ta; dims=1, init=Float8_E4M3FN(0.0f0)) + ct.store(b, pid, ct.broadcast_to(s, (16,))) + return + end, Tuple{AT, AT}) + +@test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, cumsum(ct.load(a, pid, (16,)); dims=1)) + return + end, Tuple{AT, AT}) + +# Comparisons stay allowed (cuTile Python allows them too). DLFP8Types +# implements them at the bit level, which does not compile in kernels; the +# broadcast gate never consults that, upcasting the operands to Float32 +# instead, so they lower to `ftof` + `cmpf`. +@test @filecheck begin + @check_label "entry" + code_tiled(Tuple{AT, AT, ct.TileArray{Int32,1,spec1d}}) do a, b, c + ta = ct.load(a, ct.bid(1), (16,)) + tb = ct.load(b, ct.bid(1), (16,)) + @check "ftof" + @check "ftof" + @check "cmpf" + ct.store(c, ct.bid(1), ifelse.(ta .< tb, Int32(1), Int32(0))) + return + end +end + +# The remaining comparisons, compile-only. The function barrier keeps the +# kernel closure's captured `f` a concrete singleton. +function compiles_cmp(f) + isnothing(code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ta = ct.load(a, pid, (16,)) + tb = ct.load(b, pid, (16,)) + ct.store(c, pid, ifelse.(f.(ta, tb), Int32(1), Int32(0))) + return + end, Tuple{AT, AT, ct.TileArray{Int32,1,spec1d}})) +end +@test compiles_cmp(<=) +@test compiles_cmp(==) +# `isless` forwards to `isless(::Float32, ::Float32)`, whose Base definition +# does not compile under broadcast for any float tile yet (its `isnan` guard +# trips a scalar-vs-tile shape mismatch in cmpf). The gate's upcast is what +# gets FP8 to exactly that point; this flips when the underlying issue is fixed. +@test_broken compiles_cmp(isless) + +# Explicit conversion is the sanctioned escape hatch, in both directions: it +# passes through the gate to the constructor overlays and lowers to one `ftof`. +@test @filecheck begin + @check_label "entry" + code_tiled(Tuple{AT, ct.TileArray{Float32,1,spec1d}}) do a, b + pid = ct.bid(1) + @check "ftof" + @check_not "ftof" + ct.store(b, pid, Float32.(ct.load(a, pid, (16,)))) + return + end +end +@test @filecheck begin + @check_label "entry" + code_tiled(Tuple{ct.TileArray{Float32,1,spec1d}, AT}) do a, b + pid = ct.bid(1) + @check "ftof" + @check_not "ftof" + ct.store(b, pid, Float8_E4M3FN.(ct.load(a, pid, (16,)))) + return + end +end + +# `ifelse` selects between unmodified values, so it stays available too and +# lowers to a plain `select` — no conversion in sight. +@test @filecheck begin + @check_label "entry" + code_tiled(Tuple{AT, AT, ct.TileArray{Int32,1,spec1d}, AT}) do a, b, m, c + pid = ct.bid(1) + mask = ct.load(m, pid, (16,)) .> Int32(0) + @check "select" + @check_not "ftof" + ct.store(c, pid, ifelse.(mask, ct.load(a, pid, (16,)), ct.load(b, pid, (16,)))) + return + end +end + +# Element-wise application of anything else is rejected, even a lambda whose +# body is only a cast: cuTile Python has no per-element operations on restricted +# types either. `convert(Tile{Float32}, tile)` is the supported spelling. +@test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, map(x -> Float32(x), ct.load(a, pid, (16,)))) + return + end, Tuple{AT, ct.TileArray{Float32,1,spec1d}}) + +# The gate only applies inside kernels, so host arithmetic is untouched. +@test Float8_E4M3FN(1.0f0) + Float8_E4M3FN(1.0f0) == Float8_E4M3FN(2.0f0) +@test -Float8_E4M3FN(1.0f0) == Float8_E4M3FN(-1.0f0) +@test sqrt(Float8_E4M3FN(4.0f0)) == Float8_E4M3FN(2.0f0) +@test Float8_E4M3FN(1.0f0) < Float8_E4M3FN(2.0f0) + +end + # Execution kernels are plain top-level functions, each defined next to the # test that exercises it. Kernels parametric on accumulator dtype must stay at # top level — defining them inside a testset scope boxes them into closures. @@ -78,17 +239,6 @@ function rt_e5m2(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) ct.store(b, pid, convert(ct.Tile{Float32}, convert(ct.Tile{Float8_E5M2}, tile))) return end -# FMA in FP8: load Float32, convert to FP8, multiply-add in FP8, convert back. -# Inputs whose products and sums also stay representable, so the result is exact. -function fma_e4m3(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, - c::ct.TileArray{Float32,1}, d::ct.TileArray{Float32,1}) - pid = ct.bid(1) - ta = convert(ct.Tile{Float8_E4M3FN}, ct.load(a, pid, (16,))) - tb = convert(ct.Tile{Float8_E4M3FN}, ct.load(b, pid, (16,))) - tc = convert(ct.Tile{Float8_E4M3FN}, ct.load(c, pid, (16,))) - ct.store(d, pid, convert(ct.Tile{Float32}, muladd.(ta, tb, tc))) - return -end # Non-scaled FP8 matmul with both allowed accumulator dtypes (f16 and f32). function mma_dl_fp8(A::ct.TileArray{Float8_E4M3FN,2}, B::ct.TileArray{Float8_E4M3FN,2}, C::ct.TileArray{Tacc,2}, D::ct.TileArray{Float32,2}) where {Tacc<:Union{Float16,Float32}} @@ -116,15 +266,6 @@ let a = CuArray(representable), b = CUDA.zeros(Float32, length(representable)) @test Array(b) == representable end -let av = Float32[1.0, 2.0, 0.5, 4.0, 1.5, 2.0, -1.0, -0.5, 3.0, 0.5, 1.0, 2.0, -2.0, 1.0, 0.5, 4.0], - bv = Float32[2.0, 1.0, 4.0, 0.5, 2.0, 3.0, 2.0, 4.0, 1.0, 2.0, 1.0, 0.5, 2.0, 1.0, 2.0, 1.0], - cv = Float32[0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0] - a, b, c = CuArray(av), CuArray(bv), CuArray(cv) - d = CUDA.zeros(Float32, length(av)) - @cuda backend=cuTile blocks=1 fma_e4m3(a, b, c, d) - @test Array(d) == av .* bv .+ cv -end - @testset "mma → $Tacc acc" for Tacc in (Float32, Float16) M = 16 ah = Float8_E4M3FN.(Float32.(rand(0:2, M, M)) ./ 2) diff --git a/test/extensions/Microfloats/codegen.jl b/test/extensions/Microfloats/codegen.jl index efbdf698..29810250 100644 --- a/test/extensions/Microfloats/codegen.jl +++ b/test/extensions/Microfloats/codegen.jl @@ -239,4 +239,77 @@ end end end +# Microfloats are restricted floats: storage / tensor-core operand formats +# without general arithmetic. Microfloats' scalar fallbacks would otherwise let +# the broadcast path compile into a silent ftof/op/ftof round-trip, and the +# direct tile operators into an `addf` the tileiras verifier rejects. +@testset "restricted arithmetic" begin + AT = ct.TileArray{Float8_E4M3FN,1,spec1d} + + # broadcast (upstream's Float32 round-trip fallback) + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) .+ ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + + # direct tile operator (previously an MLIR verifier failure) + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) + ct.load(b, pid, (16,))) + return + end, Tuple{AT, AT, AT}) + + # unary math via broadcast + @test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, sqrt.(ct.load(a, pid, (16,)))) + return + end, Tuple{AT, AT}) + + # tile × scalar (the mixed-arithmetic guard) + @test_throws "restricted float" code_tiled(devnull, + (a, b) -> begin + pid = ct.bid(1) + ct.store(b, pid, ct.load(a, pid, (16,)) * 2.0f0) + return + end, Tuple{AT, AT}) + + # The gate is trait-based (`is_restricted_float`, registered for the whole + # `Microfloat` supertype), so every variant is covered with no per-type + # code — not just the two FP8 types DLFP8Types also provides. + F4 = ct.TileArray{Float4_E2M1FN,1,spec1d} + @test_throws "restricted float" code_tiled(devnull, + (a, b, c) -> begin + pid = ct.bid(1) + ct.store(c, pid, ct.load(a, pid, (16,)) .* ct.load(b, pid, (16,))) + return + end, Tuple{F4, F4, F4}; bytecode_version=v"13.3") + + # Comparisons stay allowed: the gate upcasts the operands to Float32, which + # is lossless without a result to round, so they lower to `ftof` + `cmpf`. + # cuTile Python's frontend accepts them too, but has no such upcast and dies + # in tileiras — Tile IR cannot compare fp8 natively. + @test @filecheck begin + @check_label "entry" + code_tiled(Tuple{AT, AT, ct.TileArray{Int32,1,spec1d}}) do a, b, c + ta = ct.load(a, ct.bid(1), (16,)) + tb = ct.load(b, ct.bid(1), (16,)) + @check "ftof" + @check "ftof" + @check "cmpf" + ct.store(c, ct.bid(1), ifelse.(ta .< tb, Int32(1), Int32(0))) + return + end + end + + # The gate only applies inside kernels, so host arithmetic is untouched. + @test Float8_E4M3FN(1.0f0) + Float8_E4M3FN(1.0f0) == Float8_E4M3FN(2.0f0) + @test -Float8_E4M3FN(1.0f0) == Float8_E4M3FN(-1.0f0) + @test sqrt(Float8_E4M3FN(4.0f0)) == Float8_E4M3FN(2.0f0) +end + end diff --git a/test/extensions/Microfloats/device.jl b/test/extensions/Microfloats/device.jl index 97ede9a5..87653f9a 100644 --- a/test/extensions/Microfloats/device.jl +++ b/test/extensions/Microfloats/device.jl @@ -30,16 +30,6 @@ function rt_f4(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}) ct.store(b, pid, convert(ct.Tile{Float32}, convert(ct.Tile{Float4_E2M1FN}, tile))) return end -function fma_e4m3(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1}, - c::ct.TileArray{Float32,1}, d::ct.TileArray{Float32,1}) - pid = ct.bid(1) - ta = convert(ct.Tile{Float8_E4M3FN}, ct.load(a, pid, (16,))) - tb = convert(ct.Tile{Float8_E4M3FN}, ct.load(b, pid, (16,))) - tc = convert(ct.Tile{Float8_E4M3FN}, ct.load(c, pid, (16,))) - ct.store(d, pid, convert(ct.Tile{Float32}, muladd.(ta, tb, tc))) - return -end - # Standalone f32 → microfloat → f32 conversion round-trips exactly for every # microfloat type on representable inputs. FP8 (e4m3/e5m2) needs Hopper (sm_90+); # E8M0FNU and Float4_E2M1FN need Blackwell (sm_100+). E8M0 is exponent-only, so @@ -56,16 +46,6 @@ if capability(device()) >= v"9" @test Array(b) == representable8 end - # FMA in FP8: load f32, convert to FP8, multiply-add in FP8, convert back. - # Inputs whose products and sums stay representable, so the result is exact. - let av = Float32[1.0, 2.0, 0.5, 4.0, 1.5, 2.0, -1.0, -0.5, 3.0, 0.5, 1.0, 2.0, -2.0, 1.0, 0.5, 4.0], - bv = Float32[2.0, 1.0, 4.0, 0.5, 2.0, 3.0, 2.0, 4.0, 1.0, 2.0, 1.0, 0.5, 2.0, 1.0, 2.0, 1.0], - cv = Float32[0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0] - a, b, c = CuArray(av), CuArray(bv), CuArray(cv) - d = CUDA.zeros(Float32, length(av)) - @cuda backend=cuTile blocks=1 fma_e4m3(a, b, c, d) - @test Array(d) == av .* bv .+ cv - end end if capability(device()) >= v"10" # E8M0 round-trip: exponent-only, so representable values are powers of two.