Questions about the Krnl dialect

[gramalingam]

A couple of questions about the Krnl dialect.

(a) What is the purpose of the "define_loops" and "optimize_loops" instructions? Are there any examples illustrating its usage? It seems to serve as a way of naming loops and to help organize loop-transformations (may be), and it is not clear to me this is the best way to do that. E.g., I find it potentially confusing to think of a "loop" as a "Value"; it is not meant to be a run-time value, it is meant to be a static-concept relevant only while compiling.

(b) I assume "iterate" is meant to be a sequential loop/iteration? It seems useful to capture parallel iterations/loops and reduction loops explicitly in the instruction, instead of having to rely on a compiler-analysis to figure it out.

(c) MLIR's linalg.generic separates the indexing operations on tensors from the computation that is done on those values (rather than relying on a compiler to figure those out). Is there any thought/plan on doing this?

[tjingrant]

Sorry for the delayed response, if you want a quicker response, feel free to ping my work email. It's a bit of a pity that we never get to talk about Krnl Dialect in much detail. And sorry that our current focus has not been on designing/engineering it to be more like what we've imagined them to be.

For a), yes that's a very true. The goal is to have ways to refer to these loops in scheduling primitives. The point of confusion you raised is for sure valid concern, the way i see it is that krnl dialect provides a set of meta-programming constructs, and such a sword is inherently double-edged: on the one hand, it allows us to perform higher-order optimization and on the other these meta-concepts are inevitably confusing.

b). Yes, explicitly conveying the reduction semantics is very important. This is also what I got from talking to other compiler/PL folks. For sure we will need to do that.

c). I'm sitting on the fence for this one. I can see from a polyhedral optimization's point of view why this is appealing (SCoP, access pattern parsing/extraction have been a nightmare to deal with in traditional toolchain), but the facts that 1. MLIR is very disciplined, 2. ML/DL workloads have much simpler access patterns mean that these information won't be difficult to derive, should we need them down the compilation pipeline.

[gramalingam]

Hi @tjingrant : thanks for the response. With regards to (a): are there examples illustrating the requirements? Why is it not possible to use "labels" to identify instructions, or use "attributes" attached to instructions?

[tjingrant]

Hello @gramalingam , I can't give you any IR examples right now because they are all sketches and I hesitate to post them prematurely. The requirements follow from the fact that mapping between optimizations (or schedules) and the loops they act on can be on-to-many (loop collapsing) or perhaps many-to-one. So attributes wouldn't work because it assumes a one-to-one mapping between optimizations and loops.

As for approach using labels, it might work, however that would require instructions to have names, then we can refer to them using symbol reference. But I hardly see this as the case for loop instructions/operations.

[gramalingam]

Hi @tjingrant :

Popping up a higher-level, let me try to describe my high-level thoughts. I currently see that there is (a) The ONNX dialect, and (b) The Krnl dialect.

As I understand it, the Krnl dialect is meant to illustrate how one might lower from ONNX dialect to lower-level MLIR dialects or other backend-specific dialects. As an illustrative example, of course, it is fine to experiment with Krnl as you see fit.

However, I think most other users will also benefit from a Krnl-level dialect instead of an ONNX-level dialect simply because of the number of instructions in the two dialects. It is better if other users have a "smaller and simpler core dialect" to deal with when lowering to other platforms/backends etc.

Would it make sense to define a simpler "pre-Krnl" dialect for this purpose that can be shared by the community? What I am thinking of is just hiding instructions like "define_loops" and "optimize_loops" and retaining only instructions like "iterate" in this "pre-Krnl" dialect. So, the "PreKrnl.iterate" would expand into a sequence consisting of "krnl.define_loops; krnl.optimize_loops; krnl.iterate" in a trivial fashion. Lowering from ONNX to PreKrnl does not have to worry about "define_loops" or "optimize_loops". Other users targetting other backends or other optimizers don't have to worry about "define_loops" or "optimize_loops" either.

Does this make sense?

Thanks.

[tjingrant]

@gramalingam this totally makes sense to me.

In fact, I have been trying to remove these define_loops & optimize_loops operations lately because they were introduced at a time when we have very little clue about MLIR; and it's pretty clear to us now that they are fairly redundant.

Let me clean it up and get back to you.

[gramalingam]

@tjingrant , that would be great, thanks!