rebase upstream

Author: wrongtest-intellif

src/tir/schedule/primitive/layout_transformation.cc

@@ -392,8 +392,9 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref,
 
   auto iter_map = arith::DetectIterMap(
       /*indices=*/transformed_block_iters, /*input_iters=*/block_iter_dom, /*predicate=*/Bool(true),
-      /*require_bijective=*/true, &analyzer, /*simplify_trivial_iterators=*/true);
-  if (iter_map.empty()) {
+      /*check_level=*/arith::IterMapLevel::Bijective, &analyzer,
+      /*simplify_trivial_iterators=*/true);
+  if (iter_map->indices.empty()) {
     throw NotBijectiveAffineIndexMapError(self->mod, index_map);
   }
 
@@ -417,7 +418,7 @@ void TransformBlockLayout(ScheduleState self, const StmtSRef& block_sref,
   // Step 5.2: Update the block body. Use the inverse map f^{-1} to replace the original block iters
   // in the body.
 
-  auto inverse_map = arith::InverseAffineIterMap(iter_map, new_block_vars);
+  auto inverse_map = arith::InverseAffineIterMap(iter_map->indices, new_block_vars);
   // Trivial block iters will be simplified in DetectIterMap, they should be mapped to constant
   // zero.
   for (const auto& iter_var : block_ptr->iter_vars) {

tests/python/unittest/test_tir_schedule_compute_at.py

@@ -1281,13 +1281,10 @@ def grouped_channel_bias_non_perfect_tiled(
                     cc = T.axis.spatial(720, c_o * 360 + c_i)
                     Y[cc, hh, ww] = X[cc, hh, ww] + B[cc // 16]
 
-    def check_sched(debug_mask):
-        sch = tir.Schedule(grouped_channel_bias, debug_mask=debug_mask)
-        loop = sch.get_loops(sch.get_block("compute"))[0]
-        sch.compute_at(sch.get_block("init"), loop)
-        tvm.ir.assert_structural_equal(sch.mod["main"], grouped_channel_bias_non_perfect_tiled)
-
-    check_sched("all")
+    sch = tir.Schedule(grouped_channel_bias, debug_mask=debug_mask)
+    loop = sch.get_loops(sch.get_block("compute"))[0]
+    sch.compute_at(sch.get_block("init"), loop)
+    tvm.ir.assert_structural_equal(sch.mod["main"], grouped_channel_bias_non_perfect_tiled)
 
 
 def test_fail_subtree_complete_block():