Ensure loop count is a constant before trying to unroll. (apache#2797)

neo-ai · Mar 20, 2019 · d07c747 · d07c747
1 parent 300b264
commit d07c747
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/src/pass/unroll_loop.cc b/src/pass/unroll_loop.cc
@@ -78,7 +78,7 @@ class LoopUnroller : public IRMutator {
 
     if ((auto_unroll && explicit_unroll_) ||
         // unroll loops with extent = 1, no matter how many steps in body
-        (value <= auto_max_extent_ && auto_max_extent_ == 1)) {
+        (0 <= value && value <= auto_max_extent_ && auto_max_extent_ == 1)) {
       return Unroll(op);
     } else {
       if (auto_unroll) {

diff --git a/tests/python/unittest/test_pass_unroll.py b/tests/python/unittest/test_pass_unroll.py
@@ -51,7 +51,20 @@ def test_unroll_fake_loop():
     ret = tvm.ir_pass.UnrollLoop(stmt, 8, 0, 1, True)
     assert isinstance(ret.first, tvm.stmt.Store)
 
+def test_unroll_single_count_loops():
+    n = tvm.var('n')
+    A = tvm.placeholder((n,), name='A')
+    B = tvm.compute((n,), lambda *i: A(*i), name='B')
+    s = tvm.create_schedule(B.op)
+    s = s.normalize()
+    dom_map = tvm.schedule.InferBound(s)
+    stmt = tvm.schedule.ScheduleOps(s, dom_map)
+    # all parameters to UnrolLoops are default values except for
+    # auto_unroll_max_extent which has been set to 1 (default:0)
+    after_unroll_stmt = tvm.ir_pass.UnrollLoop(stmt, 0, 8, 1, True)
+    assert after_unroll_stmt == stmt
 
 if __name__ == "__main__":
     test_unroll_loop()
     test_unroll_fake_loop()
+    test_unroll_single_count_loops()