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

//===- 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/DenseMap.h"
#include "llvm/ADT/Optional.h"
#include "llvm/ADT/STLExtras.h"
#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/MemoryBuffer.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/ToolOutputFile.h"

#include <map>
#include <set>

#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::SMLoc;
using llvm::StringRef;
using llvm::Twine;

using namespace mlir;

//===----------------------------------------------------------------------===//
// Special "op aliases" substitutions.
//===----------------------------------------------------------------------===//

/// Perform substitutions of known special ops.
/// This is a poor man's way of achieving "op aliases": i.e. giving an op a
/// name.
/// This is hacky and temporary until migration to the python opdsl is complete.
static void substituteOpAliases(std::string &expressionsStr) {
  for (auto kvp : SmallVector<std::pair<std::string, std::string>>{
           {"b.create<CmpIOpSGT>(", "b.create<CmpIOp>(CmpIPredicate::sgt, "},
           {"b.create<CmpFOpOGT>(", "b.create<CmpFOp>(CmpFPredicate::OGT, "},
           {"b.create<CmpFOpOLT>(", "b.create<CmpFOp>(CmpFPredicate::OLT, "},
           {"b.create<SignExtendIOp32>(",
            "b.create<SignExtendIOp>(b.getI32Type(), "},
       }) {
    size_t pos = 0;
    while ((pos = expressionsStr.find(kvp.first, pos)) != std::string::npos) {
      expressionsStr.replace(pos, kvp.first.size(), kvp.second);
      pos += kvp.second.size();
    }
  }
}

//===----------------------------------------------------------------------===//
// 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,
    doc_str,
    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_implements_interface,
    kw_attr_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);

  /// 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();

  /// 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

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::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 '"':
      return lexString(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("attr", Token::Kind::kw_attr_def)
          .Case("def", Token::Kind::kw_def)
          .Case("ods_def", Token::Kind::kw_ods_def)
          .Case("implements_interface", Token::Kind::kw_implements_interface)
          .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);
}

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++) {
    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
  //===--------------------------------------------------------------------===//

  LogicalResult parseInteger(uint64_t &value) {
    if (!curToken.is(Token::Kind::integer))
      return emitError(curToken.getLoc(), "expected integer");
    value = curToken.getUInt64IntegerValue().getValue();
    consumeToken();
    return success();
  }

  /// 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");
      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

/// Encodes an attribute use of the form:
///
///   index-list ::= integer-literal (`,` integer-literal)*
///   attr-use ::= bare-id `[` index-list `]`
struct AttrUse {
  // Referenced attribute
  StringRef attrName;
  // Indices into the attribute
  SmallVector<uint64_t, 4> indices;
  /// Affine symbol for this usage.
  /// This is represented as an affine symbol because at the time of parsing the
  /// spec and generating the op's ODS/C++, we don't know the concrete constant
  /// value. But they should be replaced with constants read from the attribute
  /// and thus folded away for concrete op instances.
  AffineExpr symbol;

  std::string getKey() {
    SmallVector<std::string, 4> indexStrs;
    for (uint64_t index : indices)
      indexStrs.push_back(std::to_string(index));
    return llvm::formatv("{0}[{1}]", attrName, llvm::join(indexStrs, ","));
  }
};

//===----------------------------------------------------------------------===//
// 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:
  /// Creates an affine parser that parses tokens from `p`.
  ///
  /// The affine parser introduces new dimensions and symbols eagerly as new
  /// `id` are discovered. To additionally support attribute use `id`s, for a
  /// parsed `id`, the resolution mechanism proceeds as follows:
  /// 1. Try to parse `id` as an attribute use (using the `attrUseParsingHook`).
  /// 2. If unsuccessful, try to match `id` to a known dim or symbol.
  /// 3. If still unsuccessful, eagerly create a new dim or symbol and add it to
  ///    the known dims or symbols (using the `bareIdParsingHook`).
  explicit AffineParser(
      Parser &p, std::function<AffineExpr(StringRef)> bareIdParsingHook,
      std::function<llvm::Optional<AffineExpr>()> attrUseParsingHook,
      AffineDimList &dimList, AffineSymbolList &symbolList)
      : parser(p), bareIdFallback(bareIdParsingHook),
        attrUseCallback(attrUseParsingHook), 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 parseAttrUseOrBareIdExpr();
  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;
  std::function<llvm::Optional<AffineExpr>()> attrUseCallback;
  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()) {
      (void)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()) {
      (void)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()) {
      (void)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()) {
      (void)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 ((void)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 ((void)parser.emitError("missing operand of negation"), nullptr);
  return (-1) * operand;
}

AffineExpr AffineParser::parseAttrUseOrBareIdExpr() {
  if (llvm::Optional<AffineExpr> attrUse = attrUseCallback())
    return attrUse.getValue();
  return parseBareIdExpr();
}

/// 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 ((void)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 ((void)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 ((void)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 parseAttrUseOrBareIdExpr();
  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)
      (void)parser.emitError("missing right operand of binary operator");
    else
      (void)parser.emitError("missing left operand of binary operator");
    return nullptr;
  default:
    if (lhs)
      (void)parser.emitError("missing right operand of binary operator");
    else
      (void)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) {
  if (failed(parser.parseToken(lDelim,
                               "expected lDelim at start of affine expr list")))
    return {};

  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) {}

  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 {
    /// The number of operands (which includes inputs and outputs) in a
    /// comprehension.
    size_t numArgs;
    AffineDimList dims;
    SmallVector<std::unique_ptr<Expression>, 4> expressions;
    llvm::DenseMap<TensorUse, unsigned> orderedTensorArgs;
  };
  LogicalResult parseTensorUse(TensorUse &result,
                               ComprehensionParsingState &state);

  /// Parses an attribute definition.
  LogicalResult parseAttrDef();

  /// Parses an optional attribute use.
  LogicalResult parseAttrUse(AttrUse &result);

  /// 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, ArrayRef<StringRef> interfaces,
                ComprehensionParsingState &state);

  /// Print the C++ StructuredOpsInterface impl of `iterator_types`.
  void printReferenceIterators(llvm::raw_ostream &os, StringRef cppOpName,
                               ComprehensionParsingState &state);

  /// Print methods related to indexing map required attributes.
  ///
  /// Specifically, this prints the definitions for the following methods:
  ///   bool hasDynamicIndexingMaps();
  ///   LogicalResult verifyIndexingMapRequiredAttributes();
  void printIndexingMapRequiredAttrMethods(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 folders.
  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;
  };

  //===--------------------------------------------------------------------===//
  // Internal bookkeeping of attributes.
  //===--------------------------------------------------------------------===//
  struct RegisteredAttr {
    StringRef elementType;
    SmallVector<uint64_t, 4> vectorDims;
    bool isArray;
    bool isOptional;

    // Returns the function to get values at the given indices from this
    // attribute.
    llvm::Optional<std::string> getValueFn(ArrayRef<uint64_t> indices) const;
  };

  //===--------------------------------------------------------------------===//
  // Per-TC def state.
  //===--------------------------------------------------------------------===//
  /// Symbols are per TC def.
  AffineSymbolList symbols;

  /// Attribute usages in all affine expressions.
  SmallVector<AttrUse, 8> attrUses;

  /// Tensors are per TC def.
  llvm::StringMap<RegisteredTensor> registeredTensors;
  unsigned nextRegisteredTensorIndex;

  /// Attributes are per TC def.
  std::map<std::string, RegisteredAttr> registeredAttrs;

  /// A map from AttrUse to AffineExpr symbol.
  llvm::StringMap<AffineExpr> registeredAttrUseToSymbol;

  StringRef docString;

  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) {
  auto createAffineBareId = [&](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;
  };

  auto tryToParseAttrUse = [&]() -> llvm::Optional<AffineExpr> {
    if (!parser.curToken.is(Token::Kind::id))
      return llvm::None;

    StringRef attrName = parser.curToken.getSpelling();
    auto it = registeredAttrs.find(attrName.str());
    if (it == registeredAttrs.end())
      return llvm::None;

    AttrUse result;
    if (failed(parseAttrUse(result)))
      return llvm::None;

    auto symbolIt = registeredAttrUseToSymbol.find(result.getKey());
    if (symbolIt == registeredAttrUseToSymbol.end()) {
      result.symbol = getAffineSymbolExpr(symbols.size(), parser.context);
      symbols.emplace_back("<attr-use>", result.symbol);
      registeredAttrUseToSymbol[result.getKey()] = result.symbol;
      attrUses.push_back(result);
    } else {
      result.symbol = symbolIt->second;
    }

    return result.symbol;
  };

  AffineParser affineParser(parser, createAffineBareId, tryToParseAttrUse, 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();
}

/// 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)) {
    uint64_t value;
    if (failed(parser.parseInteger(value)))
      return failure();
    vectorDims.push_back(value);

    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.str(),
      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();
}

LogicalResult TCParser::parseAttrUse(AttrUse &result) {
  result.attrName = parser.curToken.getSpelling();
  if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
    return failure();

  auto it = registeredAttrs.find(result.attrName.str());
  assert(it != registeredAttrs.end());
  const RegisteredAttr &attr = it->second;

  if (!attr.vectorDims.empty() || attr.isArray) {
    // This is a vector/array attribute. Parse indices for it.
    auto indexLoc = parser.curToken.getLoc();

    if (failed(parser.parseToken(Token::Kind::l_square, "expected '['")))
      return failure();

    auto parseIndex = [&]() {
      uint64_t value;
      if (failed(parser.parseInteger(value)))
        return failure();
      result.indices.push_back(value);
      return success();
    };
    if (failed(parser.parseCommaSeparatedListUntil(
            Token::Kind::r_square, parseIndex, /*allowEmptyList=*/false)))
      return failure();

    size_t rank = attr.isArray ? 1 : attr.vectorDims.size();
    if (result.indices.size() != rank)
      return parser.emitError(indexLoc,
                              "number of indices mismatch: expected " +
                                  std::to_string(rank) + ", but found " +
                                  std::to_string(result.indices.size()));
  }

  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());
  }

  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())
      return parser.emitError("Fewer LHS definitions than RHS expressions");
    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())
    return parser.emitError("Fewer RHS expressions than LHS definitions");

  // 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)
      return parser.emitError("Unexpected multi-write to a single tensor");
    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 &&
          tensor.indexingMap.getResults() != use.indexingMap.getResults()) {
        LLVM_DEBUG(llvm::dbgs() << "\nexisting: " << tensor.indexingMap);
        (void)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 more than one definitions are less. They are shaped-only operand, which
  // are used to define reduction loops. For now, only accept exactly one
  // shaped-only operand.
  if (state.numArgs > seenDefs.size() + 1) {
    failed = true;
  } else if (state.numArgs == seenDefs.size() + 1) {
    for (auto &tensorIter : registeredTensors) {
      auto &tensor = tensorIter.getValue();
      if (tensor.indexingMap)
        continue;
      if (auto *pTensorExpr =
              dyn_cast<TensorExpr>(state.expressions[0].get())) {
        SmallVector<AffineExpr, 4> exprs;
        for (auto dim : pTensorExpr->reductionDimensions)
          exprs.push_back(getAffineDimExpr(dim, parser.context));
        tensor.indexingMap = AffineMap::get(state.dims.size(), symbols.size(),
                                            exprs, parser.context);
      }
    }
  }
  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 `}`
///
///   implements-interface ::=
///     `implements_interface` `<` bare-id (`,` bare-id)* `>` `:` tc-def
///
///   ods-def ::= `ods_def` `<` bare-id `>`
///               (implements-interface)? `:`
///               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 ods-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 '>'")))
    return failure();

  // Parse optional implements-interface header (including C++ op names)
  SmallVector<StringRef> interfaces;
  bool implementsInterface = succeeded(
      parser.parseOptionalToken(Token::Kind::kw_implements_interface));
  if (implementsInterface) {
    auto parseInterfaceString = [&]() -> LogicalResult {
      StringRef interfaceName = parser.curToken.getSpelling();
      if (failed(parser.parseToken(Token::Kind::id, "expected id")))
        return failure();
      interfaces.push_back(interfaceName);
      return success();
    };
    if (failed(parser.parseToken(Token::Kind::lt, "expected '<'")) ||
        failed(parser.parseCommaSeparatedListUntil(
            Token::Kind::gt, parseInterfaceString, /*allowEmptyList=*/false)))
      return failure();
  }

  // Parse column.
  if (failed(parser.parseToken(Token::Kind::colon, "expected ':'")))
    return failure();

  // Parse TC op name.
  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();
  }

  // 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(),
                         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());
    perComprehensionStates.back().numArgs = registeredTensors.size();
    if (failed(parseOneComprehension(cppOpName, tcName,
                                     perComprehensionStates.back())))
      return failure();
  };
  if (failed(parser.parseToken(Token::Kind::r_brace, "expected '}'")))
    return failure();

  // Print.
  auto nComprehensions = perComprehensionStates.size();
  if (nComprehensions != 1)
    return parser.emitError("only 1 comprehension supported for now, got: " +
                            llvm::Twine(nComprehensions));
  if (genODSDecl) {
    auto &state = perComprehensionStates.back();
    printODS(os, cppOpName, tcName, interfaces, state);
    os << "\n";
  }
  if (genODSImpl) {
    auto &state = perComprehensionStates.back();
    std::string extraMethods;
    llvm::raw_string_ostream ss(extraMethods);
    printReferenceIterators(ss, cppOpName, state);
    printIndexingMapRequiredAttrMethods(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, ArrayRef<StringRef> interfaces,
                        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")
                              .Case("i64", "I64")
                              .Default("");
    if (odsType.empty()) {
      (void)parser.emitError(
          "unimplemented support for attribute element type: " + elementType);
      return;
    }

    const auto &dims = attr.second.vectorDims;
    if (!dims.empty()) {
      // Vector case
      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, ", "));
    } else if (attr.second.isArray) {
      // Array case
      odsType = llvm::formatv("{0}ArrayAttr", odsType);
    } else {
      // Scalar case
      odsType = llvm::formatv("{0}Attr", 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;

  // Template for Linalg named ops' ODS definitions. Parameters:
  // {0}: ODS/C++ op name
  // {1}: assembly op mnemonic
  // {2}: op interface list
  // {3}: documentation (summary + description)
  // {4}: op attribute list
  // {5}: the number of arguments for the op region
  // {6}: builder methods taking standalone attribute parameters
  // {7}: additional methods for attributes used by indexing maps
  const char *header = R"FMT(  def {0} : LinalgStructuredBase_Op<"{1}", [
    AttrSizedOperandSegments,
    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
    SingleBlockImplicitTerminator<"YieldOp">
    /*extraInterfaces=*/{2}]> {
      {3}
      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 = [
        OpBuilder<
        (ins "ValueRange":$inputs, "ValueRange":$outputs),
        [{{
          $_state.addOperands(inputs);
          $_state.addOperands(outputs);
          $_state.addAttribute(
            "operand_segment_sizes",
            $_builder.getI32VectorAttr({{
              static_cast<int32_t>(inputs.size()),
              static_cast<int32_t>(outputs.size())}));
          createAndFillStructuredOpRegion<{0}>(
            $_builder,
            $_state,
            TypeRange(inputs),
            TypeRange(outputs));
        }]>,
        OpBuilder<
        (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
             "ValueRange":$outputs),
        [{{
          $_state.addOperands(inputs);
          $_state.addOperands(outputs);
          $_state.addTypes(resultTensorTypes);
          $_state.addAttribute(
            "operand_segment_sizes",
            $_builder.getI32VectorAttr({{
              static_cast<int32_t>(inputs.size()),
              static_cast<int32_t>(outputs.size())}));
          createAndFillStructuredOpRegion<{0}>(
            $_builder,
            $_state,
            TypeRange(inputs),
            TypeRange(outputs));
        }]>,
        OpBuilder<
        (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
        [{{
          $_state.addOperands(operands);
          $_state.addAttributes(attributes);
          $_state.addTypes(resultTensorTypes);
          (void)$_state.addRegion();
        }]>
        {6}
      ];
      let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
      let parser = [{{
        return ::parseNamedStructuredOp<{0}>(parser, result);
      }];
      let hasFolder = 1;

      let extraClassDeclaration = structuredOpsBaseDecls # [{{
        // Auto-generated.
        ArrayAttr iterator_types();
        ArrayAttr indexing_maps();
        static void regionBuilder(ImplicitLocOpBuilder &b, Block &block);
        static std::function<void(ImplicitLocOpBuilder &b, Block &)>
        getRegionBuilder() {{
          return regionBuilder;
        }

        // Generic methods.
        static unsigned getNumRegionArgs() {{ return {5}; }
        std::string getLibraryCallName() {{
          return generateLibraryCallName(getOperation());
        }

        {7}
      }];
  })FMT";

  // Generate the list of extra implemented interfaces.
  std::string interfaceNameList;
  if (!interfaces.empty()) {
    llvm::raw_string_ostream ss(interfaceNameList);
    ss << ", "; // Leading comma to concat to existing list of interfaces.
    llvm::interleaveComma(interfaces, ss);
    ss.flush();
  }

  // Generate documentation.
  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());
  }

  // Generate an additional builder that has parameters for attributes.
  std::string attrBuilder;
  if (!registeredAttrs.empty()) {
    SmallVector<std::string, 4> attrParams, attrStmts;
    for (const auto &attr : registeredAttrs) {
      llvm::StringRef name = attr.first;
      attrParams.push_back(llvm::formatv("\"Attribute\":${0}", name));
      attrStmts.push_back(
          llvm::formatv("$_state.addAttribute(\"{0}\", {0});", name));
    }
    std::string attrParamsList = llvm::join(attrParams, ", ");
    std::string attrStmtsList = llvm::join(attrStmts, "\n");

    const char *builderFmt = R"FMT(
      , OpBuilder<
      (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
           "ValueRange":$outputs, {1}),
      [{{
        $_state.addOperands(inputs);
        $_state.addOperands(outputs);
        $_state.addTypes(resultTensorTypes);
        $_state.addAttribute(
          "operand_segment_sizes",
          $_builder.getI32VectorAttr({{
            static_cast<int32_t>(inputs.size()),
            static_cast<int32_t>(outputs.size())}));
        createAndFillStructuredOpRegion<{0}>(
          $_builder,
          $_state,
          TypeRange(inputs),
          TypeRange(outputs));
        {2}
      }]>
    )FMT";
    attrBuilder =
        llvm::formatv(builderFmt, cppOpName, attrParamsList, attrStmtsList);
  }

  std::string attrMethods;
  if (!registeredAttrs.empty()) {
    attrMethods = R"(
      bool hasDynamicIndexingMaps();
      LogicalResult verifyIndexingMapRequiredAttributes();
    )";
  }

  // Finally put everything together.
  os << llvm::formatv(header, cppOpName, linalgOpName, interfaceNameList, doc,
                      attrList, state.numArgs, attrBuilder, attrMethods);
}

/// 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 *foldersFmt = R"FMT(
    LogicalResult {0}::fold(ArrayRef<Attribute>,
                            SmallVectorImpl<OpFoldResult> &) {{
      return foldMemRefCast(*this);
    }
    void {0}::getEffects(SmallVectorImpl<
        SideEffects::EffectInstance<MemoryEffects::Effect> >&effects) {{
      SmallVector<Value> inputBuffers = getInputBufferOperands();
      SmallVector<Value> outputBuffers = getOutputBufferOperands();
      getGenericEffectsImpl(effects,
        getOperation()->getResults(), inputBuffers, outputBuffers);
    })FMT";
  os << llvm::formatv(foldersFmt, cppOpName);
}

// Prints methods for querying whether the current named op has attributes that
// are used by its indexing maps and for verifying those attributes have the
// expected type.
void TCParser::printIndexingMapRequiredAttrMethods(
    llvm::raw_ostream &os, StringRef cppOpName,
    ComprehensionParsingState &state) {
  // If there are no attribute used by the whole definition, then we are done.
  if (registeredAttrs.empty())
    return;

  // Otherwise, go through each attribute and generate code to verify it's
  // valid per the spec.
  SmallVector<std::string, 4> attributes;
  for (const auto &attr : registeredAttrs) {
    if (attr.second.isOptional)
      continue;

    llvm::StringRef name = attr.first;
    llvm::StringRef elementType = attr.second.elementType;
    const auto &dims = attr.second.vectorDims;

    // Get the method call to check the element type is of the expected kind.
    std::string elemTypeCheck = llvm::StringSwitch<std::string>(elementType)
                                    .Case("f32", "isF32()")
                                    .Case("i32", "isInteger(32)")
                                    .Case("i64", "isInteger(64)")
                                    .Default("");
    if (elemTypeCheck.empty()) {
      (void)parser.emitError(
          "unimplemented support for attribute element type: " + elementType);
      return;
    }

    // Scalar case.
    if (dims.empty() && !attr.second.isArray) {
      const char *attrFmt = R"FMT(
        if (auto attr = op->getAttr("{0}")) {{
          if (!attr.getType().{1}) return op->emitError(
            "incorrect type for indexing map required attribute '{0}'");
        } else {{
          return op->emitError(
            "missing indexing map required attribute '{0}'");
        }
      )FMT";

      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck));
      continue;
    }

    // Vector case.
    if (!dims.empty()) {
      SmallVector<std::string, 4> dimStrs;
      for (uint64_t dim : dims)
        dimStrs.push_back(std::to_string(dim));

      const char *attrFmt = R"FMT(
        if (auto attr = op->getAttrOfType<DenseElementsAttr>("{0}")) {{
          if (!attr.getType().getElementType().{1}) return op->emitError(
            "incorrect element type for indexing map required attribute '{0}'");
          if (attr.getType().getShape() != ArrayRef<int64_t>{{ {2} })
            return op->emitError(
              "incorrect shape for indexing map required attribute '{0}'");
        } else {
          return op->emitError(
            "missing indexing map required attribute '{0}'");
        }
      )FMT";

      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck,
                                         llvm::join(dimStrs, ", ")));
      continue;
    }

    // Array case.
    {
      const char *attrFmt = R"FMT(
        if (auto attr = op->getAttrOfType<ArrayAttr>("{0}")) {{
          for (Attribute element : attr) {{
            if (!element.getType().{1}) return emitError(
              "incorrect element type for indexing map required attribute '{0}'");
          }
        } else {{
          return op->emitError(
            "missing indexing map required attribute '{0}'");
        }
      )FMT";

      attributes.push_back(llvm::formatv(attrFmt, name, elemTypeCheck));
    }
  }

  const char *methodFmt = R"FMT(
  bool {0}::hasDynamicIndexingMaps() {{ return true; }

  LogicalResult {0}::verifyIndexingMapRequiredAttributes() {{
    Operation *op = getOperation();
    {1}
    return success();
  }
  )FMT";

  // Print everything out.
  os << llvm::formatv(methodFmt, cppOpName, llvm::join(attributes, "\n"));
}

/// 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});
    {2}
    return Builder(context).getAffineMapArrayAttr({ {3} });
  })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. Get the list of affine maps for each input/output. The AffineExpr use
  // the common arithmetic operators on AffineExpr. These affine maps will
  // replace the `{2}` placeholder.
  std::string mapsStr;
  llvm::raw_string_ostream mapsStringStream(mapsStr);

  // Create a list of all symbols.
  SmallVector<std::string, 4> symbolReplacements;
  symbolReplacements.reserve(symbols.size());
  for (unsigned i = 0; i < symbols.size(); ++i) {
    const char *symFmt =
        "\n\tauto s{0} = getAffineSymbolExpr({0}, context); (void)s{0};";
    mapsStringStream << llvm::formatv(symFmt, i);
    symbolReplacements.push_back(llvm::formatv("s{0}", i));
  }

  // Create the affine constant expressions to replace symbols for attributes.
  for (auto attrUse : llvm::enumerate(attrUses)) {
    StringRef attrName = attrUse.value().attrName;
    auto it = registeredAttrs.find(attrName.str());
    assert(it != registeredAttrs.end() && "uses should point to valid attr!");
    llvm::Optional<std::string> getValueFn =
        it->second.getValueFn(attrUse.value().indices);
    if (!getValueFn) {
      (void)parser.emitError("unimplemented getValueFn for attribute: " +
                             attrName);
      return;
    }
    std::string cstVal = llvm::formatv("{0}(){1}", attrName, *getValueFn);
    const char *cstFmt =
        "\n\tauto cst{0} = getAffineConstantExpr({1}, context);";
    mapsStringStream << llvm::formatv(cstFmt, attrUse.index(), cstVal);

    unsigned position =
        attrUse.value().symbol.cast<AffineSymbolExpr>().getPosition();
    symbolReplacements[position] = llvm::formatv("cst{0}", attrUse.index());
  }

  // For each registered tensor, construct the affine map, replace symbols by
  // the corresponding attribute values, and simplify the affine map.
  for (auto &tensorIter : registeredTensors) {
    auto &tensor = tensorIter.getValue();
    auto indexingMap = tensor.indexingMap;
    const char *mapFmt =
        "\n\tauto map{0} = AffineMap::get({1}, {2}, {3}, context);";

    std::string exprsStr;
    llvm::raw_string_ostream exprsStringStream(exprsStr);
    exprsStringStream << "{";
    llvm::interleaveComma(indexingMap.getResults(), exprsStringStream);
    exprsStringStream << "}";
    exprsStringStream.flush();
    mapsStringStream << llvm::formatv(mapFmt, tensor.index, state.dims.size(),
                                      indexingMap.getNumSymbols(), exprsStr);

    std::string replaceSymbolList =
        llvm::formatv("{ {0} }", llvm::join(symbolReplacements, ", "));

    // Note that we use `0` as the result affine map's number of symbols. All
    // symbols representing attribute usages should be folded away. But there
    // may exist additional symbols for tensor dimension upper bounds. Linalg
    // does not handle such cases right now. This needs to be fixed once we
    // need that.
    const char *replaceFmt =
        "\n\tmap{0} = map{0}.replaceDimsAndSymbols({{}, {1}, {2}, 0);";
    mapsStringStream << llvm::formatv(replaceFmt, tensor.index,
                                      replaceSymbolList, state.dims.size());
    const char *simplifyFmt = "\n\tmap{0} = simplifyAffineMap(map{0});";
    mapsStringStream << llvm::formatv(simplifyFmt, tensor.index);
  }

  mapsStringStream.flush();

  SmallVector<std::string, 4> mapList;
  mapList.reserve(state.numArgs);
  for (auto i : llvm::seq<unsigned>(0, state.numArgs))
    mapList.push_back(llvm::formatv("map{0}", i));

  // 4. Apply format to 1. using 2. and 3.
  os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr,
                      llvm::join(mapList, ", "));
}

/// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
                                  ComprehensionParsingState &state) {
  unsigned count = state.numArgs;
  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} = b.create<{1}>({2});";
      os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->operationName,
                          subExprs);
      subExprsMap[pTensorExpr] = count;
    }
  };

  const char *regionBuilderFmt = R"FMT(
  void {0}::regionBuilder(ImplicitLocOpBuilder &b, Block &block) {
    auto args = block.getArguments();
    Value {1};
    {2}
    b.create<linalg::YieldOp>(ValueRange{ {3} });
  })FMT";

  std::string valueHandleStr;
  llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);
  std::set<unsigned> usedTensorId;
  for (const auto &iter : state.orderedTensorArgs)
    usedTensorId.insert(iter.second);
  llvm::interleaveComma(usedTensorId, valueHandleStringStream, [&](auto idx) {
    valueHandleStringStream << "_" << idx << "(args[" << 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);
    });
  expressionStringStream.flush();
  substituteOpAliases(expressionsStr);

  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();
  yieldStringStream.flush();

  os << llvm::formatv(regionBuilderFmt, cppOpName, valueHandleStr,
                      expressionsStr, yieldStr);
}

llvm::Optional<std::string>
TCParser::RegisteredAttr::getValueFn(ArrayRef<uint64_t> indices) const {
  if (isArray)
    return llvm::None;

  if (!vectorDims.empty()) {
    SmallVector<std::string, 4> indexStrs;
    for (uint64_t index : indices)
      indexStrs.push_back(std::to_string(index));
    std::string indexList = llvm::join(indexStrs, ", ");
    if (elementType == "f32")
      return llvm::formatv(".getValue<float>({ {0} })", indexList).str();
    if (elementType == "i32")
      return llvm::formatv(".getValue<int>({ {0} })", indexList).str();
    if (elementType == "i64")
      return llvm::formatv(".getValue<int64_t>({ {0} })", indexList).str();

    return llvm::None;
  }

  if (elementType == "f32")
    return std::string(".convertToFloat()");
  if (elementType == "i32" || elementType == "i64")
    return std::string("");
  return llvm::None;
}

/// 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);
  (void)parseAndEmitAllTensorComprehensions(output->os(), parser);
  output->keep();

  return 0;
}