test: validate known block size through kernel launches

Replace brittle final-IR metadata matching with synchronized 512/1024-thread launches.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: cschenjunlin

tests/kernels/test_rmsnorm_autotune.py

@@ -5,7 +5,6 @@
 
 import json
 import os
-import re
 from pathlib import Path
 
 import pytest
@@ -71,8 +70,6 @@ def test_rmsnorm_direct_specializes_known_block_size(weight_dtype, weight_dtype_
     artifact = compiled._keepalive
 
     assert "known_block_size = array<i32: 512, 1, 1>" in artifact.source_ir
-    match = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", artifact.ir)
-    assert match is not None and int(match.group(1)) == 512
     if weight_dtype == torch.float32:
         weight_copy_type = "!fly.copy_atom<!fly_rocdl.cdna3.buffer_copy<128>, 32>"
         assert artifact.source_ir.count(weight_copy_type) >= 3

tests/unit/test_kernel_known_block_size.py

@@ -1,7 +1,5 @@
 """Tests for known_block_size attribute on @flyc.kernel."""
 
-import re
-
 import pytest
 
 import flydsl.compiler as flyc
@@ -129,14 +127,6 @@ def _get_source_ir(launch_fn, *args):
     return artifact.source_ir
 
 
-def _get_compiled_ir(launch_fn, x):
-    """Call the JIT function once, then return the compiled IR string."""
-    launch_fn(x, stream=torch.cuda.current_stream())
-    assert launch_fn._mem_cache, "expected at least one cached compilation"
-    artifact = next(iter(launch_fn._mem_cache.values()))
-    return artifact.ir
-
-
 # ---------------------------------------------------------------------------
 # Tests
 # ---------------------------------------------------------------------------
@@ -168,14 +158,6 @@ def test_static_numeric_dims_are_inferred(self):
         assert "known_block_size = array<i32: 64, 1, 1>" in source_ir
         assert "known_block_size = array<i32: 32, 2, 1>" in source_ir
 
-    def test_compiled_ir_has_max_flat_workgroup_size(self):
-        compiled_ir = _get_compiled_ir(_launch_bs128_4_2, self.x)
-        # The compiled IR should report max_flat_workgroup_size >= total_threads
-        match = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", compiled_ir)
-        assert match is not None, f"max_flat_workgroup_size not found in compiled IR:\n{compiled_ir}"
-        max_wg = int(match.group(1))
-        assert max_wg >= 1024, f"max_flat_workgroup_size={max_wg} < total_threads=1024"
-
     def test_dynamic_block_size_omits_exact_attribute(self):
         source_ir = _get_source_ir(_launch_dynamic_blocks, self.x, 64, 64)
         # Check for the attribute syntax, not just the substring (which may
@@ -203,9 +185,10 @@ def test_reused_explicit_name_gets_unique_symbols(self):
         assert "gpu.func @_kn_named_block_specialization(" in source_ir
         assert "gpu.func @_kn_named_block_specialization_1(" in source_ir
 
-    def test_kernel_launches_successfully(self):
-        """Ensure the kernel actually launches without hipErrorLaunchFailure."""
-        _launch_bs64(self.x, stream=torch.cuda.current_stream())
+    @pytest.mark.parametrize("launch_fn", [_launch_bs64, _launch_bs128_4_2], ids=["64", "1024"])
+    def test_kernel_launches_successfully(self, launch_fn):
+        """Exercise backend workgroup metadata, including a 1024-thread launch."""
+        launch_fn(self.x, stream=torch.cuda.current_stream())
         torch.cuda.synchronize()  # would raise if launch failed