[mlir] Reorder shape assuming bufferization.

The previous ordering continued to use the original assuming after
replacing it which is not allowed. Now, inline the region from the old
into the new before the replacement.

Differential Revision: https://reviews.llvm.org/D90375

diff --git a/mlir/lib/Dialect/Shape/Transforms/StructuralTypeConversions.cpp b/mlir/lib/Dialect/Shape/Transforms/StructuralTypeConversions.cpp
index 61e8628..041b54b 100644 (file)
--- a/mlir/lib/Dialect/Shape/Transforms/StructuralTypeConversions.cpp
+++ b/mlir/lib/Dialect/Shape/Transforms/StructuralTypeConversions.cpp
@@ -32,10 +32,9 @@ public:
 
     auto newAssumingOp =
         rewriter.create<AssumingOp>(op.getLoc(), newResultTypes, op.witness());
-
-    rewriter.replaceOp(op, newAssumingOp.getResults());
     rewriter.inlineRegionBefore(op.doRegion(), newAssumingOp.doRegion(),
                                 newAssumingOp.doRegion().end());
+    rewriter.replaceOp(op, newAssumingOp.getResults());
 
     return success();
   }