LLVM Discussion Forums

Rewrite pattern doesn't update root with new operation

Hi,

I’m trying to build a C++ rewrite specification, passing a FuncOp to the matchAndRewrite and iterating over all underlying Operations using func.walk.

For some iterations, I need to insert an extra operation in place.
For example:
%0 = “X”(%arg0)
%1 = “Op_0”(%0)
%2 = “Y”(%1)
now becomes:
%0 = “X”(%arg0)
%1 = “Op_0”(%0)
%2 = “Op_1”(%1)
%3 = “Y”(%2)

I use setInsertionPointAfter and I call startRootUpdate/finalizeRootUpdate at the begin/end of the pattern.

When I dump the graph while still in the pattern scope (applyPartialConversion), the IR shows me the new operations (Op_1 in the example) and how Y expects %2 as its input.
However after the pass is done, the new operation is not being returned and I get an error that Y expects an operation that doesn’t exist:
%0 = “X”(%arg0)
%1 = “Op_0”(%0)
%2 = “Y”(<< UNKNOWN SSA VALUE >>)

How can I make sure the rewriter properly updates the graph with the new operations?

Thanks in advance :+1:t3:

-debug='dialect-conversion' will give you a log of what’s happening. Your conversion is likely failing here and so your changes would just be undone. Have you already looked at: https://mlir.llvm.org/docs/DialectConversion/#debugging ? In general, there’s a
lot of information here: https://mlir.llvm.org/docs/DialectConversion/

Something isn’t clear to me in what you’re describing here: how are you using func.walk? In general pattern are invoked with a driver (like applyPatternsAndFoldGreedily or applyPartialConversion) on the func op itself, and the driver will process the entire function.

Thanks for this suggestion, I ran with this flag and it looks indeed like it’s failing during the pattern application step.
The next step is to figure out why the updated ‘func’ is considered illegal: any ideas?
See also my answer to @joker-eph for my usage of the pattern.

//===-------------------------------------------===//
Legalizing operation : 'func'(0x187ce00) {
  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'func -> ()' {
    ** Insert  : 'Op_1'(0x187d1e0)

    //===-------------------------------------------===//
    Legalizing operation : 'func'(0x187ce00) {
      * Fold {
      } -> FAILURE : unable to fold

      * Pattern : 'func -> ()' {
      } -> FAILURE : pattern was already applied
    } -> FAILURE : no matched legalization pattern
    //===-------------------------------------------===//
  } -> FAILURE : operation updated in-place 'func' was illegal
} -> FAILURE : no matched legalization pattern
//===-------------------------------------------===//

That’s impossible to say without knowing what target conversion set you have in place! See addLegal..., addIllegal..., and addDynamicIllegal... from the doc.

1 Like

Hi,

I use applyPartialConversion to invoke my pattern.
While still in the scope of the pattern, my inserted operation Op_1 appears in the debug prints.
However, after inspecting the debug results as suggested by @bondhugula it seems like the resulting FuncOp is illegal and therefore the conversion is undone.

This is what the code looks like:

struct AddMyPattern : public RewritePattern {

  AddMyPattern (MLIRContext *context) : RewritePattern("func", 1, context) {}

  LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override {

    auto func = llvm::dyn_cast<FuncOp>(op);

    // Notify MLIR we're updating the function in place
    rewriter.startRootUpdate(op);

    // Iterate over the operations in the function (and its region)
    func.walk([&](Operation *thisOp) {
  
      // Insert Op_1 after Op_0
      if (thisOp->getName().getStringRef.compare("Op_0") != 0) {
  
        // Compute Op_1 attribute value
        auto result = getResult();
  
        // Insert the Op_1 after the copy of the original operation
        rewriter.setInsertionPointAfter(thisOp);
  
        // Create Op_1 template (type OpTy) at the current insertion point
        auto myNewOpTy = rewriter.create<Op_1>(thisOp->getLoc(), thisOp->getResult(0).getType(), thisOp->getResult(0), result);
  
        // Extract the pointer to the actual Op_1 operation
        SmallPtrSet<Operation *, 1> myNewOp_ptr{myNewOpTy .getOperation()};
  
        // Replace the input of the destination with the output of the new operation
        thisOp->getResult(0).replaceAllUsesExcept(myNewOpTy.getResult(), myNewOp_ptr);
      }
    });
  
    // Notify MLIR in place updates are done
    rewriter.finalizeRootUpdate(op);
  
    return success(); 
  }
};

void MyPass::runOnFunction() {

    auto &context = getContext();
    ConversionTarget target(context);

    OwningRewritePatternList patterns;
    patterns.insert<MyPattern>(&context);

    if (failed(applyPartialConversion(getFunction(), target, patterns))) {
      signalPassFailure();
    }

    patterns.clear();
}

@bondhugula I followed up on your suggestion of using dynamicallyLegal.
After defining FuncOp legal when including the new Op_1 operation, and illegal when excluding, the pattern was successfully applied.

Thank you (and @joker-eph) for your support!

It isn’t clear to me why you’re using a pattern at all this way.
You could just do you func.walk from the runOnFunction() entry point without using the pattern framework.

Using a pattern can be a good idea, but in general you wouldn’t match the function and walk it this way, instead your pattern would match on Op_0 and rewrite it.

I think I understand what you mean.
Are you saying it is weird to use a pattern if you’re going to explicitly write the func.walk yourself adding extra operations?
Rather, using the pattern framework pays off when your graph topology stays the same but only the content of the operation needs modification (e.g add an extra attribute).

In the meantime, our dialect already got changed such that the contents of Op_1 are now an attribute of Op_0: no need to insert it as an operation anymore.
But it’s always to good to learn how this framework is used properly, so any advise is welcome👍🏻

Not necessarily: you can change the topology with the patterns, but the patterns are written in a way that they match the root operation they intend to change.

You can look for example of uses of the Patterns in the repo, for example the Vector dialect to LLVM conversion here: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
Or the lowering of Vector operations to loops: https://github.com/llvm/llvm-project/blob/master/mlir/lib/Conversion/VectorToSCF/VectorToSCF.cpp