Consolidate asymmetric nounroll schedule into parameterized asymmetric schedule

The no-unroll path needs a different kernel interleaving strategy than
the unrolled path: 2-group interleaving (shared A loads interleaved
with MMA) with B loads and G2S prefetches in a separate third cluster,
rather than 4-group interleaving that folds B loads and G2S directly
into the two MMA clusters.  The 4-group pattern was designed for the
unrolled kernel where the larger loop body can absorb the extra live
values; with unroll_factor=1 the tighter loop needs the third cluster
to keep VGPR pressure in check.

Author: willghatch

examples/python/7.1_schedule.py

@@ -27,7 +27,6 @@
     get_mxfp4_dbuf_pingpong_schedule,
     get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_asymmetric_schedule,
-    get_mxfp4_asymmetric_nounroll_schedule,
     get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds,
@@ -387,14 +386,12 @@ def test_dbuf_4wave_mxfp_preshuffle_b_gemm_cpp(
     options.wave_runtime = True
     options.dump_intermediates = "build/intermediates"
     options.eliminate_epilogue = eliminate_epilogue
-    if no_unroll:
-        schedule = get_mxfp4_asymmetric_nounroll_schedule(
-            eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-        )
-    else:
-        schedule = get_mxfp4_asymmetric_schedule(
-            eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-        )
+    schedule = get_mxfp4_asymmetric_schedule(
+        eliminate_epilogue=eliminate_epilogue,
+        is_bscale_shuffled=True,
+        unroll_factor=1 if no_unroll else 2,
+        unroll_kernel=not no_unroll,
+    )
     options.print_ir_after = "all" if is_debug else []
     options = set_default_run_config(options)
     gemm = wave_compile(options, gemm, schedule)

tests/kernel/wave/asm/test_waveasm_e2e.py

@@ -1169,7 +1169,6 @@ def _dbuf_mxfp4_helper(
     from wave_lang.kernel.wave.schedules import (
         get_mxfp4_dbuf_schedule,
         get_mxfp4_asymmetric_schedule,
-        get_mxfp4_asymmetric_nounroll_schedule,
     )
     from wave_lang.kernel.wave.scheduling.schedule_enums import SchedulingType
     from wave_lang.kernel.wave.utils.run_utils import set_default_run_config
@@ -1202,14 +1201,12 @@ def _dbuf_mxfp4_helper(
         )
         options.eliminate_epilogue = eliminate_epilogue
         if use_schedule:
-            if no_unroll:
-                schedule = get_mxfp4_asymmetric_nounroll_schedule(
-                    eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-                )
-            else:
-                schedule = get_mxfp4_asymmetric_schedule(
-                    eliminate_epilogue=eliminate_epilogue, is_bscale_shuffled=True
-                )
+            schedule = get_mxfp4_asymmetric_schedule(
+                eliminate_epilogue=eliminate_epilogue,
+                is_bscale_shuffled=True,
+                unroll_factor=1 if no_unroll else 2,
+                unroll_kernel=not no_unroll,
+            )
         else:
             schedule = None
             options.schedule = SchedulingType.NONE

wave_lang/kernel/wave/schedules/__init__.py

@@ -18,7 +18,6 @@
     get_mxfp4_dbuf_mixed_pingpong_schedule,
     get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule,
     get_mxfp4_asymmetric_schedule,
-    get_mxfp4_asymmetric_nounroll_schedule,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled,
     get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds,
 )
@@ -35,7 +34,6 @@
     "get_mxfp4_dbuf_pingpong_schedule_Bshuffled",
     "get_mxfp4_dbuf_pingpong_schedule_Bshuffled_lds",
     "get_mxfp4_asymmetric_schedule",
-    "get_mxfp4_asymmetric_nounroll_schedule",
     "get_mxfp4_dbuf_mixed_pingpong_shuffle_schedule",
     "get_attention_prefetch_schedule",
 ]