Fix asan failure introduced by cl/210618122 and statement walker crash for if stateme...

PiperOrigin-RevId: 211186361

diff --git a/mlir/include/mlir/IR/StmtVisitor.h b/mlir/include/mlir/IR/StmtVisitor.h
index a9605de974f3f3b83297395440324392dcbc0674..5bd0b132bb1be822938e93993ad20b9ebb8c11ea 100644 (file)
--- a/mlir/include/mlir/IR/StmtVisitor.h
+++ b/mlir/include/mlir/IR/StmtVisitor.h
@@ -159,8 +159,9 @@ public:
     static_cast<SubClass *>(this)->visitIfStmt(ifStmt);
     static_cast<SubClass *>(this)->walk(ifStmt->getThen()->begin(),
                                         ifStmt->getThen()->end());
-    static_cast<SubClass *>(this)->walk(ifStmt->getElse()->begin(),
-                                        ifStmt->getElse()->end());
+    if (ifStmt->hasElse())
+      static_cast<SubClass *>(this)->walk(ifStmt->getElse()->begin(),
+                                          ifStmt->getElse()->end());
   }
 
   void walkIfStmtPostOrder(IfStmt *ifStmt) {

diff --git a/mlir/lib/IR/Statement.cpp b/mlir/lib/IR/Statement.cpp
index 673aa4317405bc85e6a1e2a772bb0ccebf34b1a4..2a5a2d789468ab66830e31b5183fe355ff9d2b82 100644 (file)
--- a/mlir/lib/IR/Statement.cpp
+++ b/mlir/lib/IR/Statement.cpp
@@ -366,9 +366,10 @@ IfStmt::IfStmt(Location *location, unsigned numOperands, IntegerSet *set)
 }
 
 IfStmt::~IfStmt() {
-  delete thenClause;
   if (elseClause)
     delete elseClause;
+
+  delete thenClause;
   // An IfStmt's IntegerSet 'set' should not be deleted since it is
   // allocated through MLIRContext's bump pointer allocator.
 }

diff --git a/mlir/test/IR/parser.mlir b/mlir/test/IR/parser.mlir
index 9624b75919762a6cbbc769481aa43d7dc67747c8..8bdfe8b35ce0d6224746786e724f53d8113893d4 100644 (file)
--- a/mlir/test/IR/parser.mlir
+++ b/mlir/test/IR/parser.mlir
@@ -258,6 +258,19 @@ mlfunc @ifstmt(%N: affineint) {
   return    // CHECK   return
 }           // CHECK }
 
+// CHECK-LABEL: mlfunc @simple_ifstmt(%arg0 : affineint) {
+mlfunc @simple_ifstmt(%N: affineint) {
+  %c = constant 200 : affineint // CHECK   %c200 = constant 200
+  for %i = 1 to 10 {                   // CHECK   for %i0 = 1 to 10 {
+    if @@set0(%i)[%N, %c] {     // CHECK     if @@set0(%i0)[%arg0, %c200] {
+      %x = constant 1 : i32
+       // CHECK: %c1_i32 = constant 1 : i32
+      %y = "add"(%x, %i) : (i32, affineint) -> i32 // CHECK: %0 = "add"(%c1_i32, %i0) : (i32, affineint) -> i32
+      %z = "mul"(%y, %y) : (i32, i32) -> i32 // CHECK: %1 = "mul"(%0, %0) : (i32, i32) -> i32
+    }       // CHECK     }
+  }         // CHECK   }
+  return    // CHECK   return
+}           // CHECK }
 
 // CHECK-LABEL: cfgfunc @attributes() {
 cfgfunc @attributes() {