Add muladd_scaled

Author: AntonOresten

ext/DLFP8TypesExt.jl

@@ -12,6 +12,12 @@ function ct.julia_to_tile_dtype!(table::ct.TypeTable, ::Type{Float8_E5M2})
     return ct.F8E5M2(table)
 end
 
+# Non-scaled `mma`/`matmul` (`cuda_tile.mmaf`) accepts f8e4m3fn and f8e5m2
+# operands with an f16 or f32 accumulator (f16 first/preferred), mirroring
+# cuda-tile's mmaf type table and cutile-python's `_mma_supported_dtypes`.
+ct.mma_allowed_acc_dtypes(::Type{Float8_E4M3FN}) = (Float16, Float32)
+ct.mma_allowed_acc_dtypes(::Type{Float8_E5M2})   = (Float16, Float32)
+
 # Float ↔ FP8 scalar constructor overlays (for map/convert dispatch)
 const FP8Types = (Float8_E4M3FN, Float8_E5M2)
 const StandardFloats = (Float16, ct.BFloat16, Float32, ct.TFloat32, Float64)

ext/MicrofloatsExt.jl

@@ -26,6 +26,15 @@ ct.bitwidth(::Type{T}) where {T<:Microfloats.Microfloat} = Microfloats.bitwidth(
 # nearest-even on f32→E8M0FNU (only `zero` and `positive_inf` are valid).
 ct.ftof_rounding_mode(::Type{Float8_E8M0FNU}) = ct.RoundingMode.Zero
 
+# Non-scaled `mma`/`matmul` (`cuda_tile.mmaf`) accepts f8e4m3fn and f8e5m2
+# operands with an f16 or f32 accumulator — f16 first/preferred, mirroring
+# cuda-tile's mmaf type table and cutile-python's `_mma_supported_dtypes`. The
+# other Microfloats formats (E8M0FNU, Float4_E2M1FN) are only valid as
+# scaled-mma operands/scales, so they stay absent and fall through to `mma`'s
+# unsupported-dtype error.
+ct.mma_allowed_acc_dtypes(::Type{Float8_E4M3FN}) = (Float16, Float32)
+ct.mma_allowed_acc_dtypes(::Type{Float8_E5M2})   = (Float16, Float32)
+
 # Float ↔ microfloat scalar constructor overlays (for map/convert dispatch).
 # Mirrors DLFP8TypesExt: route to `Intrinsics.ftof` so kernel-side conversions
 # lower to the FToFOp Tile IR intrinsic instead of Microfloats' Float32-fallback

src/bytecode/encodings.jl

@@ -98,6 +98,8 @@ module Opcode
     const Atan2Op = 110  # since 13.2
     const PackOp = 111   # since 13.3
     const UnpackOp = 112 # since 13.3
+    # 113 (AllocaOp) not implemented
+    const MmaFScaledOp = 114 # since 13.3
 end
 
 # Enums for operation attributes
@@ -1310,6 +1312,30 @@ function encode_MmaIOp!(cb::CodeBuilder, result_type::TypeId,
     return new_op!(cb)
 end
 
+"""
+    encode_MmaFScaledOp!(cb, result_type, lhs, rhs, acc, lhs_scale, rhs_scale) -> Value
+
+Block-scaled float matrix multiply-accumulate (`acc + (lhs ⊙ lhs_scale) @
+(rhs ⊙ rhs_scale)`) for low-precision (f8/f4) inputs. Each scale element scales
+a contiguous block of K-dimension elements in its operand. Requires Tile IR
+v13.3+.
+Opcode: 114
+"""
+function encode_MmaFScaledOp!(cb::CodeBuilder, result_type::TypeId,
+                              lhs::Value, rhs::Value, acc::Value,
+                              lhs_scale::Value, rhs_scale::Value)
+    cb.version >= v"13.3" ||
+        throw(IRError("MmaFScaledOp requires Tile IR v13.3+, got v$(cb.version)"))
+    encode_varint!(cb.buf, Opcode.MmaFScaledOp)
+    encode_typeid!(cb.buf, result_type)
+    encode_operand!(cb.buf, lhs)
+    encode_operand!(cb.buf, rhs)
+    encode_operand!(cb.buf, acc)
+    encode_operand!(cb.buf, lhs_scale)
+    encode_operand!(cb.buf, rhs_scale)
+    return new_op!(cb)
+end
+
 #=============================================================================
  Integer arithmetic operations
 =============================================================================#

src/compiler/intrinsics/core.jl

@@ -458,7 +458,10 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.mma), args)
         # the table in cuTile Python's mma implementation.
         lhs_elem === rhs_elem ||
             throw(IRError("mma: float lhs and rhs must share dtype, got lhs=$lhs_elem, rhs=$rhs_elem"))
-        allowed_acc = mma_allowed_acc_dtypes(lhs_elem)
+        # `invokelatest` so extension-defined acc-dtype tables (e.g. FP8 via the
+        # Microfloats/DLFP8Types exts) are visible from codegen's world age,
+        # mirroring `lookup_bitwidth`.
+        allowed_acc = Base.invokelatest(mma_allowed_acc_dtypes, lhs_elem)
         allowed_acc === nothing &&
             throw(IRError("mma: unsupported float input dtype $lhs_elem"))
         acc_elem in allowed_acc ||
@@ -480,6 +483,57 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.mma), args)
     CGVal(result, acc.type_id, acc.jltype, acc.shape)
 end
 
+"""
+    Intrinsics.mma_scaled(lhs, lhs_scale, rhs, rhs_scale, acc) -> typeof(acc)
+
+Block-scaled matrix-multiply-accumulate computing `(lhs ⊙ lhs_scale) * (rhs ⊙
+rhs_scale) + acc`, where each scale element multiplies a contiguous block of
+`lhs`/`rhs` elements along the K dimension. Lowers to `cuda_tile.mmaf_scaled`
+(Tile IR v13.3+).
+
+`lhs`/`rhs` are low-precision floats (`f8e4m3fn`, `f8e5m2`, or `f4e2m1fn`),
+`lhs_scale`/`rhs_scale` are `f8e8m0fnu` or `f8e4m3fn`, and `acc`/result are
+`f32`. The block size `K ÷ K_s` and the exact (operand, scale) dtype pairing are
+validated by tileiras (see its `mmaf_scaled` verifier for the supported table).
+"""
+@intrinsic mma_scaled(lhs, lhs_scale, rhs, rhs_scale, acc)
+tfunc(𝕃, ::typeof(Intrinsics.mma_scaled), @nospecialize(lhs), @nospecialize(lhs_scale),
+      @nospecialize(rhs), @nospecialize(rhs_scale), @nospecialize(acc)) = CC.widenconst(acc)
+function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.mma_scaled), args)
+    cb = ctx.cb
+
+    lhs       = emit_value!(ctx, args[1])
+    lhs_scale = emit_value!(ctx, args[2])
+    rhs       = emit_value!(ctx, args[3])
+    rhs_scale = emit_value!(ctx, args[4])
+    acc       = emit_value!(ctx, args[5])
+
+    (lhs === nothing || lhs_scale === nothing || rhs === nothing ||
+     rhs_scale === nothing || acc === nothing) &&
+        throw(IRError("Cannot resolve operands for mma_scaled()"))
+
+    lhs_elem       = eltype(CC.widenconst(lhs.jltype))
+    rhs_elem       = eltype(CC.widenconst(rhs.jltype))
+    lhs_scale_elem = eltype(CC.widenconst(lhs_scale.jltype))
+    rhs_scale_elem = eltype(CC.widenconst(rhs_scale.jltype))
+    acc_elem       = eltype(CC.widenconst(acc.jltype))
+
+    # Structural invariants (cheap, with clear errors); the operand/scale dtype
+    # pairing and block-size constraints are checked by the tileiras verifier,
+    # whose messages are already precise.
+    lhs_elem === rhs_elem ||
+        throw(IRError("mma_scaled: lhs and rhs must share dtype, got lhs=$lhs_elem, rhs=$rhs_elem"))
+    lhs_scale_elem === rhs_scale_elem ||
+        throw(IRError("mma_scaled: lhs_scale and rhs_scale must share dtype, " *
+                      "got $lhs_scale_elem and $rhs_scale_elem"))
+    acc_elem === Float32 ||
+        throw(IRError("mma_scaled: acc must be Float32, got $acc_elem"))
+
+    result = encode_MmaFScaledOp!(cb, acc.type_id, lhs.v, rhs.v, acc.v,
+                                  lhs_scale.v, rhs_scale.v)
+    CGVal(result, acc.type_id, acc.jltype, acc.shape)
+end
+
 # TODO: cuda_tile.module
 
 """

src/language/operations.jl

@@ -1144,10 +1144,21 @@ end
  Matrix multiplication
 =============================================================================#
 
-# Matrix multiply-accumulate: muladd(a, b, acc) = a * b + acc
-# Handles 1D promotion, type promotion, and batched dims (≥3D).
-# Note: SA, SB, SC type parameters required to avoid ambiguity with scalar methods during codegen
+"""
+    muladd(a::Tile, b::Tile, acc::Tile) -> Tile
+
+Matrix multiply-accumulate `a * b + acc` over tiles, lowering to
+`cuda_tile.mmaf` (float) or `cuda_tile.mmai` (`i8 × i8 → i32`).
+
+`a`/`b` are 2-D matrices `(M, K)` × `(K, N)` → `(M, N)`; a 1-D operand is
+promoted (vec-mat / mat-vec) and any trailing dimensions (≥3-D) are treated as
+broadcast batch dims, lifting `Base.muladd`'s shape rules to tiles. `acc`
+carries the result dtype, which must be one tileiras allows for the input dtype
+(f16/f32 for f16 and f8; f32 for bf16/tf32; f64 for f64; i32 for i8).
+"""
 @inline function Base.muladd(a::Tile{T1, SA}, b::Tile{T2, SB}, acc::Tile{T3, SC}) where {T1, T2, T3, SA, SB, SC}
+    # SA, SB, SC type parameters avoid ambiguity with the scalar `muladd`
+    # methods during codegen.
     _muladd(a, b, acc, Val(ndims(a)), Val(ndims(b)))
 end
 
@@ -1228,6 +1239,117 @@ end
     end
 end
 
+#=============================================================================
+ Block-scaled matrix multiply-accumulate
+=============================================================================#
+
+"""
+    muladd_scaled(a, a_scale, b, b_scale, acc) -> Tile
+
+Block-scaled matrix multiply-accumulate `(a ⊙ a_scale) * (b ⊙ b_scale) + acc`,
+lowering to `cuda_tile.mmaf_scaled` (Tile IR v13.3+, Blackwell). Each scale
+element multiplies a contiguous block of `B = K ÷ K_s` elements along the K
+dimension of its operand, so `a_scale`/`b_scale` match `a`/`b` in every
+dimension except K, where they have `K_s ≤ K` entries.
+
+`a`/`b` are low-precision floats (`f8e4m3fn`, `f8e5m2`, or `f4e2m1fn`),
+`a_scale`/`b_scale` are `f8e8m0fnu` or `f8e4m3fn`, and `acc` is `f32`. Shapes
+follow [`muladd`](@ref): 2-D `(M, K)` × `(K, N)`, mat-vec, and trailing batch
+dims; vec-mat is unsupported (it would collapse K, leaving nothing to scale).
+"""
+@inline function muladd_scaled(a::Tile{Ta, SA}, a_scale::Tile, b::Tile{Tb, SB}, b_scale::Tile,
+                               acc::Tile) where {Ta, Tb, SA, SB}
+    _muladd_scaled(a, a_scale, b, b_scale, acc, Val(ndims(a)), Val(ndims(b)))
+end
+
+# 2D × 2D: swap operands (and their scales) for row-major Tile IR, exactly as
+# `_muladd` swaps for `mma`.
+@inline function _muladd_scaled(a::Tile, a_scale::Tile, b::Tile, b_scale::Tile, acc::Tile,
+                                ::Val{2}, ::Val{2})
+    Intrinsics.mma_scaled(b, b_scale, a, a_scale, acc)
+end
+
+# Mat-vec (2D × 1D): the K-vector `b` (and its scale) gain a trailing N=1 dim;
+# `acc` becomes (M, 1); then squeeze back to (M,). K — the scaled dimension —
+# is preserved, so block scaling is well-defined.
+@inline function _muladd_scaled(a::Tile, a_scale::Tile, b::Tile, b_scale::Tile, acc::Tile,
+                                ::Val{2}, ::Val{1})
+    M, K, Ks = size(a, 1), size(b, 1), size(b_scale, 1)
+    b2d = reshape(b, (K, 1))
+    b_scale2d = reshape(b_scale, (Ks, 1))
+    acc2d = reshape(acc, (M, 1))
+    result = _muladd_scaled(a, a_scale, b2d, b_scale2d, acc2d, Val(2), Val(2))
+    reshape(result, (M,))
+end
+
+# Vec-mat (1D × 2D): promoting `a` to (M, 1) collapses K to 1, leaving no K
+# dimension to block-scale. Unsupported — reshape to 2D and supply a matching
+# K_s scale instead.
+@generated function _muladd_scaled(::Tile, ::Tile, ::Tile, ::Tile, ::Tile, ::Val{1}, ::Val{2})
+    return :(throw(ArgumentError("Scaled vec-mat is not supported (the K dimension collapses to 1, which cannot be block-scaled).")))
+end
+
+# Vec-vec (1D × 1D): not supported.
+@generated function _muladd_scaled(::Tile, ::Tile, ::Tile, ::Tile, ::Tile, ::Val{1}, ::Val{1})
+    return :(throw(ArgumentError("Scaled vector-vector multiply-accumulate is not supported.")))
+end
+
+# Batched mat-vec / vec-mat (≥3D × 1D or 1D × ≥3D): not supported.
+@generated function _muladd_scaled(::Tile, ::Tile, ::Tile, ::Tile, ::Tile, ::Val{1}, ::Val{NB}) where {NB}
+    return :(throw(ArgumentError("Batched scaled vec-mat is not supported.")))
+end
+@generated function _muladd_scaled(::Tile, ::Tile, ::Tile, ::Tile, ::Tile, ::Val{NA}, ::Val{1}) where {NA}
+    return :(throw(ArgumentError("Batched scaled mat-vec is not supported.")))
+end
+
+# Batched (≥3D × ≥3D): trailing batch dims with broadcast, mirroring `_muladd`.
+# Scales carry the same batch dims as their operands; a_scale's batch must match
+# a's batch (likewise b_scale/b), then both broadcast to the common batch shape.
+@generated function _muladd_scaled(a::Tile{Ta, SA}, a_scale::Tile{Tas, SAS},
+                                   b::Tile{Tb, SB}, b_scale::Tile{Tbs, SBS},
+                                   acc::Tile{Tc, SC},
+                                   ::Val{NA}, ::Val{NB}) where {Ta, Tas, Tb, Tbs, Tc,
+                                                                SA, SAS, SB, SBS, SC, NA, NB}
+    sa = Tuple(SA.parameters);  sas = Tuple(SAS.parameters)
+    sb = Tuple(SB.parameters);  sbs = Tuple(SBS.parameters)
+
+    # Matrix dims are first two; batch dims are trailing.
+    M = sa[1]; K = sa[2]; N = sb[2]
+    Ksa = sas[2]   # a_scale K_s (a_scale is (M, K_s, batch...))
+    Ksb = sbs[1]   # b_scale K_s (b_scale is (K_s, N, batch...))
+    a_batch  = sa[3:end];  b_batch  = sb[3:end]
+    as_batch = sas[3:end]; bs_batch = sbs[3:end]
+
+    # Broadcast batch dims (pad shorter with trailing 1s, then broadcast).
+    n_batch = max(length(a_batch), length(b_batch))
+    a_batch_padded  = (a_batch...,  ntuple(Returns(1), n_batch - length(a_batch))...)
+    b_batch_padded  = (b_batch...,  ntuple(Returns(1), n_batch - length(b_batch))...)
+    as_batch_padded = (as_batch..., ntuple(Returns(1), n_batch - length(as_batch))...)
+    bs_batch_padded = (bs_batch..., ntuple(Returns(1), n_batch - length(bs_batch))...)
+    batch_shape = map(max, a_batch_padded, b_batch_padded)
+    B_flat = prod(batch_shape)
+
+    quote
+        # Reshape + broadcast to align batch dims (still trailing).
+        a_bc  = broadcast_to(reshape(a,       $((M, K, a_batch_padded...))),  $((M, K, batch_shape...)))
+        b_bc  = broadcast_to(reshape(b,       $((K, N, b_batch_padded...))),  $((K, N, batch_shape...)))
+        as_bc = broadcast_to(reshape(a_scale, $((M, Ksa, as_batch_padded...))), $((M, Ksa, batch_shape...)))
+        bs_bc = broadcast_to(reshape(b_scale, $((Ksb, N, bs_batch_padded...))), $((Ksb, N, batch_shape...)))
+        acc_bc = broadcast_to(acc, $((M, N, batch_shape...)))
+        # Flatten batch dims to one — no permute needed since row-major Tile IR
+        # already has batch as the leading (slowest) dimension.
+        a_3d  = reshape(a_bc,  $((M, K, B_flat)))
+        b_3d  = reshape(b_bc,  $((K, N, B_flat)))
+        as_3d = reshape(as_bc, $((M, Ksa, B_flat)))
+        bs_3d = reshape(bs_bc, $((Ksb, N, B_flat)))
+        acc_3d = reshape(acc_bc, $((M, N, B_flat)))
+        # mmaf_scaled with swapped operands for row-major convention.
+        result_3d = Intrinsics.mma_scaled(b_3d, bs_3d, a_3d, as_3d, acc_3d)
+        # Unflatten batch dims.
+        reshape(result_3d, $((M, N, batch_shape...)))
+    end
+end
+
 # Matrix multiplication: A * B = muladd(A, B, zeros)
 # Note: SA, SB type parameters required to avoid ambiguity with scalar*tile methods during codegen
 #

test/extensions/DLFP8Types.jl

@@ -2,6 +2,7 @@ using CUDA
 using DLFP8Types: Float8_E4M3FN, Float8_E5M2
 
 spec1d = ct.ArraySpec{1}(16, true)
+spec2d = ct.ArraySpec{2}(16, true)
 
 @testset "codegen" begin
 
@@ -31,11 +32,24 @@ end
     end
 end
 
+# Non-scaled f8 matmul lowers to `mmaf` (f8 operands, f32 accumulator).
+@test @filecheck begin
+    @check_label "entry"
+    code_tiled(Tuple{ct.TileArray{Float8_E4M3FN,2,spec2d}, ct.TileArray{Float8_E4M3FN,2,spec2d},
+                     ct.TileArray{Float32,2,spec2d}}) do a, b, c
+        ta = ct.load(a, (1, 1), (16, 16))
+        tb = ct.load(b, (1, 1), (16, 16))
+        @check "mmaf"
+        ct.store(c, (1, 1), muladd(ta, tb, zeros(Float32, (16, 16))))
+        return
+    end
+end
+
 end
 
-# FP8 types are Blackwell-only
+# FP8 (e4m3/e5m2) conversions and matmul need Hopper (sm_90+).
 @testset "execution" begin
-if capability(device()) >= v"10"
+if capability(device()) >= v"9"
 
 # Round-trip Float32 → FP8 → Float32 on values exactly representable in
 # the target FP8 type — result must match input bit-for-bit.
@@ -82,5 +96,29 @@ 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, -
     @test Array(d) == av .* bv .+ cv
 end
 
+# Non-scaled FP8 matmul with both allowed accumulator dtypes (f16 and f32).
+function mma_dl_f32(A::ct.TileArray{Float8_E4M3FN,2}, B::ct.TileArray{Float8_E4M3FN,2},
+                    C::ct.TileArray{Float32,2}, D::ct.TileArray{Float32,2})
+    a = ct.load(A, (1, 1), (16, 16)); b = ct.load(B, (1, 1), (16, 16)); c = ct.load(C, (1, 1), (16, 16))
+    ct.store(D, (1, 1), convert(ct.Tile{Float32}, muladd(a, b, c)))
+    return
+end
+function mma_dl_f16(A::ct.TileArray{Float8_E4M3FN,2}, B::ct.TileArray{Float8_E4M3FN,2},
+                    C::ct.TileArray{Float16,2}, D::ct.TileArray{Float32,2})
+    a = ct.load(A, (1, 1), (16, 16)); b = ct.load(B, (1, 1), (16, 16)); c = ct.load(C, (1, 1), (16, 16))
+    ct.store(D, (1, 1), convert(ct.Tile{Float32}, muladd(a, b, c)))
+    return
+end
+@testset "mma → $Tacc acc" for (Tacc, kern) in ((Float32, mma_dl_f32), (Float16, mma_dl_f16))
+    M = 16
+    ah = Float8_E4M3FN.(Float32.(rand(0:2, M, M)) ./ 2)
+    bh = Float8_E4M3FN.(Float32.(rand(0:2, M, M)) ./ 2)
+    ch = Tacc.(Float32.(rand(0:2, M, M)))
+    ref = Float32.(ah) * Float32.(bh) .+ Float32.(ch)
+    D = CUDA.zeros(Float32, M, M)
+    @cuda backend=cuTile blocks=1 kern(CuArray(ah), CuArray(bh), CuArray(ch), D)
+    @test Array(D) == ref
+end
+
 end
 end