[mlir][ODS] Get rid of limitations in rewriters generator

[vinograd47]

Summary

Do not limit the number of arguments in rewriter pattern.

Introduce separate FmtStrVecObject class to handle
format of variadic std::string array.

[vinograd47]

Fixed clang-format issues.

[jpienaar]

Nice!

mlir/tools/mlir-tblgen/RewriterGen.cpp

252

I think this todo is still pending

[vinograd47]

Rebased and fixed review remarks.

[vinograd47]

@jpienaar I've fixed your remark.

[vinograd47]

Rebased

[mehdi_amini]

Do we have test for this tool already? If so can you update the test to show the effect of the change?

[vinograd47]

@mehdi_amini I've added a test for the change.

[vinograd47]

Rebased

mlir/include/mlir/TableGen/Format.h

public:
 
   FmtObject(FmtObject &&that)
       : FmtObjectBase(std::move(that)), parameters(std::move(that.parameters)) {
     adapters.reserve(that.adapters.size());
     adapters = llvm::apply_tuple(CreateAdapters(), parameters);
   }
 };
 
+class FmtStrVecObject : public FmtObjectBase {
+public:
+  using StrFormatAdapter =
+      decltype(llvm::detail::build_format_adapter(std::declval<std::string>()));
+
+  FmtStrVecObject(StringRef fmt, const FmtContext *ctx,
+                  ArrayRef<std::string> params);
+  FmtStrVecObject(FmtStrVecObject const &that) = delete;
+  FmtStrVecObject(FmtStrVecObject &&that);
+
+private:
+  SmallVector<StrFormatAdapter, 16> parameters;
+};
+
 /// Formats text by substituting placeholders in format string with replacement
 /// parameters.
 ///
 /// There are two categories of placeholders accepted, both led by a '$' sign:
 ///
 /// 1. Positional placeholder: $[0-9]+
 /// 2. Special placeholder:    $[a-zA-Z_][a-zA-Z0-9_]*
 ///
inline auto tgfmt(StringRef fmt, const FmtContext *ctx, Ts &&... vals)
   using ParamTuple = decltype(std::make_tuple(
       llvm::detail::build_format_adapter(std::forward<Ts>(vals))...));
   return FmtObject<ParamTuple>(
       fmt, ctx,
       std::make_tuple(
           llvm::detail::build_format_adapter(std::forward<Ts>(vals))...));
 }
 
+inline FmtStrVecObject tgfmt(StringRef fmt, const FmtContext *ctx,
+                             ArrayRef<std::string> params) {
+  return FmtStrVecObject(fmt, ctx, params);
+}
+
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_FORMAT_H_

mlir/lib/TableGen/Format.cpp

for (auto &repl : replacements) {
 
     if (repl.index >= adapters.size()) {
       s << repl.spec << kMarkerForNoSubst;
       continue;
     }
     adapters[repl.index]->format(s, /*Options=*/"");
   }
 }
+
+FmtStrVecObject::FmtStrVecObject(StringRef fmt, const FmtContext *ctx,
+                                 ArrayRef<std::string> params)
+    : FmtObjectBase(fmt, ctx, params.size()) {
+  parameters.reserve(params.size());
+  for (std::string p : params)
+    parameters.push_back(llvm::detail::build_format_adapter(std::move(p)));
+
+  adapters.reserve(parameters.size());
+  for (auto &p : parameters)
+    adapters.push_back(&p);
+}
+
+FmtStrVecObject::FmtStrVecObject(FmtStrVecObject &&that)
+    : FmtObjectBase(std::move(that)), parameters(std::move(that.parameters)) {
+  adapters.reserve(parameters.size());
+  for (auto &p : parameters)
+    adapters.push_back(&p);
+}

mlir/test/mlir-tblgen/rewriter-indexing.td

 // CHECK: struct test3 : public ::mlir::RewritePattern {
 // We expect ODSOperand 0 here, the attribute before the operand in BOp
 // definition shouldn't shift the counter.
 // CHECK: op1 = (*castedOp0.getODSOperands(0).begin()).getDefiningOp();
 // CHECK: rewriter.create<test::BOp>((*a.getODSResults(0).begin()).getLoc()
 def test3 : Pat<(BOp $attr, (AOp:$a $input)),
                 (BOp $attr, (AOp $input), (location $a))>;
 
+def DOp : NS_Op<"d_op", []> {
+  let arguments = (ins
+    AnyInteger:$v1,
+    AnyInteger:$v2,
+    AnyInteger:$v3,
+    AnyInteger:$v4,
+    AnyInteger:$v5,
+    AnyInteger:$v6,
+    AnyInteger:$v7,
+    AnyInteger:$v8,
+    AnyInteger:$v9,
+    AnyInteger:$v10
+  );
+
+  let results = (outs AnyInteger);
+}
+
+def NativeBuilder :
+  NativeCodeCall<[{
+    nativeCall($_builder, $_loc, $0, $1, $2, $3, $4, $5, $6, $7, $8, $9)
+  }]>;
+
+// Check Pattern with large number of DAG arguments passed to NativeCodeCall
+// CHECK: struct test4 : public ::mlir::RewritePattern {
+// CHECK: nativeCall(rewriter, odsLoc, (*v1.begin()), (*v2.begin()), (*v3.begin()), (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
+def test4 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
+                (NativeBuilder $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;

mlir/tools/mlir-tblgen/RewriterGen.cpp

 
 // Helper function to match patterns.
 void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,
                                          int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   // TODO(suderman): iterate through arguments, determine their types, output
   // names.
-  SmallVector<std::string, 8> capture(8);
-  if (tree.getNumArgs() > 8) {
-    PrintFatalError(loc,
-                    "unsupported NativeCodeCall matcher argument numbers: " +
-                        Twine(tree.getNumArgs()));
-  }
+  SmallVector<std::string, 8> capture;
+  capture.push_back(opName.str());
 
   raw_indented_ostream::DelimitedScope scope(os);
 
   os << "if(!" << opName << ") return ::mlir::failure();\n";
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     std::string argName = formatv("arg{0}_{1}", depth, i);
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       os << "Value " << argName << ";\n";
     } else {
       auto leaf = tree.getArgAsLeaf(i);
       if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {
         os << "Attribute " << argName << ";\n";
       } else if (leaf.isOperandMatcher()) {
         os << "Operation " << argName << ";\n";
       }
     }
 
-    capture[i] = std::move(argName);
+    capture.push_back(std::move(argName));
   }
 
   bool hasLocationDirective;
   std::string locToUse;
   std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
   auto fmt = tree.getNativeCodeTemplate();
-  auto nativeCodeCall = std::string(tgfmt(
-      fmt, &fmtCtx.addSubst("_loc", locToUse), opName, capture[0], capture[1],
-      capture[2], capture[3], capture[4], capture[5], capture[6], capture[7]));
+  auto nativeCodeCall =
+      std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), capture));
 
   os << "if (failed(" << nativeCodeCall << ")) return ::mlir::failure();\n";
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     auto name = tree.getArgName(i);
     if (!name.empty() && name != "_") {
-      os << formatv("{0} = {1};\n", name, capture[i]);
+      os << formatv("{0} = {1};\n", name, capture[i + 1]);
     }
   }
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
-    std::string argName = capture[i];
+    std::string argName = capture[i + 1];
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       PrintFatalError(
           loc, formatv("Matching nested tree in NativeCodecall not support for "
                        "{0} as arg {1}",
                        argName, i));
     }

 
 std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree,
                                                             int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   auto fmt = tree.getNativeCodeTemplate();
-  // TODO: replace formatv arguments with the exact specified args.
-  SmallVector<std::string, 8> attrs(8);
-  if (tree.getNumArgs() > 8) {
-    PrintFatalError(loc,
-                    "unsupported NativeCodeCall replace argument numbers: " +
-                        Twine(tree.getNumArgs()));
-  }
+
+  SmallVector<std::string, 16> attrs;
+
   bool hasLocationDirective;
   std::string locToUse;
   std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
   for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
     if (tree.isNestedDagArg(i)) {
-      attrs[i] = handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1);
+      attrs.push_back(
+          handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1));
     } else {
-      attrs[i] = handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i));
+      attrs.push_back(
+          handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)));
     }
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
                             << " replacement: " << attrs[i] << "\n");
   }
-  return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs[0],
-                           attrs[1], attrs[2], attrs[3], attrs[4], attrs[5],
-                           attrs[6], attrs[7]));
+
+  return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs));
 }
 
 int PatternEmitter::getNodeValueCount(DagNode node) {
   if (node.isOperation()) {
     // If the op is bound to a symbol in the rewrite rule, query its result
     // count from the symbol info map.
     auto symbol = node.getSymbol();
     if (!symbol.empty()) {