Add num_worker_warps entry hint (#245)

* Add `num_worker_warps` entry hint

* Update src/bytecode/writer.jl

Co-authored-by: Tim Besard <tim.besard@gmail.com>

---------

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: AntonOresten

README.md

@@ -361,6 +361,7 @@ end
 | `num_ctas` | Number of CTAs in a CGA | Powers of 2 |
 | `occupancy` | Target concurrent CTAs per SM | 1–32 |
 | `opt_level` | Optimization level | 0–3 |
+| `num_worker_warps` | Worker warps per CTA in a warp-specialized kernel | 4 or 8 |
 
 Values can be plain scalars or `ct.ByTarget(...)` for per-architecture dispatch.
 `ByTarget` maps compute capabilities to values, with an optional default:

src/bytecode/writer.jl

@@ -741,17 +741,19 @@ function encode_load_store_hints_dict!(cb::CodeBuilder, hints::LoadStoreHints)
 end
 
 """
-Kernel-level compilation hints (num_ctas, occupancy).
+Kernel-level compilation hints (num_ctas, occupancy, num_worker_warps).
 Encoded as a dictionary attribute in bytecode.
 """
 @kwdef struct EntryHints
-    num_ctas::Union{Int, Nothing} = nothing    # 1, 2, 4, 8, 16
-    occupancy::Union{Int, Nothing} = nothing   # 1-32
+    num_ctas::Union{Int, Nothing} = nothing           # 1, 2, 4, 8, 16
+    occupancy::Union{Int, Nothing} = nothing          # 1-32
+    num_worker_warps::Union{Int, Nothing} = nothing   # 4 or 8
 end
 
 function encode_entry_hints(writer::BytecodeWriter, sm_arch::Union{VersionNumber, Nothing}, hints::EntryHints)
     validate_hint(:num_ctas, hints.num_ctas)
     validate_hint(:occupancy, hints.occupancy)
+    validate_hint(:num_worker_warps, hints.num_worker_warps)
 
     # CTA clusters (num_cta_in_cga > 1) are a Blackwell feature. Older tileiras
     # versions rejected the bytecode for non-Blackwell targets; tileiras 13.3
@@ -763,10 +765,16 @@ function encode_entry_hints(writer::BytecodeWriter, sm_arch::Union{VersionNumber
             "$(format_sm_arch(sm_arch))"))
     end
 
+    if hints.num_worker_warps !== nothing && writer.version < v"13.3"
+        throw(ArgumentError(
+            "num_worker_warps requires Tile IR bytecode v13.3+, got v$(writer.version)"))
+    end
+
     # Build items list (only non-nothing values)
     items = Tuple{String, Int}[]
     isnothing(hints.num_ctas) || push!(items, ("num_cta_in_cga", hints.num_ctas))
     isnothing(hints.occupancy) || push!(items, ("occupancy", hints.occupancy))
+    isnothing(hints.num_worker_warps) || push!(items, ("num_worker_warps_per_cta", hints.num_worker_warps))
 
     # Always emit optimization hints when sm_arch is specified, even with an empty
     # dict. Python cuTile emits `optimization_hints=<sm_NNN = {}>` unconditionally

src/compiler/codegen/kernel.jl

@@ -1,7 +1,7 @@
 # kernel and argument handling
 
 """
-    emit_kernel!(writer, func_buf, sci, rettype; name, sm_arch=nothing, is_entry=true, num_ctas=nothing, occupancy=nothing, const_argtypes=nothing)
+    emit_kernel!(writer, func_buf, sci, rettype; name, sm_arch=nothing, is_entry=true, num_ctas=nothing, occupancy=nothing, num_worker_warps=nothing, const_argtypes=nothing)
 
 Compile a StructuredIRCode to Tile IR bytecode.
 
@@ -17,6 +17,7 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
                       is_entry::Bool = true,
                       num_ctas::Union{Int, Nothing} = nothing,
                       occupancy::Union{Int, Nothing} = nothing,
+                      num_worker_warps::Union{Int, Nothing} = nothing,
                       cache::CacheView,
                       const_argtypes::Union{Vector{Any}, Nothing} = nothing)
     tt = writer.type_table
@@ -82,7 +83,7 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     end
 
     # Create entry hints if provided
-    entry_hints = encode_entry_hints(writer, sm_arch, EntryHints(; num_ctas, occupancy))
+    entry_hints = encode_entry_hints(writer, sm_arch, EntryHints(; num_ctas, occupancy, num_worker_warps))
 
     # Create function-level debug attribute
     func_debug_attr = make_func_debug_attr(debug_emitter, sci; linkage_name=name)

src/compiler/driver.jl

@@ -7,13 +7,15 @@ const compile_hook = Ref{Union{Nothing,Function}}(nothing)
 =============================================================================#
 
 # Compilation options for cache sharding.
-# Hint fields (opt_level, num_ctas, occupancy) represent explicit overrides only;
-# `nothing` means "consult @compiler_options meta nodes in the IR during compilation."
+# Hint fields (opt_level, num_ctas, occupancy, num_worker_warps) represent explicit
+# overrides only; `nothing` means "consult @compiler_options meta nodes in the IR
+# during compilation."
 const CGOpts = @NamedTuple{
     sm_arch::Union{VersionNumber, Nothing},
     opt_level::Union{Int, Nothing},
     num_ctas::Union{Int, Nothing},
     occupancy::Union{Int, Nothing},
+    num_worker_warps::Union{Int, Nothing},
     bytecode_version::VersionNumber
 }
 
@@ -168,6 +170,8 @@ function emit_tile(sci::StructuredIRCode, rettype, kernel_meta::Dict{Symbol,Any}
     # Resolve hints: launch()/code_tiled() kwargs > @compiler_options meta > defaults
     resolved_num_ctas = resolve_hint(opts.num_ctas, kernel_meta, :num_ctas, opts.sm_arch)
     resolved_occupancy = resolve_hint(opts.occupancy, kernel_meta, :occupancy, opts.sm_arch)
+    resolved_num_worker_warps = resolve_hint(opts.num_worker_warps, kernel_meta,
+                                             :num_worker_warps, opts.sm_arch)
 
     # Generate Tile IR bytecode
     bytecode = write_bytecode!(1; version=opts.bytecode_version) do writer, func_buf
@@ -176,6 +180,7 @@ function emit_tile(sci::StructuredIRCode, rettype, kernel_meta::Dict{Symbol,Any}
             sm_arch = opts.sm_arch,
             num_ctas = resolved_num_ctas,
             occupancy = resolved_occupancy,
+            num_worker_warps = resolved_num_worker_warps,
             cache,
             const_argtypes
         )
@@ -307,6 +312,7 @@ function emit_tile!(cache::CacheView, mi::Core.MethodInstance,
                    opt_level=unpack_hint(key.opt_level),
                    num_ctas=unpack_hint(key.num_ctas),
                    occupancy=unpack_hint(key.occupancy),
+                   num_worker_warps=unpack_hint(key.num_worker_warps),
                    bytecode_version=unpack_version(key.bytecode_version)))
     bytecode = emit_tile(sci, rettype, kernel_meta;
                          name=sanitize_name(string(mi.def.name)),

src/compiler/reflection.jl

@@ -109,7 +109,7 @@ constant_eltype(::Type{Constant{T,V}}) where {T,V} = T
 constant_value(::Type{Constant{T,V}}) where {T,V} = V
 
 """
-    code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy)
+    code_tiled([io::IO], f, argtypes; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps)
 
 Print the CUDA Tile IR for a Julia function as a textual MLIR representation.
 Analogous to `code_llvm`/`code_native`. Calls the driver directly without
@@ -120,14 +120,15 @@ function code_tiled(io::IO, @nospecialize(f), @nospecialize(argtypes);
                     opt_level::Union{Int, Nothing}=nothing,
                     num_ctas::Union{Int, Nothing}=nothing,
                     occupancy::Union{Int, Nothing}=nothing,
+                    num_worker_warps::Union{Int, Nothing}=nothing,
                     bytecode_version::VersionNumber=cuTile.bytecode_version(),
                     debuginfo::Bool=false,
                     world::UInt=Base.get_world_counter())
     stripped, const_argtypes = process_const_argtypes(f, argtypes)
     mi = lookup_method_instance(f, stripped; world)
 
     opts = CGOpts((sm_arch=sm_arch, opt_level=opt_level, num_ctas=num_ctas, occupancy=occupancy,
-                    bytecode_version=bytecode_version))
+                    num_worker_warps=num_worker_warps, bytecode_version=bytecode_version))
     cache = CacheView{CuTileResults}(:cuTile, world)
     ir, rettype = emit_julia(cache, mi; const_argtypes)
     sci, rettype, kernel_meta = emit_structured(ir, rettype)

src/language/operations.jl

@@ -1441,7 +1441,7 @@ end
  @compiler_options macro
 =============================================================================#
 
-const _COMPILER_OPTION_NAMES = Set([:num_ctas, :occupancy, :opt_level])
+const _COMPILER_OPTION_NAMES = Set([:num_ctas, :occupancy, :opt_level, :num_worker_warps])
 
 """
     @compiler_options key=val...
@@ -1450,7 +1450,7 @@ Specify per-architecture optimization hints inside a kernel function body.
 Hints are embedded as `:meta` nodes and resolved at compile time based on
 the target `sm_arch`.
 
-Supported options: `num_ctas`, `occupancy`, `opt_level`.
+Supported options: `num_ctas`, `occupancy`, `opt_level`, `num_worker_warps`.
 
 Values can be plain scalars or `ByTarget(...)` for per-architecture dispatch.
 

src/language/types.jl

@@ -561,6 +561,7 @@ Validate a kernel optimization hint value. Throws `ArgumentError` for invalid va
 - `num_ctas`: power of 2 in [1, 16]
 - `occupancy`: integer in [1, 32]
 - `opt_level`: integer in [0, 3]
+- `num_worker_warps`: either 4 or 8
 """
 function validate_hint(key::Symbol, val)
     val === nothing && return
@@ -574,6 +575,9 @@ function validate_hint(key::Symbol, val)
     elseif key === :opt_level
         val isa Integer || throw(ArgumentError("opt_level must be an integer, got $(typeof(val))"))
         0 <= val <= 3 || throw(ArgumentError("opt_level must be between 0 and 3, got $val"))
+    elseif key === :num_worker_warps
+        val isa Integer || throw(ArgumentError("num_worker_warps must be an integer, got $(typeof(val))"))
+        val in (4, 8) || throw(ArgumentError("num_worker_warps must be either 4 or 8, got $val"))
     else
         throw(ArgumentError("unknown hint key: $key"))
     end

src/launch.jl

@@ -3,7 +3,7 @@
 # Compiles a Julia function with `TileArray` arguments to Tile IR bytecode,
 # runs `tileiras` to lower bytecode → CUBIN, loads the cubin into the active
 # CUDA context, and launches it via `cudacall`. Compilation is cached per
-# `(MethodInstance, sm_arch, opt_level, num_ctas, occupancy, bytecode_version)`.
+# `(MethodInstance, sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, bytecode_version)`.
 
 using CUDACore: CUDACore, CuArray, CuModule, CuFunction, cudacall, device, capability,
                 AbstractBackend, AbstractKernel, kernel_convert, kernel_compile, PerDevice
@@ -88,7 +88,8 @@ CUDACore.kernel_compile(::TileBackend, f::F, tt::TT=Tuple{}; kwargs...) where {F
 @inline unpack_version(x::UInt16) = VersionNumber(Int(x >> 8), Int(x & 0xff))
 
 # isbits sentinel codec for `Union{Int, Nothing}` hint fields (`opt_level`,
-# `num_ctas`, `occupancy`). `-1` is unused as a value, so we use it for `nothing`.
+# `num_ctas`, `occupancy`, `num_worker_warps`). `-1` is unused as a value, so we
+# use it for `nothing`.
 const _UNSET = -1
 @inline pack_hint(x::Union{Int, Nothing}) = x === nothing ? _UNSET : x
 @inline unpack_hint(x::Int) = x == _UNSET ? nothing : x
@@ -113,12 +114,14 @@ struct TileCacheKey
     opt_level::Int
     num_ctas::Int
     occupancy::Int
+    num_worker_warps::Int
 end
 TileCacheKey(sm_arch::VersionNumber, bytecode_version::VersionNumber,
              opt_level::Union{Int, Nothing}, num_ctas::Union{Int, Nothing},
-             occupancy::Union{Int, Nothing}) =
+             occupancy::Union{Int, Nothing}, num_worker_warps::Union{Int, Nothing}) =
     TileCacheKey(pack_version(sm_arch), pack_version(bytecode_version),
-                 pack_hint(opt_level), pack_hint(num_ctas), pack_hint(occupancy))
+                 pack_hint(opt_level), pack_hint(num_ctas), pack_hint(occupancy),
+                 pack_hint(num_worker_warps))
 
 
 #=============================================================================
@@ -409,7 +412,7 @@ function emit_binary!(cache::CacheView, mi::Core.MethodInstance,
     sm_arch = unpack_version(cache.owner.sm_arch)
 
     # Resolve opt_level here (not in emit_tile) because it's a tileiras flag, not bytecode.
-    # num_ctas/occupancy are resolved in emit_tile because they're encoded in bytecode.
+    # num_ctas/occupancy/num_worker_warps are resolved in emit_tile because they're encoded in bytecode.
     _, _, kernel_meta = res.julia_ir
     opt_level = something(resolve_hint(unpack_hint(cache.owner.opt_level),
                                        kernel_meta, :opt_level, sm_arch), 3)
@@ -516,7 +519,8 @@ end
 
 """
     cuTile.cufunction(f, tt=Tuple{}; sm_arch=nothing, opt_level=nothing,
-                      num_ctas=nothing, occupancy=nothing, name=nothing) -> TileKernel
+                      num_ctas=nothing, occupancy=nothing, num_worker_warps=nothing,
+                      name=nothing) -> TileKernel
 
 Compile `f` for the cuTile backend. `tt` is the tuple of *converted*
 argument types (i.e. after `cuTileconvert`/`Adapt.adapt(KernelAdaptor(), …)`).
@@ -533,11 +537,13 @@ function cufunction(@nospecialize(f), tt::Type{<:Tuple}=Tuple{};
                     opt_level::Union{Int, Nothing}=nothing,
                     num_ctas::Union{Int, Nothing}=nothing,
                     occupancy::Union{Int, Nothing}=nothing,
+                    num_worker_warps::Union{Int, Nothing}=nothing,
                     name::Union{String, Nothing}=nothing)
     bytecode_version = check_tile_ir_support()
     resolved_sm_arch = sm_arch !== nothing ? sm_arch : default_sm_arch()
 
-    key = TileCacheKey(resolved_sm_arch, bytecode_version, opt_level, num_ctas, occupancy)
+    key = TileCacheKey(resolved_sm_arch, bytecode_version, opt_level, num_ctas, occupancy,
+                       num_worker_warps)
 
     # Single pass over `tt.parameters`: build the unwrapped argtypes tuple
     # (Constant{T,V} → T for MI lookup) and the const_argtypes vector
@@ -685,7 +691,7 @@ end
 
 """
     launch(f, grid, args...; sm_arch=nothing, opt_level=nothing,
-           num_ctas=nothing, occupancy=nothing, name=nothing)
+           num_ctas=nothing, occupancy=nothing, num_worker_warps=nothing, name=nothing)
 
 Compile and launch a Tile IR kernel. `args` are converted via
 `cuTileconvert` (CuArray → TileArray, Type → Constant). Equivalent to
@@ -715,10 +721,11 @@ function launch(@nospecialize(f), grid, args...;
                 opt_level::Union{Int, Nothing}=nothing,
                 num_ctas::Union{Int, Nothing}=nothing,
                 occupancy::Union{Int, Nothing}=nothing,
+                num_worker_warps::Union{Int, Nothing}=nothing,
                 name::Union{String, Nothing}=nothing)
     converted = map(cuTileconvert, args)
     tt = Tuple{map(Core.Typeof, converted)...}
-    kernel = cufunction(f, tt; sm_arch, opt_level, num_ctas, occupancy, name)
+    kernel = cufunction(f, tt; sm_arch, opt_level, num_ctas, occupancy, num_worker_warps, name)
     kernel(converted...; blocks=grid)
     return nothing
 end

src/precompile.jl

@@ -44,7 +44,7 @@ import REPL
         bv = bytecode_version()
         for sm_arch in [v"8.0", v"8.6", v"8.7", v"8.9",
                         v"10.0", v"11.0", v"12.0", v"12.1"]
-            key = TileCacheKey(sm_arch, bv, nothing, nothing, nothing)
+            key = TileCacheKey(sm_arch, bv, nothing, nothing, nothing, nothing)
             compile(f, argtypes, const_argtypes, key)
         end
         return

test/device/hints.jl

@@ -71,6 +71,33 @@ end
     end
 end
 
+@testset "launch with num_worker_warps" begin
+    function vadd_kernel_worker_warps(a::ct.TileArray{Float32,1},
+                        b::ct.TileArray{Float32,1},
+                        c::ct.TileArray{Float32,1})
+        pid = ct.bid(1)
+        ta = ct.load(a, pid, (16,))
+        tb = ct.load(b, pid, (16,))
+        ct.store(c, pid, ta + tb)
+        return nothing
+    end
+
+    n = 1024
+    a = CUDA.ones(Float32, n)
+    b = CUDA.ones(Float32, n) .* 2
+    c = CUDA.zeros(Float32, n)
+
+    if cuTile.bytecode_version() >= v"13.3"
+        @cuda backend=cuTile blocks=64 num_worker_warps=8 vadd_kernel_worker_warps(a, b, c)
+        @test Array(c) ≈ ones(Float32, n) .* 3
+    else
+        @test_throws "num_worker_warps requires" @cuda backend=cuTile blocks=64 num_worker_warps=8 vadd_kernel_worker_warps(a, b, c)
+    end
+
+    # Invalid value rejected before compilation.
+    @test_throws "num_worker_warps must be either 4 or 8" @cuda backend=cuTile blocks=64 num_worker_warps=3 vadd_kernel_worker_warps(a, b, c)
+end
+
 end
 
 @testset "Load / Store Optimization Hints" begin