feat(sunmmio): add two-stage SUVM edit workflow

src/target/sunmmio/sunmmio_codegen_tiles_loop.cc

@@ -3670,18 +3670,13 @@ bool CodeGenTileLangSunMMIO::TryLowerTilesScope(const tir::ForNode *op) {
                          select->false_value, select->dtype);
     }
     if (const auto *cast = expr.as<CastNode>()) {
-      SunMMIOValue value = lower_expr(cast->value, state, preferred_dtype);
+      // An explicit cast applies after its operand has been evaluated.
+      SunMMIOValue value = lower_expr(cast->value, state, std::nullopt);
       if (IsTileLike(value)) {
         DataType dst_dtype = CanonicalizeSuvmDType(cast->dtype).with_lanes(1);
         if (value.dtype == dst_dtype) {
           return value;
         }
-        if (preferred_dtype.has_value() && is_float_like_dtype(value.dtype) &&
-            is_float_like_dtype(dst_dtype) &&
-            value.dtype ==
-                CanonicalizeSuvmDType(preferred_dtype.value()).with_lanes(1)) {
-          return value;
-        }
         SunMMIOType dst_type = MakeTileType(CanonicalizeSuvmDType(cast->dtype),
                                             ExtractStaticShape(value.type));
         return builder_->Cast(NewValueName(), value, dst_type,

testing/python/sunmmio/codegen/test_tile_ops_opt_validate.py

@@ -1,4 +1,5 @@
 import os
+import re
 
 import tilelang
 import tilelang.language as T
@@ -17,6 +18,7 @@
 # os.environ["SUNMMIO_TEST_LOG_IR"] = "1"
 
 LOOSE_OPT_ARGS = ("--verify-each",)
+STRICT_OPT_ARGS = ("--verify-each", "--suvm-to-llvm-pipeline")
 
 
 def validate_sunmmio_codegen_loose(kernel, tmp_path, *, mlir_filename, expected_tokens=()):
@@ -188,6 +190,38 @@ def main(
     return main
 
 
+@target("Sunmmio")
+def fp32_select_then_bf16_cast_test(m=32, n=32):
+    input_dtype = T.float32
+    output_dtype = T.bfloat16
+    shard_policy = T.MeshShardingPolicy()
+    tensor_shape = (m, n)
+    tensor_layout = make_zz_layout(tensor_shape, [0, 1], tensor_shape)
+
+    @T.prim_func
+    def main(
+        A: T.MeshTensor(tensor_shape, shard_policy, input_dtype, layout=tensor_layout),  # type: ignore
+        C: T.MeshTensor(tensor_shape, shard_policy, output_dtype, layout=tensor_layout),  # type: ignore
+    ):
+        with T.Kernel():
+            A_shared = T.alloc_shared(tensor_shape, input_dtype)
+            C_shared = T.alloc_shared(tensor_shape, output_dtype)
+
+            T.copy(A, A_shared)
+            for i, j in T.Tiles(A_shared, parallel=True):
+                C_shared[i, j] = T.Cast(
+                    output_dtype,
+                    T.if_then_else(
+                        A_shared[i, j] > T.float32(0),
+                        A_shared[i, j],
+                        T.float32(0),
+                    ),
+                )
+            T.copy(C_shared, C)
+
+    return main
+
+
 def test_tile_elementwise_ops_2d_codegen_validates_with_npuir_opt(tmp_path):
     src = validate_sunmmio_codegen_with_npuir_opt(
         tile_elementwise_ops_2d_test(),
@@ -244,5 +278,28 @@ def test_tile_elementwise_ops_codegen_validates_loose_with_npuir_opt(tmp_path):
     )
 
 
+def test_fp32_select_is_evaluated_before_bf16_cast(tmp_path):
+    src = validate_sunmmio_codegen_with_npuir_opt(
+        fp32_select_then_bf16_cast_test(),
+        tmp_path,
+        mlir_filename="fp32_select_then_bf16_cast_suvm.mlir",
+        expected_tokens=("suvm.tile.cmpf", "suvm.tile.select", "suvm.tile.cast"),
+        opt_args=STRICT_OPT_ARGS,
+    )
+
+    select = re.search(
+        r"(?P<result>%[\w.]+) = suvm\.tile\.select .*"
+        r"!suvm\.tile<[^>]*xf32>, !suvm\.tile<[^>]*xf32>"
+        r" -> !suvm\.tile<[^>]*xf32>",
+        src,
+    )
+    assert select, src
+    assert re.search(
+        rf"suvm\.tile\.cast {re.escape(select.group('result'))} : "
+        r"!suvm\.tile<[^>]*xf32> -> !suvm\.tile<[^>]*xbf16>",
+        src,
+    ), src
+
+
 if __name__ == "__main__":
     tilelang.testing.main()