fix bugs in broadcast_rewrite

edfd65aa · wYann · df57a6cf · edfd65aa
隐藏空白更改
内联并排

Showing with 1 addition and 1 deletion

src/pass/broadcast_rewrite.cc src/pass/broadcast_rewrite.cc +1 -1

未找到文件。
--- a/src/pass/broadcast_rewrite.cc
+++ b/src/pass/broadcast_rewrite.cc
@@ -80,7 +80,7 @@ class BroadcastVecRewriter : public IRMutator {
            }
          }
-          if (secDimFit && dimFit && forLoopFit) {
+          if (secDimFit && dimFit && forLoopFit && remainDim > 1) {
            newExtent = GetInt32Const(GetItem(forInfo.ops_, idx).as<For>()->extent) / remainDim;
            tmpBuffer = VarExpr("tmp_broadcast_" + std::to_string(broadBufferCount++) + "_local_UB", dtype);
            varList = dstInfo->var_;