[drr] Enable specifying range in replacement.

[jpienaar]

Summary

This enables creating a replacement rule where range of positional replacements
need not be spelled out, or are not known (e.g., enable having a rewrite that
forward all operands to a call generically).

[Chia-hungDuan]

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

92

I'm thinking the form like "$...", which maps all the remaining. so in this case, we don't need to write down v1 ~ v3. Will this be clear?

In addition, for me, this form looks it will appear in something like,
// Binding the remaining args of SomeOp and select the range of arguments from $arg for AnotherOp.
def : Pat<(SomeOp $a0, $a1, $arg...), (AnotherOp $arg__3...)

[jpienaar]

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

92

Yes this one just changes it in the native call, not in the DAG specification.

With ops we have an explicit and fixed set of args, today one would match them by name or elide in matching with $_. We can already use native codecall expansion without changing dag matching (so these are independent but could work together). A potential problem with the dag format is it would be confusing with variadics.

[rriddle]

Can you update the description to specifically mention that this is for NativeCodeCalls? It wasn't clear on the first read.

Can you also add documentation for this feature to the .md docs?

mlir/lib/TableGen/Format.cpp

101

Can you just pass std::isdigit as is?

109

Can you add a comment here on what this is?

181–185

interleaveComma here? (using llvm::seq, or llvm::make_range on adaptors)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

1501

This change looks unrelated.

[jpienaar]

use interleaveComma, update doc

[jpienaar]

Remove unrelated change

[jpienaar]

drop unrelated change really this time

[jpienaar]

mlir/lib/TableGen/Format.cpp

101

Depends on compiler, some yes others no it seems. Reverting to explicit lambda.

[rriddle]

mlir/lib/TableGen/Format.cpp

101

(You should be able to use llvm::isDigit then)

mlir/include/mlir/TableGen/Format.h

private:
   };
 
   llvm::SmallDenseMap<PHKind, std::string, 4, PHKindInfo> builtinSubstMap;
   llvm::StringMap<std::string> customSubstMap;
 };
 
 /// Struct representing a replacement segment for the formatted string. It can
 /// be a segment of the formatting template (for `Literal`) or a replacement
-/// parameter (for `PositionalPH` and `SpecialPH`).
+/// parameter (for `PositionalPH`, `PositionalRangePH` and `SpecialPH`).
 struct FmtReplacement {
-  enum class Type { Empty, Literal, PositionalPH, SpecialPH };
+  enum class Type {
+    Empty,
+    Literal,
+    PositionalPH,
+    PositionalRangePH,
+    SpecialPH
+  };
 
   FmtReplacement() = default;
   explicit FmtReplacement(StringRef literal)
       : type(Type::Literal), spec(literal) {}
   FmtReplacement(StringRef spec, size_t index)
       : type(Type::PositionalPH), spec(spec), index(index) {}
+  FmtReplacement(StringRef spec, size_t index, size_t end)
+      : type(Type::PositionalRangePH), spec(spec), index(index), end(end) {}
   FmtReplacement(StringRef spec, FmtContext::PHKind placeholder)
       : type(Type::SpecialPH), spec(spec), placeholder(placeholder) {}
 
   Type type = Type::Empty;
   StringRef spec;
   size_t index = 0;
+  size_t end = kUnset;
   FmtContext::PHKind placeholder = FmtContext::PHKind::None;
+
+  static constexpr size_t kUnset = -1;
 };
 
 class FmtObjectBase {
 private:
   static std::pair<FmtReplacement, StringRef> splitFmtSegment(StringRef fmt);
   static std::vector<FmtReplacement> parseFormatString(StringRef fmt);
 
 protected:
   // The parameters are stored in a std::tuple, which does not provide runtime
   // indexing capabilities.  In order to enable runtime indexing, we use this
   // structure to put the parameters into a std::vector.  Since the parameters
   // are not all the same type, we use some type-erasure by wrapping the
   // parameters in a template class that derives from a non-template superclass.
   // Essentially, we are converting a std::tuple<Derived<Ts...>> to a
   // std::vector<Base*>.
   struct CreateAdapters {
     template <typename... Ts>
-    std::vector<llvm::detail::format_adapter *> operator()(Ts &... items) {
+    std::vector<llvm::detail::format_adapter *> operator()(Ts &...items) {
       return std::vector<llvm::detail::format_adapter *>{&items...};
     }
   };
 
   StringRef fmt;
   const FmtContext *context;
   std::vector<llvm::detail::format_adapter *> adapters;
   std::vector<FmtReplacement> replacements;
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]+
+/// 1.a Positional placeholder: $[0-9]+
+/// 1.b Positional range placeholder: $[0-9]+...
 /// 2. Special placeholder:    $[a-zA-Z_][a-zA-Z0-9_]*
 ///
 /// Replacement parameters for positional placeholders are supplied as the
 /// `vals` parameter pack with 1:1 mapping. That is, $0 will be replaced by the
 /// first parameter in `vals`, $1 by the second one, and so on. Note that you
 /// can use the positional placeholders in any order and repeat any times, for
 /// example, "$2 $1 $1 $0" is accepted.
 ///
+/// Replace parameters for positional range placeholders are supplied as if
+/// positional placeholders were specified with commas separating them.
+///
 /// Replacement parameters for special placeholders are supplied using the `ctx`
 /// format context.
 ///
 /// The `fmt` is recorded as a `StringRef` inside the returned `FmtObject`.
 /// The caller needs to make sure the underlying data is available when the
 /// `FmtObject` is used.
 ///
 /// `ctx` accepts a nullptr if there is no special placeholder is used.
 ///
 /// If no substitution is provided for a placeholder or any error happens during
 /// format string parsing or replacement, the placeholder will be outputted
 /// as-is with an additional marker '<no-subst-found>', to aid debugging.
 ///
 /// To print a '$' literally, escape it with '$$'.
 ///
 /// This utility function is inspired by LLVM formatv(), with modifications
 /// specially tailored for TableGen C++ generation usage:
 ///
 /// 1. This utility use '$' instead of '{' and '}' for denoting the placeholder
 ///    because '{' and '}' are frequently used in C++ code.
 /// 2. This utility does not support format layout because it is rarely needed
 ///    in C++ code generation.
 template <typename... Ts>
-inline auto tgfmt(StringRef fmt, const FmtContext *ctx, Ts &&... vals)
+inline auto tgfmt(StringRef fmt, const FmtContext *ctx, Ts &&...vals)
     -> FmtObject<decltype(std::make_tuple(
         llvm::detail::build_format_adapter(std::forward<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))...));

mlir/lib/TableGen/Format.cpp

FmtObjectBase::splitFmtSegment(StringRef fmt) {
   // Allow escaping dollar with '$$'
   if (fmt[1] == '$') {
     return {FmtReplacement{fmt.substr(0, 1)}, fmt.substr(2)};
   }
 
   // First try to see if it's a positional placeholder, and then handle special
   // placeholders.
 
-  size_t end = fmt.find_if_not([](char c) { return std::isdigit(c); }, 1);
+  size_t end =
+      fmt.find_if_not([](char c) { return std::isdigit(c); }, /*From=*/1);
   if (end != 1) {
     // We have a positional placeholder. Parse the index.
     size_t index = 0;
     if (fmt.substr(1, end - 1).consumeInteger(0, index)) {
       llvm_unreachable("invalid replacement sequence index");
     }
 
+    if (fmt.substr(end, 3) == "...") {
+      return {
+          FmtReplacement{fmt.substr(0, end + 3), index, FmtReplacement::kUnset},
+          fmt.substr(end + 3)};
+    }
+
     if (end == StringRef::npos) {
       // All the remaining characters are part of the positional placeholder.
       return {FmtReplacement{fmt, index}, StringRef()};
     }
     return {FmtReplacement{fmt.substr(0, end), index}, fmt.substr(end)};
   }
 
   end = fmt.find_if_not([](char c) { return std::isalnum(c) || c == '_'; }, 1);
if (repl.type == FmtReplacement::Type::SpecialPH) {
         if (subst)
           s << *subst;
         else
           s << repl.spec << kMarkerForNoSubst;
       }
       continue;
     }
 
+    if (repl.type == FmtReplacement::Type::PositionalRangePH) {
+      if (repl.index >= adapters.size()) {
+        s << repl.spec << kMarkerForNoSubst;
+        continue;
+      }
+      int it = repl.index;
+      int end = repl.end == FmtReplacement::kUnset ? adapters.size() : repl.end;
+      if (it == end)
+        continue;
+      adapters[it++]->format(s, /*Options=*/"");
+      for (; it != end; ++it)
+        adapters[it]->format(s << ", ", /*Options=*/"");
+      continue;
+    }
+
     assert(repl.type == FmtReplacement::Type::PositionalPH);
 
     if (repl.index >= adapters.size()) {
       s << repl.spec << kMarkerForNoSubst;
       continue;
     }
     adapters[repl.index]->format(s, /*Options=*/"");
   }

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

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)>;
+
+// CHECK: struct test5 : public ::mlir::RewritePattern {
+// CHECK: foo(rewriter, (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
+def test5 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
+                (NativeCodeCall<[{ foo($_builder, $3...) }]> $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

gentbl_cc_library(
     ],
     tblgen = ":mlir-tblgen",
     td_file = "include/mlir/Dialect/SparseTensor/Transforms/Passes.td",
     deps = [":PassBaseTdFiles"],
 )
 
 cc_library(
     name = "SparseTensor",
-    srcs = glob(["lib/Dialect/SparseTensor/IR/*.cpp"]),
+    srcs = glob([
+        "lib/Dialect/SparseTensor/IR/*.cpp",
+    ]),
     hdrs = ["include/mlir/Dialect/SparseTensor/IR/SparseTensor.h"],
     includes = ["include"],
     deps = [
         ":IR",
         ":SideEffectInterfaces",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorOpsIncGen",
+        ":SparseTensorUtils",
+        ":StandardOps",
+        "//llvm:Support",
+    ],
+)
+
+cc_library(
+    name = "SparseTensorUtils",
+    srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
+    hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),
+    includes = ["include"],
+    deps = [
+        ":IR",
+        ":SideEffectInterfaces",
+        ":SparseTensorAttrDefsIncGen",
+        ":SparseTensorOpsIncGen",
         ":StandardOps",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob(["lib/Dialect/SparseTensor/Transforms/*.cpp"]),
deps = [
         ":StandardOpsTransforms",
         ":TensorDialect",
         ":Transforms",
         ":VectorOps",
         "//llvm:Support",
     ],
 )
 
-cc_library(
-    name = "SparseTensorUtils",
-    srcs = glob(["lib/Dialect/SparseTensor/Utils/*.cpp"]),
-    hdrs = glob(["include/mlir/Dialect/SparseTensor/Utils/*.h"]),
-    includes = ["include"],
-    deps = [
-        ":IR",
-        "//llvm:Support",
-    ],
-)
-
 td_library(
     name = "StdOpsTdFiles",
     srcs = [
         "include/mlir/Dialect/StandardOps/IR/Ops.td",
         "include/mlir/Dialect/StandardOps/IR/StandardOpsBase.td",
     ],
     includes = ["include"],
     deps = [