Generalize struct destructuring to support arbitrary nested structs

The compiler's struct destructuring was TileArray-specific. This
generalizes it so any isbits struct with non-ghost fields (Ptr, Tuple,
nested structs) is recursively flattened into kernel parameters.

Key changes:
- `should_destructure` accepts any qualifying struct, not just TileArray
- `arg_flat_values` uses path-based keys (Vector{Union{Symbol,Int}})
  instead of single field names
- Recursive `flatten_struct_params!` / `_flatten_tuple_param!` handle
  nested structs and heterogeneous tuples
- `cache_tensor_view!` accepts a path for nested TileArrays
- `make_tensor_view` resolves lazy arg refs for nested TileArrays
- Host-side `flatten()` recursively flattens any destructurable struct

This enables passing structs containing TileArrays (e.g. Broadcasted)
directly as kernel arguments.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: maleadt

ext/CUDAExt.jl

@@ -3,7 +3,7 @@ module CUDAExt
 using cuTile
 using cuTile: TileArray, Constant, ByTarget, CGOpts, CuTileResults, DEFAULT_BYTECODE_VERSION,
               emit_code, sanitize_name, constant_eltype, constant_value, is_ghost_type,
-              resolve_hint, format_sm_arch, validate_hint
+              should_destructure, resolve_hint, format_sm_arch, validate_hint
 
 using CompilerCaching: CacheView, method_instance, results
 
@@ -254,8 +254,26 @@ return their fields in order.
 
 This is used by the launch helper to splat arguments to cudacall.
 """
-flatten(x) = is_ghost_type(typeof(x)) ? () : (x,)
-flatten(arr::TileArray{T, N}) where {T, N} = (arr.ptr, arr.sizes..., arr.strides...)
+function flatten(x)
+    T = typeof(x)
+    is_ghost_type(T) && return ()
+    cuTile.should_destructure(T) || return (x,)
+    result = Any[]
+    for fi in 1:fieldcount(T)
+        fval = getfield(x, fi)
+        fT = typeof(fval)
+        if is_ghost_type(fT) || !cuTile._is_kernel_param_type(fT)
+            continue
+        elseif fval isa Tuple
+            for elem in fval
+                push!(result, flatten(elem)...)
+            end
+        else
+            push!(result, flatten(fval)...)
+        end
+    end
+    return Tuple(result)
+end
 
 """
     to_tile_arg(x)

src/compiler/codegen/kernel.jl

@@ -38,7 +38,7 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
 
     # Build parameter list, handling ghost types, const args, and struct destructuring
     param_types = TypeId[]
-    param_mapping = Tuple{Int, Union{Nothing, Symbol}}[]
+    param_mapping = Tuple{Int, Vector{Union{Symbol, Int}}}[]
 
     for (i, argtype) in enumerate(sci.argtypes)
         argtype_unwrapped = CC.widenconst(argtype)
@@ -47,28 +47,11 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
         elseif is_const_arg(i)
             continue  # const arg: no kernel parameter
         elseif should_destructure(argtype_unwrapped)
-            # Destructure TileArray into flat parameters
-            params = argtype_unwrapped.parameters
-            ndims = params[2]::Integer
-            for fi in 1:fieldcount(argtype_unwrapped)
-                fname = fieldname(argtype_unwrapped, fi)
-                ftype = fieldtype(argtype_unwrapped, fi)
-                if fname === :sizes || fname === :strides
-                    fcount = ndims
-                    elem_type = Int32
-                else
-                    fcount = flat_field_count(ftype)
-                    elem_type = ftype <: Ptr ? Ptr{params[1]} : (ftype <: Tuple ? eltype(ftype) : ftype)
-                end
-                for _ in 1:fcount
-                    push!(param_types, tile_type_for_julia!(ctx, elem_type))
-                    push!(param_mapping, (i, fname))
-                end
-            end
+            flatten_struct_params!(ctx, param_types, param_mapping, i, argtype_unwrapped, Union{Symbol, Int}[])
             ctx.arg_types[i] = argtype_unwrapped
         else
             push!(param_types, tile_type_for_julia!(ctx, argtype_unwrapped))
-            push!(param_mapping, (i, nothing))
+            push!(param_mapping, (i, Union{Symbol, Int}[]))
         end
     end
 
@@ -90,7 +73,7 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     arg_values = make_block_args!(cb, length(param_types))
 
     # Build arg_flat_values map
-    field_values = Dict{Tuple{Int, Union{Nothing, Symbol}}, Vector{Value}}()
+    field_values = Dict{Tuple{Int, Vector{Union{Symbol, Int}}}, Vector{Value}}()
     for (param_idx, val) in enumerate(arg_values)
         key = param_mapping[param_idx]
         if !haskey(field_values, key)
@@ -100,13 +83,12 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     end
 
     # Store in context and set up slot/argument mappings
-    # arg_idx is the direct index into argtypes (2, 3, ...) which matches SlotNumber/Argument
     for (key, values) in field_values
-        arg_idx, field = key
+        arg_idx, path = key
         ctx.arg_flat_values[key] = values
 
-        if field === nothing
-            # Regular argument - create concrete CGVal
+        if isempty(path) && !haskey(ctx.arg_types, arg_idx)
+            # Regular (non-destructured) argument - create concrete CGVal
             if length(values) != 1
                 throw(IRError("Expected exactly one value for argument $arg_idx, got $(length(values))"))
             end
@@ -148,8 +130,9 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     end
 
     # Create TensorViews for all TileArray arguments at kernel entry
-    for (arg_idx, _) in ctx.arg_types
-        cache_tensor_view!(ctx, arg_idx)
+    # Walk the type tree to find all nested TileArrays
+    for (arg_idx, argtype) in ctx.arg_types
+        _create_tensor_views_recursive!(ctx, arg_idx, argtype, Union{Symbol, Int}[])
     end
 
     # Create memory ordering token
@@ -166,6 +149,96 @@ function emit_kernel!(writer::BytecodeWriter, func_buf::Vector{UInt8},
     finalize_function!(func_buf, cb, writer.debug_info)
 end
 
+"""
+    _create_tensor_views_recursive!(ctx, arg_idx, T, path)
+
+Walk the type tree and create TensorViews for all nested TileArrays.
+"""
+function _create_tensor_views_recursive!(ctx::CGCtx, arg_idx::Int, @nospecialize(T), path::Vector{Union{Symbol, Int}})
+    if T <: TileArray
+        cache_tensor_view!(ctx, arg_idx, path, T)
+    elseif T <: Tuple
+        for i in 1:length(T.parameters)
+            ptype = T.parameters[i]
+            elem_path = Union{Symbol, Int}[path..., i]
+            _create_tensor_views_recursive!(ctx, arg_idx, ptype, elem_path)
+        end
+    else
+        for fi in 1:fieldcount(T)
+            ftype = fieldtype(T, fi)
+            is_ghost_type(ftype) && continue
+            fname = fieldname(T, fi)
+            field_path = Union{Symbol, Int}[path..., fname]
+            if ftype <: TileArray
+                cache_tensor_view!(ctx, arg_idx, field_path, ftype)
+            elseif should_destructure(ftype) || ftype <: Tuple
+                _create_tensor_views_recursive!(ctx, arg_idx, ftype, field_path)
+            end
+        end
+    end
+end
+
+"""
+    flatten_struct_params!(ctx, param_types, param_mapping, arg_idx, T, path)
+
+Recursively flatten a struct type into kernel parameters.
+"""
+function flatten_struct_params!(ctx, param_types, param_mapping, arg_idx, @nospecialize(T), path::Vector{Union{Symbol, Int}})
+    for fi in 1:fieldcount(T)
+        fname = fieldname(T, fi)
+        ftype = fieldtype(T, fi)
+        field_path = Union{Symbol, Int}[path..., fname]
+        if is_ghost_type(ftype)
+            continue
+        elseif ftype <: Tuple && ftype !== Tuple{}
+            _flatten_tuple_param!(ctx, param_types, param_mapping, arg_idx, ftype, field_path)
+        elseif should_destructure(ftype)
+            flatten_struct_params!(ctx, param_types, param_mapping, arg_idx, ftype, field_path)
+        else
+            # Scalar/Ptr field → 1 flat param (skip if no tile IR type)
+            type_id = tile_type_for_julia!(ctx, ftype; throw_error=false)
+            type_id === nothing && continue
+            push!(param_types, type_id)
+            push!(param_mapping, (arg_idx, field_path))
+        end
+    end
+end
+
+"""
+Flatten a Tuple field. If all elements are the same primitive type, emit N grouped
+params under field_path. Otherwise, index each element and recurse if needed.
+"""
+function _flatten_tuple_param!(ctx, param_types, param_mapping, arg_idx, @nospecialize(ftype), field_path)
+    N = length(ftype.parameters)
+    # Check if this is a homogeneous tuple of simple leaf types (NTuple{N, T})
+    et = eltype(ftype)
+    if isconcretetype(et) && !is_ghost_type(et) && !should_destructure(et) && !(et <: Tuple)
+        # Simple case: all elements are the same leaf type (e.g., NTuple{N, Int32})
+        for _ in 1:N
+            push!(param_types, tile_type_for_julia!(ctx, et))
+            push!(param_mapping, (arg_idx, field_path))
+        end
+    else
+        # Heterogeneous or complex tuple: recurse per element
+        for i in 1:N
+            elem_type = ftype.parameters[i]
+            elem_path = Union{Symbol, Int}[field_path..., i]
+            if is_ghost_type(elem_type)
+                continue
+            elseif elem_type <: Tuple && elem_type !== Tuple{}
+                _flatten_tuple_param!(ctx, param_types, param_mapping, arg_idx, elem_type, elem_path)
+            elseif should_destructure(elem_type)
+                flatten_struct_params!(ctx, param_types, param_mapping, arg_idx, elem_type, elem_path)
+            else
+                type_id = tile_type_for_julia!(ctx, elem_type; throw_error=false)
+                type_id === nothing && continue
+                push!(param_types, type_id)
+                push!(param_mapping, (arg_idx, elem_path))
+            end
+        end
+    end
+end
+
 # getfield for destructured arguments (lazy chain extension)
 function emit_getfield!(ctx::CGCtx, args, @nospecialize(result_type))
     length(args) >= 2 || return nothing
@@ -195,27 +268,28 @@ function emit_getfield!(ctx::CGCtx, args, @nospecialize(result_type))
         if field isa Symbol
             # Field access: extend chain with symbol
             new_chain = Union{Symbol, Int}[chain..., field]
-            # Check if this resolves to a scalar field (auto-materialize leaf)
-            # Don't auto-materialize tuple types - they need indexing first
             rt = CC.widenconst(result_type)
             if !(rt <: Tuple)
-                values = get_arg_flat_values(ctx, arg_idx, field)
+                # Check if this path resolves to flat values
+                values = get_arg_flat_values(ctx, arg_idx, new_chain)
                 if values !== nothing && length(values) == 1
                     # Scalar field - materialize immediately
                     type_id = tile_type_for_julia!(ctx, rt)
                     return CGVal(values[1], type_id, rt)
                 end
             end
             return arg_ref_value(arg_idx, new_chain, rt)
-        elseif field isa Integer && !isempty(chain) && chain[end] isa Symbol
-            # Tuple indexing: chain ends with field name, now indexing into it
-            # This is a leaf - materialize immediately
-            field_name = chain[end]
-            values = get_arg_flat_values(ctx, arg_idx, field_name)
+        elseif field isa Integer
+            # Tuple indexing into a field
+            # Look up the parent path (which should be a tuple field)
+            values = get_arg_flat_values(ctx, arg_idx, chain)
             if values !== nothing && 1 <= field <= length(values)
                 type_id = tile_type_for_julia!(ctx, CC.widenconst(result_type))
                 return CGVal(values[field], type_id, CC.widenconst(result_type))
             end
+            # Not a simple flat tuple — extend chain (element may be destructured)
+            new_chain = Union{Symbol, Int}[chain..., Int(field)]
+            return arg_ref_value(arg_idx, new_chain, CC.widenconst(result_type))
         end
     end
 
@@ -241,15 +315,11 @@ function emit_getindex!(ctx::CGCtx, args, @nospecialize(result_type))
     if is_arg_ref(obj_tv)
         arg_idx, chain = obj_tv.arg_ref
 
-        # If chain ends with a symbol (field name), we're indexing into a tuple field
-        # Try to materialize immediately
-        if !isempty(chain) && chain[end] isa Symbol
-            field_name = chain[end]
-            values = get_arg_flat_values(ctx, arg_idx, field_name)
-            if values !== nothing && 1 <= index <= length(values)
-                type_id = tile_type_for_julia!(ctx, CC.widenconst(result_type))
-                return CGVal(values[index], type_id, CC.widenconst(result_type))
-            end
+        # Try to materialize: the chain should point to a tuple field
+        values = get_arg_flat_values(ctx, arg_idx, chain)
+        if values !== nothing && 1 <= index <= length(values)
+            type_id = tile_type_for_julia!(ctx, CC.widenconst(result_type))
+            return CGVal(values[index], type_id, CC.widenconst(result_type))
         end
 
         # Otherwise extend the chain

src/compiler/codegen/utils.jl

@@ -151,12 +151,14 @@ mutable struct CGCtx
     slots::Dict{Int, CGVal}       # Slot number -> CGVal
     block_args::Dict{Int, CGVal}  # BlockArg id -> CGVal (for control flow)
 
-    # Destructured argument handling (for TileArray fields)
-    arg_flat_values::Dict{Tuple{Int, Union{Nothing, Symbol}}, Vector{Value}}
+    # Destructured argument handling: path-keyed flat values
+    # Key: (arg_idx, path) where path is e.g. [:ptr] or [:a, :sizes]
+    arg_flat_values::Dict{Tuple{Int, Vector{Union{Symbol, Int}}}, Vector{Value}}
     arg_types::Dict{Int, Type}
 
-    # Cached TensorViews for TileArray arguments (arg_idx -> (Value, TypeId))
-    tensor_views::Dict{Int, Tuple{Value, TypeId}}
+    # Cached TensorViews for TileArray arguments
+    # Key: arg_idx::Int for top-level, or (arg_idx, path) for nested
+    tensor_views::Dict{Any, Tuple{Value, TypeId}}
 
     # Bytecode infrastructure
     cb::CodeBuilder
@@ -188,9 +190,9 @@ function CGCtx(; cb::CodeBuilder, tt::TypeTable, sci::StructuredIRCode,
         Dict{Int, CGVal}(),
         Dict{Int, CGVal}(),
         Dict{Int, CGVal}(),
-        Dict{Tuple{Int, Union{Nothing, Symbol}}, Vector{Value}}(),
+        Dict{Tuple{Int, Vector{Union{Symbol, Int}}}, Vector{Value}}(),
         Dict{Int, Type}(),
-        Dict{Int, Tuple{Value, TypeId}}(),
+        Dict{Any, Tuple{Value, TypeId}}(),
         cb, tt, sci, token, token_type, type_cache, sm_arch, cache,
     )
 end
@@ -238,12 +240,27 @@ end
 =============================================================================#
 
 """
-    get_arg_flat_values(ctx, arg_idx, field=nothing) -> Union{Vector{Value}, Nothing}
+    get_arg_flat_values(ctx, arg_idx, path) -> Union{Vector{Value}, Nothing}
 
-Get the flat Tile IR values for an argument or its field.
+Get the flat Tile IR values for a destructured argument at the given path.
 """
-function get_arg_flat_values(ctx::CGCtx, arg_idx::Int, field::Union{Nothing, Symbol}=nothing)
-    get(ctx.arg_flat_values, (arg_idx, field), nothing)
+function get_arg_flat_values(ctx::CGCtx, arg_idx::Int, path::Vector{Union{Symbol, Int}})
+    get(ctx.arg_flat_values, (arg_idx, path), nothing)
+end
+
+# Convenience: single field name
+function get_arg_flat_values(ctx::CGCtx, arg_idx::Int, field::Symbol)
+    get_arg_flat_values(ctx, arg_idx, Union{Symbol, Int}[field])
+end
+
+# Convenience: no path = top-level
+function get_arg_flat_values(ctx::CGCtx, arg_idx::Int)
+    # Collect all values for this arg_idx across all paths
+    result = Value[]
+    for ((idx, _path), vals) in ctx.arg_flat_values
+        idx == arg_idx && append!(result, vals)
+    end
+    isempty(result) ? nothing : result
 end
 
 """
@@ -385,23 +402,61 @@ end
     should_destructure(T) -> Bool
 
 Check if a type should be destructured into flat parameters.
+Any isbits struct with non-ghost, non-primitive fields qualifies.
 """
 function should_destructure(@nospecialize(T))
     T = CC.widenconst(T)
     isstructtype(T) || return false
     is_ghost_type(T) && return false
     isprimitivetype(T) && return false
-    T <: TileArray && return true
+    T <: Tuple && return false  # Tuples are handled as flat params, not recursed
+    # Must have a concrete layout we can iterate
+    try fieldcount(T) catch; return false end
+    for fi in 1:fieldcount(T)
+        ft = fieldtype(T, fi)
+        _is_kernel_param_type(ft) && return true
+    end
+    return false
+end
+
+# Check if a field type contributes kernel parameters (recursively).
+# Only isbits types (or types containing them) can be kernel parameters.
+function _is_kernel_param_type(@nospecialize(ft))
+    is_ghost_type(ft) && return false
+    isprimitivetype(ft) && return true
+    ft <: Ptr && return true
+    if ft <: Tuple && ft !== Tuple{}
+        # Check if any tuple element is a kernel param type
+        for p in ft.parameters
+            _is_kernel_param_type(p) && return true
+        end
+        return false
+    end
+    # For structs, only recurse if isbits (prevents infinite recursion on DataType etc.)
+    isstructtype(ft) && isbitstype(ft) && should_destructure(ft) && return true
     return false
 end
 
 """
     flat_field_count(T) -> Int
 
-Count flat parameters a type expands to.
+Count flat parameters a type expands to (recursive).
 """
-flat_field_count(::Type{<:NTuple{N, T}}) where {N, T} = N
-flat_field_count(::Type) = 1
+function flat_field_count(@nospecialize(T))
+    if is_ghost_type(T)
+        return 0
+    elseif should_destructure(T)
+        count = 0
+        for fi in 1:fieldcount(T)
+            count += flat_field_count(fieldtype(T, fi))
+        end
+        return count
+    elseif T <: Tuple
+        return length(T.parameters)
+    else
+        return 1
+    end
+end
 
 #-----------------------------------------------------------------------------
 # Argument helpers