fix: cluster dims

Author: avik-pal

ext/ReactantPythonCallExt/pycall.jl

@@ -166,6 +166,8 @@ function overlayed_pycall_with_triton(
 
     grid = canonicalize_grid(grid, metadata)
 
+    # TODO: actual cluster_x/y/z
+
     return @opcall triton_call(
         pyconvert(String, compiled_kernel.asm["source"]),
         filter(x -> x isa Reactant.TracedType, args)...;
@@ -176,6 +178,9 @@ function overlayed_pycall_with_triton(
         block_x=@opcall(constant(num_warps * device_properties.warp_size)),
         block_y=@opcall(constant(1)),
         block_z=@opcall(constant(1)),
+        cluster_x=@opcall(constant(1)),
+        cluster_y=@opcall(constant(1)),
+        cluster_z=@opcall(constant(1)),
         num_ctas,
         num_warps,
         threads_per_warp=device_properties.warp_size,

src/Ops.jl

@@ -1850,6 +1850,9 @@ function triton_call(
     block_x::TracedRNumber{<:Integer},
     block_y::TracedRNumber{<:Integer},
     block_z::TracedRNumber{<:Integer},
+    cluster_x::TracedRNumber{<:Integer},
+    cluster_y::TracedRNumber{<:Integer},
+    cluster_z::TracedRNumber{<:Integer},
     num_ctas::Integer=1,
     num_warps::Integer=4,
     threads_per_warp::Integer=32,
@@ -1894,6 +1897,9 @@ function triton_call(
         block_x.mlir_data,
         block_y.mlir_data,
         block_z.mlir_data,
+        cluster_x.mlir_data,
+        cluster_y.mlir_data,
+        cluster_z.mlir_data,
         [Reactant.TracedUtils.get_mlir_data(a) for a in args];
         fn=symref,
         result_0=result_types,

src/mlir/Dialects/TritonExt.jl

@@ -20,6 +20,9 @@ function call(
     blockx::Value,
     blocky::Value,
     blockz::Value,
+    clusterx::Value,
+    clustery::Value,
+    clusterz::Value,
     inputs::Vector{Value};
     result_0::Vector{IR.Type},
     fn,
@@ -33,7 +36,9 @@ function call(
     location=Location(),
 )
     op_ty_results = IR.Type[result_0...,]
-    operands = Value[gridx, gridy, gridz, blockx, blocky, blockz, inputs...]
+    operands = Value[
+        gridx, gridy, gridz, blockx, blocky, blockz, clusterx, clustery, clusterz, inputs...
+    ]
     owned_regions = Region[]
     successors = Block[]
     attributes = NamedAttribute[namedattribute("fn", fn),]