Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[mlir][Interfaces] Clean up DestinationStyleOpInterface #67015

Merged

Conversation

matthias-springer
Copy link
Member

  • "init" operands are specified with MutableOperandRange (which gives access to the underlying OpOperand *). No more magic numbers.
  • Remove most interface methods and make them helper functions. Only getInitsMutable should be implemented.
  • Provide separate helper functions for accessing mutable/immutable operands (OpOperand/Value, in line with [mlir][IR] Change MutableOperandRange::operator[] to return an OpOperand & #66515): getInitsMutable and getInits (same naming convention as auto-generated op accessors). getInputOperands was not renamed because this function cannot return a MutableOperandRange (because the operands are not necessarily consecutive). OpOperandVector is no longer needed.
  • The new getDpsInits/getDpsInitsMutable is more efficient than the old getDpsInitOperands because no SmallVector is created. The new functions return a range of operands.
  • Fix a bug in getDpsInputOperands: out-of-bounds operands were potentially returned.

@llvmbot
Copy link
Collaborator

llvmbot commented Sep 21, 2023

@llvm/pr-subscribers-mlir-bufferization
@llvm/pr-subscribers-mlir-scf
@llvm/pr-subscribers-mlir-vector
@llvm/pr-subscribers-mlir-linalg
@llvm/pr-subscribers-mlir-tensor
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Changes
  • "init" operands are specified with MutableOperandRange (which gives access to the underlying OpOperand *). No more magic numbers.
  • Remove most interface methods and make them helper functions. Only getInitsMutable should be implemented.
  • Provide separate helper functions for accessing mutable/immutable operands (OpOperand/Value, in line with #66515): getInitsMutable and getInits (same naming convention as auto-generated op accessors). getInputOperands was not renamed because this function cannot return a MutableOperandRange (because the operands are not necessarily consecutive). OpOperandVector is no longer needed.
  • The new getDpsInits/getDpsInitsMutable is more efficient than the old getDpsInitOperands because no SmallVector is created. The new functions return a range of operands.
  • Fix a bug in getDpsInputOperands: out-of-bounds operands were potentially returned.

Patch is 80.32 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/67015.diff

34 Files Affected:

  • (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td (+1-3)
  • (modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td (+2-2)
  • (modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td (+1-8)
  • (modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td (+8-20)
  • (modified) mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td (+3-10)
  • (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+3-5)
  • (modified) mlir/include/mlir/Interfaces/DestinationStyleOpInterface.h (-5)
  • (modified) mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td (+129-226)
  • (modified) mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp (+9-10)
  • (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+32-37)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/BubbleUpExtractSlice.cpp (+2-2)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/ConstantFold.cpp (+2-2)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/DecomposeLinalgOps.cpp (+4-4)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp (+11-10)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp (+23-23)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/EliminateEmptyTensors.cpp (+5-5)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/EraseUnusedOperandsAndResults.cpp (+25-18)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Fusion.cpp (+2-2)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Generalization.cpp (+2-2)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/InlineScalarOperands.cpp (+3-2)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Loops.cpp (+8-7)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Padding.cpp (+3-3)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/SplitReduction.cpp (+9-10)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Tiling.cpp (+9-9)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp (+2-3)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp (+7-5)
  • (modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+3-3)
  • (modified) mlir/lib/Dialect/Linalg/Utils/Utils.cpp (+16-17)
  • (modified) mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp (+3-2)
  • (modified) mlir/lib/Interfaces/DestinationStyleOpInterface.cpp (+4-12)
  • (modified) mlir/test/lib/Dialect/Linalg/TestLinalgElementwiseFusion.cpp (+1-1)
  • (modified) mlir/test/lib/Dialect/Test/TestOps.td (+6-9)
  • (modified) mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml (+2-3)
  • (modified) mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp (+3-4)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 6fcba25a0f29752..9761ab12134ad28 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -264,9 +264,7 @@ def Bufferization_MaterializeInDestinationOp
       return ::llvm::cast<RankedTensorType>(getResult().getType());
     }
 
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {1, 2};  // `dest` operand
-    }
+    MutableOperandRange getDpsInitsMutable() { return getDestMutable(); }
   }];
 
   let assemblyFormat = "$source `in` $dest attr-dict `:` type($source)";
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
index 839861c2369ca1d..9ca029b489ad144 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -555,12 +555,12 @@ def LinalgStructuredInterface
         are expection. For example, in `map` output operand isn't used in
         the block.
       }],
-      /*retTy=*/"OpOperandVector",
+      /*retTy=*/"::llvm::SmallVector<OpOperand *>",
       /*methodName=*/"getOpOperandsMatchingBBargs",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        OpOperandVector result;
+        ::llvm::SmallVector<OpOperand *> result;
         result.reserve($_op->getNumOperands());
         llvm::transform(
           this->getOperation()->getOpOperands(),
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
index 4d06747a05d6350..da12e7c83b22b89 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
@@ -149,14 +149,7 @@ def Linalg_SoftmaxOp : Linalg_Op<"softmax",
     int64_t getOutputOperandRank() {
       return getOutputOperandType().getRank();
     }
-    // Method to implement DestinationStyleOpInterface.
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      std::pair<unsigned, unsigned> outputsIndexAndLength =
-        getODSOperandIndexAndLength(1);
-      return std::make_pair<int64_t, int64_t>(
-          outputsIndexAndLength.first,
-          outputsIndexAndLength.first + outputsIndexAndLength.second);
-    }
+    MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
   }];
   let hasVerifier = 1;
 }
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
index c8d579949dc4eb6..21a5e5cc47aeb5c 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -207,10 +207,8 @@ def GenericOp : LinalgStructuredBase_Op<"generic", [
     getRegionBuilder() {
       return nullptr;
     }
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      int64_t getNumOperands = this->getNumOperands();
-      return {getNumOperands - getOutputs().size(), getNumOperands};
-    }
+
+    MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
   }];
 
   let hasCanonicalizer = 1;
@@ -283,11 +281,9 @@ def MapOp : LinalgStructuredBase_Op<"map", [
     }
 
     // Implement functions necessary for DestinationStyleOpInterface.
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      int64_t getNumOperands = this->getNumOperands();
-      return {getNumOperands - 1, getNumOperands};
-    }
-    OpOperandVector getOpOperandsMatchingBBargs() {
+    MutableOperandRange getDpsInitsMutable() { return getInitMutable(); }
+
+    SmallVector<OpOperand *> getOpOperandsMatchingBBargs() {
       return getDpsInputOperands();
     }
 
@@ -381,9 +377,7 @@ def ReduceOp : LinalgStructuredBase_Op<"reduce", [
     getRegionBuilder() {
       return nullptr;
     }
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {getInits().size(), getNumOperands()};
-    }
+    MutableOperandRange getDpsInitsMutable() { return getInitsMutable(); }
   }];
 
   let hasCustomAssemblyFormat = 1;
@@ -446,10 +440,7 @@ def TransposeOp : LinalgStructuredBase_Op<"transpose", [
     }
 
     // Implement functions necessary for DestinationStyleOpInterface.
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      int64_t getNumOperands = this->getNumOperands();
-      return {getNumOperands - 1, getNumOperands};
-    }
+    MutableOperandRange getDpsInitsMutable() { return getInitMutable(); }
 
     static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
         mlir::ArrayRef<mlir::NamedAttribute>)>
@@ -517,10 +508,7 @@ def BroadcastOp : LinalgStructuredBase_Op<"broadcast", [
     }
 
     // Implement functions necessary for DestinationStyleOpInterface.
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      int64_t getNumOperands = this->getNumOperands();
-      return {getNumOperands - 1, getNumOperands};
-    }
+    MutableOperandRange getDpsInitsMutable() { return getInitMutable(); }
 
     static std::function<void(mlir::ImplicitLocOpBuilder &, mlir::Block &,
         mlir::ArrayRef<mlir::NamedAttribute>)>
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index d1c33d8b4c03c3f..86a250b77dcc8ee 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -750,9 +750,7 @@ def Tensor_InsertOp : Tensor_Op<"insert", [
   }];
 
   let extraClassDeclaration = [{
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {1, 2};  // `dest` operand
-    }
+    MutableOperandRange getDpsInitsMutable() { return getDestMutable(); }
   }];
 
   let hasFolder = 1;
@@ -892,9 +890,7 @@ def Tensor_InsertSliceOp : Tensor_OpWithOffsetSizesAndStrides<"insert_slice", [
     /// and `strides` operands.
     static unsigned getOffsetSizeAndStrideStartOperandIndex() { return 2; }
 
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {1, 2};  // `dest` operand
-    }
+    MutableOperandRange getDpsInitsMutable() { return getDestMutable(); }
   }];
 
   let hasCanonicalizer = 1;
@@ -1714,10 +1710,7 @@ class Tensor_RelayoutOp<string mnemonic, list<Trait> traits = []> :
     RankedTensorType getDestType() {
       return ::llvm::cast<RankedTensorType>(getDest().getType()); };
 
-    /// Return position for init operand. Init operand is `dest`.
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {1, 2}; // `dest` operand
-    }
+    MutableOperandRange getDpsInitsMutable() { return getDestMutable(); }
 
     /// Interface method for ConditionallySpeculatable.
     Speculation::Speculatability getSpeculatability();
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index aba85f86a7eef9c..701eefcc1e7da6a 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -1330,8 +1330,8 @@ def Vector_TransferReadOp :
     // MaskableOpInterface methods.
     bool supportsPassthru() { return true; }
 
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {0, 0};  // empty range (no init operands)
+    MutableOperandRange getDpsInitsMutable() {
+      return MutableOperandRange(getOperation(), /*start=*/0, /*length=*/0);
     }
   }];
 
@@ -1494,9 +1494,7 @@ def Vector_TransferWriteOp :
     ///  ops of other dialects.
     Value getValue() { return getVector(); }
 
-    std::pair<int64_t, int64_t> getDpsInitsPositionRange() {
-      return {1, 2};  // `source` operand
-    }
+    MutableOperandRange getDpsInitsMutable() { return getSourceMutable(); }
   }];
 
   let hasFolder = 1;
diff --git a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.h b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.h
index 9ace95c6f3d3b2f..6649371f3ed321e 100644
--- a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.h
+++ b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.h
@@ -17,11 +17,6 @@
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
-/// OpOperand vector that implicitly converts to a Value vector.
-struct OpOperandVector : public llvm::SmallVector<OpOperand *> {
-  operator SmallVector<Value>();
-};
-
 namespace detail {
 /// Verify that `op` conforms to the invariants of DestinationStyleOpInterface
 LogicalResult verifyDestinationStyleOpInterface(Operation *op);
diff --git a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
index 8558e0279e3e5bb..6b1ab10210099cc 100644
--- a/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
+++ b/mlir/include/mlir/Interfaces/DestinationStyleOpInterface.td
@@ -13,16 +13,16 @@ include "mlir/IR/OpBase.td"
 
 def DestinationStyleOpInterface : OpInterface<"DestinationStyleOpInterface"> {
   let description = [{
-    Ops that are in destination style have designated init operands, which act
+    Ops that are in destination style have designated "init" operands, which act
     as initial tensor values for the results of the operation or the init
     buffers to which the results of the op will be written.
 
     Init operands must be ranked tensors or ranked memrefs. Input operands can
     have any type. All non-init operands are DPS inputs.
 
-    It is assumed that the init operands of the op are the operands at
-    position [start, end). The positions are defined by getDpsInitsPositionRange
-    method.
+    The init operands of this op are specified by the MutableOperandRange that
+    the `getDpsInitsMutable` interface methods returns. This implies that the
+    init operands must be a consecutive range of operands.
 
     If the op has "tensor semantics", then the input operands are either ranked
     tensors or other non-tensor/memref types ("scalars"). The init operands are
@@ -50,241 +50,144 @@ def DestinationStyleOpInterface : OpInterface<"DestinationStyleOpInterface"> {
     Example of an op that is not in destination style: `%r = tensor.pad %t`.
     This op is not in destination style because `%r` and `%t` have different
     shape.
-
-    Each op that wants to implement DestinationStyleOpInterface needs to define
-    the getDpsInitsPositionRange() method.
   }];
 
   let cppNamespace = "::mlir";
 
   let methods = [
-    // This method has to be defined for every DPS op.
     InterfaceMethod<
       /*desc=*/"Return start and end indices of the init operands range.",
-      /*retTy=*/"std::pair<int64_t, int64_t>",
-      /*methodName=*/"getDpsInitsPositionRange",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/""
-    >,
-    //===------------------------------------------------------------------===//
-    // Operands handling.
-    //===------------------------------------------------------------------===//
-    // The operand list is assumed to start with the input operands and end
-    // with the init operands. Therefore, all methods to access the inputs
-    // and inits can be expressed if the number of init operands is know.
-    InterfaceMethod<
-      /*desc=*/"Return the number of inits.",
-      /*retTy=*/"int64_t",
-      /*methodName=*/"getNumDpsInits",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        return end - start;
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return the init operands.",
-      /*retTy=*/"::mlir::OpOperandVector",
-      /*methodName=*/"getDpsInitOperands",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-
-        ::mlir::OpOperandVector result;
-        result.reserve(end - start);
-        for (int i = start; i < end; ++i)
-          result.push_back(&$_op->getOpOperand(i));
-        return result;
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return the `i`-th init operand.",
-      /*retTy=*/"::mlir::OpOperand *",
-      /*methodName=*/"getDpsInitOperand",
-      /*args=*/(ins "int64_t":$i),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        assert(i >= 0 && i < $_op.getNumDpsInits());
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        return &$_op->getOpOperand(start + i);
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Set the `i`-th init operand.",
-      /*retTy=*/"void",
-      /*methodName=*/"setDpsInitOperand",
-      /*args=*/(ins "int64_t":$i, "::mlir::Value":$value),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        assert(i >= 0 && i < $_op.getNumDpsInits());
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        $_op->setOperand(start + i, value);
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return the number of inputs.",
-      /*retTy=*/"int64_t",
-      /*methodName=*/"getNumDpsInputs",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        return $_op.getNumOperands() - $_op.getNumDpsInits();
-      }]
+      /*retTy=*/"::mlir::MutableOperandRange",
+      /*methodName=*/"getDpsInitsMutable",
+      /*args=*/(ins)
     >,
-    InterfaceMethod<
-      /*desc=*/"Return the input operands.",
-      /*retTy=*/"::mlir::OpOperandVector",
-      /*methodName=*/"getDpsInputOperands",
-      /*args=*/(ins),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        int64_t numInits = end - start;
-        int64_t numOperands = $_op.getNumOperands();
-
-        ::mlir::OpOperandVector result;
-        result.reserve(numOperands - numInits);
-        for (int i = 0; i < start; ++i)
-          result.push_back(&$_op->getOpOperand(i));
-        for (int i = end; i < numOperands; ++i)
-          result.push_back(&$_op->getOpOperand(end + i));
+  ];
 
-        return result;
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/[{ Return the `i`-th input operand.  }],
-      /*retTy=*/"::mlir::OpOperand *",
-      /*methodName=*/"getDpsInputOperand",
-      /*args=*/(ins "int64_t":$i),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        assert(i >= 0 && i < getNumDpsInputs());
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        return &$_op->getOpOperand(i < start ? i : i + end - start) ;
-      }]
-    >,
-    //===------------------------------------------------------------------===//
-    // Input and DpsInit arguments handling.
-    //===------------------------------------------------------------------===//
-    InterfaceMethod<
-      /*desc=*/"Return true if `opOperand` is an input.",
-      /*retTy=*/"bool",
-      /*methodName=*/"isDpsInput",
-      /*args=*/(ins "::mlir::OpOperand *":$opOperand),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        auto operandNumber = opOperand->getOperandNumber();
-        return operandNumber < start || operandNumber >= end;
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return true if `opOperand` is an init.",
-      /*retTy=*/"bool",
-      /*methodName=*/"isDpsInit",
-      /*args=*/(ins "::mlir::OpOperand *":$opOperand),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        auto operandNumber = opOperand->getOperandNumber();
-        return operandNumber >= start && operandNumber < end;
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/[{
-        Return true if the `opOperand` is a scalar value. A scalar is defined
-        as neither a memref nor a tensor value.
-      }],
-      /*retTy=*/"bool",
-      /*methodName=*/"isScalar",
-      /*args=*/(ins "::mlir::OpOperand *":$opOperand),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
+  let extraSharedClassDeclaration = [{
+    ::mlir::OperandRange getDpsInits() {
+      return $_op.getDpsInitsMutable();
+    }
+
+    /// Return the number of DPS inits.
+    int64_t getNumDpsInits() { return $_op.getDpsInits().size(); }
+
+    /// Return the `i`-th DPS init.
+    ::mlir::OpOperand *getDpsInitOperand(int64_t i) {
+      return &$_op.getDpsInitsMutable()[i];
+    }
+
+    /// Set the `i`-th DPS init.
+    void setDpsInitOperand(int64_t i, Value value) {
+      assert(i >= 0 && i < $_op.getNumDpsInits());
+      $_op->setOperand($_op.getDpsInits().getBeginOperandIndex() + i, value);
+    }
+
+    /// Return the number of DPS inits.
+    int64_t getNumDpsInputs() {
+      return $_op->getNumOperands() - $_op.getNumDpsInits();
+    }
+
+    /// Return the DPS input operands.
+    ::llvm::SmallVector<::mlir::OpOperand *> getDpsInputOperands() {
+      ::mlir::OperandRange range = $_op.getDpsInits();
+      int64_t firstInitPos = range.getBeginOperandIndex();
+      int64_t numInits = range.size();
+      int64_t numOperands = $_op->getNumOperands();
+      ::llvm::SmallVector<::mlir::OpOperand *> result;
+      result.reserve(numOperands - numInits);
+      for (int64_t i = 0; i < firstInitPos; ++i)
+        result.push_back(&$_op->getOpOperand(i));
+      for (int64_t i = firstInitPos + numInits; i < numOperands; ++i)
+        result.push_back(&$_op->getOpOperand(i));
+      return result;
+    }
+
+    /// Return the DPS input operands.
+    ::llvm::SmallVector<::mlir::Value> getDpsInputs() {
+      return ::llvm::to_vector(::llvm::map_range($_op.getDpsInputOperands(), [](OpOperand *o) { return o->get(); }));
+    }
+
+    /// Return the `i`-th DPS input operand.
+    ::mlir::OpOperand *getDpsInputOperand(int64_t i) {
+      ::mlir::OperandRange range = $_op.getDpsInits();
+      int64_t firstInitPos = range.getBeginOperandIndex();
+      int64_t numInits = range.size();
+      assert(i >= 0 && i < $_op->getNumOperands() - numInits);
+      return &$_op->getOpOperand(
+          i < firstInitPos ? i : i + firstInitPos + numInits);
+    }
+
+    /// Return "true" if `opOperand` is an "input".
+    bool isDpsInput(::mlir::OpOperand *opOperand) {
+      assert(opOperand->getOwner() == $_op);
+      return !$_op.isDpsInit(opOperand);
+    }
+
+    /// Return "true" if `opOperand` is an "init".
+    bool isDpsInit(::mlir::OpOperand *opOperand) {
+      assert(opOperand->getOwner() == $_op);
+      ::mlir::OperandRange range = $_op.getDpsInits();
+      auto operandNumber = opOperand->getOperandNumber();
+      return operandNumber >= range.getBeginOperandIndex()
+          && operandNumber < range.getBeginOperandIndex() + range.size();
+    }
+
+    /// Return "true" if `opOperand` is a scalar value. A sclar is defined as
+    /// neither a MemRef nor a tensor value.
+    bool isScalar(::mlir::OpOperand *opOperand) {
+      assert(opOperand->getOwner() == $_op.getOperation());
+      return !::llvm::isa<MemRefType, TensorType>(opOperand->get().getType());
+    }
+
+    /// Return the OpResult that is tied to the given OpOperand.
+    ::mlir::OpResult getTiedOpResult(::mlir::OpOperand *opOperand) {
         assert(opOperand->getOwner() == $_op.getOperation());
-        return !::llvm::isa<MemRefType, TensorType>(opOperand->get().getType());
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return the OpResult that is tied to the given OpOperand.",
-      /*retTy=*/"::mlir::OpResult",
-      /*methodName=*/"getTiedOpResult",
-      /*args=*/(ins "::mlir::OpOperand *":$opOperand),
-      /*methodBody=*/"",
-      /*defaultImplementation=*/[{
-        assert(opOperand->getOwner() == $_op.getOperation());
-
-        auto [start, end] = $_op.getDpsInitsPositionRange();
-        int64_t resultIndex = opOperand->getOperandNumber() - start;
+        ::mlir::OperandRange range = $_op.getDpsInits();
+        int64_t resultIndex =
+            opOperand->getOperandNumber() - range.getBeginOperandIndex();
         assert(resultIndex >= 0 &&
                resultIndex < $_op->getNumResults() );
         return $_op->getResult(resultIndex);
-      }]
-    >,
-    InterfaceMethod<
-      /*desc=*/"Return the OpOperand that is tied to the given OpResult.",
-      /*retTy=*/"::mlir::OpOperand *",
-      /*methodName=*/"getTiedOpOperand",...
[truncated]

* "init" operands are specified with `MutableOperandRange` (which gives access to the underlying `OpOperand *`). No more magic numbers.
* Remove most interface methods and make them helper functions. Only `getInitsMutable` should be implemented.
* Provide separate helper functions for accessing mutable/immutable operands (`OpOperand`/`Value`, in line with llvm#66515): `getInitsMutable` and `getInits` (same naming convention as auto-generated op accessors). `getInputOperands` was not renamed because this function cannot return a `MutableOperandRange` (because the operands are not necessarily consecutive). `OpOperandVector` is no longer needed.
* The new `getDpsInits`/`getDpsInitsMutable` is more efficient than the old `getDpsInitOperands` because no `SmallVector` is created. The new functions return a range of operands.
* Fix a bug in `getDpsInputOperands`: out-of-bounds operands were potentially returned.

BEGIN_PUBLIC
No public commit message needed for presubmit.
END_PUBLIC
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants