[mlir][linalg] Add linalg_opdsl tool first draft.

[stellaraccident]

Summary

• Mostly imported from experimental repo as-is with cosmetic changes.
• Temporarily left out emission code (for building ops at runtime) to keep review size down.
• Documentation and lit tests added fresh.
• Sample op library that represents current Linalg named ops included.

[mehdi_amini]

I can see how having a TC-like DSL to write nice IR and emit Linalg generic operations is gonna be very nice for our end-to-end story, but it isn't clear to me why the YAML is useful / desirable as part of the publicly exported MLIR python bindings?

mlir/docs/Dialects/Linalg.md

678 ↗

(On Diff #328345)

I believe these is an important layering with a strong separation between this tool and python niceties, and the "core" of MLIR, including the Linalg dialect.
I.e. IMO the tool should not be part of the dialect documentation but have a separate page somewhere, but not under dialect on the website.

mlir/lib/Bindings/Python/CMakeLists.txt

28 ↗

(On Diff #328345)

Can you commit these changes ahead of this patch? They seem NFC and unrelated (right?)

[mehdi_amini]

In D97995#2605203, @mehdi_amini wrote:

I can see how having a TC-like DSL to write nice IR and emit Linalg generic operations is gonna be very nice for our end-to-end story, but it isn't clear to me why the YAML is useful / desirable as part of the publicly exported MLIR python bindings?

I suspect there are pieces of the design that aren't clear to me at the moment, in particular the setup looks monolithic to me (some comment inline may hint at this as well). From what I gather in terms of "independent components", the coupling between the tool and the bindings / native code isn't obvious to me. It seems like the DSL can be a "pure" python library, the tool is mostly about using the DSL and emitting YAML, and the bindings should be mostly bridging the DSL to emit IR.

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

s/build/built/ ?

Why do we need to build MLIR to be able to use this language and the tool?
Isn't the language suitable to be a pure-python library? And the tool could just use it directly without building any native code?

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py

30

The doc here makes it look like a general nice frontend for AffineExpr, isn't it suitable in mlir.ir? Or as mlir.dialects.affine?

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py

5

"a tensor-comprehension" is something that deserve to be explained / documented, I am not sure it is a widely known concept, or at least not obvious enough in the context of MLIR to be used as-is I think.

188

Thanks for providing types as comment here by the way! :)
(as non-seasoned python programmer, not having type is always a pita)

[stellaraccident]

In D97995#2605203, @mehdi_amini wrote:

I can see how having a TC-like DSL to write nice IR and emit Linalg generic operations is gonna be very nice for our end-to-end story, but it isn't clear to me why the YAML is useful / desirable as part of the publicly exported MLIR python bindings?

A couple of points. First, the way this is intertwined in the Python code right now is not great, creating a hard dependency on YAML for using any of it. I had started to, but backed off in a first patch, factoring it to an optional serializer and make the core just about export to lists/dicts and such. That is mechanics and I'll separate it some more in a future patch so that it becomes an optional export path from these definitions.

As for utility, I've been iterating on this for a few weeks, and I think the utility of YAML as an export format is high:

• Tools like this work better with a human readable intermediate form. One of the issues with the prior tool flow was that it did everything at once, with a lot of magic in the processing of the DSL and no way to inspect what it was actually declaring without looking at generated C++ code. When I started this code, I was just exporting to a repr form for lit testing and found it to lack a structure that made it easy to discern exactly what the final result was doing. What we have here produces a high "glancability" index to easily discern whether what was heard by the tool was what you intended to say. That has obvious value for testing but has been pretty valuable for developing it too.
• We actually had a non core kernel engineer hacking on some of these ops over the last week or so, and he found that a) Handwriting the YAML op specs was better than what we had before and he had made independent progress working directly on that, and b) When I showed him these forms, he immediately got it and wanted to use it as a tool to author the op specs. The YAML form in LinAlg is good for reading/diffing/at rest, but the Python form is better for authoring -- especially for those further from the compiler domain. I've got ample feedback on that point.
• Within the LLVM ecosystem, YAML is already well supported on the C++ side. Leveraging that and making our Python tools interop through an already well supported, readable interchange mechanism strictly increases value. I believe strongly that a tool like this needs to have an externalized form, and why shouldn't it be something that interops well and aids testing?
• In the domain, this is not an uncommon use: Coming, for example, from PyTorch, they often represent fanout of a bunch of math definitions through YAML intermediate forms that are code generation amenable. It's just how things are increasingly done because it is easy/readable, and if your higher level DSL needs to be replaced/overhauled, you still have an asset that you can use/extend/diff/transform.

Also, for the record, I mean it when I say that I don't believe this should ever get promoted into the build system but remain as a runtime tool and development aid for building MLIR based compilers: It is much safer to do code reviews on the distilled YAML form and not leave critical op definition work strictly to a less constrained tool like this (which will require its own debugging, development, etc). We can put the last mile serialization in some other place in the codebase (not part of the API), but it will reduce testability. Once I fix the mechanics, it'll be one pure Python file and it is better to keep it together.

[stellaraccident]

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

It uses primitives like AffineMaps and AffineExpressions as is, and in the IR generation mode (not yet committed), it is pretty tied to contexts and ops. I originally tried to keep those separate to preserve a pure python version, but I eventually decided that it was nicer to just presume the presence of the MLIR API and not try to build something abstracted against what already exists and that we support. The merit is in pulling things together, not in building another disconnected thing.

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py

30

+1 but would like to do that as a followup. The current Python API package structure needs some evolution... I struggled a good deal with what the package layout should be of the "library" side of the API, and decided that it was better to have some solid examples like this in tree, embedded in a tool, and then promote useful pieces out to the core API. It'll be easier to get a good feel of the package structure that way.

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py

188

The types are all valid, btw (MyPy requires you to add these). I just couldn't figure out an easy way to get MyPy enabled in the MLIR repo in one step (it is in the presubmit of the side repo where this was developed).

[stellaraccident]

mlir/docs/Dialects/Linalg.md

678 ↗

(On Diff #328345)

Moved the docs to Tools/LinalgOpDsl.md. I haven't touched the website building and suspect I need to update an index/nav file somewhere (either in this change or a followon).

mlir/lib/Bindings/Python/CMakeLists.txt

28 ↗

(On Diff #328345)

Will do.

[stellaraccident]

Comments.

[stellaraccident]

Remove cmake changes.

[mehdi_amini]

In D97995#2605252, @stellaraccident wrote:

In D97995#2605203, @mehdi_amini wrote:

I can see how having a TC-like DSL to write nice IR and emit Linalg generic operations is gonna be very nice for our end-to-end story, but it isn't clear to me why the YAML is useful / desirable as part of the publicly exported MLIR python bindings?

A couple of points. First, the way this is intertwined in the Python code right now is not great, creating a hard dependency on YAML for using any of it. I had started to, but backed off in a first patch, factoring it to an optional serializer and make the core just about export to lists/dicts and such. That is mechanics and I'll separate it some more in a future patch so that it becomes an optional export path from these definitions.

As for utility, I've been iterating on this for a few weeks, and I think the utility of YAML as an export format is high:
[...]

I'm not sure I followed everything you write here, so I'm not entirely sure we're focusing on the same aspect at the moment, but I may just be behind on the design and the direction as well.
You seem here to argument in favor of using YAML, but the context isn't clear to me. Since you're talking about the "prior tool flow", for which the YAML has already been committed as a replacement, I don't think there is much controversy there: using YAML as the input to the build system for Linalg seems fine to me.

That is however very different from the public python bindings, i.e. I am surprised that import mlir would publicly exposed a tool that is used to generate YAML used to build MLIR itself.
On the other hand, import mlir exposing this nice authoring language and emitting MLIR IR (linalg generic) seems very natural to me.

[mehdi_amini]

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

(missed this one since phab does not show it in the current diff)

in the IR generation mode (not yet committed), it is pretty tied to contexts and ops.

Right, as I tried to explain elsewhere, the "IR generation mode" clearly belongs to the public python API.
I'm not convinced that the *tool* that is used to generate YAML does, i.e. I'm skeptical about the setup with mlir.tools.* as is.

[stellaraccident]

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

Gotcha - so two things that might help.

1. I'd like to treat this current location as temporary. I actually got hung up putting thought into how to layer this into the public API but was having trouble seeing some reorganizations that need to happen there and decided to opt for something isolated to at least get the code landed and open up some collaborations: I couldn't quite see all of the moving parts yet and the implications on packaging and distribution are a bit subtle (also, I'd like to exercise a bit of care in evolving the "public" API).

2. I think that having an entrypoint for the tool in the API is fine/called-for. The use is ultimately not just for generating in-tree op definitions within MLIR-core but in defining special purpose ops in other projects that have a need (i.e. PyTorch has ~1500 ops, many of them representable by Linalg, and many of those we likely will end up building in as named ops). Beyond the basics that we include in MLIR-core, downstream projects will likely have their own sets of named ops, and having easy access to the tools and libraries is nice (and cheap). Ditto even for backends like IREE which will end up with special purpose ops that match specific hardware and deployment scenarios. I don't want to fragment the MLIR Python APIs in to a bunch of small, dependent pieces when it is very incremental to add some more support code and utilities. The coarseness to expect from these kinds of packages is pretty different compared to C++.

With respect to the overall organization of the mlir namespace, I want to keep this piece isolated for now and design that properly. Right now, the namespace is really not organized for much beyond being a place for shims and helpers on top of the native libraries: it needs some more care to lay it out to house more general purpose libraries at a couple of levels, and I prefer to do that explicitly. Having this stuff off to the side and out of the way will help to see the right way to organize the overall namespace since it is easier to organize things when you can see some of the pieces that need to fit.

I think I've demonstrated some good stewardship of this part of the codebase and would like some leeway to stage practical steps like this a bit incrementally.

[nicolasvasilache]

Thanks for pushing this forward @stellaraccident!

I'd characterize the current comment / discussion as the "layering needs to evolve", which is best served by having this landed and iterated upon by a distributed group of people.

Looking forward to playing with this.

mlir/docs/Dialects/Linalg.md

759 ↗

(On Diff #328345)

Note for self: do we verify early and aggressively?
In the past, late type binding has been annoying to debug.

770 ↗

(On Diff #328345)

Note for self: I assume we dump the linalg.generic body and provide a reasonable error message on failure ?

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py

120

FMI, what does this do? I don't see any setter/getter/deleter

260

nit: should the global variables exported here be grouped together at the end of the file for ease of tracking?

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py

5

@stellaraccident seems easy to define the term "tensor-comprehension" in the #### Language Guide as a line in a TensorDef that generates a single linalg.generic?
Saying it's loosely inspired from https://arxiv.org/abs/1802.04730 to represent linalg structured ops and it is evolving should be enough context for now.

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/config.py

10

+1 finding the good balance is an iterative process.

[stellaraccident]

Rebase.

[stellaraccident]

Comments.

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py

120

"local" here refers to symbols and dims as part of this build action. "global" accounts for if you created this AffineBuildState subordinate to another, which is what we do when constructing affine maps across a comprehension which share the symbol and dim space.

Added a class comment.

(I want to rework this/the API is awkward - it kind of grew out of the code you had in the original C++ TC lang parser)

[stellaraccident]

Thanks for the reviews. I generally agree with Mehdi's instincts in the fullness of time with respect to namespace organization and layering, but there are a number of things that need to be done better to strike the right balance. It is my opinion that keeping this isolated as we do that work is in the best interests of both iteration/progress and being purposeful how we expand the intended-to-be-public parts of the API. We'll get there - just takes time.

[mehdi_amini]

In general, can you please avoid landing patches while there are discussions that haven't settled?
We don't mark explicitly revision as "request changes" when we comment, but that's somehow implicit to make sure resolve these before landing a patch: i.e. I don't consider that getting Nicolas' approval in this context was enough.

[nicolasvasilache]

In D97995#2607525, @mehdi_amini wrote:

In general, can you please avoid landing patches while there are discussions that haven't settled?
We don't mark explicitly revision as "request changes" when we comment, but that's somehow implicit to make sure resolve these before landing a patch: i.e. I don't consider that getting Nicolas' approval in this context was enough.

How can we strike a better discussion balance to ensure that "done is better than perfect" and that we can improve incrementally rather than have to wait for the monolith to be shiny before unblocking others ?

[rriddle]

In D97995#2607525, @mehdi_amini wrote:

In general, can you please avoid landing patches while there are discussions that haven't settled?
We don't mark explicitly revision as "request changes" when we comment, but that's somehow implicit to make sure resolve these before landing a patch: i.e. I don't consider that getting Nicolas' approval in this context was enough.

+1. While only one reviewer is required to approve before a patch can land, it is imperative that a reasonable consensus has been reached across all reviewers. If it isn't clear if a consensus has been reached, you should wait a bit for acknowledgement (especially if the discussion is still ongoing).

[mehdi_amini]

In D97995#2607540, @nicolasvasilache wrote:

How can we strike a better discussion balance to ensure that "done is better than perfect" and that we can improve incrementally rather than have to wait for the monolith to be shiny before unblocking others ?

I'm not sure what you mean, I see your characterization as totally off here, you may need to elaborate. Note that this patch landed less than 24h after being posted, and it isn't like "a trivial" fix or a small change.
A key requirement to be able to move incrementally requires to discuss and agree on the general direction. I missed it but I haven't seen a single public discussion about any of this (other than on IREE chat room, which obviously can't really be a reference for upstream work).

Regardless, there is a cost to working upstream: it is a core part of what makes the quality of the project, I have a very strong believe in the culture that is set there. This was also one of the main motivation to develop MLIR inside LLVM.

[mehdi_amini]

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

I'd like to treat this current location as temporary

OK!
Could we put this one level higher into "mlir.linalg_opdsl"" instead? The mlir.tools namespace does not make it clearly temporary, and worse kind of "bless" new "tools" to get there.

I think that having an entrypoint for the tool in the API is fine/called-for. The use is ultimately not just for generating in-tree op definitions within MLIR-core but in defining special purpose ops in other projects that have a need

This is fair, but that still does not mean it should be exposed in the public python API. Anyone who has a need to build native named op would have to integrate with the build system, at which point there have our /bin/ available (just like they have mlir-tblgen available). So bin/linalg-dsl-to-yaml is a very reasonable tool entry point to have. And even if this tool is written in python, and is using MLIR python APIs (to reuse the DSL for example), that still does not call for a public mlir.tools as far as I can tell.

So: I don't disagree with any of your use-case, I just haven't perceived yet why exposing mlir.tools is neither needed or desirable to achieve them.

With respect to the overall organization of the mlir namespace, I want to keep this piece isolated for now and design that properly. [...] I think I've demonstrated some good stewardship of this part of the codebase and would like some leeway to stage practical steps like this a bit incrementally.

Absolutely, I'm fine with incrementally develop and improve this. However this works well when we build some shared understanding and agreement of the direction first. And right now I feel this whole track (YAML first, and DSL now) felt a short on this: looks like there has been a missed step between the downstream discussions and prototyping and the landing upstream.

[stellaraccident]

In D97995#2607558, @mehdi_amini wrote:

In D97995#2607540, @nicolasvasilache wrote:

How can we strike a better discussion balance to ensure that "done is better than perfect" and that we can improve incrementally rather than have to wait for the monolith to be shiny before unblocking others ?

I'm not sure what you mean, I see your characterization as totally off here, you may need to elaborate. Note that this patch landed less than 24h after being posted, and it isn't like "a trivial" fix or a small change.
A key requirement to be able to move incrementally requires to discuss and agree on the general direction. I missed it but I haven't seen a single public discussion about any of this (other than on IREE chat room, which obviously can't really be a reference for upstream work).

Regardless, there is a cost to working upstream: it is a core part of what makes the quality of the project, I have a very strong believe in the culture that is set there. This was also one of the main motivation to develop MLIR inside LLVM.

Probably for a longer discussion off thread, but I understand perfectly what Nicolas is referring to. I roughly slow my velocity down by 5-10x over what it would naturally be in order to stage and telegraph feature advancements that would not be controversial and would be obvious if done in a domain-aligned sub-project but have a very high consensus overhead in the general purpose MLIR project as a whole. I think there are a bunch of reasons for that, some good and some legacy, but in my experience, it is not the norm in our other MLIR derived sub projects (two of which are LLVM projects themselves). I generally feel that the core MLIR project is trying to do too many things (as evidenced by these kind of "there's no place to put the tool I need to do my job" kind of issues), especially as it straddles the line between general purpose, mid-level IR infra, and user-centric Tensor Compute domain abstractions. I think that both of those have a strong place upstream, but if I had it to do over again, I would have made them two separate projects. That still may be the right answer, either by directory structure or convention.

In this case, I made an explicit choice on the reviewer: Nicolas is the right person to be both providing the technical and roadmap vision for this component. Further, where it is placed, in the python tree (complete with caveats and tradeoffs) was literally designed and mostly implemented by me, and I have a lot of state on what can/should go where while we are expanding it. I was annoyed, given the timezones, that he was literally asleep while the debate about structure took place -- when what I was looking for was a review of the merits. I did let that annoyance come in to play when I re-read the threads this morning and decided that we were actually converged. In my defense, there is an extraordinary amount of debate about structure that goes on around here, and I have trouble determining what is drive-by vs substance from indirect stakeholders.

I'll revert this patch if this is deemed an issue worthy of more discussion.

[stellaraccident]

mlir/docs/Dialects/Linalg.md

774 ↗

(On Diff #328345)

Could we put this one level higher into "mlir.linalg_opdsl"" instead? The mlir.tools namespace does not make it clearly temporary, and worse kind of "bless" new "tools" to get there.

Yes. And I will also take two followups:

1. Start to re-organize the python tree into namespace packages. As the tree is organized now, it is actually not possible to split the namespace into different directories: we have one place we put namespaced python code, and that corresponds with the packages we produce (still in an out of tree repo because there has been no time/volunteers to move it in). As a user who has used this stuff in a variety of ways, I do question the utility of actually splitting (i.e. in to top level tools folders), but we should have the option and today we do not due to limitations of python's package setup (if you don't use namespace packages).
2. Start a thread on the organization of the python directory tree(s) and dist packages. It's been on my mind for a while, but it is essentially just me trying to drive the esoteric bits for a lot of this stuff and it has been low on my priority list to make more work for myself while also trying to hit a point where this stuff does anything useful at all (this has been discussed on discourse, and I assume that everyone is in the same boat I am and biasing towards work that enables e2e utility vs things that are merely slowing us down a bit).

[mehdi_amini]

I don't think it is worthy of a revert here. My intent is mostly to set the expectation to avoid future frictions.
Here I disagree with your assessment on the reviewer selection and approval process. MLIR is moving slower than downstream project, and that is fine. The review time on a patch here does not prevent you from continuing to make progress in a downstream fork. Note that LLVM can be moving much much slower as well.

So please for such a patch in the future: wait for at least a few days before landing upstream to give a chance to get a wider look. I'd also like to see more RFC on Discourse to set the direction and involve the community.

We have good precedent of "getting agreement on the direction and implementing incrementally", for example the C API or the Python bindings. We also have very good examples of what happens when the process is shortcut, like the original python bindings for example. ;)

[stellaraccident]

In D97995#2608158, @mehdi_amini wrote:

I don't think it is worthy of a revert here. My intent is mostly to set the expectation to avoid future frictions.

Thanks.

Here I disagree with your assessment on the reviewer selection and approval process. MLIR is moving slower than downstream project, and that is fine. The review time on a patch here does not prevent you from continuing to make progress in a downstream fork. Note that LLVM can be moving much much slower as well.

It's too bad that it will be a while before we have real f2f community meetups again. I am going to disagree, and I feel that there is an actual real dynamics and project scoping issue here, but it is the kind of thing better worked out f2f with some empathy vs cross continents while our family members are texting us asking when we are going to make them lunch :)

So please for such a patch in the future: wait for at least a few days before landing upstream to give a chance to get a wider look. I'd also like to see more RFC on Discourse to set the direction and involve the community.

We have good precedent of "getting agreement on the direction and implementing incrementally", for example the C API or the Python bindings. We also have very good examples of what happens when the process is shortcut, like the original python bindings for example. ;)

I think we should also have a good community discussion on that one at some point. I agree that the outcome is showing early signs of being good, but I would hesitate to hold that up as a model for future work. I took that up personally because of a failure to converge on a starting point and implementation plan for quite a long time -- and I simply couldn't wait any longer to have a real thing to build user-centric solutions out of. I agree that once it got underway, we got the job done... quite a bit more slowly than I would have expected, but it got there. Slow is one thing, but I think that the community bugs, though, are in the activation energy and force of will required to start things or push new directions. Many try, but very few people are being successful at picking up the initial velocity. And that's a real problem, because MLIR still has quite a way to go to be a generally useful piece of infra for the scale of problems we all want it to solve, and a lot of the gaps are in things that still need to be envisioned/invented.

Also, say what you will of the original python bindings: they served a useful purpose, and I am aware of at least one project that is still using them to good effect, and they were a key to them scaling their work to production.

[mehdi_amini]

In D97995#2608196, @stellaraccident wrote:

[...] And that's a real problem, because MLIR still has quite a way to go to be a generally useful piece of infra for the scale of problems we all want it to solve, and a lot of the gaps are in things that still need to be envisioned/invented.

My strong belief is that a project like this one can't scale or achieve its ambition without proper care and alignment in the development: we may share the same goal here, but we may also have divergent opinion on some aspects that I consider critical for a multi-parties open-source project like this one.

The energy required to do things are inherent at some point in a project like MLIR. I believe it is much worse in LLVM itself actually. Many users of LLVM have a fork and never contribute upstream, most of my work on LLVM has been organized as a pipeline: I always prototyped in a fork, with ~ 10 patches ahead, cleaning up the oldest ones and sending them for reviews, expecting the review process to take multiple days per patch. That never blocked me from progressing in my fork (even when this fork was Apple internal copy of LLVM downstream). Many features actually landed in production inside some Apple products using LLVM before trying to clean it up, refactor it, and send it to LLVM (with a discussion on llvm-dev@ as necessary).

The impression I get from this track of work reminds me of other examples I don't really want to replicate in MLIR, like how TensorFlow is developed internally: discussion, brainstorm, design, and direction are hidden from the community and when the code is there, it's already "pre-approved" and the review is really just rubber stamping.

By the way, I rather have an RFC discussed and agree upon, and skip pre-commit code reviews on the patches (but following the coding standards and being responsive on post-commit comments!). This is more in line with my experience of "moving fast incrementally" practiced in LLVM for example.

Also, say what you will of the original python bindings: they served a useful purpose, and I am aware of at least one project that is still using them to good effect, and they were a key to them scaling their work to production.

There are many projects that are "serving a purpose" and are "useful in production" that don't belong to MLIR, and I'm perfectly fine with that. But I don't have to work on them or maintain them either.
In this particular case: while these bindings are serving a purpose, they will serve the purpose of a single user and won't generalize. If a user of MLIR wants something ad-hoc to their need, then developing it out-of-tree makes sense. When it's time to upstream, it has to generalize further.

There is a huge cost to refactor and evolve things. I always believed that managing this carefully was important for the health of such a project, and I applied this in LLVM for years as much as possible. But this looks even more critical to me now that I know about Google: a soon a some code is written in MLIR and adopted inside google, the cost of improving/refactoring/evolving it is prohibitive for any Googler (but River, for now...). It can quickly spiral to the point where the only hope of evolution for MLIR will come from non-google contributors...
I have frequent reminders of this: most recently I was trying to work with the ExecutionEngine runtime and I regretted every day not pushing back more on the patches that landed there! I actually gave up for now on many of the refactoring I was trying to conduct there. This another area of the codebase that "is serving a purpose" for someone, but that may just be better not in the monorepo in this form.

[stellaraccident]

In D97995#2608307, @mehdi_amini wrote:

In D97995#2608196, @stellaraccident wrote:

[...] And that's a real problem, because MLIR still has quite a way to go to be a generally useful piece of infra for the scale of problems we all want it to solve, and a lot of the gaps are in things that still need to be envisioned/invented.

My strong belief is that a project like this one can't scale or achieve its ambition without proper care and alignment in the development: we may share the same goal here, but we may also have divergent opinion on some aspects that I consider critical for a multi-parties open-source project like this one.

You and I almost always agree on the destination, as in this case. I think there may be some practical differences on that path...

The energy required to do things are inherent at some point in a project like MLIR. I believe it is much worse in LLVM itself actually. Many users of LLVM have a fork and never contribute upstream, most of my work on LLVM has been organized as a pipeline: I always prototyped in a fork, with ~ 10 patches ahead, cleaning up the oldest ones and sending them for reviews, expecting the review process to take multiple days per patch. That never blocked me from progressing in my fork (even when this fork was Apple internal copy of LLVM downstream). Many features actually landed in production inside some Apple products using LLVM before trying to clean it up, refactor it, and send it to LLVM (with a discussion on llvm-dev@ as necessary).

When I'm working on my own or in a silo, that is my preferred operation style as well. I think our issue right now is that we are trying to synchronize the work of ~a dozen people on this aspect (not just at Google/IREE but a pretty good ragtag crew of people all trying to use this stuff for the first time across companies and projects), and you hit these clench point items that are causing a pipeline stall across individuals. Given the sprawling nature of the projects, forks, etc, it gets gross pretty quickly trying to synchronize people like that when you find a critical gap upstream.

The impression I get from this track of work reminds me of other examples I don't really want to replicate in MLIR, like how TensorFlow is developed internally: discussion, brainstorm, design, and direction are hidden from the community and when the code is there, it's already "pre-approved" and the review is really just rubber stamping.

I'll try to give more visibility on the lists next time around.

By the way, I rather have an RFC discussed and agree upon, and skip pre-commit code reviews on the patches (but following the coding standards and being responsive on post-commit comments!). This is more in line with my experience of "moving fast incrementally" practiced in LLVM for example.

Also, say what you will of the original python bindings: they served a useful purpose, and I am aware of at least one project that is still using them to good effect, and they were a key to them scaling their work to production.

There are many projects that are "serving a purpose" and are "useful in production" that don't belong to MLIR, and I'm perfectly fine with that. But I don't have to work on them or maintain them either.
In this particular case: while these bindings are serving a purpose, they will serve the purpose of a single user and won't generalize. If a user of MLIR wants something ad-hoc to their need, then developing it out-of-tree makes sense. When it's time to upstream, it has to generalize further.

There is a huge cost to refactor and evolve things. I always believed that managing this carefully was important for the health of such a project, and I applied this in LLVM for years as much as possible. But this looks even more critical to me now that I know about Google: a soon a some code is written in MLIR and adopted inside google, the cost of improving/refactoring/evolving it is prohibitive for any Googler (but River, for now...). It can quickly spiral to the point where the only hope of evolution for MLIR will come from non-google contributors...

I probably shouldn't be saying this but I'm pretty sure that the way that Google is managing at-head development of LLVM is going to just grind to a halt at some point. On my side of the fence, we are very reticent to contribute more code to Google's internal repo because of this force. New features get developed entirely on GitHub to the extent possible and then only promoted once/if mature/needed and there is no other way to do it. If it continues on the path it is for ~years more, I estimate that I'll end up being one of those non-Google contributors.

I have frequent reminders of this: most recently I was trying to work with the ExecutionEngine runtime and I regretted every day not pushing back more on the patches that landed there! I actually gave up for now on many of the refactoring I was trying to conduct there. This another area of the codebase that "is serving a purpose" for someone, but that may just be better not in the monorepo in this form.

[stellaraccident]

I took the afternoon and worked out how to open up the python directory tree/namespace and moved things into a more natural location: https://reviews.llvm.org/D98096

[stellaraccident]

Mehdi: As I mentioned privately in chat, sorry for the drama. Have a good weekend.

mlir/docs/Tools/LinalgOpDsl.md

+# linalg_opdsl tool
+
+Python based DSL for authoring Linalg op definitions and generating
+`linalg.generic` IR based on them for samples.
+
+The tool `linalg_opdsl` provides a high level DSL for constructing
+structured op definitions in a way that can be exported to built-in, named
+structured ops via the above YAML-based definitions or used interactively to
+emit corresponding `linalg.generic` IR for the composition.
+
+## Basic usage
+
+The tool is bundled with the MLIR Python bindings. To use from the CMake build
+tree, MLIR must be build with Python bindings enabled
+(`-DMLIR_BINDINGS_PYTHON_ENABLED=ON`). Then add the `python` directory in the
+build tree to your `PYTHONPATH` environment variable (i.e.
+`export PYTHONPATH=$PWD/build/python`). Optionally, use an installed MLIR
+package, if available, to avoid building.
+
+```shell
+# Dump the `core_named_ops.py` module as YAML.
+python -m python -m mlir.tools.linalg_opdsl.dump_oplib .ops.core_named_ops
+```
+
+The tool is meant for use during both development and runtime, but not as
+a build tool of the core compiler: in order to export static named op
+definitions to be built as part of the compiler, the corresponding Linalg
+dialect YAML file must be updated and reviewed. TODO: Develop a script to
+automate op updates to these files.
+
+## Language Guide
+
+The language presented here is loosely inspired from the
+[Tensor Comprehensions](https://arxiv.org/pdf/1802.04730.pdf) work, adapted to
+represent linalg structured ops.
+
+This tool is new and rapidly evolving. For language examples, refer to the
+built-in ops in the `mlir.tools.linalg_opdsl.ops` package
+(`lib/Bindings/Python/mlir/tools/linalg_opdsl/ops` in the repository).
+
+Using a matmul as an example, we will decompose the language:
+
+```python
+T1 = TV.T1
+T2 = TV.T2
+
+@linalg_structured_op
+def matmul(A=TensorDef(T1, S.M, S.K),
+           B=TensorDef(T2, S.K, S.N),
+           C=TensorDef(U, S.M, S.N, output=True)):
+  """Performs a matrix multiplacation of two 2D inputs.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
+```
+
+Here we have a simple type polymorphic contraction that takes arguments `A`
+and `B` and outputs `C`. Each is bound to a `TensorDef`, which specifies:
+
+* The symbolic element type (`T1`, `T2`, `U` above).
+* Symbolic shape expressions with symbols that are bound globally for the op (
+note that in this simple example, the shape expressions are just symbol
+references, but they are permitted to be a constrained set of affine
+expressions).
+* Usage (`output=True`).
+
+The docstring will be transferred to the op definition verbatim.
+
+Special identifying op interfaces can be declared for the op via
+`implements(interface1[, interface2...])`.
+
+## Assignments
+
+The bulk of language consists of assignment expressions of the form above.
+The iteration dimension order is determined lexically based on the order
+encountered in the expression (following operator precedence if math operators
+are used). TODO: Introduce a directive to fix the dimension bindings.
+
+Reduction dimensions are inferred to be any dimensions on the RHS that are not
+on the LHS.
+
+A number of arithmetic primitive functions are supported:
+
+* `PrimFn.add(a, b)` (also via overloading the binary `+` operator)
+* `PrimFn.exp(a)`
+* `PrimFn.log(a)`
+* `PrimFn.mul(a, b)` (also via overloading the binary `*` operator)
+* `PrimFn.max(a, b)`
+* `PrimFn.sub(a, b)` (also via overloading the binary `-` operator)
+
+Reduction functions can appear as the outer-most function on the RHS:
+
+* `ReduceFn.add` (also overloading the inplace `+=` on a LHS)
+* `ReduceFn.mul`
+* `ReduceFn.max`
+
+There are also special forms:
+
+* `cast(TypeVar, operand)`
+
+## Types
+
+All types in assignment expressions are late bound based on actual input
+and output types of constructed ops. Assignment expressions with no `cast`
+calls will generally require uniform types throughout and will fail to
+verify if violated. The presence of a `cast` allows for a limited form of
+numeric type conversion between element types that can be derived from inputs
+and outputs (and in the future, attributes). `cast` calls with a `TypeVar`
+first argument are emitted as `symbolic_cast` primitives in the YAML definition.
+
+Casting will perform `int<->float` type conversions and will perform any
+necessary extension or truncation within type family. Note that presently,
+any integer type is assumed to be signed for the purpose of determing how to
+extend or truncate. Supporting unsigned integer types is left for future work.
+
+Not all functions are applicable for all numeric types, and on mismatch, op
+verification will fail.

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/dump_oplib.py

+#!/usr/bin/which python
+# Command line tool to load an oplib module and dump all of the operations
+# it contains in some format.
+"""Loads one or more modules containing op definitions and dumps them.
+
+The dump format can be:
+
+* `--dump_format=yaml` (default)
+* `--dump_format=repr`
+
+Positional arguments are interpreted as module names (optionally, relative to
+this module). Loose module files can be specified via `--file <filepath>`.
+
+Sample usage:
+  # Dump the YAML op definitions for the core named ops (as in the dialect
+  # source tree).
+  python -m mlir.tools.linalg_opdsl.dump_oplib .ops.core_named_ops
+
+Note: YAML output is emitted in "document list" format with each operation
+as its own "document". Practically, this means that each operation (or group
+of composite ops) is emitted with a "---" preceding it, which can be useful
+for testing.
+"""
+
+import argparse
+import importlib
+
+from .lang import *
+from .lang.config import *
+from .lang.yaml_helper import *
+
+
+def create_arg_parser() -> argparse.ArgumentParser:
+  p = argparse.ArgumentParser(description="Dump an oplib in various formats")
+  p.add_argument("modules",
+                 metavar="M",
+                 type=str,
+                 nargs="*",
+                 help="Op module to dump")
+  p.add_argument("--file",
+                 metavar="F",
+                 type=str,
+                 nargs="*",
+                 help="Python op file to dump")
+  p.add_argument("--format",
+                 type=str,
+                 dest="format",
+                 default="yaml",
+                 choices=("yaml", "repr"),
+                 help="Format in which to dump")
+  return p
+
+
+def load_module_from_file(module_name, file_path):
+  spec = importlib.util.spec_from_file_location(module_name, file_path)
+  m = importlib.util.module_from_spec(spec)
+  spec.loader.exec_module(m)
+  return m
+
+
+def main(args):
+  # Load all configs.
+  configs = []
+  modules = []
+  for module_name in args.modules:
+    modules.append(
+        importlib.import_module(module_name, package="mlir.tools.linalg_opdsl"))
+  for i, file_path in enumerate(args.file or []):
+    modules.append(load_module_from_file(f"_mlir_eval_oplib{i}", file_path))
+  for m in modules:
+    for attr_name, value in m.__dict__.items():
+      # TODO: This class layering is awkward.
+      if isinstance(value, DefinedOpCallable):
+        try:
+          linalg_config = LinalgOpConfig.from_linalg_op_def(value.model)
+        except Exception as e:
+          raise ValueError(
+              f"Could not create LinalgOpConfig from {value.model}") from e
+        configs.extend(linalg_config)
+
+  # Print.
+  if args.format == "yaml":
+    print(yaml_dump_all(configs))
+  elif args.format == "repr":
+    for config in configs:
+      print(repr(config))
+
+
+if __name__ == "__main__":
+  main(create_arg_parser().parse_args())

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/__init__.py

+from .dsl import *

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/affine.py

+#  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
+"""DSL for constructing affine expressions and maps.
+
+These python wrappers allow construction of affine expressions in a more
+pythonic fashion that is later instantiated as an IR AffineExpr. Separating the
+AST from construction of the map allows for manipulations of symbols and dims
+beyond the scope of one expression.
+
+Affine expression construction:
+  >>> with _ir.Context():
+  ...   s = AffineBuildState()
+  ...   (S.K + S.M).build(s)
+  ...   (S.K * S.M).build(s)
+  ...   (S.K // S.M).build(s)
+  ...   (S.K / S.M).build(s)
+  ...   (S.K % 4).build(s)
+  ...   (D.i + D.j * 4).build(s)
+  ...   s
+  AffineExpr(s0 + s1)
+  AffineExpr(s0 * s1)
+  AffineExpr(s0 floordiv s1)
+  AffineExpr(s0 ceildiv s1)
+  AffineExpr(s0 mod 4)
+  AffineExpr(d0 + d1 * 4)
+  AffineBuildState<
+    symbols={'K': 0, 'M': 1}
+    dims={'i': 0, 'j': 1}>
+
+In the DSL, dimensions and symbols are name-uniqued instances of DimDef and
+SymbolDef. There are shortcut "expando" instances that will create a
+corresponding DimDef/SymbolDef upon accessing an attribute:
+
+Referencing a named dimension:
+
+  >>> D.i
+  Dim(i)
+  >>> D.a is D.b
+  False
+  >>> D.a is D.a
+  True
+
+Referencing a named symbol:
+
+  >>> S.foobar
+  Symbol(foobar)
+  >>> S.a is S.b
+  False
+  >>> S.a is S.a
+  True
+"""
+
+from typing import Callable, Dict, Optional, Tuple, Union
+
+from mlir import ir as _ir
+
+__all__ = [
+    "AffineBuildState",
+    "AffineExprDef",
+    "D",
+    "DimDef",
+    "S",
+    "SymbolDef",
+]
+
+# Type aliases.
+SymbolPosMap = Dict[str, int]
+
+
+class AffineBuildState:
+  """Internal state for the AffineExprDef._create impls.
+
+  Note that a "local" AffineBuildState can be created relative to a "global"
+  AffineBuildState. In that case, any affine expressions built will inherit
+  symbol and dim bindings from the global state and will update both as new
+  ones are discovered. This allows for building expressions across contexts
+  which share a common symbol and dim space.
+  """
+
+  def __init__(self,
+               *,
+               global_state: "AffineBuildState" = None,
+               allow_new_symbols: bool = True,
+               allow_new_dims: bool = True):
+    if not global_state:
+      self.all_symbols = dict()  # type: Dict[str, int]
+      self.all_dims = dict()  # type: Dict[str, int]
+    else:
+      # Alias the global dict.
+      self.all_symbols = global_state.all_symbols
+      self.all_dims = global_state.all_dims
+
+    # Map of symbols and dims in the current build.
+    self.local_symbols = dict()  # type: Dict[str, int]
+    self.local_dims = dict()  # type: Dict[str, int]
+    self.allow_new_symbols = allow_new_symbols
+    self.allow_new_dims = allow_new_dims
+
+  def get_dim(self, dimname: str) -> int:
+    """Gets the dim position given a name."""
+    pos = self.all_dims.get(dimname)
+    if pos is None:
+      if not self.allow_new_dims:
+        raise ValueError(
+            f"New dimensions not allowed in the current affine expression: "
+            f"Requested '{dimname}', Availble: {self.all_dims}")
+      pos = len(self.all_dims)
+      self.all_dims[dimname] = pos
+    self.local_dims[dimname] = pos
+    return pos
+
+  def get_symbol(self, symname: str) -> int:
+    """Geta a symbol position given a name."""
+    pos = self.all_symbols.get(symname)
+    if pos is None:
+      if not self.allow_new_symbols:
+        raise ValueError(
+            f"New symbols not allowed in the current affine expression: "
+            f"Requested '{symname}', Availble: {self.all_symbols}")
+      pos = len(self.all_symbols)
+      self.all_symbols[symname] = pos
+    self.local_symbols[symname] = pos
+    return pos
+
+  @property
+  def local_dim_count(self) -> int:
+    return len(self.local_dims)
+
+  @property
+  def local_symbol_count(self) -> int:
+    return len(self.local_symbols)
+
+  @property
+  def dim_count(self) -> int:
+    return len(self.all_dims)
+
+  @property
+  def symbol_count(self) -> int:
+    return len(self.all_symbols)
+
+  def __repr__(self):
+    lines = [f"AffineBuildState<"]
+    lines.append(f"  symbols={self.local_symbols}")
+    lines.append(f"  dims={self.local_dims}>")
+    return "\n".join(lines)
+
+
+class AffineExprDef:
+  """Base class for an affine expression being defined."""
+
+  def build(self, state: Optional[AffineBuildState] = None) -> _ir.AffineExpr:
+    """Builds the corresponding _ir.AffineExpr from the definitions.
+    """
+    state = AffineBuildState() if state is None else state
+    expr = self._create(state)
+    return expr
+
+  def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
+    raise NotImplementedError()
+
+  @staticmethod
+  def coerce_from(py_value):
+    if isinstance(py_value, int):
+      return AffineConstantExpr(py_value)
+    assert isinstance(py_value, AffineExprDef)
+    return py_value
+
+  def visit_affine_exprs(self, callback):
+    """Visits all AffineExprDefs including self."""
+    callback(self)
+
+  def __add__(lhs, rhs):
+    rhs = AffineExprDef.coerce_from(rhs)
+    return AffineBinaryExprDef(_ir.AffineAddExpr, lhs, rhs)
+
+  def __mul__(lhs, rhs):
+    rhs = AffineExprDef.coerce_from(rhs)
+    return AffineBinaryExprDef(_ir.AffineMulExpr, lhs, rhs)
+
+  def __mod__(lhs, rhs):
+    rhs = AffineExprDef.coerce_from(rhs)
+    return AffineBinaryExprDef(_ir.AffineModExpr, lhs, rhs)
+
+  def __floordiv__(lhs, rhs):
+    rhs = AffineExprDef.coerce_from(rhs)
+    return AffineBinaryExprDef(_ir.AffineFloorDivExpr, lhs, rhs)
+
+  def __truediv__(lhs, rhs):
+    # TODO: Not really a ceil div - taking liberties for the DSL.
+    rhs = AffineExprDef.coerce_from(rhs)
+    return AffineBinaryExprDef(_ir.AffineCeilDivExpr, lhs, rhs)
+
+
+class AffineConstantExpr(AffineExprDef):
+  """An affine constant being defined."""
+
+  def __init__(self, value: int):
+    assert isinstance(value, int)
+    self.value = value
+
+  def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
+    return _ir.AffineConstantExpr.get(self.value)
+
+  def __repr__(self):
+    return f"Const({self.value})"
+
+
+class AffineBinaryExprDef(AffineExprDef):
+  """An affine binary expression being defined."""
+
+  def __init__(self, ir_ctor, lhs: AffineExprDef, rhs: AffineExprDef):
+    self.ir_ctor = ir_ctor
+    self.lhs = lhs
+    self.rhs = rhs
+
+  def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
+    return self.ir_ctor.get(self.lhs._create(state), self.rhs._create(state))
+
+  def visit_affine_exprs(self, callback):
+    """Visits all AffineExprDefs including self."""
+    super().visit_affine_exprs(callback)
+    self.lhs.visit_affine_exprs(callback)
+    self.rhs.visit_affine_exprs(callback)
+
+  def __repr__(self):
+    return f"{self.ir_ctor.__name__}({repr(self.lhs)}, {repr(self.rhs)})"
+
+
+class DimDef(AffineExprDef):
+  """Represents a named dimension.
+
+  """
+  ALL_DIMS = dict()  # type: Dict[str, "DimDef"]
+  dimname: str
+
+  def __new__(cls, dimname: str):
+    existing = cls.ALL_DIMS.get(dimname)
+    if existing is not None:
+      return existing
+    new = super().__new__(cls)
+    new.dimname = dimname
+    cls.ALL_DIMS[dimname] = new
+    return new
+
+  def __repr__(self):
+    return f"Dim({self.dimname})"
+
+  def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
+    pos = state.get_dim(self.dimname)
+    return _ir.AffineDimExpr.get(position=pos)
+
+  @classmethod
+  def create_expando(cls):
+    """Create an expando class that creates unique symbols based on attr access.
+    """
+
+    class ExpandoDims:
+
+      def __getattr__(self, n):
+        return cls(n)
+
+    return ExpandoDims()
+
+
+class SymbolDef(AffineExprDef):
+  """Represents a named symbol.
+
+    >>> s1 = SymbolDef("s1")
+    >>> s1
+    Symbol(s1)
+    >>> s2 = SymbolDef("s2")
+    >>> s1 is s2
+    False
+    >>> s1 is SymbolDef("s1")
+    True
+  """
+  ALL_SYMBOLS = dict()  # type: Dict[str, "SymbolDef"]
+  symname: str
+
+  def __new__(cls, symname: str):
+    existing = cls.ALL_SYMBOLS.get(symname)
+    if existing is not None:
+      return existing
+    new = super().__new__(cls)
+    new.symname = symname
+    cls.ALL_SYMBOLS[symname] = new
+    return new
+
+  def __repr__(self):
+    return f"Symbol({self.symname})"
+
+  def _create(self, state: AffineBuildState) -> _ir.AffineExpr:
+    pos = state.get_symbol(self.symname)
+    return _ir.AffineSymbolExpr.get(position=pos)
+
+  @classmethod
+  def create_expando(cls):
+    """Create an expando class that creates unique symbols based on attr access.
+    """
+
+    class ExpandoSymbols:
+
+      def __getattr__(self, n):
+        return cls(n)
+
+    return ExpandoSymbols()
+
+
+# Global accessor for on-demand dims and symbols.
+D = DimDef.create_expando()
+S = SymbolDef.create_expando()

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/comprehension.py

+#  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
+"""Model classes representing a tensor comprehension.
+
+These classes model the language more at an AST level as evaluated. Reasoning
+about it typically involves processing this form into config objects that
+represent actual op definitions (i.e. YAML).
+"""
+
+from typing import Callable, Dict, List, Optional, Sequence, Set, Tuple, Union
+
+from mlir import ir as _ir
+
+from .affine import *
+from .scalar_expr import *
+from .types import *
+from .yaml_helper import *
+
+# Type aliases.
+AffineDimList = Dict[str, _ir.AffineExpr]
+
+
+class TensorExpression:
+  """An expression that can appear on the RHS of a comprehension."""
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    raise NotImplementedError()
+
+  def visit_affine_exprs(self, callback):
+    """Visits all affine expressions reachable by the expression."""
+    pass
+
+  def _get_all_dim_defs(self) -> Set[DimDef]:
+    """Recursively gets all DimDef affine expressions that are referenced."""
+    results = set()
+
+    def visitor(affine_expr):
+      if isinstance(affine_expr, DimDef):
+        results.add(affine_expr)
+
+    self.visit_affine_exprs(visitor)
+    return results
+
+  def collect_uses(self, uses: Set["TensorUse"]):
+    """Collects all TensorUses reachable through this expression."""
+    pass
+
+  def __add__(self, rhs: "TensorExpression") -> "TensorExpression":
+    return PrimFn.add(self, rhs)
+
+  def __mul__(self, rhs) -> "TensorExpression":
+    return PrimFn.mul(self, rhs)
+
+  def __sub__(self, rhs) -> "TensorExpression":
+    return PrimFn.sub(self, rhs)
+
+  def __hash__(self):
+    return hash(id(self))
+
+
+class TensorUse(TensorExpression):
+  """A used tensor represented by its (tensor_name, indices).
+
+  Note that forming a comprehension via direct assignment is performed through
+  __setitem__ on the TensorDef level. However, performing a reduction with
+  compound ops (+=, *=, etc) is done by doing a:
+    TensorDef.__getitem__
+    TensorUse.__iadd__
+    TensorDef.__setitem__
+  """
+
+  def __init__(self, tensor_def: "TensorDef", indices: Sequence[AffineExprDef]):
+    self.tensor_def = tensor_def
+    self.indices = tuple(indices)
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    assert self.tensor_def.tensor_name is not None
+    return ScalarArg(self.tensor_def.tensor_name).expr()
+
+  @property
+  def tensor_name(self) -> str:
+    n = self.tensor_def.tensor_name
+    assert n is not None, "TensorDef not attached"
+    return n
+
+  def visit_affine_exprs(self, callback):
+    for ind in self.indices:
+      ind.visit_affine_exprs(callback)
+
+  def collect_uses(self, uses: Set["TensorUse"]):
+    uses.add(self)
+
+  def __iadd__(self, rhs: TensorExpression) -> TensorExpression:
+    return ReduceFn.add(*self._compute_reduce_dims(rhs))(rhs)
+
+  def _compute_reduce_dims(self, rhs: TensorExpression) -> Set[DimDef]:
+    """For implicit reductions, computes default reduction dims.
+
+    Assumes that the rhs is the expression being reduced and self is being
+    reduced into. Any indices referenced on the rhs and not in self are
+    considered reduction dims and will be ordered as encountered on the rhs.
+    """
+    rhs_dims = rhs._get_all_dim_defs()
+    lhs_dims = self._get_all_dim_defs()
+    return rhs_dims - lhs_dims
+
+  def __repr__(self):
+    return f"{self.tensor_name}[{', '.join([repr(i) for i in self.indices])}]"
+
+
+class TensorDef:
+  """Bookkeeping of a single registered tensor, held in dict by name."""
+
+  def __init__(self,
+               type_var: TypeVar,
+               *shape: AffineExprDef,
+               indexing_map: Optional[_ir.AffineMap] = None,
+               output: bool = False):
+    if not isinstance(type_var, TypeVar):
+      raise ValueError(f"TensorDef requires a TypeVar. Got: {repr(type_var)}")
+    self.owner = None  # type: Optional["LinalgOpDef"]
+    self.type_var = type_var
+    self.shape = shape
+    self.indexing_map = indexing_map
+    self.output = output
+    self.tensor_name = None  # type: Optional[str]
+    self.registered_index = -1  # type: int
+
+  @property
+  def rank(self) -> int:
+    """The rank of the tensor."""
+    return len(self.shape)
+
+  def attach(self, index: int, tensor_name: str, owner: "LinalgOpDef"):
+    if self.owner:
+      raise ValueError(f"TensorDef already registered with op: {self}")
+    self.registered_index = index
+    self.tensor_name = tensor_name
+    self.owner = owner
+
+  def __getitem__(self, dims) -> TensorUse:
+    assert self.owner, "TensorDef is not attached to an op"
+    state = AffineBuildState(global_state=self.owner._affine_state,
+                             allow_new_symbols=False)
+    if not isinstance(dims, tuple):
+      dims = (dims,)  # Handle single subscript case.
+    # Special case: (None) is a 0d-scalar use.
+    if dims == (None,):
+      dims = ()
+
+    exprs = []
+    for expr_def in dims:
+      if not isinstance(expr_def, AffineExprDef):
+        raise KeyError(
+            "A TensorDef can only be subscripted by a tuple of affine dims")
+      exprs.append(expr_def)
+    return TensorUse(self, exprs)
+
+  def __setitem__(self, dims, value):
+    """Creates a new 1:1 comprehension by binding this tensor to an expression.
+
+    Note that due to the way assignment works in Python, we have to capture
+    direct assignment as a setitem on the TensorDef.
+    """
+    if not isinstance(value, TensorExpression):
+      raise ValueError(f"Only TensorExpressions can be assigned to TensorDefs. "
+                       f"Got: {repr(value)}")
+    use = self[dims]
+    comp = Comprehension((use, value))
+    self.owner.comprehensions.append(comp)
+
+  def __hash__(self):
+    return hash(id(self))
+
+  def __repr__(self):
+    output = "OUTPUT " if self.output else ""
+    return (f"{self.tensor_name}:TensorDef({output}{repr(self.type_var)}, "
+            f"shape={self.shape})")
+
+
+class Comprehension:
+  """Represents a single comprehension."""
+
+  def __init__(self, *bindings: Tuple[TensorUse, TensorExpression]):
+    self.definitions = list()  # List[TensorUse]
+    self.values = list()  # List[TensorExpression]
+
+    # Find the lhs to reduction rhs.
+    for assign, value in bindings:
+      if isinstance(value, ReduceApply):
+        if value.lhs:
+          raise ValueError(f"Reduction expression already assigns: {value}")
+        value.lhs = assign
+      self.definitions.append(assign)
+      self.values.append(value)
+
+  @property
+  def all_reduction_dims(self) -> Set[Tuple[DimDef, ...]]:
+    """Gets the reduction dims for the comprehension or None."""
+    result = set()
+    for use in self.values:
+      if isinstance(use, ReduceApply):
+        result.add(use.reduce.reduce_dims)
+      else:
+        result.add(tuple())
+    return result
+
+  def __repr__(self):
+    if len(self.definitions) > 1:
+      defs_repr = f"({', '.join(repr(d) for d in self.definitions)})"
+      values_repr = f"({', '.join(repr(v) for v in self.values)})"
+    else:
+      defs_repr = f"{repr(self.definitions[0])}"
+      values_repr = f"{repr(self.values[0])}"
+
+    return f"{defs_repr} = {values_repr}"
+
+
+class PrimFnType:
+  """Primitive operations."""
+
+  def __init__(self, prim_name: str):
+    self.prim_name = prim_name
+
+  def __call__(self, *args):
+    return PrimApply(self, args)
+
+  def reduce(self, *reduce_dims: DimDef):
+    """Shortcut to create a Reduce operation from this primitive."""
+    return ReduceFnType(self, *reduce_dims)
+
+  def __repr__(self):
+    return f"{self.prim_name}"
+
+
+class PrimFn:
+  add = PrimFnType("add")
+  exp = PrimFnType("exp")
+  log = PrimFnType("log")
+  mul = PrimFnType("mul")
+  max = PrimFnType("max")
+  sub = PrimFnType("sub")
+
+
+class ReduceFnType:
+  """A reduction operator that reduces into its LHS from its RHS."""
+
+  def __init__(self, operator: PrimFnType, *reduce_dims: DimDef):
+    """Initializes the ReduceFn with a primitive function and dims."""
+    if not isinstance(operator, PrimFnType):
+      raise ValueError(f"Reduce expected a Prim operator. Got: {operator}")
+    self.operator = operator
+    self.reduce_dims = tuple(reduce_dims)
+
+  def __call__(self, *args: TensorExpression):
+    return ReduceApply(self, args)
+
+  def __repr__(self):
+    return (f"reduce_{self.operator.prim_name}"
+            f"({', '.join(repr(d) for d in self.reduce_dims)})")
+
+
+class ReduceFn:
+  add = PrimFn.add.reduce
+  mul = PrimFn.mul.reduce
+  max = PrimFn.max.reduce
+
+
+class PrimApply(TensorExpression):
+  """Application of a primitive."""
+
+  def __init__(self, prim: PrimFnType, args: Sequence[TensorExpression]):
+    self.prim = prim
+    self.args = tuple(args)
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    return ScalarApplyFn(self.prim.prim_name,
+                         *[arg.to_scalar_expression() for arg in self.args
+                          ]).expr()
+
+  def visit_affine_exprs(self, callback):
+    for arg in self.args:
+      arg.visit_affine_exprs(callback)
+
+  def collect_uses(self, uses: Set["TensorUse"]):
+    for arg in self.args:
+      arg.collect_uses(uses)
+
+  def __repr__(self):
+    return f"{repr(self.prim)}({', '.join(repr(a) for a in self.args)})"
+
+
+class cast(TensorExpression):
+  """Casts the element type to a type (typically symbolic TypeVar)."""
+
+  def __init__(self, to_type: TypeVar, operand: TensorExpression):
+    self.to_type = to_type
+    self.operand = operand
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    return ScalarSymbolicCast(self.to_type,
+                              self.operand.to_scalar_expression()).expr()
+
+  def visit_affine_exprs(self, callback):
+    self.operand.visit_affine_exprs(callback)
+
+  def collect_uses(self, uses: Set["TensorUse"]):
+    self.operand.collect_uses(uses)
+
+  def __repr__(self):
+    return f"cast({self.to_type}, {repr(self.operand)})"
+
+
+class ReduceApply(TensorExpression):
+  """Application of a reduction.
+
+  This captures the lhs separately (initial value) separately from the rhs.
+  """
+
+  def __init__(self, reduce: ReduceFnType, args: Sequence[TensorExpression]):
+    self.reduce = reduce
+    self.lhs = None  # type: Optional[TensorUse]
+    self.args = tuple(args)
+
+  def to_scalar_expression(self) -> ScalarExpression:
+    if self.lhs is None:
+      raise ValueError(f"Cannot scalarize a ReduceApply that has not been "
+                       f"bound to its lhs: {self}")
+    full_args = [self.lhs.to_scalar_expression()
+                ] + [arg.to_scalar_expression() for arg in self.args]
+    return ScalarApplyFn(self.reduce.operator.prim_name, *full_args).expr()
+
+  def visit_affine_exprs(self, callback):
+    for ind in self.reduce.reduce_dims:
+      ind.visit_affine_exprs(callback)
+    for arg in self.args:
+      arg.visit_affine_exprs(callback)
+
+  def collect_uses(self, uses: Set["TensorUse"]):
+    for arg in self.args:
+      arg.collect_uses(uses)
+
+  def __repr__(self):
+    return f"{repr(self.reduce)}({', '.join(repr(a) for a in self.args)})"
+
+
+class OpInterfaceDef:
+  """An interface that an op implements."""
+
+  def __init__(self, cpp_name: str):
+    self.cpp_name = cpp_name
+
+
+ContractionOpInterface = OpInterfaceDef("LinalgContractionOpInterface")
+
+
+class OpMetadataDef(YAMLObject):
+  """Metadata about the op (generally not behavior impacting)."""
+  yaml_tag = "!LinalgOpMetadata"
+
+  def __init__(self, name: str, cpp_op_name: Optional[str], doc: Optional[str]):
+    self.name = name
+    self.cpp_op_name = cpp_op_name if cpp_op_name is not None else name
+    self.doc = doc
+    self.implements = []  # type: List[OpInterfaceDef]
+
+  def to_yaml_custom_dict(self):
+    d = dict(
+        name=self.name,
+        cpp_op_name=self.cpp_op_name,
+        doc=self.doc,
+    )
+    if self.implements:
+      d["implements"] = [intr.cpp_name for intr in self.implements]
+    return d
+
+
+class LinalgOpDef:
+  """Definition of a linalg op."""
+
+  def __init__(self,
+               name: str,
+               cpp_op_name: Optional[str] = None,
+               doc: Optional[str] = None):
+    self.metadata = OpMetadataDef(name=name, cpp_op_name=cpp_op_name, doc=doc)
+    self.registered_tensors = dict()  # type: Dict[str, TensorDef]
+    self.comprehensions = list()  # type: List[Comprehension]
+    self._affine_state = AffineBuildState()
+
+  @property
+  def inputs(self) -> Sequence[TensorDef]:
+    return [t for t in self.registered_tensors.values() if not t.output]
+
+  @property
+  def outputs(self) -> Sequence[TensorDef]:
+    return [t for t in self.registered_tensors.values() if t.output]
+
+  def add_tensor(self, tensor_name: str, tensor: TensorDef):
+    """Registers a tensor."""
+    if tensor_name in self.registered_tensors:
+      raise ValueError(f"Tensor {tensor_name} is already registered "
+                       f"to {self.registered_tensors['tensor_name']}")
+    tensor.attach(len(self.registered_tensors), tensor_name, self)
+    self.registered_tensors[tensor_name] = tensor
+
+  def tensor(self, name):
+    """Gets a registered tensor by name."""
+    try:
+      return self.registered_tensors[name]
+    except KeyError:
+      raise KeyError(f"Tensor {name} is not registered")
+
+  def __repr__(self):
+    lines = [
+        f"LinalgOpDef({self.metadata.name} -> {self.metadata.cpp_op_name},"
+    ]
+    for name, tensor in self.registered_tensors.items():
+      lines.append(f"  {tensor}")
+    if self.comprehensions:
+      lines[-1] += " {"
+      for comprehension in self.comprehensions:
+        lines.append(f"    {comprehension}")
+      lines.append("}")
+    return "\n".join(lines)

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/config.py

+#  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
+"""Represents configured ops as emitted for code generation.
+
+Classes in this module generally are directly serializable to YAML for use
+by the code generator.
+
+TODO: These should just be dumb containers or serialization code but they
+currently encode too many details of how the language is interpreted. Move this
+to helpers on the comprehension objects themselves.
+"""
+
+from typing import Any, Dict, Optional
+
+from mlir import ir as _ir
+
+from .comprehension import *
+from .yaml_helper import *
+
+__all__ = [
+    "LinalgStructuredOpConfig",
+    "LinalgOpConfig",
+]
+
+
+def _serialize_affine_map(affine_map: _ir.AffineMap) -> str:
+  with affine_map.context:
+    # Affine map printing/parsing is via an AffineMap attr.
+    attr = _ir.AffineMapAttr.get(affine_map)
+    return str(attr)
+
+
+class TensorUseConfig:
+  """Wrapper around a TensorUse with additional context-bound state."""
+
+  def __init__(self, tensor_use: TensorUse, indexing_map: _ir.AffineMap):
+    self.tensor_use = tensor_use
+    self.indexing_map = indexing_map
+
+  def __repr__(self):
+    return f"Use({self.tensor_use}, indexing_map={self.indexing_map})"
+
+
+class TensorDefConfig(YAMLObject):
+  """Wrapper around a TensorDef with additional context-bound state."""
+  yaml_tag = "LinalgTensorDef"
+
+  def __init__(self, tensor_def: TensorDef, shape_map: _ir.AffineMap):
+    self.tensor_def = tensor_def
+    self.shape_map = shape_map
+    self.indexing_map = None  # type: Optional[_ir.AffineMap]
+
+  def to_yaml_custom_dict(self):
+
+    def get_usage():
+      if self.tensor_def.output:
+        return "output"
+      else:
+        return "input"
+
+    return dict(
+        name=self.tensor_def.tensor_name,
+        usage=get_usage(),
+        shape=_serialize_affine_map(self.shape_map),
+        element_type_var=self.tensor_def.type_var.name,
+    )
+
+  def __repr__(self):
+    return f"Def({self.tensor_def}, shape_map={self.shape_map}, indexing_map={self.indexing_map})"
+
+
+class LinalgIndexingMapsConfig(YAMLObject):
+  """Abstracts the style of indexing maps that the op exports.
+
+  Presently only static (tied to the op name) indexing maps are supported. In
+  the future, it is expected that we will have additional variants:
+    - Dynamic based on attributes
+    - Dynamic based on operands
+  Each is expected to require a different variant of specification.
+  """
+  yaml_tag = "!LinalgIndexingMapsConfig"
+
+  def __init__(self,
+               static_indexing_maps: Optional[Sequence[_ir.AffineMap]] = None):
+    self.static_indexing_maps = static_indexing_maps
+
+  def to_yaml_custom_dict(self):
+    if self.static_indexing_maps is not None:
+      return dict(static_indexing_maps=[
+          _serialize_affine_map(m) for m in self.static_indexing_maps
+      ])
+    raise ValueError(
+        f"LinalgIndexingMapsConfig must have one type of indexing map"
+        f"(got none)")
+
+
+class LinalgStructuredOpConfig(YAMLObject):
+  """Configuration for metadata sufficient to construct a linalg single
+  contraction named op."""
+
+  yaml_tag = "!LinalgStructuredOpConfig"
+
+  def __init__(self,
+               comprehension: Comprehension,
+               context: Optional[_ir.Context] = None):
+    self.context = context if context is not None else _ir.Context()
+    self.affine_state = AffineBuildState()
+    self.writes = list()  # type: List[Tuple[TensorUse, TensorExpression]]
+    self.tensor_args = dict()  # type: Dict[TensorDef, TensorDefConfig]
+    self.uses = dict()  # type: Dict[TensorUse, TensorUseConfig]
+
+    # Compute the ordered set of writes.
+    collected_uses = set()
+    for write_use, read_use in zip(comprehension.definitions,
+                                   comprehension.values):
+      self.writes.append((write_use, read_use))
+
+    for write_use, read_use in self.writes:
+      collected_uses.add(write_use)
+      read_use.collect_uses(collected_uses)
+
+    # Need to add all definitions before uses, so process twice.
+    for use in collected_uses:
+      self.add_tensor_arg(use.tensor_def)
+    for use in collected_uses:
+      self.add_use(use)
+
+    # Now normalize all defs and uses indexing maps now that full count of
+    # dims and symbols are known.
+    for cuse in self.uses.values():
+      cuse.indexing_map = self._normalize_affine_map(cuse.indexing_map)
+    for cdef in self.tensor_args.values():
+      cdef.shape_map = self._normalize_affine_map(cdef.shape_map,
+                                                  with_dims=False)
+
+    # Now for each write use, propagate the indexing maps from the use to the
+    # tensor, ensuring that there are not conflicts.
+    for write_use, _ in self.writes:
+      write_tensor_def = self.tensor_args[write_use.tensor_def]
+      if write_tensor_def.indexing_map:
+        raise ValueError(
+            f"Unexpected multi-write to a single tensor: {write_tensor_def}")
+      write_tensor_def.indexing_map = self.uses[write_use].indexing_map
+
+    # For each read use, propagate the indexing maps from the use to the
+    # tensor, ensuring that there are not conflicts.
+    for _, read_expr in self.writes:
+      read_uses = set()  # type: Set[TensorUse]
+      read_expr.collect_uses(read_uses)
+      for read_use in read_uses:
+        read_tensor_def = self.tensor_args[read_use.tensor_def]
+        if (read_tensor_def.indexing_map and
+            read_tensor_def.indexing_map != self.uses[read_use].indexing_map):
+          raise ValueError(
+              f"Unexpected multi-read of a tensor with different accesses:"
+              f"{read_tensor_def} vs {read_use}")
+        read_tensor_def.indexing_map = self.uses[read_use].indexing_map
+
+    # Sanity check that all defs have an indexing map.
+    assert all(d.indexing_map for d in self.tensor_args.values()), (
+        f"Missing indexing map on TensorDef: {self.tensor_args}")
+
+    # Collect reduction dims and ensure all the same.
+    all_reduction_dims = set(comprehension.all_reduction_dims)
+    if len(all_reduction_dims) != 1:
+      raise ValueError(
+          f"All writes within a generic must have the same reduction "
+          f"dims. Got: {all_reduction_dims}")
+    self.reduction_dims = next(iter(all_reduction_dims))
+
+    # Generate the scalar assignments (used to build a body).
+    self.assignments = [
+        ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
+        for write_use, read_expr in self.writes
+    ]
+
+  @property
+  def ordered_tensor_args(self) -> Sequence[TensorDefConfig]:
+    return sorted(self.tensor_args.values(),
+                  key=lambda tdc: tdc.tensor_def.registered_index)
+
+  @property
+  def ordered_tensor_uses(self) -> Sequence[TensorUseConfig]:
+    return sorted(self.uses.values(),
+                  key=lambda tuc: tuc.tensor_use.tensor_def.registered_index)
+
+  @property
+  def ordered_dims(self) -> Sequence[Tuple[str, int]]:
+    """Gets the ordered list of dim bindings (symbolic name, position).
+
+    TODO: The original parser relies on parse ordering to arrive at the
+    iterator types, but that ordering is not defined on the Python side, so
+    this may be ambiguous.
+    """
+    return list(self.affine_state.all_dims.items())
+
+  @property
+  def indexing_maps(self) -> Sequence[_ir.AffineMap]:
+    return [use.indexing_map for use in self.ordered_tensor_uses]
+
+  @property
+  def iterator_types(self) -> Sequence[str]:
+
+    def get_type(symbolic_name, position):
+      for reduction_dim_expr in self.reduction_dims:
+        if reduction_dim_expr.dimname == symbolic_name:
+          return "reduction"
+      return "parallel"
+
+    return [get_type(*dim) for dim in self.ordered_dims]
+
+  def add_tensor_arg(self, tensor_def: TensorDef):
+    if tensor_def in self.tensor_args:
+      return
+    with self.context:
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_dims=False)
+      exprs = []
+      for expr in tensor_def.shape:
+        exprs.append(expr.build(state=local_state))
+      assert local_state.local_dim_count == 0
+      indexing_map = _ir.AffineMap.get(dim_count=0,
+                                       symbol_count=local_state.symbol_count,
+                                       exprs=exprs)
+
+      def_config = TensorDefConfig(tensor_def, indexing_map)
+      self.tensor_args[tensor_def] = def_config
+
+  def add_use(self, tensor_use: TensorUse):
+    if tensor_use in self.uses:
+      return
+    with self.context:
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_symbols=False)
+      exprs = []
+      for expr in tensor_use.indices:
+        exprs.append(expr.build(state=local_state))
+      assert local_state.local_symbol_count == 0
+      indexing_map = _ir.AffineMap.get(dim_count=local_state.dim_count,
+                                       symbol_count=local_state.symbol_count,
+                                       exprs=exprs)
+
+      use_config = TensorUseConfig(tensor_use, indexing_map)
+      self.uses[tensor_use] = use_config
+
+  def _normalize_affine_map(self,
+                            affine_map: _ir.AffineMap,
+                            with_dims: bool = True) -> _ir.AffineMap:
+    """Normalizes an indexing map to have the max known symbols and dims."""
+    with self.context:
+      return _ir.AffineMap.get(
+          dim_count=self.affine_state.dim_count if with_dims else 0,
+          symbol_count=self.affine_state.symbol_count,
+          exprs=list(affine_map.results))
+
+  def to_yaml_custom_dict(self):
+    self_dict = dict(
+        args=self.ordered_tensor_args,
+        # TODO: Refactor the hierarchy internally when supporting more
+        # than static (preserving this serialized form).
+        indexing_maps=LinalgIndexingMapsConfig(
+            static_indexing_maps=self.indexing_maps),
+        iterator_types=self.iterator_types,
+        assignments=self.assignments,
+    )
+    return self_dict
+
+  def __repr__(self):
+    lines = [f"LinalgGenericOpConfig(reduction_dims={self.reduction_dims},"]
+    lines.append("tensor_args=[")
+    for def_config in self.ordered_tensor_args:
+      lines.append(f"  {repr(def_config)}")
+    lines.append("], indexing_maps=[")
+    for m in self.indexing_maps:
+      lines.append(f"  {repr(m)}")
+    lines.append(f"], iterator_types=[")
+    for t in self.iterator_types:
+      lines.append(f"  {t}")
+    lines.append("])")
+    return "\n".join(lines)
+
+
+class LinalgOpConfig(YAMLObject):
+  """Container for any supported linalg op type.
+
+  This includes the concrete type by name for ease of parsing by systems
+  that ignore tags.
+  """
+  yaml_tag = "!LinalgOpConfig"
+
+  def __init__(self,
+               metadata: OpMetadataDef,
+               *,
+               structured_op: Optional[LinalgStructuredOpConfig] = None):
+    self.metadata = metadata
+    self.structured_op = structured_op
+
+  def to_yaml_custom_dict(self):
+    self_dict = dict(metadata=self.metadata,)
+    if self.structured_op:
+      self_dict["structured_op"] = self.structured_op
+    return self_dict
+
+  @staticmethod
+  def from_linalg_op_def(
+      tc_op_def: LinalgOpDef,
+      context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
+    """Expands a LinalgOpDef into corresponding Linalg configured ops."""
+    # TODO: Many LinalgOpDef patterns need to expand to multiple generics.
+    assert len(
+        tc_op_def.comprehensions) == 1, "Only one comprehension supported"
+    return [
+        LinalgOpConfig(tc_op_def.metadata,
+                       structured_op=LinalgStructuredOpConfig(
+                           tc_op_def.comprehensions[0], context)),
+    ]
+
+  def __repr__(self):
+    return (f"LinalgOpConfig(metadata={self.metadata},\n"
+            f"structured_op={self.structured_op})")

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/dsl.py

+#  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
+
+from typing import Dict, List
+
+from contextlib import contextmanager
+import functools
+import inspect
+import threading
+
+from mlir import ir
+from .comprehension import *
+
+_CONTEXT = threading.local()
+
+
+@contextmanager
+def bind_op_def(model: LinalgOpDef):
+  if hasattr(_CONTEXT, "current_op_def"):
+    raise ValueError("Cannot recursively define an operation")
+  _CONTEXT.current_op_def = model
+  try:
+    yield model
+  finally:
+    del _CONTEXT.current_op_def
+
+
+def current_op_def() -> LinalgOpDef:
+  try:
+    return _CONTEXT.current_op_def
+  except AttributeError:
+    raise ValueError(
+        "Attempt to access the current op definition being defined "
+        "but none is set. Did you mean to call this in an op definition?")
+
+
+class DefinedOpCallable:
+  """Callable that wraps any defined op function."""
+
+  def __init__(self, op_name: str, model: LinalgOpDef):
+    self.op_name = op_name
+    self.model = model
+
+  def __call__(self, *args, **kwargs):
+    # TODO: Upstream the emitter and invoke here
+    raise NotImplementedError("Linalg generic emission not yet implemented")
+
+
+def linalg_structured_op(dsl_func=None,
+                         *,
+                         op_name=None,
+                         op_class_name=None) -> DefinedOpCallable:
+  if dsl_func is None:
+    # Curry the keyword args in for delayed application.
+    return functools.partial(tc_def_op,
+                             op_name=op_name,
+                             op_class_name=op_class_name)
+  # Determine default names by introspecting the function.
+  if op_name is None:
+    op_name = dsl_func.__name__
+  if op_class_name is None:
+    # Camel case it.
+    op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
+
+  tc_model = LinalgOpDef(name=op_name,
+                         cpp_op_name=op_class_name,
+                         doc=inspect.getdoc(dsl_func))
+
+  # Extract arguments and TensorDefs from the signature.
+  dsl_func_args = list()
+  sig = inspect.signature(dsl_func)
+  for param_name, param in sig.parameters.items():
+    param_default = param.default
+    if not isinstance(param_default, TensorDef):
+      raise ValueError(f"@tc_def_op function parameters must be defaulted as "
+                       f"TensorDef(...): Found {param_name}: {param_default}")
+    dsl_func_args.append(param_default)
+    tc_model.add_tensor(param_name, param_default)
+
+  # Invoke the DSL func to finish populating the model.
+  with bind_op_def(tc_model):
+    dsl_func(*dsl_func_args)
+
+  # TODO: The returned callable should be an IR emitter but that is not
+  # upstreamed yet.
+  return DefinedOpCallable(op_name, tc_model)
+
+
+def implements(*interfaces: OpInterfaceDef):
+  current_op_def().metadata.implements.extend(interfaces)

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/scalar_expr.py

+#  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
+"""Models DAGs of scalar math expressions.
+
+Used for generating region bodies at the "math" level where they are still type
+polymorphic. This is modeled to be polymorphic by attribute name for interop
+with serialization schemes that are just plain-old-dicts.
+
+These classes are typically not user accessed and are created as a by-product
+of interpreting a comprehension DSL and model the operations to perform in the
+op body. The class hierarchy is laid out to map well to a form of YAML that
+can be easily consumed from the C++ side, not necessarily for ergonomics.
+"""
+
+from typing import Optional, Sequence
+
+from .yaml_helper import *
+from .types import *
+
+__all__ = [
+    "ScalarAssign",
+    "ScalarApplyFn",
+    "ScalarArg",
+    "ScalarExpression",
+    "ScalarSymbolicCast",
+]
+
+
+class ScalarApplyFn:
+  """A type of ScalarExpression that applies a named function to operands."""
+
+  def __init__(self, fn_name: str, *operands: "ScalarExpression"):
+    self.fn_name = fn_name
+    self.operands = operands
+
+  def expr(self) -> "ScalarExpression":
+    return ScalarExpression(scalar_apply=self)
+
+  def __repr__(self):
+    return f"ScalarApplyFn<{self.fn_name}>({', '.join(self.operands)})"
+
+
+class ScalarArg:
+  """A type of ScalarExpression that references a named argument."""
+
+  def __init__(self, arg: str):
+    self.arg = arg
+
+  def expr(self) -> "ScalarExpression":
+    return ScalarExpression(scalar_arg=self)
+
+  def __repr__(self):
+    return f"(ScalarArg({self.arg})"
+
+
+class ScalarSymbolicCast:
+  """A type of ScalarExpression that symbolically casts an operand to a TypeVar.
+  """
+
+  def __init__(self, to_type: TypeVar, operand: "ScalarExpression"):
+    self.to_type = to_type
+    self.operand = operand
+
+  def expr(self) -> "ScalarExpression":
+    return ScalarExpression(symbolic_cast=self)
+
+  def __repr__(self):
+    return f"ScalarSymbolicCast({self.to_type}, {self.operand})"
+
+
+class ScalarExpression(YAMLObject):
+  """An expression on scalar values.
+
+  Can be one of:
+    - ScalarApplyFn
+    - ScalarArg
+    - ScalarSymbolicCast
+  """
+  yaml_tag = "!ScalarExpression"
+
+  def __init__(self,
+               scalar_apply: Optional[ScalarApplyFn] = None,
+               scalar_arg: Optional[ScalarArg] = None,
+               symbolic_cast: Optional[ScalarSymbolicCast] = None):
+    if (bool(scalar_apply) + bool(scalar_arg) + bool(symbolic_cast)) != 1:
+      raise ValueError(
+          "One of 'scalar_apply', 'scalar_block_arg', 'symbolic_cast' must be "
+          "specified")
+    self.scalar_apply = scalar_apply
+    self.scalar_arg = scalar_arg
+    self.symbolic_cast = symbolic_cast
+
+  def to_yaml_custom_dict(self):
+    if self.scalar_apply:
+      return dict(scalar_apply=dict(
+          fn_name=self.scalar_apply.fn_name,
+          operands=list(self.scalar_apply.operands),
+      ))
+    elif self.scalar_arg:
+      return dict(scalar_arg=self.scalar_arg.arg)
+    elif self.symbolic_cast:
+      # Note that even though operands must be arity 1, we write it the
+      # same way as for apply because it allows handling code to be more
+      # generic vs having a special form.
+      return dict(symbolic_cast=dict(type_var=self.symbolic_cast.to_type.name,
+                                     operands=[self.symbolic_cast.operand]))
+    else:
+      raise ValueError(f"Unexpected ScalarExpression type: {self}")
+
+
+class ScalarAssign(YAMLObject):
+  """An assignment to a named argument (LHS of a comprehension)."""
+  yaml_tag = "!ScalarAssign"
+
+  def __init__(self, arg: str, value: ScalarExpression):
+    self.arg = arg
+    self.value = value
+
+  def to_yaml_custom_dict(self):
+    return dict(arg=self.arg, value=self.value)
+
+  def __repr__(self):
+    return f"ScalarAssign({self.arg}, {self.value})"

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/types.py

+#  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
+"""Facility for symbolically referencing type variables.
+
+Type variables are instances of the TypeVar class, which is uniqued by name.
+An "expando" accessor `TV` is provided that generates a named TypeVar for
+any attribute access:
+
+  >>> TV.T
+  TypeVar(T)
+  >>> TV.T is TV.U
+  False
+  >>> TV.T is TV.T
+  True
+"""
+
+from enum import Enum
+from typing import Dict
+
+__all__ = [
+    "TypeVar",
+    "TV",
+
+    # TypeVar aliases.
+    "T",
+    "U",
+    "V",
+]
+
+
+class TypeVar:
+  """A replaceable type variable.
+
+  Type variables are uniqued by name.
+  """
+  ALL_TYPEVARS = dict()  # type: Dict[str, "TypeVar"]
+
+  def __new__(cls, name: str):
+    existing = cls.ALL_TYPEVARS.get(name)
+    if existing is not None:
+      return existing
+    new = super().__new__(cls)
+    new.name = name
+    cls.ALL_TYPEVARS[name] = new
+    return new
+
+  def __repr__(self):
+    return f"TypeVar({self.name})"
+
+  @classmethod
+  def create_expando(cls):
+    """Create an expando class that creates unique type vars on attr access."""
+
+    class ExpandoTypeVars:
+
+      def __getattr__(self, n):
+        return cls(n)
+
+    return ExpandoTypeVars()
+
+
+# Expando access via TV.foo
+TV = TypeVar.create_expando()
+
+# Some common type name aliases.
+T = TV.T
+U = TV.U
+V = TV.V

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/lang/yaml_helper.py

+#  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
+"""YAML serialization is routed through here to centralize common logic."""
+
+import sys
+
+try:
+  import yaml
+except ModuleNotFoundError as e:
+  raise ModuleNotFoundError(
+      f"This tool requires PyYAML but it was not installed. "
+      f"Recommend: {sys.executable} -m pip install PyYAML") from e
+
+__all__ = [
+    "yaml_dump",
+    "yaml_dump_all",
+    "YAMLObject",
+]
+
+
+class YAMLObject(yaml.YAMLObject):
+
+  @classmethod
+  def to_yaml(cls, dumper, self):
+    """Default to a custom dictionary mapping."""
+    return dumper.represent_mapping(cls.yaml_tag, self.to_yaml_custom_dict())
+
+  def to_yaml_custom_dict(self):
+    raise NotImplementedError()
+
+  def as_linalg_yaml(self):
+    return yaml_dump(self)
+
+
+def multiline_str_representer(dumper, data):
+  if len(data.splitlines()) > 1:
+    return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
+  else:
+    return dumper.represent_scalar('tag:yaml.org,2002:str', data)
+
+
+yaml.add_representer(str, multiline_str_representer)
+
+
+def yaml_dump(data, sort_keys=False, **kwargs):
+  return yaml.dump(data, sort_keys=sort_keys, **kwargs)
+
+
+def yaml_dump_all(data, sort_keys=False, explicit_start=True, **kwargs):
+  return yaml.dump_all(data,
+                       sort_keys=sort_keys,
+                       explicit_start=explicit_start,
+                       **kwargs)

mlir/lib/Bindings/Python/mlir/tools/linalg_opdsl/ops/core_named_ops.py

+from ..lang import *
+
+T1 = TV.T1
+T2 = TV.T2
+
+Batch = S.Batch
+
+
+@linalg_structured_op
+def matmul(A=TensorDef(T1, S.M, S.K),
+           B=TensorDef(T2, S.K, S.N),
+           C=TensorDef(U, S.M, S.N, output=True)):
+  """Performs a matrix multiplacation of two 2D inputs.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
+
+
+@linalg_structured_op
+def batch_matmul(A=TensorDef(T1, Batch, S.M, S.K),
+                 B=TensorDef(T2, Batch, S.K, S.N),
+                 C=TensorDef(U, Batch, S.M, S.N, output=True)):
+  """Performs a batched matrix multiplacation of two 3D inputs.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  C[D.b, D.m, D.n] += cast(U, A[D.b, D.m, D.k]) * cast(U, B[D.b, D.k, D.n])
+
+
+@linalg_structured_op
+def matvec(A=TensorDef(T1, S.M, S.N),
+           y=TensorDef(T2, S.N),
+           x=TensorDef(U, S.M, output=True)):
+  """Performs a matrix-vector multiplication.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  x[D.m] += cast(U, A[D.m, D.n]) * cast(U, y[D.n])
+
+
+@linalg_structured_op
+def vecmat(y=TensorDef(T1, S.M),
+           A=TensorDef(T2, S.M, S.N),
+           x=TensorDef(U, S.N, output=True)):
+  """Performs a vector-matrix multiplacation.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  x[D.n] += cast(U, y[D.m]) * cast(U, A[D.m, D.n])
+
+
+@linalg_structured_op
+def dot(A=TensorDef(T1, S.M), B=TensorDef(T2, S.M), C=TensorDef(U,
+                                                                output=True)):
+  """Performs a dot product of two vectors to a scalar result.
+
+  Numeric casting is performed on the operands to the inner multiply, promoting
+  them to the same data type as the accumulator/output.
+  """
+  implements(ContractionOpInterface)
+  C[None] += cast(U, A[D.m]) * cast(U, B[D.m])

mlir/test/Bindings/Python/tools/linalg_opdsl/assignments.py

+# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+
+from mlir.tools.linalg_opdsl.lang import *
+
+# CHECK: ---
+# CHECK-LABEL: matmul
+# CHECK: assignments:
+# CHECK:  -
+# CHECK:    arg: C
+# CHECK:    value:
+# CHECK:      scalar_apply:
+# CHECK:        fn_name: add
+# CHECK:        operands:
+# CHECK:          scalar_apply:
+# CHECK:            fn_name: mul
+# CHECK:            operands:
+# CHECK:              symbolic_cast:
+# CHECK:                type_var: U
+# CHECK:                operands:
+# CHECK:                  scalar_arg: A
+# CHECK:              symbolic_cast:
+# CHECK:                type_var: U
+# CHECK:                operands:
+# CHECK:                  scalar_arg: B
+@linalg_structured_op
+def matmul(A=TensorDef(T, S.M, S.K),
+           B=TensorDef(T, S.K, S.N),
+           C=TensorDef(U, S.M, S.N, output=True)):
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])

mlir/test/Bindings/Python/tools/linalg_opdsl/doctests.py

+# RUN: %PYTHON %s
+
+import doctest
+import importlib
+
+def test_module(module_name):
+  print(f"--- Testing module: {module_name}")
+  m = importlib.import_module(module_name)
+  doctest.testmod(m, verbose=True, raise_on_error=True, report=True)
+
+
+test_module("mlir.tools.linalg_opdsl.lang.affine")
+test_module("mlir.tools.linalg_opdsl.lang.types")

mlir/test/Bindings/Python/tools/linalg_opdsl/interfaces.py

+# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+
+from mlir.tools.linalg_opdsl.lang import *
+
+# CHECK: ---
+# CHECK-LABEL: matmul
+# CHECK:      implements:
+# CHECK-NEXT: - LinalgContractionOpInterface
+@linalg_structured_op
+def matmul(A=TensorDef(T, S.M, S.K),
+           B=TensorDef(T, S.K, S.N),
+           C=TensorDef(U, S.M, S.N, output=True)):
+  implements(ContractionOpInterface)
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])

mlir/test/Bindings/Python/tools/linalg_opdsl/lit.local.cfg

+# TODO: This tool requires PyYAML, which is not yet a required build/test
+# dependency. Remove this exclusion once it is a required dep.
+
+# Since both lit and the python bindings use the same python interpreter,
+# we can just check whether yaml can be imported here and exclude if not.
+try:
+  import yaml
+except ModuleNotFoundError:
+  config.unsupported = True

mlir/test/Bindings/Python/tools/linalg_opdsl/shape_maps_iteration.py

+# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib --file %s | FileCheck %s
+
+from mlir.tools.linalg_opdsl.lang import *
+
+
+# Verify that simple case with iteration order defined lexically and reduction
+# dims auto discovered emits the right shape, indexing maps and iterator types.
+# CHECK: ---
+# CHECK-LABEL: matmul
+# CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
+# CHECK: shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
+# CHECK: shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
+# CHECK: static_indexing_maps:
+# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
+# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
+# CHECK-NEXT: - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
+# CHECK: iterator_types:
+# CHECK-NEXT: - parallel
+# CHECK-NEXT: - parallel
+# CHECK-NEXT: - reduction
+@linalg_structured_op
+def matmul(A=TensorDef(T, S.M, S.K),
+           B=TensorDef(T, S.K, S.N),
+           C=TensorDef(U, S.M, S.N, output=True)):
+  C[D.m, D.n] += cast(U, A[D.m, D.k]) * cast(U, B[D.k, D.n])
+
+
+# Verifies that assignment to a scalar (represented as [None]) is represented
+# correctly.
+# CHECK: ---
+# CHECK-LABEL: dot
+# CHECK: shape: affine_map<()[s0] -> (s0)>
+# CHECK: shape: affine_map<()[s0] -> (s0)>
+# CHECK: shape: affine_map<()[s0] -> ()>
+# CHECK: static_indexing_maps:
+# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
+# CHECK-NEXT: - affine_map<(d0)[s0] -> (d0)>
+# CHECK-NEXT: - affine_map<(d0)[s0] -> ()>
+# CHECK: iterator_types:
+# CHECK-NEXT: - reduction
+@linalg_structured_op
+def dot(A=TensorDef(T, S.M), B=TensorDef(T, S.M), C=TensorDef(U, output=True)):
+  C[None] += cast(U, A[D.m]) * cast(U, B[D.m])

mlir/test/Bindings/Python/tools/linalg_opdsl/test_core_named_ops.py

+# RUN: %PYTHON -m mlir.tools.linalg_opdsl.dump_oplib .ops.core_named_ops | FileCheck %s
+
+# Just verify that at least one known op is generated.
+# CHECK: name: matmul