[mlir] Add a postprocessing parameter in Pattern

[jcai19]

Summary

This adds a parameter SupplementalPatterns in tablegen class Pattern for
postprocessing code. For example, this can be used to ensure ops are
placed in the correct device by copying the atttributes that decide
devicement placement in Tensorflow dialect to prevent performance
regression.

[jpienaar]

This is unfortunately not safe in general, the same named attribute may mean different things etc. Now having a postprocess option where it could be opted in, that we could do.

[jcai19]

Add a postprocessing parameter in Pattern.

[jcai19]

Remove unnecessary code.

[jcai19]

In D157032#4558862, @jpienaar wrote:

This is unfortunately not safe in general, the same named attribute may mean different things etc. Now having a postprocess option where it could be opted in, that we could do.

Thanks for the information! I've made the suggested change.

[jcai19]

Update caption and summary.

[jcai19]

Update caption.

[jcai19]

@jpienaar @rriddle @antiagainst this change is required to fix a performance issue in Tensorflow. Please let me know if it makes sense to you. Thanks.

[jcai19]

Fix how ResultIndex is calculated.

[jpienaar]

(ugh sorry forgot to hit send!)

Overall looks good, just wanted to check on restrictions/ordering.

mlir/docs/DeclarativeRewrites.md

696

OOC why not just have single post-process hook that points to a function?

This does look useful for reusing though. And allow for different API's for post-processing.

Do these execute in the order mentioned here?
Is nesting of these supported? (e.g., (CopyAttrFoo (GetParent $src), $dest))

mlir/include/mlir/IR/PatternBase.td

109

Invoke additional code on ... ?

Actually what restrictions re here? Does this allow arbitrary changes to source and targets ops?

mlir/tools/mlir-tblgen/RewriterGen.cpp

1108

supplementary

[jcai19]

Update the example in the doc.

[jcai19]

mlir/docs/DeclarativeRewrites.md

696

OOC why not just have single post-process hook that points to a function?

I'm not sure how a post-process hook works, but the number of post-processing patterns we need for each rewrite varies, and a pattern may either return mlir::OpResult or mlir::Instruction. Below is an example in Tensorflow,

def CopyAttrs2: NativeCodeCallVoid<"CopyAttributes($0, $1)">;

def CopyAttrs3: NativeCodeCallVoid<"CopyAttributes($0, $1.getOwner())">;

def HashTableAndInitializeTableToV2 : Pat<
  (TF_InitializeTableFromTextFileOp:$src (TF_HashTableOp $container,
    $shared_name, $use_node_name_sharing, $key_dtype, $value_dtype), $filename,
    $key_index, $value_index, $vocab_size, $delimiter, $offset),
  (TF_InitializeTableFromTextFileV2Op:$dest1 (TF_HashTableV2Op:$dest2 $container,
    $shared_name, $use_node_name_sharing, $key_dtype, $value_dtype),
    $filename, $key_index, $value_index, $vocab_size, $delimiter, $offset),
  [], [(CopyAttrs2 $src, $dest1), (CopyAttrs3 $src, $dest2)]>;

Would it be easy to implement this example with a function?

Do these execute in the order mentioned here?

Yes patterns are executed in the specified order. I've updated the example to make it clearer.

Is nesting of these supported? (e.g., (CopyAttrFoo (GetParent $src), $dest))

Yes they are. I've added nested patterns in the example.

mlir/include/mlir/IR/PatternBase.td

109

Invoke additional code on ... ?

Done.

Actually what restrictions re here? Does this allow arbitrary changes to source and targets ops?

I'm not aware of any restrictions other than their results will never be used to replace source patterns.

[jcai19]

mlir/docs/DeclarativeRewrites.md

696

FYI I've committed this change as this is need to unblock our work in Tensorflow. Let me know if you think having a post-process hook that points to a function is a better solution and I'll address that in another patch. Thanks!

[jcai19]

Fixed a typo in the doc in commit https://reviews.llvm.org/rG3c5b4dabdc06dd380391ac965b16961610c0db77.

[mehdi_amini]

This broke the build (you can see it in the pre-merge build log on this revision as well, but you should have got some emails already)

[jcai19]

In D157032#4590167, @mehdi_amini wrote:

This broke the build (you can see it in the pre-merge build log on this revision as well, but you should have got some emails already)

Thanks Mehdi! Applogy for the inconvenience. I've fixed the test failure and relanded the change in https://reviews.llvm.org/rGd22965f0d623392fdf87e5a09a0276aa70e8dfed

mlir/docs/DeclarativeRewrites.md

 *   Matching and generating variadic operand/result ops in nested patterns.
 *   Packing and unpacking variadic operands/results during generation.
 *   [`NativeCodeCall`](#nativecodecall-transforming-the-generated-op) returning
     more than one results.
 
 ## Rule Definition
 
 The core construct for defining a rewrite rule is defined in
-[`OpBase.td`][OpBase] as
+[`PatternBase.td`][PatternBase] as
 
 ```tablegen
 class Pattern<
     dag sourcePattern, list<dag> resultPatterns,
     list<dag> additionalConstraints = [],
+    list<dag> supplementalPatterns = [],
     dag benefitsAdded = (addBenefit 0)>;
 ```
 
 A declarative rewrite rule contains two main components:
 
 *   A *source pattern*, which is used for matching a DAG of operations.
 *   One or more *result patterns*, which are used for generating DAGs of
     operations to replace the matched DAG of operations.

 
 *   Use normal `TypeConstraint`s on previous bound symbols (the first result of
     `TwoResultOp` must be a float tensor);
 *   Define new `Constraint` for previous bound symbols (the second result of
     `TwoResultOp` must has no use);
 *   Apply constraints on multiple bound symbols (`$input` and `TwoResultOp`'s
     first result must have the same element type).
 
+### Supplying additional result patterns
+
+Sometimes we need to add additional code after the result patterns, e.g. coping
+the attributes of the source op to the result ops. These can be specified via
+`SupplementalPatterns` parameter. Similar to auxiliary patterns, they are not
+for replacing results in the source pattern.
+
+For example, we can write
+
+```tablegen
+def GetOwner: NativeCodeCall<"$0.getOwner()">;
+
+def CopyAttrFoo: NativeCodeCallVoid<
+  "$1->setAttr($_builder.getStringAttr(\"foo\"), $0->getAttr(\"foo\"))">;
+
+def CopyAttrBar: NativeCodeCallVoid<
+  "$1->setAttr($_builder.getStringAttr(\"bar\"), $0->getAttr(\"bar\"))">;
+
+
+def : Pattern<
+  (ThreeResultOp:$src ...),
+  [(ZeroResultOp:$dest1 ...), (ThreeResultOp:$dest2 ...)],
+  [(CopyAttrFoo (GetOwner $src), $dest1),
+    (CopyAttrBar (GetOwner $src), (GetOwner $dest2))]>;
+```
+
+This will copy the attribute `foo` and `bar` of `ThreeResultOp` in the source
+pattern to `TwoResultOp` and `OneResultOp` in the result patterns respectively.
+The patterns are executed in the order they are specified.
+
 ### Adjusting benefits
 
 The benefit of a `Pattern` is an integer value indicating the benefit of
 matching the pattern. It determines the priorities of patterns inside the
 pattern rewrite driver. A pattern with a higher benefit is applied before one
 with a lower benefit.
 
 In DRR, a rule is set to have a benefit of the number of ops in the source

mlir/include/mlir/IR/PatternBase.td

 // Then the values generated by `FiveResultOp` will be replaced by
 //
 // * `FiveResultOp`#0: `TwoResultOp1`#1
 // * `FiveResultOp`#1: `TwoResultOp1`#0
 // * `FiveResultOp`#2: `TwoResultOp2`#0
 // * `FiveResultOp`#3: `TwoResultOp2`#1
 // * `FiveResultOp`#4: `TwoResultOp2`#1
 class Pattern<dag source, list<dag> results, list<dag> preds = [],
+  list<dag> supplemental_results = [],
   dag benefitAdded = (addBenefit 0)> {
   dag sourcePattern = source;
   // Result patterns. Each result pattern is expected to replace one result
   // of the root op in the source pattern. In the case of more result patterns
   // than needed to replace the source op, only the last N results generated
   // by the last N result pattern is used to replace a N-result source op.
   // So that the beginning result patterns can be used to generate additional
   // ops to aid building the results used for replacement.
   list<dag> resultPatterns = results;
   // Multi-entity constraints. Each constraint here involves multiple entities
   // matched in source pattern and places further constraints on them as a
   // whole.
   list<dag> constraints = preds;
+  // Optional patterns that are executed after the result patterns. Similar to
+  // auxiliary patterns, they are not used for replacement. These patterns can
+  // be used to invoke additional code after the result patterns, e.g. copy
+  // the attributes from the source op to the result ops.
+  list<dag> supplementalPatterns = supplemental_results;
   // The delta value added to the default benefit value. The default value is
   // the number of ops in the source pattern. The rule with the highest final
   // benefit value will be applied first if there are multiple rules matches.
   // This delta value can be either positive or negative.
   dag benefitDelta = benefitAdded;
 }
 
 // Form of a pattern which produces a single result.
 class Pat<dag pattern, dag result, list<dag> preds = [],
+  list<dag> supplemental_results = [],
   dag benefitAdded = (addBenefit 0)> :
-  Pattern<pattern, [result], preds, benefitAdded>;
+  Pattern<pattern, [result], preds, supplemental_results, benefitAdded>;
 
 // Native code call wrapper. This allows invoking an arbitrary C++ expression
 // to create an op operand/attribute or replace an op result.
 //
 // ## Placeholders
 //
 // If used as a DAG leaf, i.e., `(... NativeCodeCall<"...">:$arg, ...)`,
 // the wrapped expression can take special placeholders listed below:

mlir/include/mlir/TableGen/Pattern.h

public:
 
   // Returns the operator wrapper object corresponding to the given `node`'s DAG
   // operator.
   Operator &getDialectOp(DagNode node);
 
   // Returns the constraints.
   std::vector<AppliedConstraint> getConstraints() const;
 
+  // Returns the number of supplemental auxiliary patterns generated by applying
+  // this rewrite rule.
+  int getNumSupplementalPatterns() const;
+
+  // Returns the DAG tree root node of the `index`-th supplemental result
+  // pattern.
+  DagNode getSupplementalPattern(unsigned index) const;
+
   // Returns the benefit score of the pattern.
   int getBenefit() const;
 
   using IdentifierLine = std::pair<StringRef, unsigned>;
 
   // Returns the file location of the pattern (buffer identifier + line number
   // pair).
   std::vector<IdentifierLine> getLocation() const;

mlir/lib/TableGen/Pattern.cpp

for (auto *it : *listInit) {
     }
 
     ret.emplace_back(cast<llvm::DefInit>(dagInit->getOperator())->getDef(),
                      dagInit->getNameStr(), std::move(entities));
   }
   return ret;
 }
 
+int Pattern::getNumSupplementalPatterns() const {
+  auto *results = def.getValueAsListInit("supplementalPatterns");
+  return results->size();
+}
+
+DagNode Pattern::getSupplementalPattern(unsigned index) const {
+  auto *results = def.getValueAsListInit("supplementalPatterns");
+  return DagNode(cast<llvm::DagInit>(results->getElement(index)));
+}
+
 int Pattern::getBenefit() const {
   // The initial benefit value is a heuristic with number of ops in the source
   // pattern.
   int initBenefit = getSourcePattern().getNumOps();
   llvm::DagInit *delta = def.getValueAsDag("benefitDelta");
   if (delta->getNumArgs() != 1 || !isa<llvm::IntInit>(delta->getArg(0))) {
     PrintFatalError(&def,
                     "The 'addBenefit' takes and only takes one integer value");

mlir/tools/mlir-tblgen/RewriterGen.cpp

for (int i = replStartIndex; i < numResultPatterns; ++i) {
           val,
           "for (auto v: ::llvm::SmallVector<::mlir::Value, 4>{ {0} }) {{\n"
           "  tblgen_repl_values.push_back(v);\n}\n",
           "\n");
     }
     os << "\nrewriter.replaceOp(op0, tblgen_repl_values);\n";
   }
 
+  // Process supplemtal patterns.
+  int numSupplementalPatterns = pattern.getNumSupplementalPatterns();
+  for (int i = 0, offset = -numSupplementalPatterns;
+       i < numSupplementalPatterns; ++i) {
+    DagNode resultTree = pattern.getSupplementalPattern(i);
+    auto val = handleResultPattern(resultTree, offset++, 0);
+    if (resultTree.isNativeCodeCall() &&
+        resultTree.getNumReturnsOfNativeCode() == 0)
+      os << val << ";\n";
+  }
+
   LLVM_DEBUG(llvm::dbgs() << "--- done emitting rewrite logic ---\n");
 }
 
 std::string PatternEmitter::getUniqueSymbol(const Operator *op) {
   return std::string(
       formatv("tblgen_{0}_{1}", op->getCppClassName(), nextValueId++));
 }