[MLIR][PDL] Add optional attribute to enable commutativity in PDL

[srishti-pm]

Summary

Added an optional unitAttr "commutative" to the pdl.operation op. If
present, it specifies to the matcher that the operation should be
treated as commutative w.r.t. all its operands.

Co-authors: Srishti Srivastava, Prateek Gupta

Signed-off-by: Srishti Srivastava <[e-mail]>

[sfuniak]

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

280–282

It is not clear to me from this description whether, with the commutative attribute present, the matched operation will always be treated as commutative or if it is treated as commutative provided it has a suitable IsCommutative trait / attribute.

397–400

Is there a reason not to pass the UnitAttr() temporary directly to the build function? Maybe someone more familiar with LLVM style can comment.

[srishti-pm]

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

280–282

I believe there is no IsCommutative trait defined and thus the matched operation should always be treated as commutative when the commutative attribute is present. Is there something I'm missing here? I believe you have more knowledge about this implementation than me.

397–400

Yes, I think the original author to this (Prateek) did this because all the other values are also not being passed directly. But, I think someone more familiar with the style could answer this. Pinging @bondhugula to share his views here.

[rriddle]

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

280–282

Well, there is (https://github.com/llvm/llvm-project/blob/29511ec7da0a740d7dec2d0f57d858d691e73ce1/mlir/include/mlir/IR/OpDefinition.h#L1116) which we could check to see if an operation is commutative. I guess the real question is if we always want to generate a commutative matcher for commutative operations, or only when requested.

[Chia-hungDuan]

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

280–282

Just FYI, the commutative in DRR is on demand based. The cases I see don't always want the commutative property to be applied while matching

[rriddle]

Can you post an RFC on discourse discussing the desired extensions to PDL/PDLInterp/ByteCode to enable commutative matching? I know @Mogball had some concerns about some extensions to the PDLInterp/bytecode, and it would be good to discuss the full design.

[Mogball]

The implementation of commutative matching is very byte-code heavy but it's effectively just permuting every combination of operands. I think this can be implemented by just generating one pattern for every combination. Commutative patterns for an arbitrary number of operands are undesirable in both cases due to combinatorial explosion... Generally the approach is to have a canonical form for the order of operands. I wonder if this could be implemented generatively by creating patterns that reorder operands to the "canonical" form set by the commutative pattern.

[Chia-hungDuan]

In D118683#3367927, @Mogball wrote:

The implementation of commutative matching is very byte-code heavy but it's effectively just permuting every combination of operands. I think this can be implemented by just generating one pattern for every combination. Commutative patterns for an arbitrary number of operands are undesirable in both cases due to combinatorial explosion... Generally the approach is to have a canonical form for the order of operands. I wonder if this could be implemented generatively by creating patterns that reorder operands to the "canonical" form set by the commutative pattern.

Agree with Jeff. In DRR, only pair of operands are allowed to mark as commutative matching. Sorry, I didn't see many cases and not sure how often do we need the commutative property applied on more than two operands. I was wondering if an operation have more than two operands can be commutative, does that imply the design of operation could be refactored?

[srishti-pm]

In D118683#3367760, @rriddle wrote:

Can you post an RFC on discourse discussing the desired extensions to PDL/PDLInterp/ByteCode to enable commutative matching? I know @Mogball had some concerns about some extensions to the PDLInterp/bytecode, and it would be good to discuss the full design.

@rriddle, I have created it. Link: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798. Kindly add your suggestions here.

[srishti-pm]

In D118683#3367927, @Mogball wrote:

The implementation of commutative matching is very byte-code heavy but it's effectively just permuting every combination of operands. I think this can be implemented by just generating one pattern for every combination. Commutative patterns for an arbitrary number of operands are undesirable in both cases due to combinatorial explosion... Generally the approach is to have a canonical form for the order of operands. I wonder if this could be implemented generatively by creating patterns that reorder operands to the "canonical" form set by the commutative pattern.

Creating more patterns rather than this approach does seem better to me as well. I have created an RFC to track the discussion on the desired requirements. Its link: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798. Please feel free to add your comments there. Once a conclusion is reached there, I will begin implementing the same.

[srishti-pm]

This change is obsolete now, based on our concluding remark here: https://discourse.llvm.org/t/mlir-pdl-extending-pdl-pdlinterp-bytecode-to-enable-commutative-matching/60798/19?u=srishtisrivastava. I am currently working on the finalized utility and will create a revision to incorporate the same into llvm-project.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

 def PDL_OperationOp : PDL_Op<"operation", [AttrSizedOperandSegments]> {
   let summary = "Define an operation within a pattern";
   let description = [{
     `pdl.operation` operations define operation nodes within a pattern. Within
     a match sequence, i.e. when directly nested within a `pdl.pattern`, these
     operations correspond to input operations, or those that already existing
     within the MLIR module. Inside of a `pdl.rewrite`, these operations
     correspond to operations that should be created as part of the replacement
-    sequence.
+    sequence. The `commutative` attribute if present, specifies to the matcher
+    that the operation should be treated as commutative with respect to all its
+    operands.
 
     `pdl.operation`s are composed of a name, and a set of attribute, operand,
     and result type values, that map to what those that would be on a
     constructed instance of that operation. The results of a `pdl.operation` are
     a handle to the operation itself. Handles to the results of the operation
     can be extracted via `pdl.result`.
 
     Example:
let description = [{
     ```
 
     ```mlir
     // We can match the results as so:
     %result = pdl.type
     %otherResults = pdl.types
     %op = pdl.operation "foo.op" -> (%result, %otherResults : !pdl.type, !pdl.range<type>)
     ```
+
+    ```mlir
+    // If the operation is commutative:
+    %op = pdl.operation "foo.op" commutative (%operand : !pdl.value) -> (%result : !pdl.type)
+    ```
   }];
 
   let arguments = (ins OptionalAttr<StrAttr>:$name,
+                       OptionalAttr<UnitAttr>:$commutative,
                        Variadic<PDL_InstOrRangeOf<PDL_Value>>:$operands,
                        Variadic<PDL_Attribute>:$attributes,
                        StrArrayAttr:$attributeNames,
                        Variadic<PDL_InstOrRangeOf<PDL_Type>>:$types);
   let results = (outs PDL_Operation:$op);
   let assemblyFormat = [{
-    ($name^)? (`(` $operands^ `:` type($operands) `)`)?
+    ($name^)? (`commutative` $commutative^)? (`(` $operands^ `:` type($operands) `)`)?
     custom<OperationOpAttributes>($attributes, $attributeNames)
     (`->` `(` $types^ `:` type($types) `)`)? attr-dict
   }];
 
   let builders = [
     OpBuilder<(ins CArg<"Optional<StringRef>", "llvm::None">:$name,
       CArg<"ValueRange", "llvm::None">:$operandValues,
       CArg<"ArrayRef<StringRef>", "llvm::None">:$attrNames,
       CArg<"ValueRange", "llvm::None">:$attrValues,
       CArg<"ValueRange", "llvm::None">:$resultTypes), [{
       auto nameAttr = name ? $_builder.getStringAttr(*name) : StringAttr();
+      auto commutativeAttr = UnitAttr();
       build($_builder, $_state, $_builder.getType<OperationType>(), nameAttr,
-            operandValues, attrValues, $_builder.getStrArrayAttr(attrNames),
-            resultTypes);
+            commutativeAttr, operandValues, attrValues,
+            $_builder.getStrArrayAttr(attrNames), resultTypes);
     }]>,
   ];
   let extraClassDeclaration = [{
     /// Returns true if the operation type referenced supports result type
     /// inference.
     bool hasTypeInference();
   }];
 }

mlir/test/Dialect/PDL/ops.mlir

 
 pdl.pattern @attribute_with_dict : benefit(1) {
   %root = operation
   rewrite %root {
     %attr = attribute {some_unit_attr} attributes {pdl.special_attribute}
     apply_native_rewrite "NativeRewrite"(%attr : !pdl.attribute)
   }
 }
+
+// -----
+
+// Check that the optional attribute for commutative operations is working.
+pdl.pattern @operations : benefit(1) {
+  %type = pdl.type
+  %input0 = pdl.operand
+  %input1 = pdl.operand
+  %root = pdl.operation "foo.op" commutative (%input0, %input1 : !pdl.value, !pdl.value) ->(%type : !pdl.type)
+  pdl.rewrite %root with "rewriter"
+}