Reworking GetTensorizeloopmapping

Author: masahi

src/tir/schedule/analysis/analysis.cc

@@ -2096,87 +2096,43 @@ Optional<TensorizeInfo> GetTensorizeLoopMapping(const tir::ScheduleState& self,
   if (offset < 0) {
     return NullOpt;
   }
-  // We align the block and desc block's bindings from the right side
-  // block     (v0=..., v1=..., v2=...)
-  //                    ^ i_block
-  // desc_block(        v1=..., v2=...)
-  //                    ^ i_desc
-
-  std::vector<IterVarType> iter_types = GetBlockVarTypes(block_sref);
-  ICHECK(block_loops.size() == iter_types.size());
-
-  for (int i_desc = 0, i_block = offset; i_desc < n_desc_vars; ++i_desc, ++i_block) {
-    // For each block var binding, we find
-    const PrimExpr& block_bind = block->iter_values[i_block];
-    const PrimExpr& desc_bind = desc_block->iter_values[i_desc];
-    LOG(INFO) << "block bind: " << block_bind;
-    LOG(INFO) << "desc bind: " << desc_bind;
-    // Step 4.1. Find the corresponding loop of the i-th block var of block
-    const tir::ForNode* block_loop = nullptr;
-    for (int i = block_loops.size() - 1; i >= 0; --i) {
-      // Check if block_bind = block_loops[i]->loop_var + stuff-irrelevant-of-loop-vars
-      PrimExpr r = analyzer.Simplify(block_bind - block_loops[i]->loop_var);
-      const auto* int_block_extent = block_loops[i]->extent.as<IntImmNode>();
-      const auto* int_desc_extent = desc_loops[i_desc]->extent.as<IntImmNode>();
-
-      if (i_desc != n_desc_vars - 1 && iter_types[i] == IterVarType::kCommReduce) continue;
-
-      // if (int_block_extent->value == int_desc_extent->value) {
-      if (!tir::UsesVar(r, [&block_loop_vars](const tir::VarNode* var) {
-            return block_loop_vars.count(var);
-          })) {
-        block_loop = block_loops[i];
-        LOG(INFO) << "Selected " << i << " th block loop " << block_loops[i]->loop_var << ", "
-                  << block_loop->extent;
-        break;
-      } else {
-        LOG(INFO) << i << " th block loop not ok "
-                  << ", " << block_loops[i]->loop_var << ", " << block_loops[i]->extent;
-      }
-    }
-    if (block_loop == nullptr) {
-      return NullOpt;
-    }
-    // Step 4.2. Find the corresponding loop of the i-th block var of desc
-    const tir::ForNode* desc_loop = nullptr;
-    for (int i = 0, n = desc_loops.size(); i < n; ++i) {
-      // Check if desc_bind = loops[i]->loop_var + stuff-irrelevant-of-loop-vars
-      PrimExpr r = analyzer.Simplify(desc_bind - desc_loops[i]->loop_var);
-      if (!tir::UsesVar(r, [&desc_loop_vars](const tir::VarNode* var) {
-            return desc_loop_vars.count(var);
-          })) {
-        desc_loop = desc_loops[i];
-        LOG(INFO) << "Selected " << i << " th desc loop " << desc_loop->extent;
-        ;
+
+  std::vector<IterVarType> iter_types_desc;
+  for (const IterVar& iter_var : desc_block->block->iter_vars) {
+    iter_types_desc.push_back(iter_var->iter_type);
+  }
+
+  std::vector<IterVarType> iter_types_block = GetBlockVarTypes(block_sref);
+  ICHECK(block_loops.size() == iter_types_block.size());
+
+  int next_block_ind = block_loops.size() - 1;
+  for (int i_desc = n_desc_vars - 1; i_desc >= 0; --i_desc) {
+    const tir::ForNode* desc_loop = desc_loops[i_desc];
+    const auto* int_desc_extent = desc_loop->extent.as<IntImmNode>();
+
+    for (int i_block = next_block_ind; i_block >= 0; --i_block) {
+      const tir::ForNode* block_loop = block_loops[i_block];
+      const auto* int_block_extent = block_loop->extent.as<IntImmNode>();
+
+      LOG(INFO) << i_desc << ", " << i_block << ", " << iter_types_block[i_block] << ", "
+                << iter_types_desc[i_desc] << ", " << int_block_extent->value << ", "
+                << int_desc_extent->value;
+
+      if (int_block_extent->value % int_desc_extent->value != 0) continue;
+      if (iter_types_block[i_block] != iter_types_desc[i_desc]) continue;
+
+      const tir::StmtSRef& block_loop_sref = self->stmt2ref[block_loop];
+      auto it = ret->loop_map.find(block_loop_sref);
+      if (it == ret->loop_map.end()) {
+        ret->loop_map.Set(block_loop_sref, GetRef<tir::For>(desc_loop));
+	next_block_ind = i_block - 1;
+        LOG(INFO) << "Selected " << i_block << " th block loop " << block_loops[i_block]->loop_var
+                  << ", " << block_loop->extent << " for i_desc = " << i_desc;
         break;
       }
     }
-    if (desc_loop == nullptr) {
-      return NullOpt;
-    }
-    // Step 4.3. Check divisibility of loop extents
-    PrimExpr block_extent = analyzer.Simplify(block_loop->extent);
-    PrimExpr desc_extent = analyzer.Simplify(desc_loop->extent);
-    if (const auto* int_block_extent = block_extent.as<IntImmNode>()) {
-      if (const auto* int_desc_extent = desc_extent.as<IntImmNode>()) {
-        if (int_block_extent->value % int_desc_extent->value != 0) {
-          return NullOpt;
-        }
-      } else {
-        return NullOpt;
-      }
-    } else {
-      return NullOpt;
-    }
-    // Step 4.4. Maps the result of Step 4.1 to Step 4.2
-    const tir::StmtSRef& block_loop_sref = self->stmt2ref[block_loop];
-    auto it = ret->loop_map.find(block_loop_sref);
-    if (it == ret->loop_map.end()) {
-      ret->loop_map.Set(block_loop_sref, GetRef<tir::For>(desc_loop));
-    } else if ((*it).second.get() != desc_loop) {
-      return NullOpt;
-    }
   }
+
   for (int i = 0, n = desc_loops.size(); i < n; ++i) {
     ret->desc_loop_indexer.Set(GetRef<tir::For>(desc_loops[i]), Integer(i));
   }