Fix bug

iree-org · Aug 22, 2023 · 8813bbd · 8813bbd
1 parent 5439bf9
commit 8813bbd
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/compiler/src/iree/compiler/Dialect/Flow/Transforms/RaiseSpecialOps.cpp b/compiler/src/iree/compiler/Dialect/Flow/Transforms/RaiseSpecialOps.cpp
@@ -286,17 +286,16 @@ struct RaiseSpecialOpsPass : public RaiseSpecialOpsBase<RaiseSpecialOpsPass> {
     IRRewriter rewriter(&getContext());
 
     getOperation()->walk([&](linalg::GenericOp op) {
-      rewriter.setInsertionPoint(op);
-
       // Try raising to tensor.export.
+      rewriter.setInsertionPoint(op);
       FailureOr<linalg::GenericOp> maybeNewOp =
           raiseTensorExtractToInput(op, rewriter);
-      // Update op if we succeeded.
       if (succeeded(maybeNewOp)) {
-        op = maybeNewOp.value();
+        op = *maybeNewOp;
       }
 
       // Try raising to a view-like operation.
+      rewriter.setInsertionPoint(op);
       (void)tryRaiseToView(op, rewriter);
     });