[Gluon] Fix warp_specialize when the default function has no results (#7145)

Author: Mogball

python/test/gluon/test_frontend.py

@@ -321,7 +321,7 @@ def anchor(x):
 @filecheck_test
 @gluon.jit
 def test_warp_specialize():
-    # CHECK:         [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
+    # CHECK:       [[BLOCKED:#.*]] = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
     # CHECK-LABEL: test_warp_specialize
     # CHECK-NEXT:    [[A:%.*]] = tt.make_range {end = 1 : i32, start = 0 : i32}
     # CHECK-NEXT:    [[B:%.*]] = tt.make_range {end = 2 : i32, start = 0 : i32}
@@ -352,6 +352,10 @@ def test_warp_specialize():
     anchor(a)
     anchor(b)
 
+    # CHECK: ttg.warp_specialize([[A]], [[B]], [[C]])
+    # CHECK: (tensor<1xi32, [[BLOCKED]]>, tensor<2xi32, [[BLOCKED]]>, tensor<4xi32, [[BLOCKED]]>) -> ()
+    ttgl.warp_specialize((pair, c, e), warp_specialize_worker0, [warp_specialize_worker1], [4], [48])
+
 
 @gluon.jit
 def mbarrier_kernel():

python/triton/experimental/gluon/language/_semantic.py

@@ -257,7 +257,9 @@ def warp_specialize(self, args, default_partition, worker_partitions, worker_num
         default_block = builder.new_block()
         builder.set_insertion_point_to_start(default_block)
         default_results = generator.call_JitFunction(default_partition, args, kwargs={})
-        mlir_results = flatten_values_to_ir(default_results)
+        mlir_results = []
+        if default_results is not None:
+            mlir_results = flatten_values_to_ir(default_results)
         builder.create_warp_yield(mlir_results)
         result_types = [r.get_type() for r in mlir_results]
 
@@ -281,4 +283,6 @@ def warp_specialize(self, args, default_partition, worker_partitions, worker_num
 
         builder.set_insertion_point_after(ws_op.get_operation())
         mlir_results = [ws_op.get_result(i) for i in range(len(result_types))]
+        if default_results is None:
+            return
         return tuple(unflatten_ir_values(mlir_results, [r.type for r in default_results]))