[mlir] Support dialect-wide canonicalization pattern registration

[springerm]

Summary

• Add hasCanonicalizer option to Dialect.
• Initialize canonicalizer with dialect-wide canonicalization patterns.
• Add test case to TestDialect.

Dialect-wide canonicalization patterns are useful if a canonicalization pattern does not conceptually associate with any single operation, i.e., it should not be registered as part of an operation's getCanonicalizationPatterns function. E.g., this is the case for canonicalization patterns that match an op interface.

[mehdi_amini]

LGTM, but please wait for at least a day to give a change to River to have a look as well.

mlir/include/mlir/IR/Dialect.h

77

The dialect already has a getContext() member I believe

[rriddle]

LGTM

mlir/include/mlir/IR/Dialect.h

71

Comments should be ///

77

Yes, the context parameter here isn't necessary. Dialect already has a context member.

mlir/test/lib/Dialect/Test/TestDialect.cpp

295

Can you use the inline version on RewritePatternSet?

i.e: https://github.com/llvm/llvm-project/blob/d82f2a123f9c443911fc40009d2017915b850758/mlir/include/mlir/IR/PatternMatch.h#L937

[springerm]

rebase

[springerm]

address comments

[GMNGeoffrey]

So one thought that comes up for me with this patch is that while we have general ideas about what a canonicalization for an op should limit itself to (should be conceptually related to the op), I'm not sure there's restrictions in the code or clear documentation. I think this could lead to a weird scenario where just having registered a dialect will change the behavior of the canonicalize pass, even if that dialect is completely unused. This can already happen today with op canonicalizations if the canonicalization doesn't require the presence of the op it is registered on. I think that such a canonicalization is kind of obviously a bad idea (but we should maybe document that in https://mlir.llvm.org/docs/Canonicalization/). With the addition of dialect canonicalization patterns though, it seems like this scenario is more likely, especially with the stated use case of a canonicalization that operates on an interface. Since interfaces are not tied to dialects, this means someone could declare a canonicalization pattern that operates on CastOpInterface, which would make sense because it's something their dialect uses, but this could affect the behavior of other dialects. In some ways this is a powerful mechanism for declaring canonicalizations on other ops you don't own, but if we want to support that I think it should probably be more specific. I'm not sure there's much we can do in the code to enforce this, but maybe a dialect canonicalization pattern should ensure that it only matches when ops from that dialect are present (and we should document this and a similar restriction on op canonicalization patterns). Thoughts?

mlir/include/mlir/IR/Dialect.h

public:
   /// addOperation.
   bool allowsUnknownOperations() const { return unknownOpsAllowed; }
 
   /// Return true if this dialect allows for unregistered types, i.e., types
   /// prefixed with the dialect namespace but not registered with addType.
   /// These are represented with OpaqueType.
   bool allowsUnknownTypes() const { return unknownTypesAllowed; }
 
+  /// Register dialect-wide canonicalization patterns. This method should only
+  /// be used to register canonicalization patterns that do not conceptually
+  /// belong to any single operation in the dialect. (In that case, use the op's
+  /// canonicalizer.) E.g., canonicalization patterns for op interfaces should
+  /// be registered here.
+  virtual void getCanonicalizationPatterns(RewritePatternSet &results) const {}
+
   /// Registered hook to materialize a single constant operation from a given
   /// attribute value with the desired resultant type. This method should use
   /// the provided builder to create the operation without changing the
   /// insertion position. The generated operation is expected to be constant
   /// like, i.e. single result, zero operands, non side-effecting, etc. On
   /// success, this hook should return the value generated to represent the
   /// constant value. Otherwise, it should return null on failure.
   virtual Operation *materializeConstant(OpBuilder &builder, Attribute value,

mlir/include/mlir/IR/OpBase.td

class Dialect {
   // attributes.
   bit hasRegionArgAttrVerify = 0;
 
   // If this dialect overrides the hook for verifying region result attributes.
   bit hasRegionResultAttrVerify = 0;
 
   // If this dialect overrides the hook for op interface fallback.
   bit hasOperationInterfaceFallback = 0;
+
+  // If this dialect overrides the hook for canonicalization patterns.
+  bit hasCanonicalizer = 0;
 }
 
 //===----------------------------------------------------------------------===//
 // Type definitions
 //===----------------------------------------------------------------------===//
 
 // A type, carries type constraints.
 class Type<Pred condition, string descr = "",

mlir/include/mlir/TableGen/Dialect.h

public:
 
   // Returns the list of dialect (class names) that this dialect depends on.
   // These are dialects that will be loaded on construction of this dialect.
   ArrayRef<StringRef> getDependentDialects() const;
 
   // Returns the dialects extra class declaration code.
   llvm::Optional<StringRef> getExtraClassDeclaration() const;
 
+  /// Returns true if this dialect has a canonicalizer.
+  bool hasCanonicalizer() const;
+
   // Returns true if this dialect has a constant materializer.
   bool hasConstantMaterializer() const;
 
   /// Returns true if this dialect has an operation attribute verifier.
   bool hasOperationAttrVerify() const;
 
   /// Returns true if this dialect has a region argument attribute verifier.
   bool hasRegionArgAttrVerify() const;

mlir/lib/TableGen/Dialect.cpp

ArrayRef<StringRef> Dialect::getDependentDialects() const {
   return dependentDialects;
 }
 
 llvm::Optional<StringRef> Dialect::getExtraClassDeclaration() const {
   auto value = def->getValueAsString("extraClassDeclaration");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }
 
+bool Dialect::hasCanonicalizer() const {
+  return def->getValueAsBit("hasCanonicalizer");
+}
+
 bool Dialect::hasConstantMaterializer() const {
   return def->getValueAsBit("hasConstantMaterializer");
 }
 
 bool Dialect::hasOperationAttrVerify() const {
   return def->getValueAsBit("hasOperationAttrVerify");
 }
 

mlir/lib/Transforms/Canonicalizer.cpp

Canonicalizer() {
     config.enableRegionSimplification = enableRegionSimplification;
     config.maxIterations = maxIterations;
   }
 
   /// Initialize the canonicalizer by building the set of patterns used during
   /// execution.
   LogicalResult initialize(MLIRContext *context) override {
     RewritePatternSet owningPatterns(context);
+    for (auto *dialect : context->getLoadedDialects())
+      dialect->getCanonicalizationPatterns(owningPatterns);
     for (auto *op : context->getRegisteredOperations())
       op->getCanonicalizationPatterns(owningPatterns, context);
     patterns = std::move(owningPatterns);
     return success();
   }
   void runOnOperation() override {
     (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns,
                                        config);

mlir/test/Transforms/test-canonicalize.mlir

func @result_shape_per_dim(%arg0 : tensor<2x3x?xf32>, %arg1 : tensor<?x5xf32>)
   %2 = memref.dim %0#0, %c1 : tensor<?x5xf32>
   %3 = memref.dim %0#1, %c0 : tensor<2x3x?xf32>
   %4 = memref.dim %0#1, %c1 : tensor<2x3x?xf32>
   %5 = memref.dim %0#1, %c2 : tensor<2x3x?xf32>
   // CHECK-DAG: %[[D0:.+]] = memref.dim %[[ARG_1]], %[[C0]]
   // CHECK-DAG: %[[D1:.+]] = memref.dim %[[ARG_0]], %[[C2]]
   // CHECK: return %[[D0]], %[[C5]], %[[C2]], %[[C3]], %[[D1]]
   return %1, %2, %3, %4, %5 : index, index, index, index, index
-}
+}
+
+// CHECK-LABEL: test_dialect_canonicalizer
+func @test_dialect_canonicalizer() -> (i32) {
+  %0 = "test.dialect_canonicalizable"() : () -> (i32)
+  // CHECK: %[[CST:.*]] = constant 42 : i32
+  // CHECK: return %[[CST]]
+  return %0 : i32
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

 
 Optional<MutableOperandRange>
 TestBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return targetOperandsMutable();
 }
 
 //===----------------------------------------------------------------------===//
+// TestDialectCanonicalizerOp
+//===----------------------------------------------------------------------===//
+
+static LogicalResult
+dialectCanonicalizationPattern(TestDialectCanonicalizerOp op,
+                               PatternRewriter &rewriter) {
+  rewriter.replaceOpWithNewOp<ConstantOp>(op, rewriter.getI32Type(),
+                                          rewriter.getI32IntegerAttr(42));
+  return success();
+}
+
+void TestDialect::getCanonicalizationPatterns(
+    RewritePatternSet &results) const {
+  results.add(&dialectCanonicalizationPattern);
+}
+
+//===----------------------------------------------------------------------===//
 // TestFoldToCallOp
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
   using OpRewritePattern<FoldToCallOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(FoldToCallOp op,

mlir/test/lib/Dialect/Test/TestOps.td

 include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "TestInterfaces.td"
 
 def Test_Dialect : Dialect {
   let name = "test";
   let cppNamespace = "::mlir::test";
+  let hasCanonicalizer = 1;
   let hasConstantMaterializer = 1;
   let hasOperationAttrVerify = 1;
   let hasRegionArgAttrVerify = 1;
   let hasRegionResultAttrVerify = 1;
   let hasOperationInterfaceFallback = 1;
   let dependentDialects = ["::mlir::DLTIDialect"];
 
   let extraClassDeclaration = [{

 
 // An op that always fold itself.
 def TestPassthroughFold : TEST_Op<"passthrough_fold"> {
   let arguments = (ins AnyType:$op);
   let results = (outs AnyType);
   let hasFolder = 1;
 }
 
+def TestDialectCanonicalizerOp : TEST_Op<"dialect_canonicalizable"> {
+  let arguments = (ins);
+  let results = (outs I32);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns (Symbol Binding)
 
 // Test symbol binding.
 def OpSymbolBindingA : TEST_Op<"symbol_binding_a", []> {
   let arguments = (ins I32:$operand, I64Attr:$attr);
   let results = (outs I32);
 }

mlir/tools/mlir-tblgen/DialectGen.cpp

static const char *const typeParserDecl = R"(
   /// Parse a type registered to this dialect.
   ::mlir::Type parseType(::mlir::DialectAsmParser &parser) const override;
 
   /// Print a type registered to this dialect.
   void printType(::mlir::Type type,
                  ::mlir::DialectAsmPrinter &os) const override;
 )";
 
+/// The code block for the canonicalization pattern registration hook.
+static const char *const canonicalizerDecl = R"(
+  /// Register canonicalization patterns.
+  void getCanonicalizationPatterns(
+      ::mlir::RewritePatternSet &results) const override;
+)";
+
 /// The code block for the constant materializer hook.
 static const char *const constantMaterializerDecl = R"(
   /// Materialize a single constant operation from a given attribute value with
   /// the desired resultant type.
   ::mlir::Operation *materializeConstant(::mlir::OpBuilder &builder,
                                          ::mlir::Attribute value,
                                          ::mlir::Type type,
                                          ::mlir::Location loc) override;
static void emitDialectDecl(Dialect &dialect,
   // Check for any attributes/types registered to this dialect.  If there are,
   // add the hooks for parsing/printing.
   if (!dialectAttrs.empty())
     os << attrParserDecl;
   if (!dialectTypes.empty())
     os << typeParserDecl;
 
   // Add the decls for the various features of the dialect.
+  if (dialect.hasCanonicalizer())
+    os << canonicalizerDecl;
   if (dialect.hasConstantMaterializer())
     os << constantMaterializerDecl;
   if (dialect.hasOperationAttrVerify())
     os << opAttrVerifierDecl;
   if (dialect.hasRegionArgAttrVerify())
     os << regionArgAttrVerifierDecl;
   if (dialect.hasRegionResultAttrVerify())
     os << regionResultAttrVerifierDecl;