text: //===- mlir-linalg-ods-gen.cpp - Linalg ODS generation from math form -----===//
  //
  // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
  // See https://llvm.org/LICENSE.txt for license information.
  // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
  //
  //===----------------------------------------------------------------------===//
  //
  // This file contains the implementation for the Tensor Comprehension-inspired
  // parser and ODS pretty-printer for specifying Linalg "named ops" from a
  // mathematical form.
  //
  //===----------------------------------------------------------------------===//
  
  #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/Support/Casting.h"
  #include "llvm/Support/CommandLine.h"
  #include "llvm/Support/FormatVariadic.h"
  #include "llvm/Support/ToolOutputFile.h"
  
  #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"));
  
  static llvm::cl::opt<std::string>
      outputFilename("o", llvm::cl::desc("Output filename"),
                     llvm::cl::value_desc("filename"), llvm::cl::init("-"));
  
  static llvm::cl::opt<bool>
      genODSDecl("gen-ods-decl", llvm::cl::desc("Emit the ODS ops declarations."),
                 llvm::cl::cat(ODSGenCat));
  
  static llvm::cl::opt<bool>
      genODSImpl("gen-impl", llvm::cl::desc("Emit the ops implementations"),
                 llvm::cl::init(false), llvm::cl::cat(ODSGenCat));
  
  static llvm::cl::opt<bool> testEmitIncludeTdHeader(
      "test-emit-include-td-header",
      llvm::cl::desc("Include LinalgStructuredOps.td for end-to-end "
                     "tblgen testing."),
      llvm::cl::init(false), llvm::cl::cat(ODSGenCat));
  
  using llvm::SetVector;
  using llvm::SMLoc;
  using llvm::StringRef;
  using llvm::Twine;
  
  using namespace mlir;
  
  //===----------------------------------------------------------------------===//
  // Lexer
  //===----------------------------------------------------------------------===//
  
  namespace {
  /// This class represents a specific token in the input format.
  class Token {
  public:
    enum class Kind {
      // Markers.
      eof,
      error,
  
      // Tokens with no info.
      colon,
      comma,
      equal,
      gt,
      l_brace,
      l_paren,
      lt,
      minus,
      plus,
      r_brace,
      r_paren,
      semicolon,
      star,
  
      // Keywords.
      kw_def,
      FIRST_KEYWORD = kw_def,
      kw_ods_def,
      kw_floordiv,
      kw_ceildiv,
      kw_mod,
      LAST_KEYWORD = kw_mod,
  
      // String valued tokens.
      id,
      integer,
    };
  
    Token(Kind kind, StringRef spelling) : kind(kind), spelling(spelling) {}
  
    /// Return the bytes that make up this token.
    StringRef getSpelling() const { return spelling; }
  
    /// Return the kind of this token.
    Kind getKind() const { return kind; }
  
    /// Return a location for this token.
    llvm::SMLoc getLoc() const {
      return llvm::SMLoc::getFromPointer(spelling.data());
    }
  
    /// Return if this token is a keyword.
    bool isKeyword() const {
      return kind >= Kind::FIRST_KEYWORD && kind <= Kind::LAST_KEYWORD;
    }
    bool is(Kind k) const { return kind == k; }
    bool isNot(Kind k) const { return kind != k; }
  
    Optional<uint64_t> getUInt64IntegerValue() const {
      bool isHex = spelling.size() > 1 && spelling[1] == 'x';
  
      uint64_t result = 0;
      if (spelling.getAsInteger(isHex ? 0 : 10, result))
        return None;
      return result;
    }
  
  private:
    /// Discriminator that indicates the kind of token this is.
    Kind kind;
  
    /// A reference to the entire token contents; this is always a pointer into
    /// a memory buffer owned by the source manager.
    StringRef spelling;
  };
  
  /// This class implements a simple lexer.
  class Lexer {
  public:
    Lexer(llvm::SourceMgr &mgr);
  
    /// 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);
  
  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();
  
    /// Lex an identifier.
    Token lexIdentifier(const char *tokStart);
  
    // Lex an integer.
    Token lexInteger(const char *tokStart);
  
    // Skip a comment line, starting with a '//'.
    void skipComment();
  
    llvm::SourceMgr &srcMgr;
    StringRef curBuffer;
    const char *curPtr;
  };
  } // end anonymous namespace
  
  Lexer::Lexer(llvm::SourceMgr &mgr) : srcMgr(mgr) {
    curBuffer = srcMgr.getMemoryBuffer(mgr.getMainFileID())->getBuffer();
    curPtr = curBuffer.begin();
  }
  
  Token Lexer::emitError(llvm::SMLoc loc, const Twine &msg) {
    srcMgr.PrintMessage(loc, llvm::SourceMgr::DK_Error, msg);
    return formToken(Token::Kind::error, loc.getPointer());
  }
  Token Lexer::emitError(const char *loc, const Twine &msg) {
    return emitError(llvm::SMLoc::getFromPointer(loc), msg);
  }
  
  int Lexer::getNextChar() {
    char curChar = *curPtr++;
    switch (curChar) {
    default:
      return (unsigned char)curChar;
    case 0: {
      // A nul character in the stream is either the end of the current buffer
      // or a random nul in the file. Disambiguate that here.
      if (curPtr - 1 != curBuffer.end())
        return 0;
  
      // Otherwise, return end of file.
      --curPtr;
      return EOF;
    }
    case '\n':
    case '\r':
      // Handle the newline character by ignoring it and incrementing the line
      // count. However, be careful about 'dos style' files with \n\r in them.
      // Only treat a \n\r or \r\n as a single line.
      if ((*curPtr == '\n' || (*curPtr == '\r')) && *curPtr != curChar)
        ++curPtr;
      return '\n';
    }
  }
  
  Token Lexer::lexToken() {
    while (true) {
      const char *tokStart = curPtr;
  
      // This always consumes at least one character.
      int curChar = getNextChar();
      switch (curChar) {
      default:
        // Handle identifiers: [a-zA-Z_]
        if (isalpha(curChar) || curChar == '_')
          return lexIdentifier(tokStart);
  
        // Handle integers: [0-9]
        if (isdigit(curChar))
          return lexInteger(tokStart);
  
        // Unknown character, emit an error.
        return emitError(tokStart, "unexpected character");
  
      case EOF:
        // Return EOF denoting the end of lexing.
        return formToken(Token::Kind::eof, tokStart);
  
      // Lex punctuation.
      case ':':
        return formToken(Token::Kind::colon, tokStart);
      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::r_brace, tokStart);
      case ')':
        return formToken(Token::Kind::r_paren, 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 '/':
        if (*curPtr == '/') {
          skipComment();
          continue;
        }
        // Unknown character, emit an error.
        return emitError(tokStart, "unexpected character: not a comment");
  
      // Ignore whitespace characters.
      case 0:
      case ' ':
      case '\t':
      case '\n':
        return lexToken();
      }
    }
  }
  
  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("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);
  }
  
  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);
  }
  
  /// Skip a comment line, starting with a '//'.
  void Lexer::skipComment() {
    // Advance over the second '/' in a '//' comment.
    assert(*curPtr == '/');
    ++curPtr;
  
    while (true) {
      switch (*curPtr++) {
      case '\n':
      case '\r':
        // Newline is end of comment.
        return;
      case 0:
        // If this is the end of the buffer, end the comment.
        if (curPtr - 1 == curBuffer.end()) {
          --curPtr;
          return;
        }
        LLVM_FALLTHROUGH;
      default:
        // Skip over other characters.
        break;
      }
    }
  }
  
  namespace {
  
  class Parser {
  public:
    Parser(llvm::SourceMgr &mgr, MLIRContext *ctx)
        : lexer(mgr), curToken(lexer.lexToken()), context(ctx) {}
  
    //===--------------------------------------------------------------------===//
    // Lexer Utilities
    //===--------------------------------------------------------------------===//
  
    /// 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();
    }
    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");
        consumeToken(rightToken);
        return success();
      }
  
      if (failed(parseCommaSeparatedList(parseElement)) ||
          failed(
              parseToken(rightToken, "expected ',' or right-terminating token")))
        return failure();
  
      return success();
    }
  
    Lexer lexer;
    Token curToken;
    MLIRContext *context;
  };
  } // namespace
  
  //===----------------------------------------------------------------------===//
  // Affine parsing.
  //===----------------------------------------------------------------------===//
  
  namespace {
  
  /// Lower precedence ops (all at the same precedence level). LNoOp is false in
  /// the boolean sense.
  enum AffineLowPrecOp {
    /// Null value.
    LNoOp,
    Add,
    Sub
  };
  
  /// Higher precedence ops - all at the same precedence level. HNoOp is false
  /// in the boolean sense.
  enum AffineHighPrecOp {
    /// Null value.
    HNoOp,
    Mul,
    FloorDiv,
    CeilDiv,
    Mod
  };
  
  using AffineDimList = SmallVector<std::pair<StringRef, AffineExpr>, 4>;
  using AffineSymbolList = SmallVector<std::pair<StringRef, AffineExpr>, 4>;
  
  /// This is a specialized parser for affine expressions.
  class AffineParser {
  public:
    explicit AffineParser(Parser &p,
                          std::function<AffineExpr(StringRef)> bareIdParsingHook,
                          AffineDimList &dimList, AffineSymbolList &symbolList)
        : parser(p), bareIdFallback(bareIdParsingHook), dims(dimList),
          symbols(symbolList) {}
  
    /// Parse a comma-separated list of affine exprs.
    SmallVector<AffineExpr, 4>
    parseAffineExprs(Token::Kind lDelim = Token::Kind::l_paren,
                     Token::Kind rDelim = Token::Kind::r_paren);
  
    /// Parse a single affine expr.`.
    AffineExpr parseAffineExpr();
  
  private:
    // Binary affine op parsing.
    AffineLowPrecOp consumeIfLowPrecOp();
    AffineHighPrecOp consumeIfHighPrecOp();
  
    // AffineExpr parsing.
    AffineExpr parseParentheticalExpr();
    AffineExpr parseNegateExpression(AffineExpr lhs);
    AffineExpr parseIntegerExpr();
    AffineExpr parseBareIdExpr();
  
    AffineExpr getAffineBinaryOpExpr(AffineHighPrecOp op, AffineExpr lhs,
                                     AffineExpr rhs, SMLoc opLoc);
    AffineExpr getAffineBinaryOpExpr(AffineLowPrecOp op, AffineExpr lhs,
                                     AffineExpr rhs);
    AffineExpr parseAffineOperandExpr(AffineExpr lhs);
    AffineExpr parseAffineLowPrecOpExpr(AffineExpr llhs, AffineLowPrecOp llhsOp);
    AffineExpr parseAffineHighPrecOpExpr(AffineExpr llhs, AffineHighPrecOp llhsOp,
                                         SMLoc llhsOpLoc);
  
    Parser &parser;
    std::function<AffineExpr(StringRef)> bareIdFallback;
    AffineDimList &dims;
    AffineSymbolList &symbols;
  };
  } // end anonymous namespace
  
  /// Create an affine binary high precedence op expression (mul's, div's, mod).
  /// opLoc is the location of the op token to be used to report errors
  /// for non-conforming expressions.
  AffineExpr AffineParser::getAffineBinaryOpExpr(AffineHighPrecOp op,
                                                 AffineExpr lhs, AffineExpr rhs,
                                                 SMLoc opLoc) {
    switch (op) {
    case Mul:
      if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant()) {
        parser.emitError(opLoc,
                         "non-affine expression: at least one of the multiply "
                         "operands has to be either a constant or symbolic");
        return nullptr;
      }
      return lhs * rhs;
    case FloorDiv:
      if (!rhs.isSymbolicOrConstant()) {
        parser.emitError(opLoc,
                         "non-affine expression: right operand of floordiv "
                         "has to be either a constant or symbolic");
        return nullptr;
      }
      return lhs.floorDiv(rhs);
    case CeilDiv:
      if (!rhs.isSymbolicOrConstant()) {
        parser.emitError(opLoc, "non-affine expression: right operand of ceildiv "
                                "has to be either a constant or symbolic");
        return nullptr;
      }
      return lhs.ceilDiv(rhs);
    case Mod:
      if (!rhs.isSymbolicOrConstant()) {
        parser.emitError(opLoc, "non-affine expression: right operand of mod "
                                "has to be either a constant or symbolic");
        return nullptr;
      }
      return lhs % rhs;
    case HNoOp:
      llvm_unreachable("can't create affine expression for null high prec op");
      return nullptr;
    }
    llvm_unreachable("Unknown AffineHighPrecOp");
  }
  
  /// Create an affine binary low precedence op expression (add, sub).
  AffineExpr AffineParser::getAffineBinaryOpExpr(AffineLowPrecOp op,
                                                 AffineExpr lhs, AffineExpr rhs) {
    switch (op) {
    case AffineLowPrecOp::Add:
      return lhs + rhs;
    case AffineLowPrecOp::Sub:
      return lhs - rhs;
    case AffineLowPrecOp::LNoOp:
      llvm_unreachable("can't create affine expression for null low prec op");
      return nullptr;
    }
    llvm_unreachable("Unknown AffineLowPrecOp");
  }
  
  /// Consume this token if it is a lower precedence affine op (there are only
  /// two precedence levels).
  AffineLowPrecOp AffineParser::consumeIfLowPrecOp() {
    switch (parser.curToken.getKind()) {
    case Token::Kind::plus:
      parser.consumeToken();
      return AffineLowPrecOp::Add;
    case Token::Kind::minus:
      parser.consumeToken();
      return AffineLowPrecOp::Sub;
    default:
      return AffineLowPrecOp::LNoOp;
    }
  }
  
  /// Consume this token if it is a higher precedence affine op (there are only
  /// two precedence levels)
  AffineHighPrecOp AffineParser::consumeIfHighPrecOp() {
    switch (parser.curToken.getKind()) {
    case Token::Kind::star:
      parser.consumeToken(Token::Kind::star);
      return Mul;
    case Token::Kind::kw_floordiv:
      parser.consumeToken(Token::Kind::kw_floordiv);
      return FloorDiv;
    case Token::Kind::kw_ceildiv:
      parser.consumeToken(Token::Kind::kw_ceildiv);
      return CeilDiv;
    case Token::Kind::kw_mod:
      parser.consumeToken(Token::Kind::kw_mod);
      return Mod;
    default:
      return HNoOp;
    }
  }
  
  /// Parse a high precedence op expression list: mul, div, and mod are high
  /// precedence binary ops, i.e., parse a
  ///   expr_1 op_1 expr_2 op_2 ... expr_n
  /// where op_1, op_2 are all a AffineHighPrecOp (mul, div, mod).
  /// All affine binary ops are left associative.
  /// Given llhs, returns (llhs llhsOp lhs) op rhs, or (lhs op rhs) if llhs is
  /// null. If no rhs can be found, returns (llhs llhsOp lhs) or lhs if llhs is
  /// null. llhsOpLoc is the location of the llhsOp token that will be used to
  /// report an error for non-conforming expressions.
  AffineExpr AffineParser::parseAffineHighPrecOpExpr(AffineExpr llhs,
                                                     AffineHighPrecOp llhsOp,
                                                     SMLoc llhsOpLoc) {
    AffineExpr lhs = parseAffineOperandExpr(llhs);
    if (!lhs)
      return nullptr;
  
    // Found an LHS. Parse the remaining expression.
    auto opLoc = parser.curToken.getLoc();
    if (AffineHighPrecOp op = consumeIfHighPrecOp()) {
      if (llhs) {
        AffineExpr expr = getAffineBinaryOpExpr(llhsOp, llhs, lhs, opLoc);
        if (!expr)
          return nullptr;
        return parseAffineHighPrecOpExpr(expr, op, opLoc);
      }
      // No LLHS, get RHS
      return parseAffineHighPrecOpExpr(lhs, op, opLoc);
    }
  
    // This is the last operand in this expression.
    if (llhs)
      return getAffineBinaryOpExpr(llhsOp, llhs, lhs, llhsOpLoc);
  
    // No llhs, 'lhs' itself is the expression.
    return lhs;
  }
  
  /// Parse an affine expression inside parentheses.
  ///
  ///   affine-expr ::= `(` affine-expr `)`
  AffineExpr AffineParser::parseParentheticalExpr() {
    if (failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
      return nullptr;
    if (parser.curToken.is(Token::Kind::r_paren))
      return (parser.emitError("no expression inside parentheses"), nullptr);
  
    auto expr = parseAffineExpr();
    if (!expr)
      return nullptr;
    if (failed(parser.parseToken(Token::Kind::r_paren, "expected ')'")))
      return nullptr;
  
    return expr;
  }
  
  /// Parse the negation expression.
  ///
  ///   affine-expr ::= `-` affine-expr
  AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
    if (failed(parser.parseToken(Token::Kind::minus, "expected '-'")))
      return nullptr;
  
    AffineExpr operand = parseAffineOperandExpr(lhs);
    // Since negation has the highest precedence of all ops (including high
    // precedence ops) but lower than parentheses, we are only going to use
    // parseAffineOperandExpr instead of parseAffineExpr here.
    if (!operand)
      // Extra error message although parseAffineOperandExpr would have
      // complained. Leads to a better diagnostic.
      return (parser.emitError("missing operand of negation"), nullptr);
    return (-1) * operand;
  }
  
  /// Parse a bare id that may appear in an affine expression.
  ///
  ///   affine-expr ::= bare-id
  AffineExpr AffineParser::parseBareIdExpr() {
    if (parser.curToken.isNot(Token::Kind::id))
      return (parser.emitError("expected id"), nullptr);
  
    StringRef sRef = parser.curToken.getSpelling();
    for (auto &list : {dims, symbols}) {
      for (auto entry : list) {
        if (entry.first == sRef) {
          parser.consumeToken(Token::Kind::id);
          return entry.second;
        }
      }
    }
  
    // Not found, check fallback path.
    AffineExpr expr = bareIdFallback(sRef);
    if (expr) {
      parser.consumeToken(Token::Kind::id);
      return expr;
    }
  
    return (parser.emitError("use of undeclared id"), nullptr);
  }
  
  /// Parse a positive integral constant appearing in an affine expression.
  ///
  ///   affine-expr ::= integer-literal
  AffineExpr AffineParser::parseIntegerExpr() {
    auto val = parser.curToken.getUInt64IntegerValue();
    if (!val.hasValue() || (int64_t)val.getValue() < 0)
      return (parser.emitError("constant too large for index"), nullptr);
  
    parser.consumeToken(Token::Kind::integer);
    return getAffineConstantExpr((int64_t)val.getValue(), parser.context);
  }
  
  /// Parses an expression that can be a valid operand of an affine expression.
  /// lhs: if non-null, lhs is an affine expression that is the lhs of a binary
  /// operator, the rhs of which is being parsed. This is used to determine
  /// whether an error should be emitted for a missing right operand.
  //  Eg: for an expression without parentheses (like i + j + k + l), each
  //  of the four identifiers is an operand. For i + j*k + l, j*k is not an
  //  operand expression, it's an op expression and will be parsed via
  //  parseAffineHighPrecOpExpression(). However, for i + (j*k) + -l, (j*k) and
  //  -l are valid operands that will be parsed by this function.
  AffineExpr AffineParser::parseAffineOperandExpr(AffineExpr lhs) {
    switch (parser.curToken.getKind()) {
    case Token::Kind::id:
      return parseBareIdExpr();
    case Token::Kind::integer:
      return parseIntegerExpr();
    case Token::Kind::l_paren:
      return parseParentheticalExpr();
    case Token::Kind::minus:
      return parseNegateExpression(lhs);
    case Token::Kind::kw_ceildiv:
    case Token::Kind::kw_floordiv:
    case Token::Kind::kw_mod:
    case Token::Kind::plus:
    case Token::Kind::star:
      if (lhs)
        parser.emitError("missing right operand of binary operator");
      else
        parser.emitError("missing left operand of binary operator");
      return nullptr;
    default:
      if (lhs)
        parser.emitError("missing right operand of binary operator");
      else
        parser.emitError("expected affine expression");
      return nullptr;
    }
  }
  
  /// Parse affine expressions that are bare-id's, integer constants,
  /// parenthetical affine expressions, and affine op expressions that are a
  /// composition of those.
  ///
  /// All binary op's associate from left to right.
  ///
  /// {add, sub} have lower precedence than {mul, div, and mod}.
  ///
  /// Add, sub'are themselves at the same precedence level. Mul, floordiv,
  /// ceildiv, and mod are at the same higher precedence level. Negation has
  /// higher precedence than any binary op.
  ///
  /// llhs: the affine expression appearing on the left of the one being parsed.
  /// This function will return ((llhs llhsOp lhs) op rhs) if llhs is non null,
  /// and lhs op rhs otherwise; if there is no rhs, llhs llhsOp lhs is returned
  /// if llhs is non-null; otherwise lhs is returned. This is to deal with left
  /// associativity.
  ///
  /// Eg: when the expression is e1 + e2*e3 + e4, with e1 as llhs, this function
  /// will return the affine expr equivalent of (e1 + (e2*e3)) + e4, where
  /// (e2*e3) will be parsed using parseAffineHighPrecOpExpr().
  AffineExpr AffineParser::parseAffineLowPrecOpExpr(AffineExpr llhs,
                                                    AffineLowPrecOp llhsOp) {
    AffineExpr lhs;
    if (!(lhs = parseAffineOperandExpr(llhs)))
      return nullptr;
  
    // Found an LHS. Deal with the ops.
    if (AffineLowPrecOp lOp = consumeIfLowPrecOp()) {
      if (llhs) {
        AffineExpr sum = getAffineBinaryOpExpr(llhsOp, llhs, lhs);
        return parseAffineLowPrecOpExpr(sum, lOp);
      }
      // No LLHS, get RHS and form the expression.
      return parseAffineLowPrecOpExpr(lhs, lOp);
    }
    auto opLoc = parser.curToken.getLoc();
    if (AffineHighPrecOp hOp = consumeIfHighPrecOp()) {
      // We have a higher precedence op here. Get the rhs operand for the llhs
      // through parseAffineHighPrecOpExpr.
      AffineExpr highRes = parseAffineHighPrecOpExpr(lhs, hOp, opLoc);
      if (!highRes)
        return nullptr;
  
      // If llhs is null, the product forms the first operand of the yet to be
      // found expression. If non-null, the op to associate with llhs is llhsOp.
      AffineExpr expr =
          llhs ? getAffineBinaryOpExpr(llhsOp, llhs, highRes) : highRes;
  
      // Recurse for subsequent low prec op's after the affine high prec op
      // expression.
      if (AffineLowPrecOp nextOp = consumeIfLowPrecOp())
        return parseAffineLowPrecOpExpr(expr, nextOp);
      return expr;
    }
    // Last operand in the expression list.
    if (llhs)
      return getAffineBinaryOpExpr(llhsOp, llhs, lhs);
    // No llhs, 'lhs' itself is the expression.
    return lhs;
  }
  
  /// Parse an affine expression.
  ///  affine-expr ::= `(` affine-expr `)`
  ///                | `-` affine-expr
  ///                | affine-expr `+` affine-expr
  ///                | affine-expr `-` affine-expr
  ///                | affine-expr `*` affine-expr
  ///                | affine-expr `floordiv` affine-expr
  ///                | affine-expr `ceildiv` affine-expr
  ///                | affine-expr `mod` affine-expr
  ///                | bare-id
  ///                | integer-literal
  ///
  /// Additional conditions are checked depending on the production. For eg.,
  /// one of the operands for `*` has to be either constant/symbolic; the second
  /// operand for floordiv, ceildiv, and mod has to be a positive integer.
  AffineExpr AffineParser::parseAffineExpr() {
    return parseAffineLowPrecOpExpr(nullptr, AffineLowPrecOp::LNoOp);
  }
  
  SmallVector<AffineExpr, 4> AffineParser::parseAffineExprs(Token::Kind lDelim,
                                                            Token::Kind rDelim) {
    parser.parseToken(lDelim, "expected lDelim at start of affine expr list");
  
    SmallVector<AffineExpr, 4> exprs;
    auto parseElt = [&]() -> LogicalResult {
      auto elt = parseAffineExpr();
      exprs.push_back(elt);
      return elt ? success() : failure();
    };
  
    if (failed(parser.parseCommaSeparatedListUntil(rDelim, parseElt,
                                                   /*allowEmptyList=*/true)))
      llvm_unreachable("Failed AffineExpr parsing");
  
    return exprs;
  }
  
  //===----------------------------------------------------------------------===//
  // TC parsing.
  //===----------------------------------------------------------------------===//
  
  namespace {
  
  /// Base class for expressions involved in TC parsing.
  struct Expression {
    enum class Kind {
      Uninitialized = 0,
      TensorExpr = 1,
      TensorUse = 2,
    };
  
    explicit Expression(Kind k = Kind::Uninitialized) : kind(k) {}
    virtual ~Expression() = default;
  
    operator bool() const { return kind != Kind::Uninitialized; }
  
    Kind kind;
  };
  
  /// Encodes a tensor use of the form:
  ///
  ///   affine-expr-list ::= affine-expr (`,` affine-expr)*
  ///   tensor-use ::= bare-id `(` `)`
  ///                | bare-id `(` affine-expr-list `)`
  ///
  /// The affine-expr-list is stored as an AffineMap.
  struct TensorUse : public Expression {
    TensorUse() : TensorUse("", AffineMap()) {}
    TensorUse(StringRef name, AffineMap map)
        : Expression(Kind::TensorUse), tensorId(name), indexingMap(map) {}
    TensorUse(const TensorUse &use) = default;
  
    static bool classof(const Expression *e) {
      return e->kind == Kind::TensorUse;
    }
  
    bool operator==(const TensorUse &other) const {
      return tensorId == other.tensorId && indexingMap == other.indexingMap;
    }
  
    /// Visitation function. Performs preorder or postorder traversal depending on
    /// `PreOrder` and applies `callback` on each node.
    template <typename Lambda, bool PreOrder>
    void visit(Lambda callback) const;
  
    StringRef tensorId;
    AffineMap indexingMap;
  };
  
  /// Encodes 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 `)`
  ///
  /// Underlying op-arg are stored by unique_ptr to base class.
  struct TensorExpr : public Expression {
    TensorExpr(StringRef name,
               SmallVectorImpl<std::unique_ptr<Expression>> &&exprs,
               ArrayRef<unsigned> reductionDims)
        : Expression(Kind::TensorExpr), operationName(name),
          expressions(std::move(exprs)),
          reductionDimensions(reductionDims.begin(), reductionDims.end()) {}
  
    static bool classof(const Expression *e) {
      return e->kind == Kind::TensorExpr;
    }
  
    bool operator==(const TensorExpr &other) const {
      if (operationName != other.operationName)
        return false;
      if (expressions.size() != other.expressions.size())
        return false;
      for (unsigned i = 0, e = expressions.size(); i < e; ++i)
        if (*expressions[i] != *other.expressions[i])
          return false;
      for (unsigned i = 0, e = reductionDimensions.size(); i < e; ++i)
        if (reductionDimensions[i] != other.reductionDimensions[i])
          return false;
      return true;
    }
  
    /// Visitation function. Performs preorder or postorder traversal depending on
    /// `PreOrder` and applies `callback` on each node.
    template <typename Lambda, bool PreOrder>
    void visit(Lambda callback) const;
  
    StringRef operationName;
    SmallVector<std::unique_ptr<Expression>, 4> expressions;
    SetVector<unsigned> reductionDimensions;
  };
  
  /// This is a specialized parser for a TCDef.
  /// This maintains the dims it finds in an eager fashion.
  class TCParser {
    enum class EagerDiscoveryMode { None = 0, Symbols, Dimensions };
  
  public:
    explicit TCParser(Parser &p);
  
    /// Uses the AffineParser to parse the affine exprs used in a tensor
    /// definition. If `discoveryMode` is set to Symbols (resp. Dimensions), new
    /// symbols (resp. dimensions) are added eagerly. Otherwise, an error is
    /// emitted on new identifiers.
    SmallVector<AffineExpr, 4>
    parseAffineExprs(EagerDiscoveryMode discoveryMode, AffineDimList &dims,
                     Token::Kind lDelim = Token::Kind::l_paren,
                     Token::Kind rDelim = Token::Kind::r_paren);
  
    /// Parse the information for a tensor def.
    /// All the affine-expr must be dimensionless (i.e. contain only expressions
    /// involving symbols and constants), but can otherwise contain arbitrary
    /// affine expressions.
    LogicalResult parseTensorDef(bool isOutput);
  
    /// Parses a tensor use.
    struct ComprehensionParsingState {
      AffineDimList dims;
      SmallVector<std::unique_ptr<Expression>, 4> expressions;
      llvm::DenseMap<TensorUse, unsigned> orderedTensorArgs;
    };
    LogicalResult parseTensorUse(TensorUse &result,
                                 ComprehensionParsingState &state);
  
    /// 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,
                                        ComprehensionParsingState &state);
  
    /// Parse and print the information for a TC def.
    /// When `gen-ods-decl` is used, this prints the ODS declaration for the TC.
    /// When `gen-impl` is used, this prints the C++ implementation for the extra
    /// methods defined in ODS (`iterator_types`, `indexing_maps` and
    /// `regionBuilder`).
    LogicalResult parseAndEmitODSDef(llvm::raw_ostream &os);
  
    /// Print the ODS class that defines a new `cppOpName` for a `linalgOpName`.
    void printODS(llvm::raw_ostream &os, StringRef cppOpName,
                  StringRef linalgOpName, ComprehensionParsingState &state);
  
    /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
    void printReferenceIterators(llvm::raw_ostream &os, StringRef cppOpName,
                                 ComprehensionParsingState &state);
  
    /// Print the C++ StructuredOpsInterface impl of `indexing_maps`.
    void printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef cppOpName,
                                    ComprehensionParsingState &state);
  
    /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
    void printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                            ComprehensionParsingState &state);
  
    /// Print the C++ impl for named ops canonicalizers and fodlers.
    void printCanonicalizersAndFolders(llvm::raw_ostream &os,
                                       StringRef cppOpName);
  
  private:
    //===--------------------------------------------------------------------===//
    // Internal bookkeeping of tensors.
    //===--------------------------------------------------------------------===//
    struct RegisteredTensor {
      StringRef type;
      AffineMap shape;
      bool isOutput;
      AffineMap indexingMap;
      unsigned index;
    };
  
    //===--------------------------------------------------------------------===//
    // Per-TC def state.
    //===--------------------------------------------------------------------===//
    /// Symbols are per TC def.
    AffineSymbolList symbols;
    /// Tensors are per TC def.
    llvm::StringMap<RegisteredTensor> registeredTensors;
    unsigned nextRegisteredTensorIndex;
  
    Parser &parser;
  };
  } // namespace
  
  namespace llvm {
  
  template <>
  struct DenseMapInfo<TensorUse> {
    static TensorUse getEmptyKey() { return TensorUse("", AffineMap()); }
    static TensorUse getTombstoneKey() {
      return TensorUse(DenseMapInfo<StringRef>::getTombstoneKey(),
                       DenseMapInfo<AffineMap>::getTombstoneKey());
    }
    static unsigned getHashValue(const TensorUse &val) {
      return ::llvm::hash_value(val.tensorId); // don't care about collisions.
    }
    static bool isEqual(const TensorUse &LHS, const TensorUse &RHS) {
      return LHS == RHS;
    }
  };
  
  } // namespace llvm
  
  //===----------------------------------------------------------------------===//
  // Visitation functions.
  //===----------------------------------------------------------------------===//
  
  template <typename Lambda, bool PreOrder>
  void visit(const Expression &expr, Lambda callback) {
    switch (expr.kind) {
    default:
      llvm_unreachable("Unexpected kind");
    case Expression::Kind::TensorExpr:
      static_cast<const TensorExpr &>(expr).visit<Lambda, PreOrder>(callback);
      break;
    case Expression::Kind::TensorUse:
      static_cast<const TensorUse &>(expr).visit<Lambda, PreOrder>(callback);
      break;
    }
  }
  
  template <typename Lambda>
  void visitPreorder(const Expression &expr, Lambda callback) {
    visit<Lambda, false>(expr, callback);
  }
  
  template <typename Lambda>
  void visitPostorder(Expression &expr, Lambda callback) {
    visit<Lambda, true>(expr, callback);
  }
  
  template <typename Lambda, bool PreOrder>
  void TensorExpr::visit(Lambda callback) const {
    if (!PreOrder)
      callback(*this);
    for (auto &e : expressions)
      ::visit<Lambda, PreOrder>(*e, callback);
    if (PreOrder)
      callback(*this);
  }
  
  template <typename Lambda, bool PreOrder>
  void TensorUse::visit(Lambda callback) const {
    callback(*this);
  }
  
  //===----------------------------------------------------------------------===//
  // TC parsing functions.
  //===----------------------------------------------------------------------===//
  TCParser::TCParser(Parser &p)
      : symbols(), registeredTensors(), nextRegisteredTensorIndex(0), parser(p) {}
  
  /// Uses the AffineParser to parse the affine exprs used in a tensor
  /// definition. All identifiers are interpreted as symbols, new symbols are
  /// added eagerly.
  SmallVector<AffineExpr, 4>
  TCParser::parseAffineExprs(EagerDiscoveryMode discoveryMode,
                             AffineDimList &dims, Token::Kind lDelim,
                             Token::Kind rDelim) {
    AffineParser affineParser(
        parser,
        [&](StringRef sRef) {
          AffineExpr expr;
          if (discoveryMode == EagerDiscoveryMode::Symbols) {
            expr = getAffineSymbolExpr(symbols.size(), parser.context);
            symbols.emplace_back(sRef, expr);
          } else if (discoveryMode == EagerDiscoveryMode::Dimensions) {
            expr = getAffineDimExpr(dims.size(), parser.context);
            dims.emplace_back(sRef, expr);
          }
          return expr;
        },
        dims, symbols);
    return affineParser.parseAffineExprs(lDelim, rDelim);
  }
  
  /// Parse the information for a tensor def of the form:
  ///
  ///   affine-expr-list ::= affine-expr (`,` affine-expr )*
  ///   tensor-typedef ::= type `(` `)`
  ///                    | type `(` affine-expr-list `)`
  ///   tensor-def ::= bare-id `:` tensor-typedef
  LogicalResult TCParser::parseTensorDef(bool isOutput) {
    StringRef tensorId = parser.curToken.getSpelling();
    if (failed(parser.parseToken(Token::Kind::id, "expected an id")) ||
        failed(parser.parseToken(Token::Kind::colon, "expected colon")))
      return failure();
  
    StringRef tensorType = parser.curToken.getSpelling();
    if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
      return failure();
  
    AffineDimList emptyDims;
    auto exprs = parseAffineExprs(EagerDiscoveryMode::Symbols, emptyDims);
    assert(emptyDims.empty() && "Unexpected dimension in tensor def");
    AffineMap map =
        AffineMap::get(/*dimCount=*/0, symbols.size(), exprs, parser.context);
  
    auto iterBoolPair = registeredTensors.try_emplace(
        tensorId, RegisteredTensor{tensorType, map, isOutput, AffineMap(),
                                   nextRegisteredTensorIndex++});
    (void)iterBoolPair;
    assert(iterBoolPair.second && "Could not emplace tensor registration");
    LLVM_DEBUG(llvm::dbgs() << "Recorded: " << tensorId << " "
                            << "with typeString: " << tensorType << " "
                            << "and shape: " << map << "\n");
  
    return success();
  }
  
  /// Parses a tensor use of the form:
  ///
  ///   affine-expr-list ::= affine-expr (`,` affine-expr)*
  ///   tensor-use ::= bare-id `(` `)`
  ///                | bare-id `(` affine-expr-list `)`
  LogicalResult TCParser::parseTensorUse(TensorUse &result,
                                         ComprehensionParsingState &state) {
    StringRef tensorId = parser.curToken.getSpelling();
    if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
      return failure();
  
    auto exprs = parseAffineExprs(EagerDiscoveryMode::Dimensions, state.dims);
    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();
  }
  
  /// 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 `)`
  LogicalResult TCParser::parseExpression(TensorUse currentDefinition,
                                          std::unique_ptr<Expression> &result,
                                          ComprehensionParsingState &state) {
    StringRef opOrTensor = parser.curToken.getSpelling();
    if (registeredTensors.count(opOrTensor) > 0) {
      TensorUse use;
      auto res = parseTensorUse(use, state);
      if (failed(res))
        return res;
      result = std::make_unique<TensorUse>(use);
      return success();
    }
  
    if (failed(parser.parseToken(Token::Kind::id, "expected an operation")))
      return failure();
  
    // This is an op.
    SmallVector<unsigned, 4> reductionDims;
    SmallVector<std::unique_ptr<Expression>, 4> expressions;
  
    // Check if it has a reduction set, discover dimensions eagerly.
    if (parser.curToken.is(Token::Kind::lt)) {
      auto iters = parseAffineExprs(EagerDiscoveryMode::Dimensions, state.dims,
                                    Token::Kind::lt, Token::Kind::gt);
      for (auto iter : iters)
        reductionDims.push_back(iter.cast<AffineDimExpr>().getPosition());
    }
  
    // If this op is a reduction, it's first argument is the `currentDefinition`
    // tensor use.
    if (!reductionDims.empty())
      expressions.push_back(std::make_unique<TensorUse>(currentDefinition));
    LLVM_DEBUG(llvm::dbgs() << "op: " << opOrTensor << "\n");
  
    auto parseExpr = [&]() -> LogicalResult {
      std::unique_ptr<Expression> e;
      if (failed(parseExpression(currentDefinition, e, state)))
        return failure();
      expressions.push_back(std::move(e));
      return success();
    };
    if (failed(parser.parseToken(Token::Kind::l_paren, "expected '('")) ||
        failed(parser.parseCommaSeparatedListUntil(
            Token::Kind::r_paren, parseExpr, /*allowEmptyList=*/true)))
      return failure();
  
    result = std::make_unique<TensorExpr>(opOrTensor, std::move(expressions),
                                          reductionDims);
  
    return success();
  }
  
  //===----------------------------------------------------------------------===//
  // Parse and Emit functions.
  //===----------------------------------------------------------------------===//
  
  /// Parse the information for a single comprehension.
  ///
  ///   tensor-def-list ::= tensor-def (`,` tensor-def)*
  ///   tensor-expr-list ::= tensor-expr (`,` tensor-expr)*
  ///   comprehension ::= tensor-def-list `=` tensor-expr-list `;`
  LogicalResult
  TCParser::parseOneComprehension(StringRef cppOpName, StringRef linalgOpName,
                                  ComprehensionParsingState &state) {
    // 1. Parse LHS of `=`, these become the definitions that appear as the output
    // tensors or read/write buffers.
    SmallVector<TensorUse, 4> definitions;
    auto parseUse = [&]() -> LogicalResult {
      TensorUse use;
      if (failed(parseTensorUse(use, state)))
        return failure();
      definitions.push_back(use);
      return success();
    };
    if (failed(parser.parseCommaSeparatedListUntil(Token::Kind::equal, parseUse,
                                                   /*allowEmptyList=*/true)))
      return failure();
  
    // 2. Parse RHS of `=`, this becomes the expressions from which we emit
    // computations.
    unsigned idx = 0;
    auto parseExpr = [&]() -> LogicalResult {
      std::unique_ptr<Expression> expr;
      if (idx >= definitions.size()) {
        parser.emitError("Fewer LHS definitions than RHS expressions");
        return failure();
      }
      if (failed(parseExpression(definitions[idx++], expr, state)))
        return failure();
      state.expressions.push_back(std::move(expr));
      return success();
    };
    if (failed(parser.parseCommaSeparatedListUntil(
            Token::Kind::semicolon, parseExpr, /*allowEmptyList=*/true)))
      return failure();
    if (idx != definitions.size()) {
      parser.emitError("Fewer RHS expressions than LHS definitions");
      return failure();
    }
  
    // 3. Postprocess.
    // 3.a. Normalize all maps to the proper state.dims and symbols counts.
    SmallVector<TensorUse, 4> allUses;
    allUses.reserve(registeredTensors.size());
    for (auto &def : definitions)
      allUses.push_back(def);
    for (auto &pExpr : state.expressions)
      visitPostorder(*pExpr, [&](const Expression &e) {
        if (auto *use = dyn_cast<TensorUse>(&e))
          allUses.push_back(*use);
      });
    for (auto &use : allUses)
      use.indexingMap =
          AffineMap::get(state.dims.size(), symbols.size(),
                         use.indexingMap.getResults(), parser.context);
  
    // 3.b. Traverse definitions
    llvm::DenseSet<StringRef> seenDefs;
    for (auto &def : definitions) {
      if (seenDefs.count(def.tensorId) > 0) {
        parser.emitError("Unexpected multi-write to a single tensor");
        return failure();
      }
      seenDefs.insert(def.tensorId);
      auto tensorIter = registeredTensors.find(def.tensorId);
      assert(tensorIter != registeredTensors.end() && "unregistered tensor");
      auto &tensor = tensorIter->getValue();
      tensor.indexingMap = def.indexingMap;
      state.orderedTensorArgs[def] = tensor.index;
    }
  
    bool failed = false;
    for (auto &pExpr : state.expressions)
      visitPostorder(*pExpr, [&](const Expression &e) {
        auto *pUse = dyn_cast<TensorUse>(&e);
        if (failed || !pUse)
          return;
        auto &use = *pUse;
        LLVM_DEBUG(llvm::dbgs()
                   << "\nuse: " << use.tensorId << " map: " << use.indexingMap);
        auto tensorIter = registeredTensors.find(use.tensorId);
        assert(tensorIter != registeredTensors.end() && "unregistered tensor");
        auto &tensor = tensorIter->getValue();
        if (tensor.indexingMap && state.orderedTensorArgs.count(use) == 0) {
          LLVM_DEBUG(llvm::dbgs() << "\nexisting: " << tensor.indexingMap);
          parser.emitError(
              "Unexpected multi-read of a tensor with different accesses");
          failed = true;
          return;
        }
        seenDefs.insert(use.tensorId);
        tensor.indexingMap = use.indexingMap;
        state.orderedTensorArgs[use] = tensor.index;
      });
    if (failed)
      return failure();
  
    return success();
  }
  
  /// Parse and print the information for a ODS def.
  ///
  ///   tensor-def-list ::= tensor-def (`,` tensor-def )*
  ///
  ///   comprehension-list ::= comprehension comprehension*
  ///
  ///   tc-def ::= `def` bare-id `(`tensor-def-list`)` `->` `(` tensor-def-list`)`
  ///     `{` 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) {
    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");
    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();
  
    // 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(),
                           parser.context);
    }
  
    if (failed(parser.parseToken(Token::Kind::l_brace, "expected '{'")))
      return failure();
  
    SmallVector<ComprehensionParsingState, 4> perComprehensionStates;
    while (parser.curToken.isNot(Token::Kind::r_brace)) {
      perComprehensionStates.push_back(ComprehensionParsingState());
      if (failed(parseOneComprehension(cppOpName, tcName,
                                       perComprehensionStates.back())))
        return failure();
    };
    parser.parseToken(Token::Kind::r_brace, "expected '}'");
  
    // Print.
    auto nComprehensions = perComprehensionStates.size();
    if (nComprehensions != 1) {
      parser.emitError("only 1 comprehension supported for now, got: " +
                       llvm::Twine(nComprehensions));
      return failure();
    }
    if (genODSDecl) {
      auto &state = perComprehensionStates.back();
      printODS(os, cppOpName, tcName, state);
      os << "\n";
    }
    if (genODSImpl) {
      auto &state = perComprehensionStates.back();
      std::string extraMethods;
      llvm::raw_string_ostream ss(extraMethods);
      printReferenceIterators(ss, cppOpName, state);
      printReferenceIndexingMaps(ss, cppOpName, state);
      printRegionBuilder(ss, cppOpName, state);
      printCanonicalizersAndFolders(ss, cppOpName);
      ss.flush();
      os << extraMethods << "\n";
    }
  
    return success();
  }
  
  //===----------------------------------------------------------------------===//
  // 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) {
    const char *header = R"FMT(  def {0} : LinalgStructuredBase_Op<"{1}", [
      NamedStructuredOpTrait,
      AttrSizedOperandSegments,
      SingleBlockImplicitTerminator<"YieldOp">]> {
        let arguments = (ins Variadic<AnyShaped>:$inputs,
                             Variadic<AnyMemRef>:$output_buffers,
                             Variadic<AnyRankedTensor>:$init_tensors);
        let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
        let regions = (region AnyRegion:$region);
  
        let skipDefaultBuilders = 1;
        let builders = [ OpBuilder<
          "ValueRange inputs, ValueRange outputBuffers",
          [{{
            $_state.addOperands(inputs);
            $_state.addOperands(outputBuffers);
            $_state.addAttribute(
              "operand_segment_sizes",
              $_builder.getI32VectorAttr({{
                static_cast<int32_t>(inputs.size()),
                static_cast<int32_t>(outputBuffers.size()),
                static_cast<int32_t>(0)}));
            buildNamedStructuredOpRegionAndAttributes<{0}>(
              $_builder,
              $_state,
              TypeRange(inputs),
              TypeRange(outputBuffers),
              TypeRange(),
              TypeRange());
          }]>, OpBuilder<
          "TypeRange resultTensorTypes, ValueRange inputs, "
          "ValueRange outputBuffers, ValueRange initTensors",
          [{{
            $_state.addOperands(inputs);
            $_state.addOperands(outputBuffers);
            $_state.addOperands(initTensors);
            $_state.addTypes(resultTensorTypes);
            $_state.addAttribute(
              "operand_segment_sizes",
              $_builder.getI32VectorAttr({{
                static_cast<int32_t>(inputs.size()),
                static_cast<int32_t>(outputBuffers.size()),
                static_cast<int32_t>(initTensors.size())}));
            buildNamedStructuredOpRegionAndAttributes<{0}>(
              $_builder,
              $_state,
              TypeRange(inputs),
              TypeRange(outputBuffers),
              TypeRange(initTensors),
              resultTensorTypes);
          }]>, OpBuilder<
          "TypeRange resultTensorTypes, ValueRange operands, ArrayRef<NamedAttribute> attributes = {{}",
          [{{
            $_state.addOperands(operands);
            $_state.addAttributes(attributes);
            $_state.addTypes(resultTensorTypes);
            (void)$_state.addRegion();
          }]>
        ];
        let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
        let parser = [{{ return ::parseNamedStructuredOp<{0}>(parser, result); }];
        let verifier = [{{ return ::verifyNamedStructuredOp(*this); }];
        let hasFolder = 1;
        let hasCanonicalizer = 1;
  
        let extraClassDeclaration = [{{
          // Auto-generated.
          ArrayAttr iterator_types();
          ArrayAttr indexing_maps();
          static void regionBuilder(Block &block);
  
          // Generic methods.
          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++;
    }
  
    os << llvm::formatv(header, cppOpName, linalgOpName, nInputs, nOutputs,
                        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(
      ArrayAttr {0}::iterator_types() {
        return Builder(getContext()).getStrArrayAttr(SmallVector<StringRef, 8>{{ {1} });
      })FMT";
  
    std::string iteratorsStr;
    llvm::raw_string_ostream ss(iteratorsStr);
    unsigned pos = 0;
    llvm::interleaveComma(
        state.dims, ss, [&](std::pair<StringRef, AffineExpr> p) {
          bool reduction = false;
          for (auto &expr : state.expressions) {
            visitPostorder(*expr, [&](const Expression &e) {
              if (auto *pTensorExpr = dyn_cast<TensorExpr>(&e)) {
                if (pTensorExpr->reductionDimensions.count(pos) > 0)
                  reduction = true;
              }
            });
            if (reduction)
              break;
          }
          ss << (reduction ? "getReductionIteratorTypeName()"
                           : "getParallelIteratorTypeName()");
          pos++;
        });
    ss.flush();
  
    os << llvm::formatv(referenceReferenceIteratorsFmt, cppOpName, iteratorsStr);
  }
  
  void TCParser::printCanonicalizersAndFolders(llvm::raw_ostream &os,
                                               StringRef cppOpName) {
    const char *canonicalizersAndFoldersFmt = R"FMT(
      void {0}::getCanonicalizationPatterns(
          OwningRewritePatternList &results,
          MLIRContext *context) {{
        results.insert<EraseDeadLinalgOp>();
        results.insert<FoldTensorCastOp>();
      }
      LogicalResult {0}::fold(ArrayRef<Attribute>,
                              SmallVectorImpl<OpFoldResult> &) {{
        return foldMemRefCast(*this);
      })FMT";
    os << llvm::formatv(canonicalizersAndFoldersFmt, cppOpName);
  }
  
  /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
  void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
                                            StringRef cppOpName,
                                            ComprehensionParsingState &state) {
    // 1. Generic string template for specifying reference indexing maps.
    const char *referenceIndexingMapsFmt =
        R"FMT(
    // This is temporary until we transition out of manually specified ops that
    // should be auto-generated with linalg-ods-gen.
    ArrayAttr {0}::indexing_maps() {
      MLIRContext *context = getContext();
      AffineExpr {1};
      bindDims(context, {1});
      return Builder(context).getAffineMapArrayAttr({ {2} });
    })FMT";
  
    // 2. Print a comma-separated list of identifiers for the AffineExpr in
    // `state.dims`. These will replace the `{1}` placeholder in both
    // `AffineExpr {1}` and `bindDims(context, {1})` ensuring the AffineExpr
    // identifiers are bound in the right order to the proper AffineDimExpr.
    std::string dimsStr;
    llvm::raw_string_ostream ss(dimsStr);
    llvm::interleaveComma(
        state.dims, ss,
        [&](std::pair<StringRef, AffineExpr> p) { ss << p.second; });
    ss.flush();
  
    // 3. Print a comma-separated list of AffineMap constructors that use the
    // identifiers from 1. The AffineExpr use the common arithmetic operators on
    // AffineExpr. These AffineMap constructors will replace the `{2}` placeholder
    // in return `SmallVector<AffineMap, 8>{{ {2} };`.
    std::string mapsStr;
    llvm::raw_string_ostream mapsStringStream(mapsStr);
    SmallVector<TensorUse, 4> orderedUses(state.orderedTensorArgs.size());
    for (const auto &it : state.orderedTensorArgs)
      orderedUses[it.second] = it.first;
    llvm::interleaveComma(orderedUses, mapsStringStream, [&](TensorUse u) {
      assert(u.indexingMap);
      const char *mapFmt = "\n\tAffineMap::get({0}, 0, {1}, context)";
      if (u.indexingMap.isEmpty()) {
        mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), "context");
        return;
      }
  
      std::string exprsStr;
      llvm::raw_string_ostream exprsStringStream(exprsStr);
      exprsStringStream << "{";
      llvm::interleaveComma(u.indexingMap.getResults(), exprsStringStream);
      exprsStringStream << "}";
      exprsStringStream.flush();
  
      mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), exprsStr);
    });
    mapsStringStream.flush();
  
    // 4. Apply format to 1. using 2. and 3.
    os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr);
  }
  
  /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
  void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                                    ComprehensionParsingState &state) {
    unsigned count = state.orderedTensorArgs.size();
    llvm::DenseMap<const TensorExpr *, unsigned> subExprsMap;
    std::function<void(llvm::raw_ostream & os, const Expression &)> printExpr;
    printExpr = [&](llvm::raw_ostream &os, const Expression &e) -> void {
      if (auto *pUse = dyn_cast<TensorUse>(&e)) {
        os << "_" << state.orderedTensorArgs.find(*pUse)->second;
        return;
      }
      auto *pTensorExpr = cast<TensorExpr>(&e);
      if (subExprsMap.count(pTensorExpr) > 0) {
        os << "_" << subExprsMap[pTensorExpr];
      } else {
        std::string subExprs;
        llvm::raw_string_ostream subExprsStringStream(subExprs);
        llvm::interleaveComma(pTensorExpr->expressions, subExprsStringStream,
                              [&](const std::unique_ptr<Expression> &e) {
                                printExpr(subExprsStringStream, *e);
                              });
        subExprsStringStream.flush();
        const char *tensorExprFmt = "\n    Value _{0} = {1}({2});";
        os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->operationName,
                            subExprs);
        subExprsMap[pTensorExpr] = count;
      }
    };
  
    const char *regionBuilderFmt = R"FMT(
    void {0}::regionBuilder(Block &block) {
      using namespace edsc;
      using namespace intrinsics;
      auto args = block.getArguments();
      Value {1};
      {2}
      (linalg_yield(ValueRange{ {3} }));
    })FMT";
  
    unsigned idx = 0;
    std::string valueHandleStr;
    llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);
    llvm::interleaveComma(
        state.orderedTensorArgs, valueHandleStringStream, [&](auto) {
          valueHandleStringStream << "_" << idx << "(args[" << idx << "])";
          idx++;
        });
  
    std::string expressionsStr;
    llvm::raw_string_ostream expressionStringStream(expressionsStr);
    for (auto &expr : state.expressions)
      visitPostorder(*expr, [&](const Expression &e) {
        if (e.kind == Expression::Kind::TensorExpr)
          printExpr(expressionStringStream, e);
      });
  
    std::string yieldStr;
    llvm::raw_string_ostream yieldStringStream(yieldStr);
    llvm::interleaveComma(state.expressions, yieldStringStream,
                          [&](const std::unique_ptr<Expression> &e) {
                            printExpr(yieldStringStream, *e);
                          });
  
    valueHandleStringStream.flush();
    expressionStringStream.flush();
    yieldStringStream.flush();
  
    os << llvm::formatv(regionBuilderFmt, cppOpName, valueHandleStr,
                        expressionsStr, yieldStr);
  }
  
  /// Iterate over each Tensor Comprehension def.
  LogicalResult parseAndEmitAllTensorComprehensions(llvm::raw_ostream &os,
                                                    Parser &parser) {
    while (parser.curToken.getKind() != Token::Kind::eof) {
      TCParser tcParser(parser);
      if (failed(tcParser.parseAndEmitODSDef(os)))
        return failure();
    }
    return success();
  }
  
  int main(int argc, char **argv) {
    llvm::cl::ParseCommandLineOptions(argc, argv, "Linalg ODS Gen");
  
    // Set up the input file.
    std::string errorMessage;
    std::unique_ptr<llvm::MemoryBuffer> file =
        mlir::openInputFile(inputFilename, &errorMessage);
    if (!file) {
      llvm::errs() << errorMessage << "\n";
      return 1;
    }
  
    std::unique_ptr<llvm::ToolOutputFile> output =
        openOutputFile(outputFilename, &errorMessage);
    if (!output) {
      llvm::errs() << errorMessage << "\n";
      exit(1);
    }
  
    // Include the proper Linalg header for end-to-end tblgen testing without
    // resorting to non-portable shell manipulations.
    if (testEmitIncludeTdHeader)
      output->os() << "include \"mlir/Dialect/Linalg/IR/LinalgStructuredOps.td\"";
  
    MLIRContext context;
    llvm::SourceMgr mgr;
    mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    Parser parser(mgr, &context);
    parseAndEmitAllTensorComprehensions(output->os(), parser);
    output->keep();
  
    return 0;
  }
  