[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 ↗

(On Diff #547939)

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 ↗

(On Diff #547939)

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

1099

supplementary

[jcai19]

Update the example in the doc.

[jcai19]

mlir/docs/DeclarativeRewrites.md

696 ↗

(On Diff #547939)

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 ↗

(On Diff #550153)

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 ↗

(On Diff #547939)

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/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";
   }
 
   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++));
 }
 
 std::string PatternEmitter::handleResultPattern(DagNode resultTree,
if (!tail.returnType && (usePartialResults || depth > 0 || resultIndex < 0)) {
     // here given that it's easier for developers to write compared to
     // aggregate-parameter builders.
     createSeparateLocalVarsForOpArgs(tree, childNodeNames);
 
     os.scope().os << formatv("{0} = rewriter.create<{1}>({2}", valuePackName,
                              resultOp.getQualCppClassName(), locToUse);
     supplyValuesForOpArgs(tree, childNodeNames, depth);
     os << "\n  );\n}\n";
+    os << formatv("for (auto attr: op0->getAttrs()) {\n"
+                  "  {0}->setAttr(attr.getName(), attr.getValue());\n"
+                  "}\n",
+                  valuePackName);
+
     return resultValue;
   }
 
   // If we are provided explicit return types, use them to build the op.
   // However, if depth == 0 and resultIndex >= 0, it means we are replacing
   // the values generated from the source pattern root op. Then we must use the
   // source pattern's value types to determine the value type of the generated
   // op here.
if (numResults != 0) {
         os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n"
                       "  tblgen_types.push_back(v.getType());\n}\n",
                       resultIndex + i);
     }
   }
   os << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, "
                 "tblgen_values, tblgen_attrs);\n",
                 valuePackName, resultOp.getQualCppClassName(), locToUse);
+  os << formatv("for (auto attr: op0->getAttrs()) {\n"
+                "  {0}->setAttr(attr.getName(), attr.getValue());\n"
+                "}\n",
+                valuePackName);
   os.unindent() << "}\n";
   return resultValue;
 }
 
 void PatternEmitter::createSeparateLocalVarsForOpArgs(
     DagNode node, ChildNodeIndexNameMap &childNodeNames) {
   Operator &resultOp = node.getDialectOp(opMap);