Fix @device_code_tiled for kernels with Constant arguments (#127)

The compile_hook in emit_ir passed unwrapped types to code_tiled
because launch does MI lookup with Constant stripped to its element
type. This caused @device_code_tiled to fail with "Unsupported Julia
type for Tile IR" for any kernel using Constant args (including the
tiled broadcast kernels).

Reconstruct Constant{T,V} types from const_argtypes in the hook so
code_tiled can recover const-seeded arguments.

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

Author: maleadt

src/compiler/interface.jl

@@ -376,11 +376,24 @@ IR where constant values are folded in.
 function emit_ir(cache::CacheView, mi::Core.MethodInstance;
                  const_argtypes::Union{Vector{Any}, Nothing}=nothing)
     # Invoke compile hook if set (for @device_code_* reflection)
-    # Pass (f, tt) tuple to enable direct use with reflection utilities
+    # Pass (f, tt) tuple to enable direct use with reflection utilities.
+    # Reconstruct Constant{T,V} types from const_argtypes so that code_tiled
+    # can recover const-seeded arguments (MI specTypes have them unwrapped).
     if compile_hook[] !== nothing
         ftype = mi.specTypes.parameters[1]
         f = isdefined(ftype, :instance) ? ftype.instance : ftype
-        tt = Tuple{mi.specTypes.parameters[2:end]...}
+        arg_types = collect(Any, mi.specTypes.parameters[2:end])
+        if const_argtypes !== nothing
+            # const_argtypes is [Const(f), arg2, ...]; arg_types omits f,
+            # so arg_types[i] corresponds to const_argtypes[i+1].
+            for i in eachindex(arg_types)
+                if const_argtypes[i+1] isa CC.Const
+                    val = const_argtypes[i+1].val
+                    arg_types[i] = Constant{typeof(val), val}
+                end
+            end
+        end
+        tt = Tuple{arg_types...}
         compile_hook[](f, tt)
     end
 

test/execution/core.jl

@@ -197,6 +197,27 @@ end # invalidations
         ct.@device_code_typed io=buf ct.launch(reflect_vadd, cld(n, 16), a, b, c)
         String(take!(buf))
     end
+
+    # @device_code_tiled with Constant arguments
+    function reflect_const_vadd(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1},
+                                c::ct.TileArray{Float32,1}, tile_size::Int)
+        pid = ct.bid(1)
+        tile_a = ct.load(a, pid, (tile_size,))
+        tile_b = ct.load(b, pid, (tile_size,))
+        ct.store(c, pid, tile_a + tile_b)
+        return
+    end
+
+    c2 = CUDA.zeros(Float32, n)
+    @test @filecheck begin
+        @check "entry @reflect_const_vadd"
+        @check "load_view"
+        @check "addf"
+        @check "store_view"
+        ct.@device_code_tiled ct.launch(reflect_const_vadd, cld(n, 16), a, b, c2,
+                                        ct.Constant(16))
+    end
+    @test Array(c2) ≈ Array(a) + Array(b)
 end
 
 @testset "assert" begin