fix auto separate bug (PaddlePaddle#108)

thisjiang · Mar 29, 2020 · c062900 · c062900
1 parent 13b63ea
commit c062900
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/cinn/optim/transform_polyfor_to_for.cc b/cinn/optim/transform_polyfor_to_for.cc
@@ -191,8 +191,7 @@ struct ForAutoSeparateMutator : ir::IRMutator<Expr*> {
 
 Expr* PolyForAutoSeparateHelper(Expr* expr) {
   ForAutoSeparateMutator mutator;
-  auto* res = mutator(expr);
-  if (res) return res;
+  return mutator(expr);
 }
 
 struct ForAutoSeparateMutatorMain : public ir::IRMutator<Expr*> {
@@ -214,7 +213,10 @@ struct ForAutoSeparateMutatorMain : public ir::IRMutator<Expr*> {
     if (res) Visit(res);
   }
 
-  void Visit(Expr* expr) { ir::IRMutator<>::Visit(expr, expr); }
+  void Visit(Expr* expr) {
+    CHECK(expr);
+    ir::IRMutator<>::Visit(expr, expr);
+  }
 };
 
 Expr PlusOneWithMinMax(Expr expr) {