[MLIR] Pre-Order WalkAndApplyPatterns Driver

[MatthiasReumann]

Hey community!

I'm currently working on a project where pre-order has some benefits over post-order traversal. Unfortunately, as of now, MLIR only includes the walkAndApplyPatterns function which iterates the IR post-order.

I'm aware of the existence of applyPatternsGreedily. However, I'm (1) not sure if it can be configured as a simple pre-order iteration¹ and (2) if this would even advisable due to the overhead added because a simple walk would suffice.

Consequently, I've implemented it myself. Please keep in mind that I'm not a expert in LLVM and MLIR - so suggestions for improvements welcome:

/**
 * @brief A pre-order version of the walkAndApplyPatterns driver.
 * @param op The operation to walk. Does not visit the op itself.
 * @param patterns A set of patterns to apply.
 * @link Adapted from https://mlir.llvm.org/doxygen/WalkPatternRewriteDriver_8cpp_source.html
 */
void walkPreOrderAndApplyPatterns(Operation* op,
                                  const FrozenRewritePatternSet& patterns) {
  MLIRContext* ctx = op->getContext();
  PatternRewriter rewriter(ctx);

  PatternApplicator applicator(patterns);
  applicator.applyDefaultCostModel();

  ctx->executeAction<PreOrderWalkDriverAction>(
      [&] {
        op->walk<WalkOrder::PreOrder>([&](Operation* x) {
          std::ignore = applicator.matchAndRewrite(x, rewriter);
        });
      },
      {op});
}

Nonetheless, my question is: Is there any particular reason not to have a pre-order walk driver included in MLIR?

Footnotes

1. At least I didn't manage to do it with the help of the documentation. ↩

[llvmbot]

@llvm/issue-subscribers-mlir

Author: matthias (MatthiasReumann)

Hey community!

I'm currently working on a project where pre-order has some benefits over post-order traversal. Unfortunately, as of now, MLIR only includes the walkAndApplyPatterns function which iterates the IR post-order.

I'm aware of the existence of applyPatternsGreedily. However, I'm (1) not sure if it can be configured as a simple pre-order iteration¹ and (2) if this would even advisable due to the overhead added because a simple walk would suffice.

Consequently, I've implemented it myself. Please keep in mind that I'm not a expert in LLVM and MLIR - so suggestions for improvements welcome:

/**
 * @<!-- -->brief A pre-order version of the walkAndApplyPatterns driver.
 * @<!-- -->param op The operation to walk. Does not visit the op itself.
 * @<!-- -->param patterns A set of patterns to apply.
 * @<!-- -->link Adapted from https://mlir.llvm.org/doxygen/WalkPatternRewriteDriver_8cpp_source.html
 */
void walkPreOrderAndApplyPatterns(Operation* op,
                                  const FrozenRewritePatternSet&amp; patterns) {
  MLIRContext* ctx = op-&gt;getContext();
  PatternRewriter rewriter(ctx);

  PatternApplicator applicator(patterns);
  applicator.applyDefaultCostModel();

  ctx-&gt;executeAction&lt;PreOrderWalkDriverAction&gt;(
      [&amp;] {
        op-&gt;walk&lt;WalkOrder::PreOrder&gt;([&amp;](Operation* x) {
          std::ignore = applicator.matchAndRewrite(x, rewriter);
        });
      },
      {op});
}

Nonetheless, my question is: Is there any particular reason not to have a pre-order walk driver included in MLIR?

Footnotes

1. At least I didn't manage to do it with the help of the documentation. ↩

[joker-eph]

This probably could be a configuration flag on this API: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h

[MatthiasReumann]

This probably could be a configuration flag on this API: https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Transforms/WalkPatternRewriteDriver.h

Sounds reasonable. Would you mind if I start working on this in the near future?

The above implementation should probably be refactored to a pre-order worklist and extensive API checks, right?

[joker-eph]

Yes, feel free to send me a PR whenever.