[mlir][linalg] Add docstring support for named op spec

[antiagainst]

Summary

Depends on D94335

[antiagainst]

Rebase

[hanchung]

mlir/docs/Dialects/Linalg.md

611

typo: attached

612

typo: lengthy

[antiagainst]

Fix typos

mlir/docs/Dialects/Linalg.md

 
 At this time, the following restrictions are imposed on the syntax and
 semantics:
 
 1.  Each def may only contain a single comprehension but each comprehension may
     perform multiple updates.
 2.  Each tensor may only be used with a single indexing expression.
 
+A `"""`-wrapped doc string can be attached to the named op. It should contain a
+oneliner for summary first, followed by lengthy description.
+
 The following specification may be used to define a named `batchmatmul` op:
 
 ```
-def batchmatmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
+def batchmatmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N))
+"""Batch matrix-multiply operation.
+
+This operation performs batch matrix-multiply over ...
+"""
+{
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
 }
 ```
 
 When `mlir-linalg-ods-gen -gen-ods-decl=1` is called, the following ODS is
 produced:
 
 ```

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

 // IMPL: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
 //
 ods_def<Test5Op>:
 def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
      attr(strides: 2xi32) {
   O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
     I(n, h * strides[0] + kh, w * strides[1] + kw, c), K(f, kh, kw, c)));
 }
+
+// ODS-LABEL: def Test6Op
+// ODS:       let summary = [{ My magic op. }];
+// ODS-NEXT:  let description = [{
+// ODS-NEXT:    It has two inputs.
+// ODS-NEXT:    It has one output.
+// ODS-NEXT:  }];
+//
+ods_def<Test6Op>:
+def test6(A: f32(M, K), B: f32(K)) -> (C: f32(M))
+"""
+My magic op.
+
+It has two inputs.
+It has one output.
+"""
+{
+  C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
+}

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

 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.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/Debug.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");
public:
   enum class Kind {
     // Markers.
     eof,
     error,
 
     // Tokens with no info.
     colon,
     comma,
+    doc_str,
     equal,
     gt,
     l_brace,
     l_paren,
     l_square,
     lt,
     minus,
     plus,
private:
   int getNextChar();
 
   /// Lex an identifier.
   Token lexIdentifier(const char *tokStart);
 
   // Lex an integer.
   Token lexInteger(const char *tokStart);
 
+  // Lex a string.
+  Token lexString(const char *tokStart);
+
   // Skip a comment line, starting with a '//'.
   void skipComment();
 
   llvm::SourceMgr &srcMgr;
   StringRef curBuffer;
   const char *curPtr;
 };
 } // end anonymous namespace
while (true) {
     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 '"':
+      return lexString(tokStart);
     case '/':
       if (*curPtr == '/') {
         skipComment();
         continue;
       }
       // Unknown character, emit an error.
       return emitError(tokStart, "unexpected character: not a comment");
 
Token Lexer::lexInteger(const char *tokStart) {
   // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
   while (isdigit(*curPtr))
     ++curPtr;
 
   StringRef str(tokStart, curPtr - tokStart);
   return Token(Token::Kind::integer, str);
 }
 
+Token Lexer::lexString(const char *tokStart) {
+  assert(curPtr[-1] == '"');
+
+  if (*curPtr == '"' && *(curPtr + 1) == '"') {
+    curPtr += 2;
+    while (true) {
+      switch (*curPtr++) {
+      case '"':
+        if (*curPtr == '"' && *(curPtr + 1) == '"') {
+          Token token(Token::Kind::doc_str,
+                      StringRef(tokStart + 3, curPtr - tokStart - 4));
+          curPtr += 2;
+          return token;
+        }
+        continue;
+      case 0:
+        // If this is a random nul character in the middle of the doc string,
+        // just include it.  If it is the end of file, then it is an error.
+        if (curPtr - 1 != curBuffer.end())
+          continue;
+        return emitError(curPtr - 1, "expected '\"\"\"' to end doc string");
+      default:
+        continue;
+      }
+    }
+  }
+
+  return emitError(curPtr - 1, "expected '\"\"\"' to start doc string");
+}
+
 /// Skip a comment line, starting with a '//'.
 void Lexer::skipComment() {
   // Advance over the second '/' in a '//' comment.
   assert(*curPtr == '/');
   ++curPtr;
 
   while (true) {
     switch (*curPtr++) {
private:
 
   /// Tensors are per TC def.
   llvm::StringMap<RegisteredTensor> registeredTensors;
   unsigned nextRegisteredTensorIndex;
 
   /// Attributes are per TC def.
   std::map<std::string, RegisteredAttr> registeredAttrs;
 
+  StringRef docString;
+
   Parser &parser;
 };
 } // namespace
 
 namespace llvm {
 
 template <>
 struct DenseMapInfo<TensorUse> {
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();
   }
 
+  // Parse optional doc string
+  if (parser.curToken.is(Token::Kind::doc_str)) {
+    docString = parser.curToken.getSpelling();
+    parser.consumeToken();
+    LLVM_DEBUG(llvm::dbgs()
+               << "parsed doc string: '''" << docString << "'''\n");
+  }
+
   // 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(),
void TCParser::printODS(llvm::raw_ostream &os, StringRef cppOpName,
   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">]> {
+      {2}
       let arguments = (ins
         Variadic<AnyShaped>:$inputs,
-        Variadic<AnyShaped>:$outputs{4}
+        Variadic<AnyShaped>:$outputs{3}
       );
       let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
       let regions = (region AnyRegion:$region);
 
       let skipDefaultBuilders = 1;
       let builders = [ OpBuilderDAG<
         (ins "ValueRange":$inputs, "ValueRange":$outputs),
         [{{
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 {5}; }
+        static unsigned getNumRegionArgs() {{ return {4}; }
         std::string getLibraryCallName() {{
           return generateLibraryCallName(getOperation());
         }
       }];
   })FMT";
 
-  unsigned nInputs = 0, nOutputs = 0;
-  for (auto &t : registeredTensors) {
-    if (t.getValue().isOutput)
-      nOutputs++;
-    else
-      nInputs++;
+  std::string doc;
+
+  if (!docString.empty()) {
+    const char *docFmt = R"FMT(
+      let summary = [{ {0} }];
+      let description = [{
+        {1}
+      }];
+    )FMT";
+
+    StringRef summary, description;
+    std::tie(summary, description) = docString.trim().split('\n');
+    doc = llvm::formatv(docFmt, summary.trim(), description.trim());
   }
 
-  os << llvm::formatv(header, cppOpName, linalgOpName, nInputs, nOutputs,
-                      attrList, state.orderedTensorArgs.size());
+  os << llvm::formatv(header, cppOpName, linalgOpName, doc, 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(