[mlir][linalg] Support parsing attributes in named op spec

[antiagainst]

Summary

With this, now we can specify a list of attributes on named ops
generated from the spec. The format is defined as

attr-id ::= bare-id (`?`)?
attr-typedef ::= type (`[` `]`)?
attr-def ::= attr-id `:` attr-typedef

tc-attr-def ::= `attr` `(` attr-def-list `)`
tc-def ::= `def` bare-id
  `(`tensor-def-list`)` `->` `(` tensor-def-list`)`
  (tc-attr-def)?

For example,

ods_def<SomeCppOp>
def some_op(...) -> (...)
attr(
  f32_attr: f32,
  i32_attr: i32,
  array_attr : f32[],
  optional_attr? : f32
)

where ? means optional attribute and [] means array type.

[mehdi_amini]

Not specific to this revision, but is there a documentation for the .tc? (I don't see docs being updated here)

[antiagainst]

In D94240#2484860, @mehdi_amini wrote:

Not specific to this revision, but is there a documentation for the .tc? (I don't see docs being updated here)

Good question. The closet I can find is https://mlir.llvm.org/docs/Dialects/Linalg/#named-payload-ops-specification, which talks about it but doesn't use explicit grammars there. I need more patches to fully utilize these attributes, I'll make sure update the docs afterwards. :)

[antiagainst]

Add support for vector attributes

[antiagainst]

Minor fixes

[antiagainst]

Minor fixes

[hanchung]

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

84

Let's follow other test case -- add // for this blank line.

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

30–33

nit: I think we would add a blank line between LLVM project/subproject headers and System #includes.

164

Remove one space after the first sentence.

[nicolasvasilache]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

31

if you use StringMap you can drop this include.

1069

why not use StringMap like 2 lines above ?

[antiagainst]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

31

StringMap iteration is not guaranteed to be deterministic. I need to print the whole map as ODS attributes so that can cause a problem for tests.

1069

Explained before. StringMap does not provide deterministic iteration.

[mehdi_amini]

Sorry I had to revert in 110775809ad1 as it broke the gcc-5 build (which has been broken for hours by another change that just got reverted, so I didn't want to lose more coverage right now and need to get it back green ASAP).

[antiagainst]

In D94240#2491177, @mehdi_amini wrote:

Sorry I had to revert in 110775809ad1 as it broke the gcc-5 build (which has been broken for hours by another change that just got reverted, so I didn't want to lose more coverage right now and need to get it back green ASAP).

No problem! Thanks for the explanation!

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-gen.tc

 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test3Op> :
 def test3(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
 }
+
+// Test attribute definitions
+// ODS-LABEL: def Test4Op
+// ODS: F32ArrayAttr:$array_attr,
+// ODS: F32:$f32_attr,
+// ODS: RankedF32ElementsAttr<[4]>:$fvec_attr,
+// ODS: I32:$i32_attr,
+// ODS: RankedI32ElementsAttr<[5, 6]>:$ivec_attr,
+// ODS: OptionalAttr<F32>:$optional_attr
+//
+ods_def<Test4Op> :
+def test4(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
+attr(
+  f32_attr: f32,
+  i32_attr: i32,
+  fvec_attr: 4xf32,
+  ivec_attr: 5x6xi32,
+  array_attr : f32[],
+  optional_attr? : f32
+) {
+  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
+}

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/ToolOutputFile.h"
 
+#include <map>
+
 #define DEBUG_TYPE "linalg-ods-gen"
 
 static llvm::cl::OptionCategory ODSGenCat("Linalg ODS Gen");
 
 // Commandline options
 static llvm::cl::opt<std::string>
     inputFilename(llvm::cl::Positional, llvm::cl::desc("<input file>"),
                   llvm::cl::init("-"), llvm::cl::value_desc("filename"));
enum class Kind {
 
     // Tokens with no info.
     colon,
     comma,
     equal,
     gt,
     l_brace,
     l_paren,
+    l_square,
     lt,
     minus,
     plus,
+    question,
     r_brace,
     r_paren,
+    r_square,
     semicolon,
     star,
 
     // Keywords.
     kw_def,
     FIRST_KEYWORD = kw_def,
     kw_ods_def,
+    kw_attr_def,
     kw_floordiv,
     kw_ceildiv,
     kw_mod,
     LAST_KEYWORD = kw_mod,
 
     // String valued tokens.
     id,
     integer,
public:
 
   /// Lex the next token and return it.
   Token lexToken();
 
   /// Emit an error to the lexer with the given location and message.
   Token emitError(llvm::SMLoc loc, const Twine &msg);
   Token emitError(const char *loc, const Twine &msg);
 
+  /// Change the position of the lexer cursor. The next token we lex will start
+  /// at the designated point in the input.
+  void resetPointer(const char *newPtr) { curPtr = newPtr; }
+
 private:
   Token formToken(Token::Kind kind, const char *tokStart) {
     return Token(kind, StringRef(tokStart, curPtr - tokStart));
   }
 
   /// Return the next character in the stream.
   int getNextChar();
 
while (true) {
     case ',':
       return formToken(Token::Kind::comma, tokStart);
     case '=':
       return formToken(Token::Kind::equal, tokStart);
     case '{':
       return formToken(Token::Kind::l_brace, tokStart);
     case '(':
       return formToken(Token::Kind::l_paren, tokStart);
+    case '[':
+      return formToken(Token::Kind::l_square, tokStart);
     case '}':
       return formToken(Token::Kind::r_brace, tokStart);
     case ')':
       return formToken(Token::Kind::r_paren, tokStart);
+    case ']':
+      return formToken(Token::Kind::r_square, tokStart);
     case '<':
       return formToken(Token::Kind::lt, tokStart);
     case '>':
       return formToken(Token::Kind::gt, tokStart);
     case '+':
       return formToken(Token::Kind::plus, tokStart);
     case '-':
       return formToken(Token::Kind::minus, tokStart);
     case ';':
       return formToken(Token::Kind::semicolon, tokStart);
     case '*':
       return formToken(Token::Kind::star, tokStart);
+    case '?':
+      return formToken(Token::Kind::question, tokStart);
     case '/':
       if (*curPtr == '/') {
         skipComment();
         continue;
       }
       // Unknown character, emit an error.
       return emitError(tokStart, "unexpected character: not a comment");
 

 Token Lexer::lexIdentifier(const char *tokStart) {
   // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
   while (isalnum(*curPtr) || *curPtr == '_' || *curPtr == '-')
     ++curPtr;
 
   // Check to see if this identifier is a keyword.
   StringRef str(tokStart, curPtr - tokStart);
   Token::Kind kind = StringSwitch<Token::Kind>(str)
+                         .Case("attr", Token::Kind::kw_attr_def)
                          .Case("def", Token::Kind::kw_def)
                          .Case("ods_def", Token::Kind::kw_ods_def)
                          .Case("floordiv", Token::Kind::kw_floordiv)
                          .Case("ceildiv", Token::Kind::kw_ceildiv)
                          .Case("mod", Token::Kind::kw_mod)
                          .Default(Token::Kind::id);
 
   return Token(kind, str);
public:
 
   /// Advance the current lexer onto the next token.
   void consumeToken() {
     assert(curToken.getKind() != Token::Kind::eof &&
            curToken.getKind() != Token::Kind::error &&
            "shouldn't advance past EOF or errors");
     curToken = lexer.lexToken();
   }
+
   void consumeToken(Token::Kind kind) {
     assert(curToken.getKind() == kind && "unexpected token");
     curToken = lexer.lexToken();
   }
+
   LogicalResult parseToken(Token::Kind kind, const Twine &msg) {
     if (curToken.getKind() != kind)
       return emitError(curToken.getLoc(), msg);
     consumeToken();
     return success();
   }
+
+  /// Parses an optional token and returns failure if failed to parse.
+  LogicalResult parseOptionalToken(Token::Kind kind) {
+    return success(consumeIf(kind));
+  }
+
   LogicalResult emitError(llvm::SMLoc loc, const Twine &msg) {
     lexer.emitError(loc, msg);
     return failure();
   }
+
   LogicalResult emitError(const Twine &msg) {
     return emitError(curToken.getLoc(), msg);
   }
+
   bool consumeIf(Token::Kind kind) {
     if (curToken.isNot(kind))
       return false;
     consumeToken(kind);
     return true;
   }
+
   LogicalResult
   parseCommaSeparatedList(llvm::function_ref<ParseResult()> parseElement) {
     // Non-empty case starts with an element.
     if (parseElement())
       return failure();
 
     // Otherwise we have a list of comma separated elements.
     while (consumeIf(Token::Kind::comma)) {
       if (parseElement())
         return failure();
     }
     return success();
   }
+
   LogicalResult
   parseCommaSeparatedListUntil(Token::Kind rightToken,
                                llvm::function_ref<ParseResult()> parseElement,
                                bool allowEmptyList) {
     // Handle the empty case.
     if (curToken.is(rightToken)) {
       if (!allowEmptyList)
         return emitError("expected list element");
public:
   struct ComprehensionParsingState {
     AffineDimList dims;
     SmallVector<std::unique_ptr<Expression>, 4> expressions;
     llvm::DenseMap<TensorUse, unsigned> orderedTensorArgs;
   };
   LogicalResult parseTensorUse(TensorUse &result,
                                ComprehensionParsingState &state);
 
+  LogicalResult parseAttrDef();
+
   /// Parses a tensor expression.
   LogicalResult parseExpression(TensorUse currentDefinition,
                                 std::unique_ptr<Expression> &result,
                                 ComprehensionParsingState &state);
 
   /// Parse a single comprehension.
   LogicalResult parseOneComprehension(StringRef cppOpName,
                                       StringRef linalgOpName,
struct RegisteredTensor {
     StringRef type;
     AffineMap shape;
     bool isOutput;
     AffineMap indexingMap;
     unsigned index;
   };
 
   //===--------------------------------------------------------------------===//
+  // Internal bookkeeping of attributes.
+  //===--------------------------------------------------------------------===//
+  struct RegisteredAttr {
+    StringRef elementType;
+    SmallVector<uint64_t, 4> vectorDims;
+    bool isArray;
+    bool isOptional;
+  };
+
+  //===--------------------------------------------------------------------===//
   // Per-TC def state.
   //===--------------------------------------------------------------------===//
   /// Symbols are per TC def.
   AffineSymbolList symbols;
+
   /// Tensors are per TC def.
   llvm::StringMap<RegisteredTensor> registeredTensors;
   unsigned nextRegisteredTensorIndex;
 
+  /// Attributes are per TC def.
+  std::map<std::string, RegisteredAttr> registeredAttrs;
+
   Parser &parser;
 };
 } // namespace
 
 namespace llvm {
 
 template <>
 struct DenseMapInfo<TensorUse> {
AffineMap map =
       AffineMap::get(state.dims.size(), symbols.size(), exprs, parser.context);
   LLVM_DEBUG(llvm::dbgs() << "Use of tensor: " << tensorId << " map: " << map
                           << "\n");
 
   result = TensorUse(tensorId, map);
   return success();
 }
 
+/// Parse the information for an attribute def of the form:
+///
+///   affine-expr-list ::= affine-expr (`,` affine-expr )*
+///   attr-id ::= bare-id (`?`)?
+///   dim-list ::= (integer-literal 'x')+
+///   attr-typedef ::= dim-list? type (`[` `]`)?
+///   attr-def ::= attr-id `:` attr-typedef
+LogicalResult TCParser::parseAttrDef() {
+  auto attrLoc = parser.curToken.getLoc();
+  StringRef attrName = parser.curToken.getSpelling();
+  if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
+    return failure();
+  bool isOptional = succeeded(parser.parseOptionalToken(Token::Kind::question));
+  if (failed(parser.parseToken(Token::Kind::colon, "expected colon")))
+    return failure();
+
+  // Parse the attribute's type. We don't expect the type to be arbitrary
+  // complex, so just use this ad-hoc handling here.
+
+  // Parse potential dimension list
+  SmallVector<uint64_t, 4> vectorDims;
+  while (parser.curToken.is(Token::Kind::integer)) {
+    vectorDims.push_back(parser.curToken.getUInt64IntegerValue().getValue());
+    parser.consumeToken();
+
+    StringRef spelling = parser.curToken.getSpelling();
+    if (spelling[0] != 'x')
+      return parser.emitError(parser.curToken.getLoc(),
+                              "expected 'x' in dimension list");
+
+    // If we had a prefix of 'x', lex the next token immediately after the 'x'.
+    if (spelling.size() != 1)
+      parser.lexer.resetPointer(spelling.data() + 1);
+
+    parser.consumeToken();
+  }
+
+  StringRef elementType = parser.curToken.getSpelling();
+  if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
+    return failure();
+
+  bool isArray = false;
+  auto arrayLoc = parser.curToken.getLoc();
+  if (succeeded(parser.parseOptionalToken(Token::Kind::l_square))) {
+    isArray = true;
+    if (failed(parser.parseToken(Token::Kind::r_square, "expected ']'")))
+      return failure();
+  }
+
+  if (!vectorDims.empty() && isArray)
+    return parser.emitError(arrayLoc, "unsupported vector array attribute");
+
+  auto iterBoolPair = registeredAttrs.emplace(
+      attrName, RegisteredAttr{elementType, vectorDims, isArray, isOptional});
+  if (!iterBoolPair.second)
+    return parser.emitError(attrLoc,
+                            "Failed to register attribute '" + attrName + "'");
+
+  LLVM_DEBUG(llvm::dbgs() << "Recorded: " << (isOptional ? "[optional]" : "")
+                          << " " << attrName << " "
+                          << "with type: " << elementType
+                          << (isArray ? "[]" : "") << "\n");
+
+  return success();
+}
+
 /// Parses a tensor expression of the form:
 ///
 ///   op-spec ::= bare-id `<` reduction-dims-list `>`
 ///             | bare-id
 ///   op-arg ::= tensor-expr
 ///            | tensor-use
 ///   op-arg-list ::= op-arg (`,` op-arg)*
 ///   tensor-expr ::= op-spec `(` op-arg-list `)`
if (failed)
     return failure();
 
   return success();
 }
 
 /// Parse and print the information for a ODS def.
 ///
 ///   tensor-def-list ::= tensor-def (`,` tensor-def )*
+///   attr-def-list ::= attr-def (`,` attr-def )*
 ///
 ///   comprehension-list ::= comprehension comprehension*
 ///
+///   tc-attr-def ::= `attr` `(` attr-def-list `)`
 ///   tc-def ::= `def` bare-id `(`tensor-def-list`)` `->` `(` tensor-def-list`)`
+///     (tc-attr-def)?
 ///     `{` comprehension-list `}`
 ///
 ///   ods-def ::= `ods_def` `<` bare-id `>` `:` tc-def
 ///
 /// All the affine-expr in a `tensor-typedef` must be dimensionless (i.e.
 /// contain only expressions involving symbols and constants), but can
 /// otherwise contain arbitrary affine expressions.
 LogicalResult TCParser::parseAndEmitODSDef(llvm::raw_ostream &os) {
+  // Parse def header (including C++ op name)
   if (failed(parser.parseToken(Token::Kind::kw_ods_def,
                                "expected 'ods_def' to define a TC ODS")) ||
       failed(parser.parseToken(Token::Kind::lt, "expected '<'")))
     return failure();
   StringRef cppOpName = parser.curToken.getSpelling();
   LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing ODS: " << cppOpName << "\n");
 
   if (failed(parser.parseToken(Token::Kind::id, "expected id")) ||
       failed(parser.parseToken(Token::Kind::gt, "expected '>'")) ||
       failed(parser.parseToken(Token::Kind::colon, "expected ':'")))
     return failure();
+
   if (failed(parser.parseToken(Token::Kind::kw_def,
                                "expected 'def' to define a TC")))
     return failure();
 
   StringRef tcName = parser.curToken.getSpelling();
   LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing TC: " << tcName << "\n");
+
+  // Parse input/output tensor definitions
   if (failed(parser.parseToken(Token::Kind::id, "expected id")) ||
       failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
     return failure();
 
   auto parseInputDef = [&]() -> LogicalResult {
     return parseTensorDef(/*isOutput=*/false);
   };
   if (failed(parser.parseCommaSeparatedListUntil(
           Token::Kind::r_paren, parseInputDef, /*allowEmptyList=*/false)))
     return failure();
 
   if (failed(parser.parseToken(Token::Kind::minus, "expected '-'")) ||
       failed(parser.parseToken(Token::Kind::gt, "expected '>'")) ||
       failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
     return failure();
   auto parseOutputDef = [&]() -> LogicalResult {
     return parseTensorDef(/*isOutput=*/true);
   };
   if (failed(parser.parseCommaSeparatedListUntil(
           Token::Kind::r_paren, parseOutputDef, /*allowEmptyList=*/false)))
     return failure();
 
+  // Parse optional attribute definitions
+  if (succeeded(parser.parseOptionalToken(Token::Kind::kw_attr_def))) {
+    if (failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
+      return failure();
+    if (failed(parser.parseCommaSeparatedListUntil(
+            Token::Kind::r_paren, std::bind(&TCParser::parseAttrDef, this),
+            /*allowEmptyList=*/false)))
+      return failure();
+  }
+
   // Since we don't declare symbols separately, we discover them eagerly: each
   // newly encountered id in a tensor shape expression is treated as a new
   // symbolic. At this point, all tensors have been parsed and all the symbols
   // that could be discovered eagerly are now known. Resize all AffineMaps to
   // normalize the number of eagerly discovered symbols.
   for (auto &tensor : registeredTensors) {
     auto &map = tensor.getValue().shape;
     map = AffineMap::get(/*dimCount=*/0, symbols.size(), map.getResults(),

 //===----------------------------------------------------------------------===//
 // Printing functions
 //===----------------------------------------------------------------------===//
 
 /// Print the ODS class that defines a new `cppOpName` for a `linalgOpName`.
 void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName,
                         StringRef linalgOpName,
                         ComprehensionParsingState &state) {
+  SmallVector<std::string, 4> attributes;
+  for (const auto &attr : registeredAttrs) {
+    llvm::StringRef name = attr.first;
+
+    llvm::StringRef elementType = attr.second.elementType;
+    std::string odsType = llvm::StringSwitch<std::string>(elementType)
+                              .Case("f32", "F32")
+                              .Case("i32", "I32")
+                              .Default("");
+    if (odsType.empty()) {
+      parser.emitError("unimplemented support for attribute element type: " +
+                       elementType);
+      return;
+    }
+
+    const auto &dims = attr.second.vectorDims;
+    if (!dims.empty()) {
+      SmallVector<std::string, 4> dimStrs;
+      for (uint64_t dim : dims)
+        dimStrs.push_back(std::to_string(dim));
+      odsType = llvm::formatv("Ranked{0}ElementsAttr<[{1}]>", odsType,
+                              llvm::join(dimStrs, ", "));
+    }
+
+    assert(dims.empty() || !attr.second.isArray);
+    if (attr.second.isArray)
+      odsType = llvm::formatv("{0}ArrayAttr", odsType);
+
+    if (attr.second.isOptional)
+      odsType = llvm::formatv("OptionalAttr<{0}>", odsType);
+
+    attributes.push_back(llvm::formatv("{0}:${1}", odsType, name));
+  }
+
+  std::string attrList = llvm::join(attributes, ",\n");
+  if (!attrList.empty())
+    attrList = ",\n" + attrList;
+
   const char *header = R"FMT(  def {0} : LinalgStructuredBase_Op<"{1}", [
     AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     SingleBlockImplicitTerminator<"YieldOp">]> {
-      let arguments = (ins Variadic<AnyShaped>:$inputs,
-                           Variadic<AnyShaped>:$outputs);
+      let arguments = (ins
+        Variadic<AnyShaped>:$inputs,
+        Variadic<AnyShaped>:$outputs{4}
+      );
       let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
       let regions = (region AnyRegion:$region);
 
       let skipDefaultBuilders = 1;
       let builders = [ OpBuilderDAG<
         (ins "ValueRange":$inputs, "ValueRange":$outputs),
         [{{
           $_state.addOperands(inputs);
SingleBlockImplicitTerminator<"YieldOp">]> {
       let extraClassDeclaration = [{{
         // Auto-generated.
         ArrayAttr iterator_types();
         ArrayAttr indexing_maps();
         static void regionBuilder(Block &block);
         static std::function<void(Block &)> getRegionBuilder() {{ return regionBuilder; }
 
         // Generic methods.
-        static unsigned getNumRegionArgs() {{ return {4}; }
+        static unsigned getNumRegionArgs() {{ return {5}; }
         std::string getLibraryCallName() {{
           return generateLibraryCallName(getOperation());
         }
       }];
   })FMT";
 
   unsigned nInputs = 0, nOutputs = 0;
   for (auto &t : registeredTensors) {
     if (t.getValue().isOutput)
       nOutputs++;
     else
       nInputs++;
   }
 
   os << llvm::formatv(header, cppOpName, linalgOpName, nInputs, nOutputs,
-                      state.orderedTensorArgs.size());
+                      attrList, state.orderedTensorArgs.size());
 }
 
 /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
 void TCParser::printReferenceIterators(llvm::raw_ostream &os,
                                        StringRef cppOpName,
                                        ComprehensionParsingState &state) {
   const char *referenceReferenceIteratorsFmt =
       R"FMT(