[Mlir-commits] [mlir] [NFC] Make AggregateOpInterface part of mlir:: instead of linalg:: (PR #70089)

[llvmlistbot at llvm.org]

================
@@ -0,0 +1,57 @@
+//===- AggregatedOpInterface.td ----------------------*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_AGGREGATEDOPINTERFACE
+#define MLIR_AGGREGATEDOPINTERFACE
+
+include "mlir/IR/OpBase.td"
+
+def AggregatedOpInterface : OpInterface<"AggregatedOpInterface"> {
+  let description = [{
+    This Interface is particularly useful in cases where we have an operation
+    that can be lowered into a sequence of simpler operations, thus essentially
+    decomposing an operation into a set of one or many simpler operations.
+    The operation being decomposed need to implement this Interface by implementing
+    the method `decomposeOperation` and return the set of values which would replace
+    the uses of the operation being decomposed.
+    Eg:
+        Assume there is an operation `CustomOp_Mul_Add` that takes in an input tensor
+        and a constant. It basically performs element-wise multiplication of the input
+        tensor with the given constant, and then performs element-wise addition of the
+        intermediate resulting tensor with the given constant.
+        `CustomOp_Mul_Add` can thus essentially be decomposed by implementing this
+        Interface.
+        `linalg::SoftmaxOp` is one such operation which makes use of this Interface
+        for implementing its decomposition.
+  }];
----------------
MaheshRavishankar wrote:

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.

https://github.com/llvm/llvm-project/pull/70089