Fixed more review comments

Change-Id: If6133dd822f33a8d32f3ddd8b8ce22b92490694e
apache · Jun 9, 2021 · 2b61414 · 2b61414
1 parent 53b99b1
commit 2b61414
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/python/tvm/tir/transform/inject_rolling_buffer.py b/python/tvm/tir/transform/inject_rolling_buffer.py
@@ -46,7 +46,7 @@ def InjectRollingBuffer():
     buffer_to_attrs = defaultdict(list)
     rolling_buffers = set()
     rolling_buffer_to_info = dict()
-    iter_vars = list()
+    for_loops = list()
     hoist_buffer_to_for = defaultdict(list)
 
     RollingBufferInfo = namedtuple(
@@ -56,7 +56,7 @@ def InjectRollingBuffer():
     def _pre_visit(stmt):
         if isinstance(stmt, tvm.tir.For):
             # Manage the stack of iter_vars
-            iter_vars.append(stmt)
+            for_loops.append(stmt)
 
         elif isinstance(stmt, tvm.tir.AttrStmt):
             if isinstance(stmt.node, tvm.tir.Buffer):
@@ -118,7 +118,7 @@ def _pre_visit(stmt):
                 # to be the rolling axis
                 roll_iter_var = None
                 roll_axis = -1
-                for loop in iter_vars:
+                for loop in for_loops:
                     iter_var = loop.loop_var
                     if iter_var in bound_iter_vars:
                         roll_iter_var = iter_var
@@ -148,7 +148,7 @@ def _pre_visit(stmt):
     def _post_visit(stmt):
         if isinstance(stmt, tvm.tir.For):
             # Manage the stack of iter_vars
-            iter_vars.pop()
+            for_loops.pop()
             # If the loop corresponds to an iter_var that needs a BufferRealize
             # hoisting to its scope, perform the hoisting
             if stmt.loop_var in hoist_buffer_to_for:

diff --git a/src/te/operation/compute_op.cc b/src/te/operation/compute_op.cc
@@ -484,9 +484,8 @@ ComputeLoopNest ComputeLoopNest::Create(const BaseComputeOpNode* self, const Sta
     }
     ret.init_nest = MakeLoopNest(stage, dom_map, begin_loop, true, skip_iter, &(ret.init_vmap),
                                  debug_keep_trivial_loop);
-    bool skip_ivar_domain = !stage->rolling_buffer;
     ret.init_predicates =
-        MakeBoundCheck(stage, dom_map, ret.init_vmap, skip_ivar_domain, skip_iter);
+        MakeBoundCheck(stage, dom_map, ret.init_vmap, !stage->rolling_buffer, skip_iter);
     for (auto& e : ret.init_predicates) {
       e = likely(e);
     }