hack to tensorize loop mapping to make conv2d work

Author: masahi

src/tir/schedule/analysis/analysis.cc

@@ -2069,16 +2069,24 @@ Optional<TensorizeInfo> GetTensorizeLoopMapping(const tir::ScheduleState& self,
     // 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 = 0, n = block_loops.size(); i < n; ++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);
-      if (!tir::UsesVar(r, [&block_loop_vars](const tir::VarNode* 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 == 0 && int_block_extent->value == int_desc_extent->value) || !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) {
@@ -2093,6 +2101,7 @@ Optional<TensorizeInfo> GetTensorizeLoopMapping(const tir::ScheduleState& self,
             return desc_loop_vars.count(var);
           })) {
         desc_loop = desc_loops[i];
+	LOG(INFO) << "Selected " << i << " th desc loop " << desc_loop->extent;;
         break;
       }
     }