[mlir][linalg] Use attributes in named ops' indexing maps

[antiagainst]

Summary

This commit adds support for parsing attribute uses in indexing
maps. These attribute uses are represented as affine symbols in
the resultant indexing maps because we can only know their
concrete value (which are coming from op attributes and are
constants) for specific op instances. The indxing_maps()
calls are synthesized to read these attributes and create affine
constants to replace the placeholder affine symbols and simplify.

Depends on D94240

[nicolasvasilache]

There are a few typos in the commit message.

[antiagainst]

Update commit message

[antiagainst]

Add doc

[hanchung]

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

375

nit: return emitError...?

1948

Do we want to add a message to this assert, like assert(... && "expected ...")?

[nicolasvasilache]

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

474

s/red/read

520

Could we please rephrase as:

The affine parser introduces new dims and symbols eagerly as they are discovered.
With the additional support for attr uses, not all ids are eagerly discovered.
For a parsed `id`, the resolution mechanism proceeds as follows:
1. Try to parse `id` as an attribute use.
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.

+/- what you find appropriate?

1964

replace symbols by the corresponding attribute value.

[antiagainst]

Address comments

mlir/docs/Dialects/Linalg.md

1.  Reduction dimensions are specified with angle bracket notation on the
     operation they apply to (e.g. `std_add<k>` specifies that `k` is a reduction
     dimension). In TC, a reduction is specified with `op=` operator and the
     reduction dimensions are inferred.
 1.  The parallel and reduction dimension are ordered by the textual program
     order. For instance, in the comprehension `O(i, j) = std_add<k, l>(...)`,
     `i` (resp. `j`) is a parallel iterator encoded by affine dimension of
     position `0` (resp. `1`); `k` (resp. `l`) is a reduction iterator encoded by
     an affine dimension of position `2` (resp. `3`).
+1.  A list of attributes can be defined for the op with the format of `attr(
+    strides: 2xi32)` and referenced in comprehension like `strides[0]`. These
+    attribute uses will be parsed as affine symbols to generate op definition
+    and implementation. For a concrete op instance, the runtime constant values
+    from the attributes will be used to replace the affine symbols and simplify
+    the indexing maps.
 
 These decisions and syntax are subject to evolution and change. In particular,
 op-specific attributes, dynamic ranks, some form of templating, shape
 calculation function specification, etc. may be added in the future.
 
 At this time, the following restrictions are imposed on the syntax and
 semantics:
 

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

 // RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 | FileCheck %s --check-prefix=ODS
 // RUN: mlir-linalg-ods-gen %s -gen-impl=1 | FileCheck %s --check-prefix=IMPL
 
 // ODS-LABEL: def Test1Op : LinalgStructuredBase_Op<"test1", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test1Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test1Op::indexing_maps() {
-//       IMPL:  AffineMap::get(2, 0, {d0, d1}, context),
-//  IMPL-NEXT:  AffineMap::get(2, 0, {d1}, context),
-//  IMPL-NEXT:  AffineMap::get(2, 0, {d0}, context) });
+//       IMPL: auto s0 = getAffineSymbolExpr(0, context); (void)s0;
+//  IMPL-NEXT: auto s1 = getAffineSymbolExpr(1, context); (void)s1;
+//  IMPL-NEXT: auto map0 = AffineMap::get(2, 2, {d0, d1}, context);
+//  IMPL-NEXT: map0 = map0.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
+//  IMPL-NEXT: map0 = simplifyAffineMap(map0);
+//  IMPL-NEXT: auto map1 = AffineMap::get(2, 2, {d1}, context);
+//  IMPL-NEXT: map1 = map1.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
+//  IMPL-NEXT: map1 = simplifyAffineMap(map1);
+//  IMPL-NEXT: auto map2 = AffineMap::get(2, 2, {d0}, context);
+//  IMPL-NEXT: map2 = map2.replaceDimsAndSymbols({}, { s0, s1 }, 2, 0);
+//  IMPL-NEXT: map2 = simplifyAffineMap(map2);
+//  IMPL-NEXT: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
 //
 //       IMPL:  void Test1Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test1Op> :
 def test1(A: f32(M, K), B: f32(K)) -> (C: f32(M)) {
   C(m) = std_addf<k>(std_mulf(A(m, k), B(k)));
 }
 
 // ODS-LABEL: def Test2Op : LinalgStructuredBase_Op<"test2", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test2Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test2Op::indexing_maps() {
-//       IMPL:  AffineMap::get(3, 0, {d0, d2}, context),
-//  IMPL-NEXT:  AffineMap::get(3, 0, {d2, d1}, context),
-//  IMPL-NEXT:  AffineMap::get(3, 0, {d0, d1}, context) });
+//       IMPL:  AffineMap::get(3, 3, {d0, d2}, context)
+//       IMPL:  AffineMap::get(3, 3, {d2, d1}, context)
+//       IMPL:  AffineMap::get(3, 3, {d0, d1}, context)
 //
 //       IMPL:  Test2Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test2Op> :
 def test2(A: f32(M, K), B: f32(K, N)) -> (C: f32(M, N)) {
   C(m, n) = std_addf<k>(std_mulf(A(m, k), B(k, n)));
 }
 
 // ODS-LABEL: def Test3Op : LinalgStructuredBase_Op<"test3", [
 //  ODS-NEXT:   AttrSizedOperandSegments
 //  ODS-NEXT:   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
 //  ODS-NEXT:   SingleBlockImplicitTerminator<"YieldOp">
 //
 // IMPL-LABEL:  ArrayAttr Test3Op::iterator_types() {
 //       IMPL:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
 //
 //       IMPL:  ArrayAttr Test3Op::indexing_maps() {
-//       IMPL:  AffineMap::get(4, 0, {d0, d1, d3}, context),
-//  IMPL-NEXT:  AffineMap::get(4, 0, {d3, d2}, context),
-//  IMPL-NEXT:  AffineMap::get(4, 0, {d0, d1, d2}, context) });
+//       IMPL:  AffineMap::get(4, 4, {d0, d1, d3}, context)
+//       IMPL:  AffineMap::get(4, 4, {d3, d2}, context)
+//       IMPL:  AffineMap::get(4, 4, {d0, d1, d2}, context)
 //
 //       IMPL:  Test3Op::regionBuilder(Block &block) {
 //       IMPL:  Value [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
 //       IMPL:  Value [[d:.*]] = std_mulf([[a]], [[b]]);
 //       IMPL:  Value [[e:.*]] = std_addf([[c]], [[d]]);
 //       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
 //
 ods_def<Test3Op> :
attr(
   i32_attr: i32,
   fvec_attr: 4xf32,
   ivec_attr: 5x6xi32,
   array_attr : f32[],
   optional_attr? : f32
 ) {
   C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
 }
+
+// Test attribute usage in affine expressions
+// IMPL-LABEL: ArrayAttr Test5Op::indexing_maps() {
+// IMPL: auto cst0 = getAffineConstantExpr(strides().getValue<int>({ 0 }), context);
+// IMPL: auto cst1 = getAffineConstantExpr(strides().getValue<int>({ 1 }), context);
+// IMPL: auto map0 = AffineMap::get(7, 9, {d0, d1 * s7 + d4, d2 * s8 + d5, d6}, context);
+// IMPL: map0 = map0.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map0 = simplifyAffineMap(map0);
+// IMPL: auto map1 = AffineMap::get(7, 9, {d3, d4, d5, d6}, context);
+// IMPL: map1 = map1.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map1 = simplifyAffineMap(map1);
+// IMPL: auto map2 = AffineMap::get(7, 7, {d0, d1, d2, d3}, context);
+// IMPL: map2 = map2.replaceDimsAndSymbols({}, { s0, s1, s2, s3, s4, s5, s6, cst0, cst1 }, 7, 0);
+// IMPL: map2 = simplifyAffineMap(map2);
+// IMPL: return {{.+}}.getAffineMapArrayAttr({ map0, map1, map2 });
+//
+ods_def<Test5Op>:
+def test5(I: f32(N, H, W, C), K: f32(F, KH, KW, C)) -> (O: f32(N, H, W, F))
+     attr(strides: 2xi32) {
+  O(n, h, w, f) = std_addf<kh, kw>(std_mulf(
+    I(n, h * strides[0] + kh, w * strides[1] + kw, c), K(f, kh, kw, c)));
+}

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

 
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/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/FormatVariadic.h"
 #include "llvm/Support/ToolOutputFile.h"
 

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

 };
 
 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),
+  /// 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,
+      : 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::parseNegateExpression(AffineExpr lhs) {
   // 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;
 }
 
+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 (parser.emitError("expected id"), nullptr);
 
   StringRef sRef = parser.curToken.getSpelling();

 //  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();
+    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:
public:
   struct ComprehensionParsingState {
     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,
private:
   //===--------------------------------------------------------------------===//
   // 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.
+    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;
 
   Parser &parser;

 
 /// 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);
+  auto createAffineBareId = [&](StringRef sRef) {
+  };
+
+  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;
+
+    // We create a new symbol for each attribute usage without reuse. This is
+    // fine given these symbols will be replaced with constants and folded away
+    // for concrete op instances.
+    result.symbol = getAffineSymbolExpr(symbols.size(), parser.context);
+    // Merely for taking the index. We don't reuse anyway.
+    symbols.emplace_back("<attr-use>", result.symbol);
+
+    attrUses.push_back(result);
+
+    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 `)`
if (failed(parser.parseToken(Token::Kind::colon, "expected colon")))
     return failure();
 
   // Parse the attribute's type. We don't expect the type to be arbitrary
   // complex, so just use this ad-hoc handling here.
 
   // Parse potential dimension list
   SmallVector<uint64_t, 4> vectorDims;
   while (parser.curToken.is(Token::Kind::integer)) {
-    vectorDims.push_back(parser.curToken.getUInt64IntegerValue().getValue());
-    parser.consumeToken();
+    uint64_t value;
+    parser.parseInteger(value);
+    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)
LogicalResult TCParser::parseAttrDef() {
   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 `)`
void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os,
   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} });
+    {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. 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} };`.
+  // 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);
+
   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");
+
+  // 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!");
+    std::string getValueFn = it->second.getValueFn(attrUse.value().indices);
+    if (getValueFn.empty()) {
+      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 tensor use, construct the affine map, replace symbols by the
+  // corresponding attribute values, and simplify the affine map.
+  for (auto tensorUse : llvm::enumerate(orderedUses)) {
+    auto indexingMap = tensorUse.value().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(u.indexingMap.getResults(), exprsStringStream);
+    llvm::interleaveComma(indexingMap.getResults(), exprsStringStream);
     exprsStringStream << "}";
     exprsStringStream.flush();
+    mapsStringStream << llvm::formatv(mapFmt, tensorUse.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, tensorUse.index(),
+                                      replaceSymbolList, state.dims.size());
+    const char *simplifyFmt = "\n\tmap{0} = simplifyAffineMap(map{0});";
+    mapsStringStream << llvm::formatv(simplifyFmt, tensorUse.index());
+  }
 
-    mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), exprsStr);
-  });
   mapsStringStream.flush();
 
+  SmallVector<std::string, 4> mapList;
+  mapList.reserve(orderedUses.size());
+  for (unsigned i = 0; i < orderedUses.size(); ++i)
+    mapList.push_back(llvm::formatv("map{0}", i));
+
   // 4. Apply format to 1. using 2. and 3.
-  os << llvm::formatv(referenceIndexingMapsFmt, cppOpName, dimsStr, mapsStr);
+  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.orderedTensorArgs.size();
   llvm::DenseMap<const TensorExpr *, unsigned> subExprsMap;
   std::function<void(llvm::raw_ostream & os, const Expression &)> printExpr;
void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef cppOpName,
   valueHandleStringStream.flush();
   expressionStringStream.flush();
   yieldStringStream.flush();
 
   os << llvm::formatv(regionBuilderFmt, cppOpName, valueHandleStr,
                       expressionsStr, yieldStr);
 }
 
+std::string
+TCParser::RegisteredAttr::getValueFn(ArrayRef<uint64_t> indices) const {
+  if (isArray)
+    return "";
+
+  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);
+    if (elementType == "i32")
+      return llvm::formatv("getValue<int>({ {0} })", indexList);
+
+    return "";
+  }
+
+  if (elementType == "f32")
+    return "getValue().convertToFloat()";
+  if (elementType == "i32")
+    return "getInt()";
+  return "";
+}
+
 /// 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();
   }