[Mlir-commits] [mlir] [llvm] [mlir][Interfaces] Add `SubsetOpInterface` and `SubsetExtractionOpInterface` (PR #70617)

[Nicolas Vasilache]

================
@@ -1,34 +1,113 @@
-//===-- SubsetInsertionOpInterface.td - Tensor Subsets -----*- tablegen -*-===//
+//===-- SubsetOpInterface.td - Tensor Subsets --------------*- 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 SUBSET_INSERTION_OP_INTERFACE
-#define SUBSET_INSERTION_OP_INTERFACE
+#ifndef SUBSET_OP_INTERFACE
+#define SUBSET_OP_INTERFACE
 
 include "mlir/IR/OpBase.td"
 
-def SubsetInsertionOpInterface : OpInterface<"SubsetInsertionOpInterface"> {
+def SubsetOpInterface : OpInterface<"SubsetOpInterface"> {
   let description = [{
-    This interface can be implemented by ops that insert a source tensor into
-    a destination tensor.
-
-    The elements in the destination tensor that are overwritten by this
-    insertion are called the "subset". How the subset is defined is up to the
-    op. E.g., "tensor.insert_slice" defines the subset via a hyperrectangular
-    slice. A scatter operation could define the subset via a list of indices.
-
-    Ops that deal with tensor subsets come in two flavours:
-    - Insertion flavor: Ops that insert a source tensor into a destination
-      tensor at the specified subset. Such ops usually return a new destination
-      tensor and implement the `DestinationStyleOpInterface`. Insertion ops can
-      implement the `SubsetInsertionOpInterface`. Example: "tensor.insert_slice"
-    - Extraction flavor: Ops that define a tensor subset. They extract a
-      specified subset from a tensor. There is currently no op interface for
-      such ops. Example: "tensor.extract_slice"
+    This interface can be implemented by ops that operate on tensor subsets. A
+    "subset" is a part of a tensor. This interface describes the subset that
+    an implementing op operates on. Only the specified subset may be accessed by
+    the op.
+
+    Subset ops come in two flavours and ops that implement the
+    `SubsetOpInterface` must also implement one of the respective interfaces.
+    - Insertion flavor: Ops that insert a source value into a destination
+      tensor at the specified subset. Such ops return an updated destination
+      tensor and usually implement the `DestinationStyleOpInterface`. Insertion
+      ops must implement the `SubsetInsertionOpInterface`.
+    - Extraction flavor: Ops that extract at a subset. Extraction ops must
+      implement the `SubsetExtractionOpInterface`.
+
+    How the subset is specified is up to the implementing op. E.g.:
+    - `tensor.extract_slice/insert_slice` describe the subset as a
+      hyperrectangular slice.
+    - `tensor.gather/scatter` describe the subset as list of indices. (Not
+      implemented yet.)
+
+    Note: This interface does not expose any interface methods to get a
+    description of the accessed subset. That is because there is currently no
+    efficient way to describe arbitrary subsets. This interface merely provides
+    interface methods to check if two subsets are equivalent or disjoint.
+  }];
+
+  let cppNamespace = "::mlir";
+  let methods = [
+      InterfaceMethod<
+        /*desc=*/[{
+          Return "true" if this op and the given candidate subset op operate on
+          an equivalent subset. Return "false" is the two subsets are disjoint
+          or cannot be proven to be equivalent.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"operatesOnEquivalentSubset",
+        /*args=*/(ins
+            "::mlir::SubsetOpInterface":$candidate,
+            "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn)
+      >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return "true" if this op and the given candidate subset op operate on
+          disjoint subsets. Return "false" is the two subsets are equivalent,
+          overlapping or cannot be proven to be disjoint.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"operatesOnDisjointSubset",
+        /*args=*/(ins
+            "::mlir::SubsetOpInterface":$candidate,
+            "::llvm::function_ref<bool(Value, Value)>":$equivalenceFn)
+      >,
+  ];
+
+  let verify = [{
+    return ::mlir::detail::verifySubsetOpInterface(
+        ::mlir::cast<::mlir::SubsetOpInterface>($_op));
+  }];
+}
+
+def SubsetExtractionOpInterface
+    : OpInterface<"SubsetExtractionOpInterface", [SubsetOpInterface]> {
+  let description = [{
+    This interface can be implemented by ops that extract a value from
+    a source tensor at a specified subset. The elements in the source tensor
+    that are read by this extraction are called "subset".
+
+    Extraction ops must have a single result value.
+  }];
+
+  let cppNamespace = "::mlir";
+  let methods = [
+      InterfaceMethod<
+        /*desc=*/[{
+          Return the source tensor operand.
+        }],
+        /*retType=*/"::mlir::OpOperand &",
+        /*methodName=*/"getSourceOperand",
+        /*args=*/(ins)
+      >,
+  ];
+
+  let verify = [{
+    return ::mlir::detail::verifySubsetExtractionOpInterface(
+        ::mlir::cast<::mlir::SubsetExtractionOpInterface>($_op));
+  }];
+}
+
+def SubsetInsertionOpInterface
+    : OpInterface<"SubsetInsertionOpInterface", [SubsetOpInterface]> {
+  let description = [{
+    This interface can be implemented by ops that insert a source value into
+    a destination tensor at a specified subset. The elements in the destination
+    tensor that are overwritten by this insertion are called "subset". The
+    updates destination tensor is returned.
----------------
nicolasvasilache wrote:

nit: updated

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