Block arithmetic on restricted floats

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

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

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

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)

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))))

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`.

src/language/operations.jl

@@ -1386,16 +1386,19 @@ 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
 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}

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}

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
 
 #=========================================================================