[mlir][OpDSL] Add support for adding canonicalization patterns.

[gysit]

Summary

Extend OpDSL with a defines method that can set the hasCanonicalizer flag for an OpDSL operation. If the flag is set via defines(Canonicalizer) the operation needs to implement the getCanonicalizationPatterns method. The revision specifies the flag for linalg.fill_tensor and adds an empty FillTensorOp::getCanonicalizationPatterns implementation.

This revision is a preparation step to replace linalg.fill by its OpDSL counterpart linalg.fill_tensor. The two are only functionally equivalent if both specify the same canonicalization patterns. The revision is thus a prerequisite for the linalg.fill replacement.

Depends On D120725

[gysit]

Rebase.

[gysit]

Rebase.

mlir/docs/Dialects/Linalg/OpDSL.md

def matmul(A=TensorDef(T1, S.M, S.K),
            B=TensorDef(T2, S.K, S.N),
            C=TensorDef(U, S.M, S.N, output=True)):
   """Performs a matrix multiplication of two 2D inputs.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   domain(D.m, D.n, D.k)
+  defines(Canonicalizer)
   implements(ContractionOpInterface)
   C[D.m, D.n] += TypeFn.cast_signed(
       U, A[D.m, D.k]) * TypeFn.cast_signed(U, B[D.k, D.n])
 ```
 
 Here we have a simple type polymorphic contraction that takes arguments `A` and
 `B` and outputs `C`. Each is bound to a `TensorDef`, which specifies:
 
 *   The symbolic element type (`T1`, `T2`, `U` above).
 *   Symbolic shape expressions with symbols that are bound globally for the op (
     note that in this simple example, the shape expressions are just symbol
     references, but they are permitted to be a constrained set of affine
     expressions).
 *   Usage (`output=True`).
 
 The docstring will be transferred to the op definition verbatim.
 
 An explicit iteration domain dimension order can be declared for the op via
 `domain(D.d0[, D.d1...])`.
 
 Special identifying op interfaces can be declared for the op via
 `implements(interface1[, interface2...])`.
 
+Extra method definitions can be declared for the op via
+`defines(definition1[, definition2...])`.
+
 ## Parameters
 
 Structured operations take two types of runtime parameters namely scalars and
 tensors. While scalars are inputs only, a tensor may be marked as an output.
 Assignment expressions index the tensor parameters to access the individual
 elements, while scalars can be accessed directly.
 
 The following example demonstrates the use of the two parameter types:

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

metadata: !LinalgOpMetadata
   doc: |-
     Fills the output tensor with the given value.
 
     Works for arbitrary ranked output tensors since the operation performs scalar
     accesses only and is thus rank polymorphic. Numeric casting is performed on
     the value operand, promoting it to the same data type as the output.
   implements:
   - LinalgFillOpInterface
+  defines:
+  - hasCanonicalizer
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: value
     kind: scalar
     type_var: T1
   - !LinalgOperandDefConfig
     name: O

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 void FillOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                          MLIRContext *context) {
   results
       .add<FoldFillWithPad, FoldFillWithTensorReshape<tensor::CollapseShapeOp>,
            FoldFillWithTensorReshape<tensor::ExpandShapeOp>,
            FoldInsertPadIntoFill>(context);
 }
 
+// TODO: Add the FillOp patterns when transitioning to the OpDSL FillOp.
+void FillTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                               MLIRContext *context) {}
+
 //===----------------------------------------------------------------------===//
 // GenericOps
 //===----------------------------------------------------------------------===//
 void GenericOp::build(
     OpBuilder &builder, OperationState &result, TypeRange resultTensorTypes,
     ValueRange inputs, ValueRange outputs, ArrayRef<AffineMap> indexingMaps,
     ArrayRef<StringRef> iteratorTypes, StringRef doc, StringRef libraryCall,
     function_ref<void(OpBuilder &, Location, ValueRange)> bodyBuild,

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

def __init__(self, cpp_name: str):
     self.cpp_name = cpp_name
 
 
 ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
 ConvolutionOpInterface = OpInterfaceDef("LinalgConvolutionOpInterface")
 FillOpInterface = OpInterfaceDef("LinalgFillOpInterface")
 
 
+class OpDefinitionDef:
+  """A method that an op implements."""
+
+  def __init__(self, def_name: str):
+    self.def_name = def_name
+
+
+Canonicalizer = OpDefinitionDef("hasCanonicalizer")
+
+
 class OpMetadataDef(YAMLObject):
   """Metadata about the op (generally not behavior impacting)."""
   yaml_tag = "!LinalgOpMetadata"
 
   def __init__(self, name: str, cpp_class_name: Optional[str],
                doc: Optional[str]):
     self.name = name
     self.cpp_class_name = cpp_class_name if cpp_class_name is not None else name
     self.doc = doc
     self.implements = []  # type: List[OpInterfaceDef]
+    self.defines = []  # type: List[OpDefinitionsDef]
 
   def to_yaml_custom_dict(self):
     d = dict(
         name=self.name,
         cpp_class_name=self.cpp_class_name,
         doc=self.doc,
     )
     if self.implements:
       d["implements"] = [intr.cpp_name for intr in self.implements]
+    if self.defines:
+      d["defines"] = [defi.def_name for defi in self.defines]
     return d
 
 
 class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,

mlir/python/mlir/dialects/linalg/opdsl/lang/dsl.py

def linalg_structured_op(dsl_func=None,
   with bind_op_def(op_def):
     dsl_func(*dsl_func_args)
 
   # TODO: The returned callable should be an IR emitter but that is not
   # upstreamed yet.
   return DefinedOpCallable(op_name, op_def)
 
 
+def domain(*dimensions: DimDef):
+  if any(not isinstance(d, DimDef) for d in dimensions):
+    raise ValueError(f"Expected dimensions of type DimDef but got {dimensions}")
+  current_op_def().domain.extend(dimensions)
+
+
 def implements(*interfaces: OpInterfaceDef):
+  if any(not isinstance(intr, OpInterfaceDef) for intr in interfaces):
+    raise ValueError(
+        f"Expected interfaces of type OpInterfaceDef but got {interfaces}")
   current_op_def().metadata.implements.extend(interfaces)
 
 
-def domain(*dimensions: DimDef):
-  if current_op_def().domain:
-    raise ValueError(f"Expected only one set of domain dimensions per operator")
-  if any(not isinstance(dim, DimDef) for dim in dimensions):
-    raise ValueError(f"Expected dimensions of type DimDef but got {dimensions}")
-  current_op_def().domain.extend(dimensions)
+def defines(*definitions: OpDefinitionDef):
+  if any(not isinstance(defi, OpDefinitionDef) for defi in definitions):
+    raise ValueError(
+        f"Expected definitions of type OpDefinitionDef but got {definitions}")
+  current_op_def().metadata.defines.extend(definitions)

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

 def fill_tensor(value=ScalarDef(T1), O=TensorDef(U, output=True)):
   """Fills the output tensor with the given value.
 
   Works for arbitrary ranked output tensors since the operation performs scalar
   accesses only and is thus rank polymorphic. Numeric casting is performed on
   the value operand, promoting it to the same data type as the output.
   """
   implements(FillOpInterface)
+  defines(Canonicalizer)
   O[None] = TypeFn.cast_signed(U, value)
 
 
 @linalg_structured_op
 def fill_rng_2d(
     min=ScalarDef(F64),
     max=ScalarDef(F64),
     seed=ScalarDef(I32),

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

 # IMPL-LABEL:  void Test4Op::regionBuilder(ImplicitLocOpBuilder &b,
 #  IMPL-NEXT:    Block &block, ArrayRef<NamedAttribute> attrs)
 #       IMPL:  UnaryFn unary_funVal = UnaryFn::exp
 #       IMPL:  BinaryFn binary_funVal = BinaryFn::add
 
 #       IMPL:  Value [[VAL0:[a-z0-9]+]] = helper.buildUnaryFn(unary_funVal, block.getArgument(0))
 #  IMPL-NEXT:  Value [[VAL1:[a-z0-9]+]] = helper.buildBinaryFn(binary_funVal, [[VAL0]], block.getArgument(0))
 #  IMPL-NEXT:  yields.push_back([[VAL1]])
+
+# @linalg_structured_op
+# def test5(value=ScalarDef(T1), O=TensorDef(U, output=True)):
+#   """Title.
+
+#   Detailed description.
+#   """
+#   implements(FillOpInterface)
+#   defines(Canonicalizer)
+#   O[None] = TypeFn.cast(U, value)
+
+--- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+  name: test5
+  cpp_class_name: Test5Op
+  doc: |-
+    Title.
+
+    Detailed description.
+  implements:
+  - LinalgFillOpInterface
+  defines:
+  - hasCanonicalizer
+structured_op: !LinalgStructuredOpConfig
+  args:
+  - !LinalgOperandDefConfig
+    name: value
+    kind: scalar
+    type_var: T1
+  - !LinalgOperandDefConfig
+    name: O
+    kind: output_tensor
+    type_var: U
+    shape_map: affine_map<() -> ()>
+  indexing_maps: !LinalgIndexingMapsConfig
+    static_indexing_maps:
+    - affine_map<() -> ()>
+    - affine_map<() -> ()>
+  iterator_types: []
+  assignments:
+  - !ScalarAssign
+    arg: O
+    value: !ScalarExpression
+      scalar_fn:
+        kind: type
+        fn_name: cast
+        type_var: U
+        operands:
+        - !ScalarExpression
+          scalar_arg: value
+
+# ODS-LABEL:  def Test5Op : LinalgStructuredBase_Op<"test5"
+#  ODS-NEXT:  /*extraInterfaces=*/[LinalgFillOpInterface])>
+
+#       ODS:  let hasCanonicalizer = 1;

mlir/test/python/dialects/linalg/opdsl/metadata.py

 # RUN: %PYTHON -m mlir.dialects.linalg.opdsl.dump_oplib --file %s | FileCheck %s
 
 from mlir.dialects.linalg.opdsl.lang import *
 
 
 # CHECK: ---
 # CHECK-LABEL: matmul
 # CHECK:      implements:
 # CHECK-NEXT: - LinalgContractionOpInterface
+# CHECK:      defines:
+# CHECK-NEXT: - hasCanonicalizer
 @linalg_structured_op
 def matmul(
     A=TensorDef(T, S.M, S.K),
     B=TensorDef(T, S.K, S.N),
     C=TensorDef(U, S.M, S.N, output=True)):
   implements(ContractionOpInterface)
+  defines(Canonicalizer)
   C[D.m, D.n] += TypeFn.cast_signed(U, A[D.m, D.k]) * TypeFn.cast_signed(
       U, B[D.k, D.n])

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

struct LinalgYAMLContext {
   MLIRContext *mlirContext;
 };
 
 struct LinalgOpMetadata {
   std::string name;
   std::string cppClassName;
   Optional<std::string> doc;
   SmallVector<std::string> implements;
+  SmallVector<std::string> defines;
 };
 
 struct SerializedAffineMap {
   AffineMapAttr affineMapAttr;
 
   AffineMap affineMap() { return affineMapAttr.getValue(); }
 };
 

 /// Metadata about the op (name, C++ name, and documentation).
 template <>
 struct MappingTraits<LinalgOpMetadata> {
   static void mapping(IO &io, LinalgOpMetadata &info) {
     io.mapRequired("name", info.name);
     io.mapRequired("cpp_class_name", info.cppClassName);
     io.mapOptional("doc", info.doc);
     io.mapOptional("implements", info.implements);
+    io.mapOptional("defines", info.defines);
   }
 };
 
 /// How the ops indexing maps are produced. Must be one of:
 ///   - static_indexing_maps: A static list of AffineMaps, possibly with
 ///     some symbols that bind to attributes of the op. Each indexing map must
 ///     be normalized over the same list of dimensions, and its symbols must
 ///     match the symbols for argument shapes.

 
 // Template for Linalg named ops' ODS definitions. Parameters:
 // {0}: ODS/C++ op name
 // {1}: assembly op mnemonic
 // {2}: op interface list
 // {3}: documentation (summary + description)
 // {4}: op attribute list
 // {5}: builder methods taking standalone attribute parameters
-// {6}: additional methods for attributes used by indexing maps
+// {6}: additional method defintions
+// {7}: additional methods for attributes used by indexing maps
 static const char structuredOpOdsHeaderFormat[] = R"FMT(
 //===----------------------------------------------------------------------===//
 // Op definition for {0}
 //===----------------------------------------------------------------------===//
 
 def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([AttrSizedOperandSegments],
   /*extraInterfaces=*/[{2}])> {
     {3}
let builders = [
         $_state.addAttributes(attributes);
         $_state.addTypes(resultTensorTypes);
         (void)$_state.addRegion();
       }]>
       {5}
     ];
     let hasCustomAssemblyFormat = 1;
     let hasFolder = 1;
+    {6}
 
     let extraClassDeclaration = structuredOpsBaseDecls # [{{
       // Auto-generated.
       ArrayAttr iterator_types();
       ArrayAttr indexing_maps();
       static void regionBuilder(ImplicitLocOpBuilder &b,
                                 Block &block, ArrayRef<NamedAttribute> attrs);
       static std::function<void(ImplicitLocOpBuilder &,
                                 Block &, ArrayRef<NamedAttribute>)>
       getRegionBuilder() {{
         return regionBuilder;
       }
 
       // Generic methods.
       static unsigned getNumRegionArgs();
       std::string getLibraryCallName();
-      {6}
+      {7}
     }];
 }
 )FMT";
 
 // Builder method taking attribute parameters. Parameters:
 // {0}: Class name
 // {1}: Comma interleaved attribute parameters
 // {2}: Attribute initialization
)FMT";
     StringRef summary, description;
     std::tie(summary, description) =
         StringRef(*opConfig.metadata->doc).trim().split('\n');
     doc = llvm::formatv(structuredOpDocFmt, summary.trim(), description.trim());
   }
 
   interfaceNameList = interleaveToString(opConfig.metadata->implements, ", ");
 
+  std::string definitionList;
+  for (const std::string &definition : opConfig.metadata->defines) {
+    static const char definitionFmt[] = "let {0} = 1;\n";
+    definitionList.append(llvm::formatv(definitionFmt, definition));
+  }
+
   if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
         return isAttribute(arg.kind);
       })) {
     SmallVector<std::string> attrDefs;
     SmallVector<std::string> attrParams;
     SmallVector<std::string> attrStmts;
     for (LinalgOperandDef &arg : opConfig.structuredOp->args) {
       static const char paramFmt[] = "\"Attribute\":${0}";
if (llvm::any_of(opConfig.structuredOp->args, [](LinalgOperandDef &arg) {
     attrBuilder = llvm::formatv(
         structuredOpBuilderFormat, opConfig.metadata->cppClassName,
         llvm::join(attrParams, ", "), llvm::join(attrStmts, "\n"));
   }
 
   os << llvm::formatv(structuredOpOdsHeaderFormat,
                       opConfig.metadata->cppClassName, opConfig.metadata->name,
                       interfaceNameList, doc, attrList, attrBuilder,
-                      attrMethods);
+                      definitionList, attrMethods);
 
   return success();
 }
 
 static LogicalResult
 generateNamedGenericOpDefns(LinalgOpConfig &opConfig,
                             GenerationContext &genContext) {
   if (!genContext.shouldGenerateDefns())