From 08d68a89bb4fe5036336baf175877911026206f1 Mon Sep 17 00:00:00 2001 From: lightzhan-intellif Date: Thu, 19 Jan 2023 01:51:01 +0000 Subject: [PATCH] Fix the crash of the pass RemoveNoOp. --- src/tir/transforms/remove_no_op.cc | 5 +++++ .../unittest/test_tir_transform_remove_no_op.py | 14 ++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/src/tir/transforms/remove_no_op.cc b/src/tir/transforms/remove_no_op.cc index 430c1f41bfaf..d35cf8b8d602 100644 --- a/src/tir/transforms/remove_no_op.cc +++ b/src/tir/transforms/remove_no_op.cc @@ -119,6 +119,11 @@ class NoOpRemover : public arith::IRMutatorWithAnalyzer { Stmt VisitStmt_(const IfThenElseNode* op) final { Stmt stmt = Parent::VisitStmt_(op); op = stmt.as(); + // Sometimes the condition can be statically determined, + // in which the type of the `stmt` will not be IfThenElseNode. + if (!op) { + return stmt; + } if (op->else_case) { bool no_op_else = is_no_op(op->else_case.value()); bool no_op_then = is_no_op(op->then_case); diff --git a/tests/python/unittest/test_tir_transform_remove_no_op.py b/tests/python/unittest/test_tir_transform_remove_no_op.py index ce37329b7ed3..06d9289aa795 100644 --- a/tests/python/unittest/test_tir_transform_remove_no_op.py +++ b/tests/python/unittest/test_tir_transform_remove_no_op.py @@ -603,5 +603,19 @@ def expected(A: T.Buffer[16, "int32"], C: T.Buffer[1, "int32"]): C[0] = C[0] + B[i] +class TestCertainConditon(BaseBeforeAfter): + """The conditon of the If-Else node is certain. + This would cause `Segmentation fault` error before.""" + + def before(): + if True: + T.evaluate(0) + else: + T.evaluate(0) + + def expected(): + T.evaluate(0) + + if __name__ == "__main__": tvm.testing.main()