[mlir][Linalg] Create a tool to generate named Linalg ops from a Tensor Comprehensions-like specification.

[nicolasvasilache]

Summary

This revision adds a tool that generates the ODS and C++ implementation for "named" Linalg ops according to the [RFC discussion](https://llvm.discourse.group/t/rfc-declarative-named-ops-in-the-linalg-dialect/745).

While the mechanisms and language aspects are by no means set in stone, this revision allows connecting the pieces end-to-end from a mathematical-like specification.

Some implementation details and short-term decisions taken for the purpose of bootstrapping and that are not set in stone include:
1. using a "[Tensor Comprehension](https://arxiv.org/abs/1802.04730)-inspired" syntax
2. implicit and eager discovery of dims and symbols when parsing
3. using EDSC ops to specify the computation (e.g. std_addf, std_mul_f, ...)

A followup revision will connect this tool to tablegen mechanisms and allow the emission of named Linalg ops that automatically lower to various loop forms and run end to end.

For the following "Tensor Comprehension-inspired" string:
```

def batch_matmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
}

```

With -gen-ods-decl=1, this emits (modulo formatting):
```
  def batch_matmulOp : LinalgNamedStructured_Op<"batch_matmul", [
    NInputs<2>,
    NOutputs<1>,
    NamedStructuredOpTraits]> {
      let arguments = (ins Variadic<LinalgOperand>:$views);
      let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
      let extraClassDeclaration = [{
        llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
        llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
        void regionBuilder(ArrayRef<BlockArgument> args);
      }];
      let hasFolder = 1;
  }
```

With -gen-ods-impl, this emits (modulo formatting):
```
  llvm::Optional<SmallVector<StringRef, 8>> batch_matmul::referenceIterators() {
      return SmallVector<StringRef, 8>{ getParallelIteratorTypeName(),
                                        getParallelIteratorTypeName(),
                                        getParallelIteratorTypeName(),
                                        getReductionIteratorTypeName() };
  }
  llvm::Optional<SmallVector<AffineMap, 8>> batch_matmul::referenceIndexingMaps()
  {
    MLIRContext *context = getContext();
    AffineExpr d0, d1, d2, d3;
    bindDims(context, d0, d1, d2, d3);
    return SmallVector<AffineMap, 8>{
        AffineMap::get(4, 0, {d0, d1, d3}),
        AffineMap::get(4, 0, {d3, d2}),
        AffineMap::get(4, 0, {d0, d1, d2}) };
  }
  void batch_matmul::regionBuilder(ArrayRef<BlockArgument> args) {
    using namespace edsc;
    using namespace intrinsics;
    ValueHandle _0(args[0]), _1(args[1]), _2(args[2]);

    ValueHandle _4 = std_mulf(_0, _1);
    ValueHandle _5 = std_addf(_2, _4);
    (linalg_yield(ValueRange{ _5 }));
  }
```

[silvas]

Do you intend for this to be "approaching production quality code" and reviewed as such or still proof-of-concept level?

[silvas]

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

24

This actually looks like it could be reasonably parsed with a custom linalg_named_op_gen.def op? Or is there a dependency issue with using MLIR for this as this needs to happen during build time?

[nicolasvasilache]

Format

[nicolasvasilache]

@silvas I'd hope closer to "approaching production quality", it is missing comments though which will make it easier to read.
Note that almost everything above l. 1000 in mlir-linalg-ods-gen.cpp is borrowed from other places.
For some reason Token, Lexer and core Parser are kept hidden within MLIR and I would very much like to expose them and avoid the copy-pasta (@rriddle what's your take on this?).

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

24

This needs to generate ODS which in turn defines new ops.
I think what you are referring to would be the linalg.generic form which has verbosity issues as well as does not let you refer to isa/cast/dyn_cast<MatmulOp>()or let you matchAndRewrite easily.

[nicolasvasilache]

Refactorings, cleanups and reformat.

[nicolasvasilache]

@silvas refactored so that things are better layereed.
Code above the following code block at line ~800 ish is taken from other places in MLIR and should be refactored out once lexer/parser is exposed.

//===----------------------------------------------------------------------===//
// TC parsing.
//===----------------------------------------------------------------------===//

[nicolasvasilache]

Add a test line to pipe the generated ODS through mlir-tblgen.

[mehdi_amini]

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

2

Missing license header?

[silvas]

First round of comments.

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

828

Is there another diff that includes this?

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

2

add test that exercises multiple comprehensions in the body

6

layering-wise would prefer to not test this here. If needed, we can add a separate test elsewhere that does this .td -> .inc file check. Strictly speaking what ends up in the .inc file is not really the concern of this component, only the contents of the .td file.

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

975

some spaces needed around first tensor-def-list

978

I don't see affine-expr or tensor-typedef mentioned locally in this comment. move this to the appropriate comment?

1371

should this be a diagnostic?

1380

How can we emit ODS before we finish processing the whole tc-def production?

1411

maybe rename to "parseAndEmitTCDef"? Also probably rename processOneComprehension to parseAndEmitOneComprehension to be consistent with that.

1429

typo in the "expected" string.

1439–1440

Can you make this comment a bit easier to understand. What is an "eagerly discovered symbol" and how does this "normalize" it?

1445–1446

Instead of the ternary, use

static AffineMap get(unsigned dimCount, unsigned symbolCount,
                      ArrayRef<AffineExpr> results, MLIRContext *context);

1455

comma separated comprehensions seems to contradict the grammar?

1656

auto here obscures things IMO

1662

auto here obscures things IMO

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

828

This is eagerly included to allow the test to pipe through mlir-tblgen and verify the ODS is well-formed (as was suggested by @ftynse, see other comment). The companion revision https://reviews.llvm.org/D76456 does the plumbing assuming a parser exists and shows how to make this run end-to-end.

In the current form this is non-functional and only exists for the purpose of verifying well-formedness and avoiding a giant diff when things can be (reasonably well) separated.

If you have strong objections against this interim state, I would rather drop the piping through tablegen rather than merge revisions (but @ftynse may have his own objections to this).

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

2

Despite the parser accepting it, there is no support for that atm and some eperiment + design is required here.
Emitting an error for now.

6

This was suggested by @ftynse to show the ODS is valid and how it connects to tblgen by mirroring this test: https://github.com/llvm/llvm-project/blob/master/mlir/test/mlir-tblgen/llvm-intrinsics.td#L11.

I am fine either way, would just like consensus on this before reverting back to the previous state.
Please reopen if you feel strongly about this.
@ftynse any strong opinion?

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

1380

Thanks!

1455

you're right, thanks!

[nicolasvasilache]

Address review comments + refactor ComprehensionParserState.

[ftynse]

It would be great to share some parts with the main parser, for example affine expression parsing. I think we can pretty much have parseAffineExpr(StringRef) declared in a private header and use it here, possibly with some semantic post-checks on the expression not involving, e.g., SSA values.

mlir/include/mlir/IR/AffineExpr.h

222

Nit: AffineExpr is a value-type, can't we just pass it by-value ?

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

87

This makes it look like it's a token start... How about FIRST_KEYWORD = kw_def, LAST_KEYWORD=kw_select?

136

Copy-pasta comment, this is not "operation assembly"

289

Missed select keyword. We could have some macro magic to make sure modifying the list of tokens also handles them in the lexer.

295

The code in getUInt64IntegerValue seems to support hex integers, but this clearly does not.

372

Nit: llvm::function_ref if you don't store the argument

546

Nit: the operand in consumeToken here and below is redundant, the case-expression just above ensures that the token is of the right kind.

841

This should be trivial to implement without resorting to virtual functions, dispatching on kind and using static_cast.

850

Nit: tensor-id is not defined

856

Nit: = default would also work

860

If you use LLVM-style type system, you would normally want to avoid virtual functions...

874

Nit: given that PreOrder is a boolean template parameter, I am not sure what "perfoms PreOrder traversal` means when the parameter is false. Post-order? In-order? Compilation error?

894

Would MutableArrayRef work instead of hardcoding SmallVector with a given size?

911

Do you care about the order of reduction dimensions?

929

Why SetVector? In TC, we wouldn't care about the order of reduction dimensions.

941–942

And what if discovery mode != symbols ?

951

Nit: something went wrong with formatting here: | ran away to the right. I personally prefer something like

foo ::= token token
        continuation line of the same rule
      | another rule

975

Nit: this comment repeats the comment on struct TensorExpr. I am worried about it getting out of sync if the syntax evolves. My recommendation would be to only keep the syntax in a single comment (preferably, the implementation of this method), and just refer to that from the other comments.

982

Nit: why pass by-pointer rather than by-reference?

[ftynse]

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

397

"proper token" is unclear as error message

928

Do you actually need pointers? Can't we just store Expressions as is, eventually with appropriate move semantics to avoid extra copies?

990

Why does a parsing function accept an _output_ stream?

995

Plz document what does it print

1051

Do you actually need this? I only see DenseMap<TensorExpr *>, which should be using a generic pointer-based map implementation.

1054

How about DenseMapInfo<StringRef>::getTombstoneKey() instead?

1067

Same as above, AffineMap also has a DenseMapInfo if I'm not mistaken.

1084

Nit: document how the visitation behaves if the callback mutates the visited object

1127

Nit: would emplace_back work?

1136

Have you considered storing tensors in an llvm::StringMap indexed by name instead of doing linear lookups every time?

1148

Naming nit: isa is widely used for downcasting, this is just a lookup; prefer is.

1173

Would emplace_back work?

1193

Could you just have a default message expected %tokenname% instead of having a similar string everywhere

1197

It looks like it would parse just about any id. "expected a type id" sounds a bit misleading because "type id" is not a production rule, and there's no additional check on the id somehow being a type.

1202

Nit: add a description in the assertion. Also, are we sure this can never happen?

1281

Ultra-nit: we tend to use single quotes rather than backticks in error messages

1283

Nit: /*allowEmptyList=*/true

1315

/*allowEmptyList=*/true

1323

This may crash if you have less LHS declarations than RHS definitions.

1328

/*allowEmptyList=*/true

1333

dimCount and symbolCount make the comment look outdated, is it?

1353

Did you check that indexings were different?

1366

Nit: I'd use early return here

1374

[Not for this commit]: I would rather have the parser accept the correct syntax, and have a separate check that implements "semantic" rules.

1415

/*allowEmptyList=*/true

1426

/*allowEmptyList=*/true

1431

typo: "symbolicc"

1491

Why is the result optional?

1561

Nit: could we use more meaningful names than ss2?

1627

C++14 supports auto for lambda arguments

1646

Alternatively, you could use ss.str() instead of valueHandleStr below.
Also, consider better names than ss, ss2, ss3. One ss is acceptable in a short function, but here it's really tricky to keep in mind which stream is associated with which string.

[ftynse]

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

6

I'm not sure I understand what is the concern here? The ODS check verifies the content of the produced .td file, _not_ the result of feeding that .td file to mlir-tblgen -gen-op-defs, which is indeed a separate concern. The IMPL check verifies the implementations of methods that are declared in the .td file and there is simply no other place where we can verify them.

The staging here is:
1a. mlir-linalg-ods-gen -gen-ods-decl %this_file% > ods.td
1b. mlir-linalg-ods-gen -gen-impl %this_file% > impl.cc
2a. mlir-tblgen -gen-op-decl ods.td > ods.h
2b. mlir-tblgen -gen-op-decl ods.td > ods.cc

3. include impl.cc and ods.cc into the implementation file; and ods.h into the header file.

@nicolasvasilache the test you referenced also has RUN lines making sure mlir-tblgen can consume what the first stage produces. Consider adding them here as well. This could help detect cases of ODS syntax change (the simple syntactic test passes, but not the piping check). That's why there is only a trivial check to make sure FileCheck eats something.

[nicolasvasilache]

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

6

I'm not sure I understand what is the concern here? 
...
Consider adding them here as well.

That's precisely what the concern was IIUC, piping through mlir-tblgen (see previous snapshot that I updated improperly https://reviews.llvm.org/D77067?id=254251).

Restored that part of the test.

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

289

Rather than continue duplicating here, MLIR should expose the lexer and parser and everyone's life will be better.

295

Yes, I am trimming liberally until MLIR exposes its lexer and parser at which point all this can disappear.

860

why ? https://llvm.org/docs/HowToSetUpLLVMStyleRTTI.html#basic-setup shows it's perfectly fine to use abstract base classes and LLVM RTTI.

911

you should otherwise your computation is non-deterministic

928

It's this or uniqu'ing, underlying storage, placement new etc etc.
I went for the simple solution.
When we have strong data that we need to scale much more we can revisit.

929

reductions loops don't commute in FP land

1136

I need 2 extra maps and really don't anticipate a single named op to ever to a point where this would matter.
Of course if proven otherwise I'm happy to reconsider.

1193

I'm reluctant to invest more in duplicating something that should be exposed by core in a later NFC revision.

1374

Agreed, there are a few other things for follwups too, thanks!

1431

it's a faster symbolcc

1491

this is what the ODS currently is because of manual "named ops", will be cleaned later.

[nicolasvasilache]

Address review comments.

[silvas]

meta-point: @ftynse let's not review the core parser code at the top of the file, as Nicolas says that they are just copypasta from the .mlir parser and won't be in the final patch.

Otherwise, thanks @ftynse for helping with the review :)

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

6

Ah, okay. Sorry for the confusion! When I saw C++ code I was assuming it was emitted by mlir-tblgen gen-op-def. But I see now that there is a mlir-linalg-tblgen -gen-impl that emits C++ as well. Sorry for the noise!!!

[nicolasvasilache]

Thanks for your details reviews @silvas @ftynse !
Anything else ?

[silvas]

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

6

Actually, when rereading I see that we do indeed invoke mlir-tblgen -gen-op-decls. I specifically object to the TBLGEN check prefixes here. I consider it a bug to do that (although i see the precendent in llvm-intrinsics.td, but I would have raised the same objection there), since it violates the layering: somebody updating mlir-tblgen shouldn't be able to break this test.

Consider the implications of what is being checked now in this test...

// TBLGEN-LABEL: linalg::batchmatmulOp declarations
^ could be broken by a change in a *comment* in the generated file :x
//       TBLGEN: class batchmatmulOpOperandAdaptor {
^ could be broken by adding a common base class to the operand adaptor classes, or a change in naming convention for the adaptor classes
//       TBLGEN: class batchmatmulOp : public Op<
^ could be changed by a change in base classes or naming convention.

Note that none of those changes I've indicated would actually break any actual use of this code. So this test is just artificially constraining the implementation of mlir-tblgen for no real value. And even if you strip it down, all you would really be testing is def batchmatmulOp results in a class batchmatmulOp in the output, which is already tested in many places, such as, say, https://github.com/llvm/llvm-project/blob/master/mlir/test/mlir-tblgen/op-decl.td

We need to be courteous to the maintainers of other components and give them the flexibility to adjust the implementations of their components.

[silvas]

@nicolasvasilache any progress on reusing the MLIR parser? I consider that refactoring as blocking for submitting this patch. I don't want us to have a custom parser copypasted here that somebody has to clean up later without a strong reason.

[ftynse]

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

6

I agree that this specific test is over-constraining for mlir-tblgen implementation. What I intended to test in intrinsicgen, and what I would like to see replicated here, is that the tablegen input produced by intrinsicgen, or my mlir-linalg-ods-gen, can be consumed by mlir-tblgen at all. Basically, we don't need to check for any output if we can find a way to check that mlir-tblgen exited with code 0 on the produced file. FileChecking the class name is just a workaround.

If we don't do this check, we risk ending up in a situation where all of the existing tests pass (mlir-tblgen still generates expected C++ from ODS, and mlir-linalg-ods-gen still generates the strings expected by its test, just those strings are no longer valid ODS), but the pipeline fails. And given mlir-tblgen's tendency to assert or crash on improperly structured yet valid TableGen, it would be annoying to debug.

[ftynse]

meta-point: @ftynse let's not review the core parser code at the top of the file, as Nicolas says that they are just copypasta from the .mlir parser and won't be in the final patch.

@silvas I wouldn't review it if it was actual copy-pasta. It is an incomplete and modified copy, which is therefore likely to have some weird behavior or be able to get into an irreversible state where the original code wouldn't get.

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

860

Because you pay the runtime overhead price for two abstractions serving essentially the same goal. Why?

911

I suppose you mean the IR you produce does not have a deterministic order of dimensions, which makes it hard to check. TC semantics says that all dimensions are interchangable, so their textual order should not matter. If it does, we should discuss the semantics and avoid branding this input as TC-like.

928

I don't think vectors of unique_ptr are simpler than vectors of values. This is an extra abstraction with associated cognitive overhead. This is also one extra dynamic allocation per element, as opposed to occasional allocations in the vector, and no strong reason to maintain the pointer as unique or to auto-deallocate (other than you forced the allocation in the first place). There is no actual uniquing of expression, neither is there underlying storage or placement new, you seem to be mistaking this with how types/attributes are handled in MLIR.

1136

Well, you currently have two extra vectors. I just don't see why prefer using a vector of pairs and implementing a search for _every one of them_ is better than using a dedicated container with accessor immediately available.

[nicolasvasilache]

@silvasean I consider that refactoring as blocking for submitting this patch. I don't want us to have a custom parser copypasted here that somebody has to clean up later without a strong reason.
I have been following precedent here, see https://reviews.llvm.org/D73405 which also introduces its own tokenizer / lexer / parser.

As far as I understand it, MLIR has been pretty opinionated about not wanting to expose its tokenizer / lexer / parser: I tried to have them exposed in the past but objections have been along the line of "it's very easy code to write anyway".
I would strongly prefer we revisit that but IMO it would be unfortunate that work is blocked on this refactoring.

Does this help mitigate your position?

[silvas]

Does this help mitigate your position?

Yes. I take back my request to break it out. I buy Chris' statement "I don’t think that splitting this out and pretending it is reusable is a good idea - too much of it is specific to decisions in the MLIR syntax".

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

6

Ah, ok. Then you can just remove the | FileCheck. The RUN line checks that the program has exit code 0, which won't be the case if mlir-tblgen runs into a syntax or processing error.

[ftynse]

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

6

Perfect, let's do this!

[silvas]

LGTM. Let's do this!

[nicolasvasilache]

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

6

Updated the test to get the minimal checkable thing: the class name.

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

860

Marking as done, this is part of the more global reply on downcasting, uptr etc.

911

Added a documentation section in Linalg.md.

928

unique_ptr + abstract base class: basically because I use downcasting.

vector<Expression> will slice unless derived classes have sizeof == 0 (i.e. there is an underlying pointer payload).

An option is to implement a similar arena + pImpl to what MLIR does for the "by-value" abstractions.
I consider this to be unnecessarily complex for my use case (parser that runs at compiler compile time):

vector<unique_ptr<...>> is a standard and simple way to solve the slicing and ownership issue, its performance drawback are not relevant at this time IMO.

1136

fair enough, done, thanks!

[nicolasvasilache]

Address review comments.

[nicolasvasilache]

Addressed

[ftynse]

Please fix the Windows build problem before landing. It looks like the pre-merge testing has such build now so you can use it for the initial check.

mlir/docs/Dialects/Linalg.md

473 ↗

(On Diff #256354)

Nit: angle bracket notation

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

6

This won't work on Windows. Consider adding a -test-emit-additional-includes flag to mlir-linalg-ods-gen and use it here instead of trying shell magic.

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

823

llvm_unreachable ?

[nicolasvasilache]

Address last review comments.

[nicolasvasilache]

Format

[nicolasvasilache]

Doc

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

let description = [{
     future.
   }];
 
   let verifier = [{ return ::verify(*this); }];
 
   let hasFolder = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Named Linalg ops, implemented as a declarative configurations of generic ops.
+//===----------------------------------------------------------------------===//
+
+def NamedStructuredOpTraits : NativeOpTrait<"linalg::NamedStructuredOpTraits">;
+
+class LinalgNamedStructured_Op<string mnemonic, list<OpTrait> props>
+  : Op<Linalg_Dialect, mnemonic,
+       !listconcat(props, [StructuredOpTraits, LinalgStructuredInterface])> {
+  string spec = ?;
+  let assemblyFormat = "`(` operands `)` attr-dict `:` "
+    "functional-type(operands, results)";
+}
+
 #endif // LINALG_STRUCTURED_OPS

mlir/include/mlir/IR/AffineExpr.h

 //===- AffineExpr.h - MLIR Affine Expr Class --------------------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //

 /// products expression, 'localExprs' is expected to have the AffineExpr
 /// for it, and is substituted into. The ArrayRef 'eq' is expected to be in the
 /// format [dims, symbols, locals, constant term].
 AffineExpr getAffineExprFromFlatForm(ArrayRef<int64_t> flatExprs,
                                      unsigned numDims, unsigned numSymbols,
                                      ArrayRef<AffineExpr> localExprs,
                                      MLIRContext *context);
 
-raw_ostream &operator<<(raw_ostream &os, AffineExpr &expr);
+raw_ostream &operator<<(raw_ostream &os, AffineExpr expr);
 
 template <typename U> bool AffineExpr::isa() const {
   if (std::is_same<U, AffineBinaryOpExpr>::value)
     return getKind() <= AffineExprKind::LAST_AFFINE_BINARY_OP;
   if (std::is_same<U, AffineDimExpr>::value)
     return getKind() == AffineExprKind::DimId;
   if (std::is_same<U, AffineSymbolExpr>::value)
     return getKind() == AffineExprKind::SymbolId;

mlir/lib/IR/AffineExpr.cpp

 //===- AffineExpr.cpp - MLIR Affine Expr Classes --------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
return uniquer.get<AffineBinaryOpExprStorage>(
       /*initFn=*/{}, static_cast<unsigned>(AffineExprKind::Mod), *this, other);
 }
 
 AffineExpr AffineExpr::compose(AffineMap map) const {
   SmallVector<AffineExpr, 8> dimReplacements(map.getResults().begin(),
                                              map.getResults().end());
   return replaceDimsAndSymbols(dimReplacements, {});
 }
-raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr &expr) {
+raw_ostream &mlir::operator<<(raw_ostream &os, AffineExpr expr) {
   expr.print(os);
   return os;
 }
 
 /// Constructs an affine expression from a flat ArrayRef. If there are local
 /// identifiers (neither dimensional nor symbolic) that appear in the sum of
 /// products expression, `localExprs` is expected to have the AffineExpr
 /// for it, and is substituted into. The ArrayRef `flatExprs` is expected to be

mlir/test/CMakeLists.txt

configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/Unit/lit.cfg.py
   )
 
 set(MLIR_TEST_DEPENDS
   FileCheck count not
   MLIRUnitTests
   mlir-cpu-runner
   mlir-edsc-builder-api-test
+  mlir-linalg-ods-gen
   mlir-opt
   mlir-sdbm-api-test
   mlir-tblgen
   mlir-translate
   cblas
   cblas_interface
   mlir_runner_utils
   mlir_c_runner_utils

mlir/test/lit.cfg.py

 # Configuration file for the 'lit' test runner.
 
 # name: The name of this test suite.
 config.name = 'MLIR'
 
 config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)
 
 # suffixes: A list of file extensions to treat as test files.
-config.suffixes = ['.td', '.mlir', '.toy', '.ll']
+config.suffixes = ['.td', '.mlir', '.toy', '.ll', '.tc']
 
 # test_source_root: The root path where tests are located.
 config.test_source_root = os.path.dirname(__file__)
 
 # test_exec_root: The root path where tests should be run.
 config.test_exec_root = os.path.join(config.mlir_obj_root, 'test')
 
 config.substitutions.append(('%PATH%', config.environment['PATH']))

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
+
+// RUN: mlir-linalg-ods-gen %s -gen-ods-decl=1 -o %t-ods.td \
+// RUN: && echo 'include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.td" ' $(cat %t-ods.td) \
+// RUN: | mlir-tblgen -gen-op-decls -I %S/../../include \
+// RUN: | FileCheck %s --check-prefix=TBLGEN
+
+
+// ODS-LABEL: def matvecOp : LinalgNamedStructured_Op<"matvec", [
+//  ODS-NEXT:   NInputs<2>,
+//  ODS-NEXT:   NOutputs<1>,
+//  ODS-NEXT:   NamedStructuredOpTraits]>
+//
+// IMPL-LABEL:  matvec::referenceIterators() {
+//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+//
+//       IMPL:  matvec::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(2, 0, {d0, d1}),
+//  IMPL-NEXT:  AffineMap::get(2, 0, {d1}),
+//  IMPL-NEXT:  AffineMap::get(2, 0, {d0}) };
+//
+//       IMPL:  matvec::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
+//       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
+//       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
+//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
+//
+def matvec(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 matmulOp : LinalgNamedStructured_Op<"matmul", [
+//  ODS-NEXT:   NInputs<2>,
+//  ODS-NEXT:   NOutputs<1>,
+//  ODS-NEXT:   NamedStructuredOpTraits]>
+//
+// IMPL-LABEL:  matmul::referenceIterators() {
+//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+//
+//       IMPL:  matmul::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(3, 0, {d0, d2}),
+//  IMPL-NEXT:  AffineMap::get(3, 0, {d2, d1}),
+//  IMPL-NEXT:  AffineMap::get(3, 0, {d0, d1}) };
+//
+//       IMPL:  matmul::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
+//       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
+//       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
+//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
+//
+def matmul(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 batchmatmulOp : LinalgNamedStructured_Op<"batchmatmul", [
+//  ODS-NEXT:   NInputs<2>,
+//  ODS-NEXT:   NOutputs<1>,
+//  ODS-NEXT:   NamedStructuredOpTraits]>
+//
+// IMPL-LABEL:  batchmatmul::referenceIterators() {
+//  IMPL-NEXT:  { {{.*}}Parallel{{.*}}, {{.*}}Parallel{{.*}}, {{.*}}Reduction{{.*}} }
+//
+//       IMPL:  batchmatmul::referenceIndexingMaps() {
+//       IMPL:  AffineMap::get(4, 0, {d0, d1, d3}),
+//  IMPL-NEXT:  AffineMap::get(4, 0, {d3, d2}),
+//  IMPL-NEXT:  AffineMap::get(4, 0, {d0, d1, d2}) };
+//
+//       IMPL:  batchmatmul::regionBuilder(ArrayRef<BlockArgument> args) {
+//       IMPL:  ValueHandle [[a:.*]](args[0]), [[b:.*]](args[1]), [[c:.*]](args[2]);
+//       IMPL:  ValueHandle [[d:.*]] = std_mulf([[a]], [[b]]);
+//       IMPL:  ValueHandle [[e:.*]] = std_addf([[c]], [[d]]);
+//       IMPL:  (linalg_yield(ValueRange{ [[e]] }));
+//
+// TBLGEN-LABEL: linalg::batchmatmulOp declarations
+//       TBLGEN: class batchmatmulOpOperandAdaptor {
+//       TBLGEN: class batchmatmulOp : public Op<
+def batchmatmul(A: f32(Batch, M, K), B: f32(K, N)) -> (C: f32(Batch, M, N)) {
+  C(b, m, n) = std_addf<k>(std_mulf(A(b, m, k), B(k, n)));
+}

mlir/tools/CMakeLists.txt

 add_subdirectory(mlir-cuda-runner)
 add_subdirectory(mlir-cpu-runner)
+add_subdirectory(mlir-linalg-ods-gen)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
 add_subdirectory(mlir-shlib)

mlir/tools/mlir-linalg-ods-gen/CMakeLists.txt

+add_llvm_tool(mlir-linalg-ods-gen
+  mlir-linalg-ods-gen.cpp
+)
+llvm_update_compile_flags(mlir-linalg-ods-gen)
+target_link_libraries(mlir-linalg-ods-gen PRIVATE
+  MLIRParser
+  MLIRSupport
+  LLVMCore
+  LLVMSupport
+  )

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 "mlir/Support/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+#define DEBUG_TYPE "linalg-ods-gen"
+
+static llvm::cl::OptionCategory IntrinsicGenCat("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(IntrinsicGenCat));
+
+static llvm::cl::opt<bool>
+    genODSImpl("gen-impl", llvm::cl::desc("Emit the ops implementations"),
+               llvm::cl::init(false), llvm::cl::cat(IntrinsicGenCat));
+
+using llvm::SetVector;
+using llvm::SMLoc;
+using llvm::StringRef;
+using llvm::Twine;
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Lexer
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class represents a specific token in the input format.
+class Token {
+public:
+  enum class Kind {
+    // Markers.
+    eof,
+    error,
+
+    // Tokens with no info.
+    colon,
+    comma,
+    equal,
+    gt,
+    l_brace,
+    l_paren,
+    lt,
+    minus,
+    plus,
+    r_brace,
+    r_paren,
+    semicolon,
+    star,
+
+    // Keywords.
+    kw_def,
+    FIRST_KEYWORD = kw_def,
+    kw_floordiv,
+    kw_ceildiv,
+    kw_mod,
+    LAST_KEYWORD = kw_mod,
+
+    // String valued tokens.
+    id,
+    integer,
+  };
+
+  Token(Kind kind, StringRef spelling) : kind(kind), spelling(spelling) {}
+
+  /// Return the bytes that make up this token.
+  StringRef getSpelling() const { return spelling; }
+
+  /// Return the kind of this token.
+  Kind getKind() const { return kind; }
+
+  /// Return a location for this token.
+  llvm::SMLoc getLoc() const {
+    return llvm::SMLoc::getFromPointer(spelling.data());
+  }
+
+  /// Return if this token is a keyword.
+  bool isKeyword() const {
+    return kind >= Kind::FIRST_KEYWORD && kind <= Kind::LAST_KEYWORD;
+  }
+  bool is(Kind k) const { return kind == k; }
+  bool isNot(Kind k) const { return kind != k; }
+
+  Optional<uint64_t> getUInt64IntegerValue() const {
+    bool isHex = spelling.size() > 1 && spelling[1] == 'x';
+
+    uint64_t result = 0;
+    if (spelling.getAsInteger(isHex ? 0 : 10, result))
+      return None;
+    return result;
+  }
+
+private:
+  /// Discriminator that indicates the kind of token this is.
+  Kind kind;
+
+  /// A reference to the entire token contents; this is always a pointer into
+  /// a memory buffer owned by the source manager.
+  StringRef spelling;
+};
+
+/// This class implements a simple lexer.
+class Lexer {
+public:
+  Lexer(llvm::SourceMgr &mgr);
+
+  /// Lex the next token and return it.
+  Token lexToken();
+
+  /// Emit an error to the lexer with the given location and message.
+  Token emitError(llvm::SMLoc loc, const Twine &msg);
+  Token emitError(const char *loc, const Twine &msg);
+
+private:
+  Token formToken(Token::Kind kind, const char *tokStart) {
+    return Token(kind, StringRef(tokStart, curPtr - tokStart));
+  }
+
+  /// Return the next character in the stream.
+  int getNextChar();
+
+  /// Lex an identifier.
+  Token lexIdentifier(const char *tokStart);
+
+  // Lex an integer.
+  Token lexInteger(const char *tokStart);
+
+  // Skip a comment line, starting with a '//'.
+  void skipComment();
+
+  llvm::SourceMgr &srcMgr;
+  StringRef curBuffer;
+  const char *curPtr;
+};
+} // end anonymous namespace
+
+Lexer::Lexer(llvm::SourceMgr &mgr) : srcMgr(mgr) {
+  curBuffer = srcMgr.getMemoryBuffer(mgr.getMainFileID())->getBuffer();
+  curPtr = curBuffer.begin();
+}
+
+Token Lexer::emitError(llvm::SMLoc loc, const Twine &msg) {
+  srcMgr.PrintMessage(loc, llvm::SourceMgr::DK_Error, msg);
+  return formToken(Token::Kind::error, loc.getPointer());
+}
+Token Lexer::emitError(const char *loc, const Twine &msg) {
+  return emitError(llvm::SMLoc::getFromPointer(loc), msg);
+}
+
+int Lexer::getNextChar() {
+  char curChar = *curPtr++;
+  switch (curChar) {
+  default:
+    return (unsigned char)curChar;
+  case 0: {
+    // A nul character in the stream is either the end of the current buffer
+    // or a random nul in the file. Disambiguate that here.
+    if (curPtr - 1 != curBuffer.end())
+      return 0;
+
+    // Otherwise, return end of file.
+    --curPtr;
+    return EOF;
+  }
+  case '\n':
+  case '\r':
+    // Handle the newline character by ignoring it and incrementing the line
+    // count. However, be careful about 'dos style' files with \n\r in them.
+    // Only treat a \n\r or \r\n as a single line.
+    if ((*curPtr == '\n' || (*curPtr == '\r')) && *curPtr != curChar)
+      ++curPtr;
+    return '\n';
+  }
+}
+
+Token Lexer::lexToken() {
+  while (true) {
+    const char *tokStart = curPtr;
+
+    // This always consumes at least one character.
+    int curChar = getNextChar();
+    switch (curChar) {
+    default:
+      // Handle identifiers: [a-zA-Z_]
+      if (isalpha(curChar) || curChar == '_')
+        return lexIdentifier(tokStart);
+
+      // Handle integers: [0-9]
+      if (isdigit(curChar))
+        return lexInteger(tokStart);
+
+      // Unknown character, emit an error.
+      return emitError(tokStart, "unexpected character");
+
+    case EOF:
+      // Return EOF denoting the end of lexing.
+      return formToken(Token::Kind::eof, tokStart);
+
+    // Lex punctuation.
+    case ':':
+      return formToken(Token::Kind::colon, tokStart);
+    case ',':
+      return formToken(Token::Kind::comma, tokStart);
+    case '=':
+      return formToken(Token::Kind::equal, tokStart);
+    case '{':
+      return formToken(Token::Kind::l_brace, tokStart);
+    case '(':
+      return formToken(Token::Kind::l_paren, tokStart);
+    case '}':
+      return formToken(Token::Kind::r_brace, tokStart);
+    case ')':
+      return formToken(Token::Kind::r_paren, tokStart);
+    case '<':
+      return formToken(Token::Kind::lt, tokStart);
+    case '>':
+      return formToken(Token::Kind::gt, tokStart);
+    case '+':
+      return formToken(Token::Kind::plus, tokStart);
+    case '-':
+      return formToken(Token::Kind::minus, tokStart);
+    case ';':
+      return formToken(Token::Kind::semicolon, tokStart);
+    case '*':
+      return formToken(Token::Kind::star, tokStart);
+    case '/':
+      if (*curPtr == '/') {
+        skipComment();
+        continue;
+      }
+      // Unknown character, emit an error.
+      return emitError(tokStart, "unexpected character: not a comment");
+
+    // Ignore whitespace characters.
+    case 0:
+    case ' ':
+    case '\t':
+    case '\n':
+      return lexToken();
+    }
+  }
+}
+
+Token Lexer::lexIdentifier(const char *tokStart) {
+  // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
+  while (isalnum(*curPtr) || *curPtr == '_' || *curPtr == '-')
+    ++curPtr;
+
+  // Check to see if this identifier is a keyword.
+  StringRef str(tokStart, curPtr - tokStart);
+  Token::Kind kind = llvm::StringSwitch<Token::Kind>(str)
+                         .Case("def", Token::Kind::kw_def)
+                         .Case("floordiv", Token::Kind::kw_floordiv)
+                         .Case("ceildiv", Token::Kind::kw_ceildiv)
+                         .Case("mod", Token::Kind::kw_mod)
+                         .Default(Token::Kind::id);
+
+  return Token(kind, str);
+}
+
+Token Lexer::lexInteger(const char *tokStart) {
+  // Match the rest of the identifier regex: [0-9a-zA-Z_\-]*
+  while (isdigit(*curPtr))
+    ++curPtr;
+
+  StringRef str(tokStart, curPtr - tokStart);
+  return Token(Token::Kind::integer, str);
+}
+
+/// Skip a comment line, starting with a '//'.
+void Lexer::skipComment() {
+  // Advance over the second '/' in a '//' comment.
+  assert(*curPtr == '/');
+  ++curPtr;
+
+  while (true) {
+    switch (*curPtr++) {
+    case '\n':
+    case '\r':
+      // Newline is end of comment.
+      return;
+    case 0:
+      // If this is the end of the buffer, end the comment.
+      if (curPtr - 1 == curBuffer.end()) {
+        --curPtr;
+        return;
+      }
+      LLVM_FALLTHROUGH;
+    default:
+      // Skip over other characters.
+      break;
+    }
+  }
+}
+
+namespace {
+
+class Parser {
+public:
+  Parser(llvm::SourceMgr &mgr, MLIRContext *ctx)
+      : lexer(mgr), curToken(lexer.lexToken()), context(ctx) {}
+
+  //===--------------------------------------------------------------------===//
+  // Lexer Utilities
+  //===--------------------------------------------------------------------===//
+
+  /// Advance the current lexer onto the next token.
+  void consumeToken() {
+    assert(curToken.getKind() != Token::Kind::eof &&
+           curToken.getKind() != Token::Kind::error &&
+           "shouldn't advance past EOF or errors");
+    curToken = lexer.lexToken();
+  }
+  void consumeToken(Token::Kind kind) {
+    assert(curToken.getKind() == kind && "unexpected token");
+    curToken = lexer.lexToken();
+  }
+  LogicalResult parseToken(Token::Kind kind, const Twine &msg) {
+    if (curToken.getKind() != kind)
+      return emitError(curToken.getLoc(), msg);
+    consumeToken();
+    return success();
+  }
+  LogicalResult emitError(llvm::SMLoc loc, const Twine &msg) {
+    lexer.emitError(loc, msg);
+    return failure();
+  }
+  LogicalResult emitError(const Twine &msg) {
+    return emitError(curToken.getLoc(), msg);
+  }
+  bool consumeIf(Token::Kind kind) {
+    if (curToken.isNot(kind))
+      return false;
+    consumeToken(kind);
+    return true;
+  }
+  LogicalResult
+  parseCommaSeparatedList(llvm::function_ref<ParseResult()> parseElement) {
+    // Non-empty case starts with an element.
+    if (parseElement())
+      return failure();
+
+    // Otherwise we have a list of comma separated elements.
+    while (consumeIf(Token::Kind::comma)) {
+      if (parseElement())
+        return failure();
+    }
+    return success();
+  }
+  LogicalResult
+  parseCommaSeparatedListUntil(Token::Kind rightToken,
+                               llvm::function_ref<ParseResult()> parseElement,
+                               bool allowEmptyList) {
+    // Handle the empty case.
+    if (curToken.is(rightToken)) {
+      if (!allowEmptyList)
+        return emitError("expected list element");
+      consumeToken(rightToken);
+      return success();
+    }
+
+    if (failed(parseCommaSeparatedList(parseElement)) ||
+        failed(
+            parseToken(rightToken, "expected ',' or right-terminating token")))
+      return failure();
+
+    return success();
+  }
+
+  Lexer lexer;
+  Token curToken;
+  MLIRContext *context;
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Affine parsing.
+//===----------------------------------------------------------------------===//
+
+namespace {
+
+/// Lower precedence ops (all at the same precedence level). LNoOp is false in
+/// the boolean sense.
+enum AffineLowPrecOp {
+  /// Null value.
+  LNoOp,
+  Add,
+  Sub
+};
+
+/// Higher precedence ops - all at the same precedence level. HNoOp is false
+/// in the boolean sense.
+enum AffineHighPrecOp {
+  /// Null value.
+  HNoOp,
+  Mul,
+  FloorDiv,
+  CeilDiv,
+  Mod
+};
+
+using AffineDimList = SmallVector<std::pair<StringRef, AffineExpr>, 4>;
+using AffineSymbolList = SmallVector<std::pair<StringRef, AffineExpr>, 4>;
+
+/// This is a specialized parser for affine expressions.
+class AffineParser {
+public:
+  explicit AffineParser(Parser &p,
+                        std::function<AffineExpr(StringRef)> bareIdParsingHook,
+                        AffineDimList &dimList, AffineSymbolList &symbolList)
+      : parser(p), bareIdFallback(bareIdParsingHook), dims(dimList),
+        symbols(symbolList) {}
+
+  /// Parse a comma-separated list of affine exprs.
+  SmallVector<AffineExpr, 4>
+  parseAffineExprs(Token::Kind lDelim = Token::Kind::l_paren,
+                   Token::Kind rDelim = Token::Kind::r_paren);
+
+  /// Parse a single affine expr.`.
+  AffineExpr parseAffineExpr();
+
+private:
+  // Binary affine op parsing.
+  AffineLowPrecOp consumeIfLowPrecOp();
+  AffineHighPrecOp consumeIfHighPrecOp();
+
+  // AffineExpr parsing.
+  AffineExpr parseParentheticalExpr();
+  AffineExpr parseNegateExpression(AffineExpr lhs);
+  AffineExpr parseIntegerExpr();
+  AffineExpr parseBareIdExpr();
+
+  AffineExpr getAffineBinaryOpExpr(AffineHighPrecOp op, AffineExpr lhs,
+                                   AffineExpr rhs, SMLoc opLoc);
+  AffineExpr getAffineBinaryOpExpr(AffineLowPrecOp op, AffineExpr lhs,
+                                   AffineExpr rhs);
+  AffineExpr parseAffineOperandExpr(AffineExpr lhs);
+  AffineExpr parseAffineLowPrecOpExpr(AffineExpr llhs, AffineLowPrecOp llhsOp);
+  AffineExpr parseAffineHighPrecOpExpr(AffineExpr llhs, AffineHighPrecOp llhsOp,
+                                       SMLoc llhsOpLoc);
+
+  Parser &parser;
+  std::function<AffineExpr(StringRef)> bareIdFallback;
+  AffineDimList &dims;
+  AffineSymbolList &symbols;
+};
+} // end anonymous namespace
+
+/// Create an affine binary high precedence op expression (mul's, div's, mod).
+/// opLoc is the location of the op token to be used to report errors
+/// for non-conforming expressions.
+AffineExpr AffineParser::getAffineBinaryOpExpr(AffineHighPrecOp op,
+                                               AffineExpr lhs, AffineExpr rhs,
+                                               SMLoc opLoc) {
+  switch (op) {
+  case Mul:
+    if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant()) {
+      parser.emitError(opLoc,
+                       "non-affine expression: at least one of the multiply "
+                       "operands has to be either a constant or symbolic");
+      return nullptr;
+    }
+    return lhs * rhs;
+  case FloorDiv:
+    if (!rhs.isSymbolicOrConstant()) {
+      parser.emitError(opLoc,
+                       "non-affine expression: right operand of floordiv "
+                       "has to be either a constant or symbolic");
+      return nullptr;
+    }
+    return lhs.floorDiv(rhs);
+  case CeilDiv:
+    if (!rhs.isSymbolicOrConstant()) {
+      parser.emitError(opLoc, "non-affine expression: right operand of ceildiv "
+                              "has to be either a constant or symbolic");
+      return nullptr;
+    }
+    return lhs.ceilDiv(rhs);
+  case Mod:
+    if (!rhs.isSymbolicOrConstant()) {
+      parser.emitError(opLoc, "non-affine expression: right operand of mod "
+                              "has to be either a constant or symbolic");
+      return nullptr;
+    }
+    return lhs % rhs;
+  case HNoOp:
+    llvm_unreachable("can't create affine expression for null high prec op");
+    return nullptr;
+  }
+  llvm_unreachable("Unknown AffineHighPrecOp");
+}
+
+/// Create an affine binary low precedence op expression (add, sub).
+AffineExpr AffineParser::getAffineBinaryOpExpr(AffineLowPrecOp op,
+                                               AffineExpr lhs, AffineExpr rhs) {
+  switch (op) {
+  case AffineLowPrecOp::Add:
+    return lhs + rhs;
+  case AffineLowPrecOp::Sub:
+    return lhs - rhs;
+  case AffineLowPrecOp::LNoOp:
+    llvm_unreachable("can't create affine expression for null low prec op");
+    return nullptr;
+  }
+  llvm_unreachable("Unknown AffineLowPrecOp");
+}
+
+/// Consume this token if it is a lower precedence affine op (there are only
+/// two precedence levels).
+AffineLowPrecOp AffineParser::consumeIfLowPrecOp() {
+  switch (parser.curToken.getKind()) {
+  case Token::Kind::plus:
+    parser.consumeToken();
+    return AffineLowPrecOp::Add;
+  case Token::Kind::minus:
+    parser.consumeToken();
+    return AffineLowPrecOp::Sub;
+  default:
+    return AffineLowPrecOp::LNoOp;
+  }
+}
+
+/// Consume this token if it is a higher precedence affine op (there are only
+/// two precedence levels)
+AffineHighPrecOp AffineParser::consumeIfHighPrecOp() {
+  switch (parser.curToken.getKind()) {
+  case Token::Kind::star:
+    parser.consumeToken(Token::Kind::star);
+    return Mul;
+  case Token::Kind::kw_floordiv:
+    parser.consumeToken(Token::Kind::kw_floordiv);
+    return FloorDiv;
+  case Token::Kind::kw_ceildiv:
+    parser.consumeToken(Token::Kind::kw_ceildiv);
+    return CeilDiv;
+  case Token::Kind::kw_mod:
+    parser.consumeToken(Token::Kind::kw_mod);
+    return Mod;
+  default:
+    return HNoOp;
+  }
+}
+
+/// Parse a high precedence op expression list: mul, div, and mod are high
+/// precedence binary ops, i.e., parse a
+///   expr_1 op_1 expr_2 op_2 ... expr_n
+/// where op_1, op_2 are all a AffineHighPrecOp (mul, div, mod).
+/// All affine binary ops are left associative.
+/// Given llhs, returns (llhs llhsOp lhs) op rhs, or (lhs op rhs) if llhs is
+/// null. If no rhs can be found, returns (llhs llhsOp lhs) or lhs if llhs is
+/// null. llhsOpLoc is the location of the llhsOp token that will be used to
+/// report an error for non-conforming expressions.
+AffineExpr AffineParser::parseAffineHighPrecOpExpr(AffineExpr llhs,
+                                                   AffineHighPrecOp llhsOp,
+                                                   SMLoc llhsOpLoc) {
+  AffineExpr lhs = parseAffineOperandExpr(llhs);
+  if (!lhs)
+    return nullptr;
+
+  // Found an LHS. Parse the remaining expression.
+  auto opLoc = parser.curToken.getLoc();
+  if (AffineHighPrecOp op = consumeIfHighPrecOp()) {
+    if (llhs) {
+      AffineExpr expr = getAffineBinaryOpExpr(llhsOp, llhs, lhs, opLoc);
+      if (!expr)
+        return nullptr;
+      return parseAffineHighPrecOpExpr(expr, op, opLoc);
+    }
+    // No LLHS, get RHS
+    return parseAffineHighPrecOpExpr(lhs, op, opLoc);
+  }
+
+  // This is the last operand in this expression.
+  if (llhs)
+    return getAffineBinaryOpExpr(llhsOp, llhs, lhs, llhsOpLoc);
+
+  // No llhs, 'lhs' itself is the expression.
+  return lhs;
+}
+
+/// Parse an affine expression inside parentheses.
+///
+///   affine-expr ::= `(` affine-expr `)`
+AffineExpr AffineParser::parseParentheticalExpr() {
+  if (failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
+    return nullptr;
+  if (parser.curToken.is(Token::Kind::r_paren))
+    return (parser.emitError("no expression inside parentheses"), nullptr);
+
+  auto expr = parseAffineExpr();
+  if (!expr)
+    return nullptr;
+  if (failed(parser.parseToken(Token::Kind::r_paren, "expected ')'")))
+    return nullptr;
+
+  return expr;
+}
+
+/// Parse the negation expression.
+///
+///   affine-expr ::= `-` affine-expr
+AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
+  if (failed(parser.parseToken(Token::Kind::minus, "expected '-'")))
+    return nullptr;
+
+  AffineExpr operand = parseAffineOperandExpr(lhs);
+  // Since negation has the highest precedence of all ops (including high
+  // precedence ops) but lower than parentheses, we are only going to use
+  // parseAffineOperandExpr instead of parseAffineExpr here.
+  if (!operand)
+    // Extra error message although parseAffineOperandExpr would have
+    // complained. Leads to a better diagnostic.
+    return (parser.emitError("missing operand of negation"), nullptr);
+  return (-1) * operand;
+}
+
+/// Parse a bare id that may appear in an affine expression.
+///
+///   affine-expr ::= bare-id
+AffineExpr AffineParser::parseBareIdExpr() {
+  if (parser.curToken.isNot(Token::Kind::id))
+    return (parser.emitError("expected id"), nullptr);
+
+  StringRef sRef = parser.curToken.getSpelling();
+  for (auto &list : {dims, symbols}) {
+    for (auto entry : list) {
+      if (entry.first == sRef) {
+        parser.consumeToken(Token::Kind::id);
+        return entry.second;
+      }
+    }
+  }
+
+  // Not found, check fallback path.
+  AffineExpr expr = bareIdFallback(sRef);
+  if (expr) {
+    parser.consumeToken(Token::Kind::id);
+    return expr;
+  }
+
+  return (parser.emitError("use of undeclared id"), nullptr);
+}
+
+/// Parse a positive integral constant appearing in an affine expression.
+///
+///   affine-expr ::= integer-literal
+AffineExpr AffineParser::parseIntegerExpr() {
+  auto val = parser.curToken.getUInt64IntegerValue();
+  if (!val.hasValue() || (int64_t)val.getValue() < 0)
+    return (parser.emitError("constant too large for index"), nullptr);
+
+  parser.consumeToken(Token::Kind::integer);
+  return getAffineConstantExpr((int64_t)val.getValue(), parser.context);
+}
+
+/// Parses an expression that can be a valid operand of an affine expression.
+/// lhs: if non-null, lhs is an affine expression that is the lhs of a binary
+/// operator, the rhs of which is being parsed. This is used to determine
+/// whether an error should be emitted for a missing right operand.
+//  Eg: for an expression without parentheses (like i + j + k + l), each
+//  of the four identifiers is an operand. For i + j*k + l, j*k is not an
+//  operand expression, it's an op expression and will be parsed via
+//  parseAffineHighPrecOpExpression(). However, for i + (j*k) + -l, (j*k) and
+//  -l are valid operands that will be parsed by this function.
+AffineExpr AffineParser::parseAffineOperandExpr(AffineExpr lhs) {
+  switch (parser.curToken.getKind()) {
+  case Token::Kind::id:
+    return parseBareIdExpr();
+  case Token::Kind::integer:
+    return parseIntegerExpr();
+  case Token::Kind::l_paren:
+    return parseParentheticalExpr();
+  case Token::Kind::minus:
+    return parseNegateExpression(lhs);
+  case Token::Kind::kw_ceildiv:
+  case Token::Kind::kw_floordiv:
+  case Token::Kind::kw_mod:
+  case Token::Kind::plus:
+  case Token::Kind::star:
+    if (lhs)
+      parser.emitError("missing right operand of binary operator");
+    else
+      parser.emitError("missing left operand of binary operator");
+    return nullptr;
+  default:
+    if (lhs)
+      parser.emitError("missing right operand of binary operator");
+    else
+      parser.emitError("expected affine expression");
+    return nullptr;
+  }
+}
+
+/// Parse affine expressions that are bare-id's, integer constants,
+/// parenthetical affine expressions, and affine op expressions that are a
+/// composition of those.
+///
+/// All binary op's associate from left to right.
+///
+/// {add, sub} have lower precedence than {mul, div, and mod}.
+///
+/// Add, sub'are themselves at the same precedence level. Mul, floordiv,
+/// ceildiv, and mod are at the same higher precedence level. Negation has
+/// higher precedence than any binary op.
+///
+/// llhs: the affine expression appearing on the left of the one being parsed.
+/// This function will return ((llhs llhsOp lhs) op rhs) if llhs is non null,
+/// and lhs op rhs otherwise; if there is no rhs, llhs llhsOp lhs is returned
+/// if llhs is non-null; otherwise lhs is returned. This is to deal with left
+/// associativity.
+///
+/// Eg: when the expression is e1 + e2*e3 + e4, with e1 as llhs, this function
+/// will return the affine expr equivalent of (e1 + (e2*e3)) + e4, where
+/// (e2*e3) will be parsed using parseAffineHighPrecOpExpr().
+AffineExpr AffineParser::parseAffineLowPrecOpExpr(AffineExpr llhs,
+                                                  AffineLowPrecOp llhsOp) {
+  AffineExpr lhs;
+  if (!(lhs = parseAffineOperandExpr(llhs)))
+    return nullptr;
+
+  // Found an LHS. Deal with the ops.
+  if (AffineLowPrecOp lOp = consumeIfLowPrecOp()) {
+    if (llhs) {
+      AffineExpr sum = getAffineBinaryOpExpr(llhsOp, llhs, lhs);
+      return parseAffineLowPrecOpExpr(sum, lOp);
+    }
+    // No LLHS, get RHS and form the expression.
+    return parseAffineLowPrecOpExpr(lhs, lOp);
+  }
+  auto opLoc = parser.curToken.getLoc();
+  if (AffineHighPrecOp hOp = consumeIfHighPrecOp()) {
+    // We have a higher precedence op here. Get the rhs operand for the llhs
+    // through parseAffineHighPrecOpExpr.
+    AffineExpr highRes = parseAffineHighPrecOpExpr(lhs, hOp, opLoc);
+    if (!highRes)
+      return nullptr;
+
+    // If llhs is null, the product forms the first operand of the yet to be
+    // found expression. If non-null, the op to associate with llhs is llhsOp.
+    AffineExpr expr =
+        llhs ? getAffineBinaryOpExpr(llhsOp, llhs, highRes) : highRes;
+
+    // Recurse for subsequent low prec op's after the affine high prec op
+    // expression.
+    if (AffineLowPrecOp nextOp = consumeIfLowPrecOp())
+      return parseAffineLowPrecOpExpr(expr, nextOp);
+    return expr;
+  }
+  // Last operand in the expression list.
+  if (llhs)
+    return getAffineBinaryOpExpr(llhsOp, llhs, lhs);
+  // No llhs, 'lhs' itself is the expression.
+  return lhs;
+}
+
+/// Parse an affine expression.
+///  affine-expr ::= `(` affine-expr `)`
+///                | `-` affine-expr
+///                | affine-expr `+` affine-expr
+///                | affine-expr `-` affine-expr
+///                | affine-expr `*` affine-expr
+///                | affine-expr `floordiv` affine-expr
+///                | affine-expr `ceildiv` affine-expr
+///                | affine-expr `mod` affine-expr
+///                | bare-id
+///                | integer-literal
+///
+/// Additional conditions are checked depending on the production. For eg.,
+/// one of the operands for `*` has to be either constant/symbolic; the second
+/// operand for floordiv, ceildiv, and mod has to be a positive integer.
+AffineExpr AffineParser::parseAffineExpr() {
+  return parseAffineLowPrecOpExpr(nullptr, AffineLowPrecOp::LNoOp);
+}
+
+SmallVector<AffineExpr, 4> AffineParser::parseAffineExprs(Token::Kind lDelim,
+                                                          Token::Kind rDelim) {
+  parser.parseToken(lDelim, "expected lDelim at start of affine expr list");
+
+  SmallVector<AffineExpr, 4> exprs;
+  auto parseElt = [&]() -> LogicalResult {
+    auto elt = parseAffineExpr();
+    exprs.push_back(elt);
+    return elt ? success() : failure();
+  };
+
+  if (failed(parser.parseCommaSeparatedListUntil(rDelim, parseElt,
+                                                 /*allowEmptyList=*/true)))
+    assert(false);
+
+  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() = 0;
+
+  bool operator==(const Expression &e) const;
+  operator bool() const { return kind != Kind::Uninitialized; }
+
+  Kind kind;
+};
+
+/// Encodes a tensor use of the form:
+///
+///   affine-expr-list ::= affine-expr (`,` affine-expr)*
+///   tensor-use ::= bare-id `(` `)`
+///                | bare-id `(` affine-expr-list `)`
+///
+/// The affine-expr-list is stored as an AffineMap.
+struct TensorUse : public Expression {
+  TensorUse() : TensorUse("", AffineMap()) {}
+  TensorUse(StringRef name, AffineMap map)
+      : Expression(Kind::TensorUse), tensorId(name), indexingMap(map) {}
+  TensorUse(const TensorUse &use) = default;
+
+  static bool classof(const Expression *e) {
+    return e->kind == Kind::TensorUse;
+  }
+
+  bool operator==(const TensorUse &other) const {
+    return tensorId == other.tensorId && indexingMap == other.indexingMap;
+  }
+
+  /// Visitation function. Performs preorder or postorder traversal depending on
+  /// `PreOrder` and applies `callback` on each node.
+  template <typename Lambda, bool PreOrder> void visit(Lambda callback) const;
+
+  StringRef tensorId;
+  AffineMap indexingMap;
+};
+
+/// Encodes a tensor expression of the form:
+///
+///   op-spec ::= bare-id `<` reduction-dims-list `>`
+///             | bare-id
+///   op-arg ::= tensor-expr
+///            | tensor-use
+///   op-arg-list ::= op-arg (`,` op-arg)*
+///   tensor-expr ::= op-spec `(` op-arg-list `)`
+///
+/// Underlying op-arg are stored by unique_ptr to base class.
+struct TensorExpr : public Expression {
+  TensorExpr(StringRef name,
+             SmallVectorImpl<std::unique_ptr<Expression>> &&exprs,
+             ArrayRef<unsigned> reductionDims)
+      : Expression(Kind::TensorExpr), opId(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 (opId != other.opId)
+      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 opId;
+  SmallVector<std::unique_ptr<Expression>, 4> expressions;
+  SetVector<unsigned> reductionDimensions;
+};
+
+Expression::~Expression() {}
+
+bool Expression::operator==(const Expression &e) const {
+  if (this->kind != e.kind)
+    return false;
+  if (e.kind == Expression::Kind::TensorUse)
+    return static_cast<const TensorUse &>(*this) ==
+           static_cast<const TensorUse &>(e);
+  if (e.kind == Expression::Kind::TensorExpr)
+    return static_cast<const TensorExpr &>(*this) ==
+           static_cast<const TensorExpr &>(e);
+  llvm_unreachable("Unexpected case");
+}
+
+/// This is a specialized parser for a TCDef.
+/// This maintains the dims it finds in an eager fashion.
+class TCParser {
+  enum class EagerDiscoveryMode { None = 0, Symbols, Dimensions };
+
+public:
+  explicit TCParser(Parser &p);
+
+  /// Uses the AffineParser to parse the affine exprs used in a tensor
+  /// definition. If `discoveryMode` is set to Symbols (resp. Dimensions), new
+  /// symbols (resp. dimensions) are added eagerly. Otherwise, an error is
+  /// emitted on new identifiers.
+  SmallVector<AffineExpr, 4>
+  parseAffineExprs(EagerDiscoveryMode discoveryMode, AffineDimList &dims,
+                   Token::Kind lDelim = Token::Kind::l_paren,
+                   Token::Kind rDelim = Token::Kind::r_paren);
+
+  /// Parse the information for a tensor def.
+  /// All the affine-expr must be dimensionless (i.e. contain only expressions
+  /// involving symbols and constants), but can otherwise contain arbitrary
+  /// affine expressions.
+  LogicalResult parseTensorDef(bool isOutput);
+
+  /// Parses a tensor use.
+  struct ComprehensionParsingState {
+    AffineDimList dims;
+    SmallVector<std::unique_ptr<Expression>, 4> expressions;
+    llvm::DenseMap<TensorUse, unsigned> orderedTensorArgs;
+  };
+  LogicalResult parseTensorUse(TensorUse &result,
+                               ComprehensionParsingState &state);
+
+  /// Parses a tensor expression.
+  LogicalResult parseExpression(TensorUse currentDefinition,
+                                std::unique_ptr<Expression> &result,
+                                ComprehensionParsingState &state);
+
+  /// Parse a single comprehension.
+  LogicalResult parseOneComprehension(StringRef cppOpName,
+                                      StringRef linalgOpName,
+                                      ComprehensionParsingState &state);
+
+  /// Parse and print the information for a TC def.
+  /// When `gen-ods-decl` is used, this prints the ODS declaration for the TC.
+  /// When `gen-impl` is used, this prints the C++ implementation for the extra
+  /// methods defined in ODS (referenceIterators, referenceIndexingMaps and
+  /// regionBuilder).
+  LogicalResult parseAndEmitTCDef(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);
+
+  /// Print the C++ StructuredOpsInterface impl of `referenceIterators`.
+  void printReferenceIterators(llvm::raw_ostream &os, StringRef opId,
+                               ComprehensionParsingState &state);
+
+  /// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
+  void printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef opId,
+                                  ComprehensionParsingState &state);
+
+  /// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
+  void printRegionBuilder(llvm::raw_ostream &os, StringRef opId,
+                          ComprehensionParsingState &state);
+
+private:
+  //===--------------------------------------------------------------------===//
+  // Internal bookkeeping of tensors.
+  //===--------------------------------------------------------------------===//
+  struct RegisteredTensor {
+    StringRef name;
+    StringRef type;
+    AffineMap shape;
+    bool isOutput;
+    AffineMap indexingMap;
+  };
+  void registerTensor(StringRef tensorId, StringRef tensorType, AffineMap map,
+                      bool isOutput);
+  RegisteredTensor &getRegisteredTensor(StringRef id);
+  unsigned getRegisteredTensorIndex(StringRef id);
+  bool isRegisteredTensor(StringRef id);
+
+  //===--------------------------------------------------------------------===//
+  // Per-TC def state.
+  //===--------------------------------------------------------------------===//
+  /// Symbols are per TC def.
+  AffineSymbolList symbols;
+  /// Tensors are per TC def.
+  SmallVector<RegisteredTensor, 4> tensors;
+
+  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);
+}
+
+//===----------------------------------------------------------------------===//
+// Internal bookkeeping of tensors.
+//===----------------------------------------------------------------------===//
+void TCParser::registerTensor(StringRef tensorId, StringRef tensorType,
+                              AffineMap map, bool isOutput) {
+  tensors.emplace_back(
+      RegisteredTensor{tensorId, tensorType, map, isOutput, AffineMap()});
+  LLVM_DEBUG(llvm::dbgs() << "Recorded: " << tensorId << " "
+                          << "with typeString: " << tensorType << " "
+                          << "and shape: " << map << "\n");
+}
+
+TCParser::RegisteredTensor &TCParser::getRegisteredTensor(StringRef id) {
+  for (unsigned i = 0, e = tensors.size(); i < e; ++i)
+    if (tensors[i].name == id)
+      return tensors[i];
+  llvm_unreachable("Unregistered tensor");
+}
+
+unsigned TCParser::getRegisteredTensorIndex(StringRef id) {
+  for (unsigned i = 0, e = tensors.size(); i < e; ++i)
+    if (tensors[i].name == id)
+      return i;
+  llvm_unreachable("Unregistered tensor");
+}
+
+bool TCParser::isRegisteredTensor(StringRef id) {
+  for (unsigned i = 0, e = tensors.size(); i < e; ++i)
+    if (tensors[i].name == id)
+      return true;
+  return false;
+}
+
+//===----------------------------------------------------------------------===//
+// TC parsing functions.
+//===----------------------------------------------------------------------===//
+TCParser::TCParser(Parser &p) : parser(p) {}
+
+/// Uses the AffineParser to parse the affine exprs used in a tensor
+/// definition. All identifiers are interpreted as symbols, new symbols are
+/// added eagerly.
+SmallVector<AffineExpr, 4>
+TCParser::parseAffineExprs(EagerDiscoveryMode discoveryMode,
+                           AffineDimList &dims, Token::Kind lDelim,
+                           Token::Kind rDelim) {
+  AffineParser affineParser(
+      parser,
+      [&](StringRef sRef) {
+        AffineExpr expr;
+        if (discoveryMode == EagerDiscoveryMode::Symbols) {
+          expr = getAffineSymbolExpr(symbols.size(), parser.context);
+          symbols.emplace_back(sRef, expr);
+        } else if (discoveryMode == EagerDiscoveryMode::Dimensions) {
+          expr = getAffineDimExpr(dims.size(), parser.context);
+          dims.emplace_back(sRef, expr);
+        }
+        return expr;
+      },
+      dims, symbols);
+  return affineParser.parseAffineExprs(lDelim, rDelim);
+}
+
+/// Parse the information for a tensor def of the form:
+///
+///   affine-expr-list ::= affine-expr (`,` affine-expr )*
+///   tensor-typedef ::= type `(` `)`
+///                    | type `(` affine-expr-list `)`
+///   tensor-def ::= bare-id `:` tensor-typedef
+LogicalResult TCParser::parseTensorDef(bool isOutput) {
+  StringRef tensorId = parser.curToken.getSpelling();
+  if (failed(parser.parseToken(Token::Kind::id, "expected an id")) ||
+      failed(parser.parseToken(Token::Kind::colon, "expected colon")))
+    return failure();
+
+  StringRef tensorType = parser.curToken.getSpelling();
+  if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
+    return failure();
+
+  AffineDimList emptyDims;
+  auto exprs = parseAffineExprs(EagerDiscoveryMode::Symbols, emptyDims);
+  assert(emptyDims.empty() && "Unexpected dimension in tensor def");
+  AffineMap map =
+      AffineMap::get(/*dimCount=*/0, symbols.size(), exprs, parser.context);
+
+  registerTensor(tensorId, tensorType, map, isOutput);
+
+  return success();
+}
+
+/// Parses a tensor use of the form:
+///
+///   affine-expr-list ::= affine-expr (`,` affine-expr)*
+///   tensor-use ::= bare-id `(` `)`
+///                | bare-id `(` affine-expr-list `)`
+LogicalResult TCParser::parseTensorUse(TensorUse &result,
+                                       ComprehensionParsingState &state) {
+  StringRef tensorId = parser.curToken.getSpelling();
+  if (failed(parser.parseToken(Token::Kind::id, "expected an id")))
+    return failure();
+
+  auto exprs = parseAffineExprs(EagerDiscoveryMode::Dimensions, state.dims);
+  AffineMap map =
+      AffineMap::get(state.dims.size(), symbols.size(), exprs, parser.context);
+  LLVM_DEBUG(llvm::dbgs() << "Use of tensor: " << tensorId << " map: " << map
+                          << "\n");
+
+  result = TensorUse(tensorId, map);
+  return success();
+}
+
+/// Parses a tensor expression of the form:
+///
+///   op-spec ::= bare-id `<` reduction-dims-list `>`
+///             | bare-id
+///   op-arg ::= tensor-expr
+///            | tensor-use
+///   op-arg-list ::= op-arg (`,` op-arg)*
+///   tensor-expr ::= op-spec `(` op-arg-list `)`
+LogicalResult TCParser::parseExpression(TensorUse currentDefinition,
+                                        std::unique_ptr<Expression> &result,
+                                        ComprehensionParsingState &state) {
+  StringRef opOrTensor = parser.curToken.getSpelling();
+  if (isRegisteredTensor(opOrTensor)) {
+    TensorUse use;
+    auto res = parseTensorUse(use, state);
+    if (failed(res))
+      return res;
+    result = std::make_unique<TensorUse>(use);
+    return success();
+  }
+
+  if (failed(parser.parseToken(Token::Kind::id, "expected an operation")))
+    return failure();
+
+  // This is an op.
+  SmallVector<unsigned, 4> reductionDims;
+  SmallVector<std::unique_ptr<Expression>, 4> expressions;
+
+  // Check if it has a reduction set, discover dimensions eagerly.
+  if (parser.curToken.is(Token::Kind::lt)) {
+    auto iters = parseAffineExprs(EagerDiscoveryMode::Dimensions, state.dims,
+                                  Token::Kind::lt, Token::Kind::gt);
+    for (auto iter : iters)
+      reductionDims.push_back(iter.cast<AffineDimExpr>().getPosition());
+  }
+
+  // If this op is a reduction, it's first argument is the `currentDefinition`
+  // tensor use.
+  if (!reductionDims.empty())
+    expressions.push_back(std::make_unique<TensorUse>(currentDefinition));
+  LLVM_DEBUG(llvm::dbgs() << "op: " << opOrTensor << "\n");
+
+  auto parseExpr = [&]() -> LogicalResult {
+    std::unique_ptr<Expression> e;
+    if (failed(parseExpression(currentDefinition, e, state)))
+      return failure();
+    expressions.push_back(std::move(e));
+    return success();
+  };
+  if (failed(parser.parseToken(Token::Kind::l_paren, "expected '('")) ||
+      failed(parser.parseCommaSeparatedListUntil(
+          Token::Kind::r_paren, parseExpr, /*allowEmptyList=*/true)))
+    return failure();
+
+  result = std::make_unique<TensorExpr>(opOrTensor, std::move(expressions),
+                                        reductionDims);
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Parse and Emit functions.
+//===----------------------------------------------------------------------===//
+
+/// Parse the information for a single comprehension.
+///
+///   tensor-def-list ::= tensor-def (`,` tensor-def)*
+///   tensor-expr-list ::= tensor-expr (`,` tensor-expr)*
+///   comprehension ::= tensor-def-list `=` tensor-expr-list `;`
+LogicalResult
+TCParser::parseOneComprehension(StringRef cppOpName, StringRef linalgOpName,
+                                ComprehensionParsingState &state) {
+  // 1. Parse LHS of `=`, these become the definitions that appear as the output
+  // tensors or read/write buffers.
+  SmallVector<TensorUse, 4> definitions;
+  auto parseUse = [&]() -> LogicalResult {
+    TensorUse use;
+    if (failed(parseTensorUse(use, state)))
+      return failure();
+    definitions.push_back(use);
+    return success();
+  };
+  if (failed(parser.parseCommaSeparatedListUntil(Token::Kind::equal, parseUse,
+                                                 /*allowEmptyList=*/true)))
+    return failure();
+
+  // 2. Parse RHS of `=`, this becomes the expressions from which we emit
+  // computations.
+  unsigned idx = 0;
+  auto parseExpr = [&]() -> LogicalResult {
+    std::unique_ptr<Expression> expr;
+    if (idx >= definitions.size()) {
+      parser.emitError("Fewer LHS definitions than RHS expressions");
+      return failure();
+    }
+    if (failed(parseExpression(definitions[idx++], expr, state)))
+      return failure();
+    state.expressions.push_back(std::move(expr));
+    return success();
+  };
+  if (failed(parser.parseCommaSeparatedListUntil(
+          Token::Kind::semicolon, parseExpr, /*allowEmptyList=*/true)))
+    return failure();
+  if (idx != definitions.size()) {
+    parser.emitError("Fewer RHS expressions than LHS definitions");
+    return failure();
+  }
+
+  // 3. Postprocess.
+  // 3.a. Normalize all maps to the proper state.dims and symbols counts.
+  SmallVector<TensorUse, 4> allUses;
+  allUses.reserve(tensors.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) {
+    auto &tensor = getRegisteredTensor(def.tensorId);
+    if (seenDefs.count(def.tensorId) > 0) {
+      parser.emitError("Unexpected multi-write to a single tensor");
+      return failure();
+    }
+    seenDefs.insert(def.tensorId);
+    tensor.indexingMap = def.indexingMap;
+    state.orderedTensorArgs[def] = getRegisteredTensorIndex(def.tensorId);
+  }
+
+  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;
+      auto &tensor = getRegisteredTensor(use.tensorId);
+      LLVM_DEBUG(llvm::dbgs()
+                 << "\nuse: " << use.tensorId << " map: " << use.indexingMap);
+      if (tensor.indexingMap && state.orderedTensorArgs.count(use) == 0) {
+        LLVM_DEBUG(llvm::dbgs() << "\nexisting: " << tensor.indexingMap);
+        parser.emitError(
+            "Unexpected multi-read of a tensor with different accesses");
+        failed = true;
+        return;
+      }
+      seenDefs.insert(use.tensorId);
+      tensor.indexingMap = use.indexingMap;
+      state.orderedTensorArgs[use] = getRegisteredTensorIndex(use.tensorId);
+    });
+  if (failed)
+    return failure();
+
+  return success();
+}
+
+/// Parse and print the information for a TC def.
+///
+///   tensor-def-list ::= tensor-def (`,` tensor-def )*
+///
+///   comprehension-list ::= comprehension comprehension*
+///
+///   tc-def ::= `def` bare-id `(`tensor-def-list`)` `->` `(` tensor-def-list`)`
+///     `{` comprehension-list `}`
+///
+/// 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::parseAndEmitTCDef(llvm::raw_ostream &os) {
+  if (failed(parser.parseToken(Token::Kind::kw_def,
+                               "expected 'def' to define a TC")))
+    return failure();
+
+  StringRef tcName = parser.curToken.getSpelling();
+  LLVM_DEBUG(llvm::dbgs() << "\n\nStart parsing tc: " << tcName << "\n");
+  if (failed(parser.parseToken(Token::Kind::id, "expected id")) ||
+      failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
+    return failure();
+
+  auto parseInputDef = [&]() -> LogicalResult {
+    return parseTensorDef(/*isOutput=*/false);
+  };
+  if (failed(parser.parseCommaSeparatedListUntil(
+          Token::Kind::r_paren, parseInputDef, /*allowEmptyList=*/false)))
+    return failure();
+
+  if (failed(parser.parseToken(Token::Kind::minus, "expected '-'")) ||
+      failed(parser.parseToken(Token::Kind::gt, "expected '>'")) ||
+      failed(parser.parseToken(Token::Kind::l_paren, "expected '('")))
+    return failure();
+  auto parseOutputDef = [&]() -> LogicalResult {
+    return parseTensorDef(/*isOutput=*/true);
+  };
+  if (failed(parser.parseCommaSeparatedListUntil(
+          Token::Kind::r_paren, parseOutputDef, /*allowEmptyList=*/false)))
+    return failure();
+
+  // Since we don't declare symbols separately, we discover them eagerly: each
+  // newly encountered id in a tensor shape expression is treated as a new
+  // symbolic. At this point, all tensors have been parsed and all the symbols
+  // that could be discovered eagerly are now known. Resize all AffineMaps to
+  // normalize the number of eagerly discovered symbols.
+  for (auto &tensor : tensors) {
+    auto &map = tensor.shape;
+    map = AffineMap::get(/*dimCount=*/0, symbols.size(), map.getResults(),
+                         parser.context);
+  }
+
+  if (failed(parser.parseToken(Token::Kind::l_brace, "expected '{'")))
+    return failure();
+
+  SmallVector<ComprehensionParsingState, 4> perComprehensionStates;
+  while (parser.curToken.isNot(Token::Kind::r_brace)) {
+    perComprehensionStates.push_back(ComprehensionParsingState());
+    if (failed(parseOneComprehension(tcName, tcName,
+                                     perComprehensionStates.back())))
+      return failure();
+  };
+  parser.parseToken(Token::Kind::r_brace, "expected '}'");
+
+  // Print.
+  auto nComprehensions = perComprehensionStates.size();
+  if (nComprehensions != 1) {
+    parser.emitError("only 1 comprehension supported for now, got: " +
+                     llvm::Twine(nComprehensions));
+    return failure();
+  }
+  if (genODSDecl) {
+    printODS(os, tcName, tcName);
+    os << "\n";
+  }
+  if (genODSImpl) {
+    auto &state = perComprehensionStates.back();
+    std::string extraMethods;
+    llvm::raw_string_ostream ss(extraMethods);
+    printReferenceIterators(ss, tcName, state);
+    printReferenceIndexingMaps(ss, tcName, state);
+    printRegionBuilder(ss, tcName, state);
+    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) {
+  const char *header = R"FMT(  def {0}Op : LinalgNamedStructured_Op<"{1}", [
+    NInputs<{2}>,
+    NOutputs<{3}>,
+    NamedStructuredOpTraits]> {
+      let arguments = (ins Variadic<LinalgOperand>:$views);
+      let results = (outs Variadic<AnyRankedTensor>:$output_tensors);
+      let extraClassDeclaration = [{{
+        llvm::Optional<SmallVector<StringRef, 8>> referenceIterators();
+        llvm::Optional<SmallVector<AffineMap, 8>> referenceIndexingMaps();
+        void regionBuilder(ArrayRef<BlockArgument> args);
+      }];
+      let hasFolder = 1;
+  })FMT";
+
+  unsigned nInputs = 0, nOutputs = 0;
+  for (auto &t : tensors) {
+    if (t.isOutput)
+      nOutputs++;
+    else
+      nInputs++;
+  }
+
+  os << llvm::formatv(header, cppOpName, linalgOpName, nInputs, nOutputs);
+}
+
+/// Print the C++ StructuredOpsInterface impl of `referenceIterators`.
+void TCParser::printReferenceIterators(llvm::raw_ostream &os, StringRef opId,
+                                       ComprehensionParsingState &state) {
+  const char *referenceReferenceIteratorsFmt =
+      R"FMT(
+    llvm::Optional<SmallVector<StringRef, 8>> {0}::referenceIterators() {
+      return SmallVector<StringRef, 8>{{ {1} };
+    })FMT";
+
+  std::string iteratorsStr;
+  llvm::raw_string_ostream ss(iteratorsStr);
+  unsigned pos = 0;
+  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, opId, iteratorsStr);
+}
+
+/// Print the C++ StructuredOpsInterface impl of `referenceIndexingMaps`.
+void TCParser::printReferenceIndexingMaps(llvm::raw_ostream &os, StringRef opId,
+                                          ComprehensionParsingState &state) {
+  const char *referenceIndexingMapsFmt =
+      R"FMT(
+  llvm::Optional<SmallVector<AffineMap, 8>> {0}::referenceIndexingMaps() {
+    MLIRContext *context = getContext();
+    AffineExpr {1};
+    bindDims(context, {1});
+    return SmallVector<AffineMap, 8>{{ {2} };
+  })FMT";
+
+  std::string dimsStr;
+  llvm::raw_string_ostream ss(dimsStr);
+  interleaveComma(state.dims, ss,
+                  [&](std::pair<StringRef, AffineExpr> p) { ss << p.second; });
+  ss.flush();
+
+  std::string mapsStr;
+  llvm::raw_string_ostream mapsStringStream(mapsStr);
+  SmallVector<TensorUse, 4> orderedUses(state.orderedTensorArgs.size());
+  for (auto it : state.orderedTensorArgs)
+    orderedUses[it.second] = it.first;
+  interleaveComma(orderedUses, mapsStringStream, [&](TensorUse u) {
+    assert(u.indexingMap);
+    const char *mapFmt = "\n\tAffineMap::get({0}, 0, {1})";
+    if (u.indexingMap.isEmpty()) {
+      mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), "context");
+      return;
+    }
+
+    std::string exprsStr;
+    llvm::raw_string_ostream exprsStringStream(exprsStr);
+    exprsStringStream << "{";
+    interleaveComma(u.indexingMap.getResults(), exprsStringStream);
+    exprsStringStream << "}";
+    exprsStringStream.flush();
+
+    mapsStringStream << llvm::formatv(mapFmt, state.dims.size(), exprsStr);
+  });
+  mapsStringStream.flush();
+
+  os << llvm::formatv(referenceIndexingMapsFmt, opId, dimsStr, mapsStr);
+}
+
+/// Print the C++ StructuredOpsInterface impl of `regionBuilder`.
+void TCParser::printRegionBuilder(llvm::raw_ostream &os, StringRef opId,
+                                  ComprehensionParsingState &state) {
+  unsigned count = state.orderedTensorArgs.size();
+  llvm::DenseMap<const TensorExpr *, unsigned> subExprsMap;
+  std::function<void(llvm::raw_ostream & os, const Expression &)> printExpr;
+  printExpr = [&](llvm::raw_ostream &os, const Expression &e) -> void {
+    if (auto *pUse = dyn_cast<TensorUse>(&e)) {
+      os << "_" << state.orderedTensorArgs.find(*pUse)->second;
+      return;
+    }
+    auto *pTensorExpr = cast<TensorExpr>(&e);
+    if (subExprsMap.count(pTensorExpr) > 0) {
+      os << "_" << subExprsMap[pTensorExpr];
+    } else {
+      std::string subExprs;
+      llvm::raw_string_ostream subExprsStringStream(subExprs);
+      interleaveComma(pTensorExpr->expressions, subExprsStringStream,
+                      [&](const std::unique_ptr<Expression> &e) {
+                        printExpr(subExprsStringStream, *e);
+                      });
+      subExprsStringStream.flush();
+      const char *tensorExprFmt = "\n    ValueHandle _{0} = {1}({2});";
+      os << llvm::formatv(tensorExprFmt, ++count, pTensorExpr->opId, subExprs);
+      subExprsMap[pTensorExpr] = count;
+    }
+  };
+
+  const char *regionBuilderFmt = R"FMT(
+  void {0}::regionBuilder(ArrayRef<BlockArgument> args) {
+    using namespace edsc;
+    using namespace intrinsics;
+    ValueHandle {1};
+    {2}
+    (linalg_yield(ValueRange{ {3} }));
+  })FMT";
+
+  unsigned idx = 0;
+  std::string valueHandleStr;
+  llvm::raw_string_ostream valueHandleStringStream(valueHandleStr);
+  interleaveComma(state.orderedTensorArgs, valueHandleStringStream, [&](auto) {
+    valueHandleStringStream << "_" << idx << "(args[" << idx << "])";
+    idx++;
+  });
+
+  std::string expressionsStr;
+  llvm::raw_string_ostream expressionStringStream(expressionsStr);
+  for (auto &expr : state.expressions)
+    visitPostorder(*expr, [&](const Expression &e) {
+      if (e.kind == Expression::Kind::TensorExpr)
+        printExpr(expressionStringStream, e);
+    });
+
+  std::string yieldStr;
+  llvm::raw_string_ostream yieldStringStream(yieldStr);
+  interleaveComma(state.expressions, yieldStringStream,
+                  [&](const std::unique_ptr<Expression> &e) {
+                    printExpr(yieldStringStream, *e);
+                  });
+
+  valueHandleStringStream.flush();
+  expressionStringStream.flush();
+  yieldStringStream.flush();
+
+  os << llvm::formatv(regionBuilderFmt, opId, valueHandleStr, expressionsStr,
+                      yieldStr);
+}
+
+/// Iterate over each Tensor Comprehension def.
+LogicalResult parseAndEmitAllTensorComprehensions(llvm::raw_ostream &os,
+                                                  Parser &parser) {
+  while (parser.curToken.getKind() != Token::Kind::eof) {
+    TCParser tcParser(parser);
+    if (failed(tcParser.parseAndEmitTCDef(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);
+  }
+
+  MLIRContext context;
+  llvm::SourceMgr mgr;
+  mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
+  Parser parser(mgr, &context);
+  parseAndEmitAllTensorComprehensions(output->os(), parser);
+  output->keep();
+
+  return 0;
+}