Take zero extent loops as NoOp and remove it and add unittest for the same (#3724)

diff --git a/src/pass/remove_no_op.cc b/src/pass/remove_no_op.cc
index c367bd6..c7797f4 100644 (file)
--- a/src/pass/remove_no_op.cc
+++ b/src/pass/remove_no_op.cc
@@ -70,6 +70,9 @@ class NoOpRemover : public IRMutator {
   Stmt Mutate_(const For* op, const Stmt& s) final {
     Stmt stmt = IRMutator::Mutate_(op, s);
     op = stmt.as<For>();
+    if (is_zero(op->extent)) {
+      return Evaluate::make(0);
+    }
     return is_no_op(op->body) ? MakeEvaluate({op->min, op->extent}) : stmt;
   }
   Stmt Mutate_(const Allocate* op, const Stmt& s) final {

diff --git a/tests/python/unittest/test_pass_remove_no_op.py b/tests/python/unittest/test_pass_remove_no_op.py
index b6fb171..d1ea450 100644 (file)
--- a/tests/python/unittest/test_pass_remove_no_op.py
+++ b/tests/python/unittest/test_pass_remove_no_op.py
@@ -39,7 +39,10 @@ def test_remove_no_op():
                            i + 1)
     stmt2 = tvm.make.Block(stmt, store)
     assert(tvm.ir_pass.RemoveNoOp(stmt2) == store)
-
+    # remove zero extent loop
+    stmt3 = tvm.make.For(i, 0, 0, 0, 0, store)
+    ret = tvm.ir_pass.RemoveNoOp(stmt3)
+    assert(isinstance(ret, tvm.stmt.Evaluate))
 
 if __name__ == "__main__":
     test_remove_no_op()