Fix CancelNestedInv pass

Before, the passes tried to delete the innerInvOp but this failed as
they still used the innerInnerUnitary. Now, the innerInnerUnitary is
cloned and the operation can be deleted (happens within the
replaceOp(...)).

Author: Ectras

mlir/lib/Dialect/Flux/IR/Modifiers/InvOp.cpp

@@ -29,7 +29,7 @@ using namespace mlir::flux;
 namespace {
 
 /**
- * @brief Cancel nested inverse modifiers, i.e., `inv(inv(x)) = x`.
+ * @brief Cancel nested inverse modifiers, i.e., `inv(inv(x)) => x`.
  */
 struct CancelNestedInv final : OpRewritePattern<InvOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -44,8 +44,28 @@ struct CancelNestedInv final : OpRewritePattern<InvOp> {
 
     // Remove both inverse operations
     auto innerInnerUnitary = innerInvOp.getBodyUnitary();
-    rewriter.replaceOp(op, innerInnerUnitary.getOperation());
-    rewriter.eraseOp(innerInvOp);
+    auto* clonedOp = rewriter.clone(*innerInnerUnitary.getOperation());
+    rewriter.replaceOp(op, clonedOp->getResults());
+
+    return success();
+  }
+};
+
+/**
+ * @brief Reorder inverse and control modifiers, i.e., `inv(ctrl(x)) =>
+ * ctrl(inv(x))`.
+ */
+struct ReorderInvCtrl final : OpRewritePattern<InvOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(InvOp op,
+                                PatternRewriter& rewriter) const override {
+    auto ctrlOp = llvm::dyn_cast<CtrlOp>(op.getBodyUnitary().getOperation());
+    if (!ctrlOp) {
+      return failure();
+    }
+
+    rewriter.replaceOp(op, ctrlOp);
 
     return success();
   }
@@ -203,5 +223,5 @@ LogicalResult InvOp::verify() {
 
 void InvOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                         MLIRContext* context) {
-  results.add<CancelNestedInv>(context);
+  results.add<CancelNestedInv, ReorderInvCtrl>(context);
 }

mlir/lib/Dialect/Quartz/IR/Modifiers/InvOp.cpp

@@ -27,7 +27,7 @@ using namespace mlir::quartz;
 namespace {
 
 /**
- * @brief Cancel nested inverse modifiers.
+ * @brief Cancel nested inverse modifiers, i.e., `inv(inv(x)) => x`.
  */
 struct CancelNestedInv final : OpRewritePattern<InvOp> {
   using OpRewritePattern::OpRewritePattern;
@@ -40,8 +40,8 @@ struct CancelNestedInv final : OpRewritePattern<InvOp> {
     }
 
     auto innerInnerUnitary = innerInvOp.getBodyUnitary();
-    rewriter.replaceOp(invOp, innerInnerUnitary.getOperation());
-    rewriter.eraseOp(innerInvOp);
+    auto* clonedOp = rewriter.clone(*innerInnerUnitary.getOperation());
+    rewriter.replaceOp(invOp, clonedOp->getResults());
 
     return success();
   }