[NFC] Make AggregateOpInterface part of mlir:: instead of linalg:: #70089
Conversation
|
@llvm/pr-subscribers-mlir-linalg @llvm/pr-subscribers-mlir Author: Abhishek Varma (Abhishek-Varma) Changes-- Currently, AggregateOpInterface is part of mlir::linalg:: namespace Signed-off-by: Abhishek Varma <abhvarma@amd.com> Full diff: https://github.com/llvm/llvm-project/pull/70089.diff 2 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 44e82f452b3cef1..99b2520b5750da9 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -817,34 +817,4 @@ def LinalgStructuredInterface
let verifyWithRegions = 1;
}
-def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
- let description = [{
- Interface for decomposing aggregated operations into a sequence of simpler
- ops.
- }];
- let cppNamespace = "::mlir::linalg";
- let methods = [
- InterfaceMethod<
- /*desc=*/[{
- Method to decompose the operation into simpler operations.
-
- On success, this method returns one `Value` per result in the
- original operation.
- The order of the returned values must match the order of the
- original values.
- In other words, the returned vector can be used directly with
- `RewriterBase::replaceOp(this, returnedValues)`.
- }],
- /*retType=*/"FailureOr<SmallVector<Value>>",
- /*methodName=*/"decomposeOperation",
- /*args=*/(ins
- "OpBuilder &":$b),
- /*methodBody=*/"",
- /*defaultImplementation=*/[{
- return {};
- }]
- >
- ];
-}
-
#endif // LINALG_IR_LINALGINTERFACES
diff --git a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
index 4c52d803e114762..fabe609d2190f4f 100644
--- a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
+++ b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
@@ -206,5 +206,35 @@ def DestinationStyleOpInterface : OpInterface<"DestinationStyleOpInterface"> {
let verifyWithRegions = 1;
}
+def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
+ let description = [{
+ Interface for decomposing aggregated operations into a sequence of simpler
+ ops.
+ }];
+ let cppNamespace = "::mlir";
+ let methods = [
+ InterfaceMethod<
+ /*desc=*/[{
+ Method to decompose the operation into simpler operations.
+
+ On success, this method returns one `Value` per result in the
+ original operation.
+ The order of the returned values must match the order of the
+ original values.
+ In other words, the returned vector can be used directly with
+ `RewriterBase::replaceOp(this, returnedValues)`.
+ }],
+ /*retType=*/"FailureOr<SmallVector<Value>>",
+ /*methodName=*/"decomposeOperation",
+ /*args=*/(ins
+ "OpBuilder &":$b),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return {};
+ }]
+ >
+ ];
+}
+
#endif // MLIR_DESTINATIONSTYLEOPINTERFACE
|
Abhishek-Varma
force-pushed
the
nfc_aggregate_op_interface
branch
3 times, most recently
from
c88b335
to
1f208cf
Compare
|
The change looks fine but what is your use case? Also, we may want to also move the |
Hi @dcaballe So, the use case occurs in IREE where we have operations from Also regarding the |
@qcolombet see comment above, but I think its more about, this doesnt need to be in Linalg, so avoiding unnecessary dependency on Linalg by moving this out. |
| def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> { | ||
| let description = [{ | ||
| Interface for decomposing aggregated operations into a sequence of simpler | ||
| ops. |
There was a problem hiding this comment.
An "aggregated operations" isn't a terminology that I'm familiar with, so not a "standard MLIR term" I suspect.
So please provide a more extensive documentation for the interface, with examples preferably.
We should know what this interface really does, what is "decomposing", when it is used and why.
There was a problem hiding this comment.
Hi @joker-eph - I've added a description as requested. Can you please take a look and let me know if it sounds better?
-- Currently, AggregateOpInterface is part of mlir::linalg:: namespace so this commit makes it part of a generic mlir:: namespace. Signed-off-by: Abhishek Varma <abhishek@nod-labs.com>
Abhishek-Varma
force-pushed
the
nfc_aggregate_op_interface
branch
from
1f208cf
to
e953a8d
Compare
| Interface. | ||
| `linalg::SoftmaxOp` is one such operation which makes use of this Interface | ||
| for implementing its decomposition. | ||
| }]; |
There was a problem hiding this comment.
Thanks for adding a doc!
This seems pretty ambiguous and under-defined to me: the whole concept of MLIR and progressive lowering is that almost every operation will be decomposed, gradually, through many layer of abstractions: so most ops (above LLVM) are fitting this description.
I still don't know when I would use this interface and ask to "decompose" such an operation, what is the use-case?
(also the Aggregate naming confused me because Aggregate is how LLVM named "struct" types).
There was a problem hiding this comment.
also the Aggregate naming confused me because Aggregate is how LLVM named "struct" types).
Oh okay, I understand. We can change the naming here perhaps, in that case.
I still don't know when I would use this interface and ask to "decompose" such an operation, what is the use-case?
I understand your point. I'll try to explain the use-case which necessitates this NFC change.
So, the idea is to perform a similar operation like linalg::Softmax for a few other Linalg-Like operations in IREE by implementing the already defined upstream interface used by linalg::Softmax here on top-of-main and then instead of having separate individual passes for their decomposition, we only have one pass which would do the needful.
I'm not sure if we're discussing about removing this interface altogether from top-of-main now or we shouldn't move it from linalg:: namespace to a more generic namespace - since the Interface already exists.
Can you confirm once? @joker-eph
Also, perhaps @MaheshRavishankar can add a better context here if I missed any pointers pertaining to the use-case.
There was a problem hiding this comment.
For context, when this interface was first introduced the concept we wanted to capture is instructions that can be represented with a sequence of simpler instructions within the same dialect.
There was a problem hiding this comment.
The same thing could be achieved with providing proper pattern rewriters, but we found that it was easier to discover when these particular patterns were part of the related "instruction" itself.
At the time we kept this interface within linalg because we would have expected that using that more broadly would warrant a proper RFC.
There was a problem hiding this comment.
See https://reviews.llvm.org/D154363 for more context
There was a problem hiding this comment.
Echo-ing @qcolombet comments. All operations eventually get decomposed, you are right. As Quentin mentioned you could have a pattern rewrite that decomposes an operation into smaller parts. We might maybe need a better term for it, but this interface is meant to group together a set of operations that exhibit a common behavior. Some examples are Softmax (that is in tree), FlashAttention and few others (that are actually out-of-tree in IREE). These are all "concepts" that are common in typical ML programmer land, and maintaining these as an "aggregate operation" in the compiler on the input better represents the programmer intent. You could do several transformations on these aggregate operations themselves (like tiling), but at some point in the compilation flow, it is better to lower them into the decomposed sequence of operations (typically a DAG of linalg-like operations) so that the rest of the compilation flow works on the decomposed sequence (like reductions or matmuls, etc). All these operations seem to follow a similar compilation flow, and the goal is to have some interfaces that these operations can implement so that the compiler just works on these interfaces. There is an in-tree example for now, softmax. We can implement the same interface on FlashAttention etc. that are not in tree as of now. (We plan to move those in tree as well if that helps, just like Softmax was moved into tree by Quentin)
So if you have a better name we could use that. (Naming is hard :) ). Also note that this is just moving an interface that is already in tree, but within Linalg, out of Linalg dialect. There doesnt seem to be a reason for this interface to live in Linalg dialect.
There was a problem hiding this comment.
Also note that this is just moving an interface that is already in tree, but within Linalg, out of Linalg dialect. There doesnt seem to be a reason for this interface to live in Linalg dialect.
I assume our messages crossed and I expanded on this above?
There was a problem hiding this comment.
I have a hard time seeing how this is something that composes across dialects or can just be reused generically. That seems quite "specific" to a specific pass at one given level of abstraction (and this is what makes it pretty contextual and not generic as I see it right now), and that's why it seems fine to me where it is.
For example would I implement this interface for the Arith dialect by splitting my
fmaintoadd+mul? This may be something I want for my use-case, but if I do this it'll be automatically picked up by the transformation when you apply it at the linalg level.
Yeah thats not the intent, but there is no way to codify it AFAICS. It is hard to predict where all these gets used. If there are issues with this, I am happy to go with one of two options
- We keep the interface in Linalg...
- We move the interface out of core into IREE... Since there we have more control of which ops implement the interface, we can prevent unintended uses of it.
There was a problem hiding this comment.
I'm fine with these 2 options, but none of them is consistent with "land this PR" right now.
There was a problem hiding this comment.
For example would I implement this interface for the Arith dialect by splitting my
fmaintoadd+mul? This may be something I want for my use-case, but if I do this it'll be automatically picked up by the transformation when you apply it at the linalg level.
That's a good point and I believe we wouldn't want that thing to be picked up at the linalg level. That being said, this is transformation specific, meaning we could have the transformation pick the level it wants or all levels if it is what the specific use case is. E.g., if op.implementThisInterface and op.isDialect(linalg). I'm not saying this is what we want (nor if it is even possible), but assuming it is, then that could be a way to have the interface being more controllable.
Not arguing either way for this PR (accept or reject), ultimately what I wanted to achieve with the original PR (admittedly without giving it much thoughts) was to do a better job at describing / discovering what is supported by the conversion patterns and what should be expanded before conversion. The idea of this interface was to capture what needed to be expanded.
could you give some feedback, then how to proceed then. This interface is not being added. It is in the Linalg dialect (where it really does not belong). So this PR is only moving it. |
I thought you concluded in the thread above where you wrote:
I'm confused by your comment there? Right now the "service" this interface provide is contextual: that is for a given op, the implementation you want from the will differ based on the pass, and this as-is can't be a generic interface, but looks like a pass-specific interface. |
so... leave it where it is, and drop this PR? IMO this is in the wrong place under Linalg... Id rather move it out of there, and be very prescriptive of which ops in tree to add this interface to (add a big red warning flag saying this is to be used with caution). But I am not tied to this. I am happy to leave it in Linalg. I dont want to move it to IREE because this has upstream dependencies, i.e. Softmax op in Linalg dialect (which also doesnt belong in Linalg, but thats a different story). |
|
While we resolve this, @Abhishek-Varma instead of moving this to |
-- Currently, AggregateOpInterface is part of mlir::linalg:: namespace
so this commit makes it part of a generic mlir:: namespace.
Signed-off-by: Abhishek Varma abhishek@nod-labs.com