[LinalgExt] Add gather operation (1/5) (iree-org#20460)

Adds `iree_linalg_ext.gather` operation which is the converse of
`iree_linalg_ext.scatter`. Both operations share similar semantics, but
while scatter writes values into `original`, gather reads values from
`source` based on `indices`.

---------

Signed-off-by: Ian Wood <[email protected]>

Author: IanWood1

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.cpp

@@ -133,29 +133,35 @@ static bool isSmallerThan(ArrayRef<int64_t> sourceShape,
                       });
 }
 
-//===----------------------------------------------------------------------===//
-// ScatterOp
-//===----------------------------------------------------------------------===//
-
-LogicalResult ScatterOp::verify() {
-  Operation *op = getOperation();
-  if (getInputs().size() != 2) {
-    return op->emitOpError("expected two input operands");
+/// Helper function to verify both `scatter` and `gather`. Since both ops share
+/// the same sementics, we can use the same function to verify them. Note: this
+/// is written from the perspective of `scatter` op. For gather, `updateType`
+/// maps to the type of the output and `originalType` maps to the type of the
+/// `source`.
+template <typename OpTy>
+static LogicalResult
+verifyGatherScatter(OpTy op, int64_t sliceRank, ShapedType originalType,
+                    ShapedType updateType, StringRef originalName,
+                    StringRef updateName) {
+  static_assert(llvm::is_one_of<OpTy, GatherOp, ScatterOp>::value,
+                "applies to only gather or scatter operations");
+  if (op.getInputs().size() != 2) {
+    return op.emitOpError("expected two input operands");
   }
-  if (getOutputs().size() != 1) {
-    return op->emitOpError("expected one output operand");
+  if (op.getOutputs().size() != 1) {
+    return op.emitOpError("expected one output operand");
   }
 
-  auto indicesType = getIndicesType();
+  auto indicesType = op.getIndicesType();
   if (indicesType.getRank() < 1 ||
       !isa<IntegerType>(indicesType.getElementType())) {
     return op->emitOpError("expected indices to be of rank 1 or greater and of "
                            "integer element type");
   }
 
-  ArrayRef<int64_t> dimMap = getDimensionMap();
+  ArrayRef<int64_t> dimMap = op.getDimensionMap();
   if (failed(isPermSequence(
-          [&]() { return this->emitOpError("dimension map is invalid."); },
+          [&]() { return op->emitOpError("dimension map is invalid."); },
           dimMap))) {
     return failure();
   }
@@ -164,23 +170,24 @@ LogicalResult ScatterOp::verify() {
     return op->emitOpError("dimension map must have at least one element");
   }
 
-  const size_t indexDepth = getIndexDepth();
-  auto originalType = getOriginalType();
-  auto updateType = getUpdateType();
+  const size_t indexDepth = op.getIndexDepth();
   const auto originalSliceRank = originalType.getRank() - indexDepth;
   if (originalSliceRank < 0) {
-    return op->emitOpError(
-        "expected original rank to be greater or equal to index depth");
+    return op->emitOpError("expected " + originalName +
+                           " rank to be greater or equal to index depth");
   }
   if (updateType.getRank() < originalSliceRank) {
-    return op->emitOpError(
-        "expected update to be at least the rank of non indexed original dims");
+    return op->emitOpError("expected " + updateName +
+                           " to be at least the rank of non indexed " +
+                           originalName + " dims");
   }
   const size_t batchRank = updateType.getRank() - originalSliceRank;
 
   if (updateType.getRank() - batchRank != originalSliceRank) {
-    return op->emitOpError("expected rank of update value - batch rank to be "
-                           "equal to rank of original value - index depth");
+    return op->emitOpError("expected rank of " + updateName +
+                           " value - batch rank to be "
+                           "equal to rank of " +
+                           originalName + " value - index depth");
   }
 
   if ((indicesType.getRank() != batchRank || indexDepth != 1) &&
@@ -196,8 +203,8 @@ LogicalResult ScatterOp::verify() {
         llvm::mismatch(indicesType.getShape().take_front(batchRank),
                        updateType.getShape().take_front(batchRank));
     if (indicesIt != indicesType.getShape().take_front(batchRank).end()) {
-      return op->emitOpError(
-                 "mismatch in shape of indices and update value at dim#")
+      return op->emitOpError("mismatch in shape of indices and " + updateName +
+                             " value at dim#")
              << (indicesIt - indicesType.getShape().begin());
     }
   }
@@ -208,7 +215,7 @@ LogicalResult ScatterOp::verify() {
   }
 
   {
-    for (auto idx : llvm::seq<int64_t>(0, getUpdateSliceRank())) {
+    for (auto idx : llvm::seq<int64_t>(0, sliceRank)) {
       int64_t updateDim = idx + batchRank;
       int64_t origDim = idx + indexDepth;
       if (originalType.isDynamicDim(origDim) ||
@@ -217,14 +224,14 @@ LogicalResult ScatterOp::verify() {
       }
       if (originalType.getDimSize(origDim) !=
           updateType.getDimSize(updateDim)) {
-        return op->emitOpError("shape of update value dim#")
-               << (updateDim) << " must match original value at dim#"
-               << (origDim);
+        return op->emitOpError("shape of " + updateName + " value dim#")
+               << (updateDim)
+               << " must match " + originalName + " value at dim#" << (origDim);
       }
     }
   }
 
-  Region &region = this->getRegion();
+  Region &region = op.getRegion();
   Block *body = &region.front();
   if (body->getNumArguments() != 2) {
     return op->emitOpError("expected region to have two arguments");
@@ -238,12 +245,12 @@ LogicalResult ScatterOp::verify() {
   }
   if (arg0Type != updateType.getElementType()) {
     return op->emitOpError("mismatch in argument 0 of region ")
-           << arg0Type << " and element type of update value "
+           << arg0Type << " and element type of " + updateName + " value "
            << updateType.getElementType();
   }
   if (arg1Type != originalType.getElementType()) {
     return op->emitOpError("mismatch in argument 1 of region ")
-           << arg1Type << " and element type of original value "
+           << arg1Type << " and element type of " + originalName + " value "
            << originalType.getElementType();
   }
   if (arg0Type != arg1Type) {
@@ -262,6 +269,15 @@ LogicalResult ScatterOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// ScatterOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult ScatterOp::verify() {
+  return verifyGatherScatter(*this, getUpdateSliceRank(), getOriginalType(),
+                             getUpdateType(), "original", "update");
+}
+
 LogicalResult
 ScatterOp::reifyResultShapes(OpBuilder &b,
                              ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
@@ -285,6 +301,22 @@ SmallVector<AffineMap> ScatterOp::getIndexingMapsForResults() {
   return {AffineMap(nullptr)};
 }
 
+//===----------------------------------------------------------------------===//
+// GatherOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GatherOp::verify() {
+  return verifyGatherScatter(*this, getOutputSliceRank(), getSourceType(),
+                             getOutputType(), "source", "output");
+}
+
+LogicalResult
+GatherOp::reifyResultShapes(OpBuilder &b,
+                            ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+  return cast<LinalgExtOp>(getOperation())
+      .reifyResultShapes(b, reifiedReturnShapes);
+}
+
 //===----------------------------------------------------------------------===//
 // SortOp
 //===----------------------------------------------------------------------===//
@@ -1950,6 +1982,7 @@ LogicalResult IREE::LinalgExt::IndexOp::verify() {
   }
 
 DEFINE_OP_GET_EFFECTS(ScatterOp)
+DEFINE_OP_GET_EFFECTS(GatherOp)
 DEFINE_OP_GET_EFFECTS(SortOp)
 DEFINE_OP_GET_EFFECTS(FftOp)
 DEFINE_OP_GET_EFFECTS(ScanOp)

compiler/src/iree/compiler/Dialect/LinalgExt/IR/LinalgExtOps.td

@@ -223,6 +223,95 @@ def IREELinalgExt_ScatterOp : IREELinalgExt_Op<"scatter",
   }];
 }
 
+def IREELinalgExt_GatherOp : IREELinalgExt_Op<"gather",
+    [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
+  let summary = "Gather operator";
+  let description = [{
+    Takes two inputs (`source` and `indices`) and outputs value (`output`).
+    The operation returns the value at the slices specified by `indices` by
+    combining the gathered value from `source` with the value in `output`
+    using the computation specified in `region`. The `region` specifies a binary
+    operation of signature `(T, T) -> T`, where `T` is the element-type of
+    `source` & `output`. The first argument is from `source` and the second is
+    from `output`.
+
+    The size of the `dimension_map` attribute is used to determine how many
+    indices are used to index into `source`, i.e. `index_depth`. The
+    `dimension_map` attribute describes which index value maps to which dimension
+    in the destination.
+
+    This operation preforms the opposite operation of `iree_linalg_ext.scatter`.
+    Instead of scattering `updates` into `original`, it gathers the values from
+    `source` into `output` using the indices in `indices`. See the documentation
+    on `iree_linalg_ext.scatter` for more details regarding the indexing/shape
+    semantics.
+  }];
+  let arguments = (ins
+      Variadic<AnyRankedTensorOrMemRefType>:$inputs,
+      Variadic<AnyRankedTensorOrMemRefType>:$outputs,
+      DenseI64ArrayAttr:$dimension_map
+  );
+  let results = (outs Variadic<AnyRankedTensor>:$results);
+  let regions = (region AnyRegion:$region);
+  let assemblyFormat = [{
+    attr-dict `dimension_map` `=` $dimension_map
+    (`ins` `(` $inputs^ `:` type($inputs) `)`)?
+    `outs` `(` $outputs `:` type($outputs) `)`
+    $region (`->` type($results)^)?
+  }];
+
+  let extraClassDeclaration = extraLinalgExtOpClassDeclaration # [{
+    static constexpr unsigned kSourceOpNum = 0;
+    static constexpr unsigned kIndicesOpNum = 1;
+    static constexpr unsigned kResultOpNum = 2;
+
+    /// Utility to get the number of indices used to index into `source`.
+    int64_t getIndexDepth() {
+      return getDimensionMap().size();
+    }
+
+    /// Utility to get rank of the portion of `output` that is contiguous.
+    int64_t getOutputSliceRank() {
+      return getSourceType().getRank() - getIndexDepth();
+    }
+
+    /// Utility to get the rank of the portion of `indices` that represents the
+    /// batch dimensions.
+    int64_t getBatchRank() {
+      return getOutputType().getRank() - getOutputSliceRank();
+    }
+
+    Value getSource(){
+      return getOperand(kSourceOpNum);
+    }
+
+    ShapedType getSourceType(){
+      return cast<ShapedType>(getSource().getType());
+    }
+
+    Value getIndices(){
+      return getOperand(kIndicesOpNum);
+    }
+
+    ShapedType getIndicesType(){
+      return cast<ShapedType>(getIndices().getType());
+    }
+
+    Value getOutput(){
+      return getOperand(kResultOpNum);
+    }
+
+    ShapedType getOutputType(){
+      return cast<ShapedType>(getOutput().getType());
+    }
+
+    /// For DPS interface.
+    MutableOperandRange getDpsInitsMutable() {
+      return getOutputsMutable();
+    }
+  }];
+}
+
 def IREELinalgExt_SortOp : IREELinalgExt_Op<"sort",
     [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      DeclareOpInterfaceMethods<TilingInterface,

compiler/src/iree/compiler/Dialect/LinalgExt/IR/test/invalid.mlir

@@ -421,6 +421,70 @@ func.func @scatter_index_depth_too_small(
 
 // -----
 
+func.func @gather_output_too_large(
+    %source : tensor<10xf32>, %idx : tensor<1xi32>,
+    %output : tensor<2xf32>) -> tensor<2xf32> {
+  // expected-error @below {{'iree_linalg_ext.gather' op mismatch in shape of indices and output value at dim#0}}
+  %0 = iree_linalg_ext.gather
+    dimension_map = [0]
+    ins(%source, %idx : tensor<10xf32>, tensor<1xi32>)
+    outs(%output : tensor<2xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32):
+      iree_linalg_ext.yield %arg0 : f32
+  } -> tensor<2xf32>
+  return %0 : tensor<2xf32>
+}
+
+// -----
+
+func.func @gather_mismatch_output_and_source(
+    %source : tensor<10x10xf32>, %idx : tensor<2xi32>,
+    %output : tensor<1xf32>) -> tensor<1xf32> {
+  // expected-error @below {{'iree_linalg_ext.gather' op shape of output value dim#0 must match source value at dim#1}}
+  %0 = iree_linalg_ext.gather
+    dimension_map = [0]
+    ins(%source, %idx : tensor<10x10xf32>, tensor<2xi32>)
+    outs(%output : tensor<1xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32):
+      iree_linalg_ext.yield %arg0 : f32
+  } -> tensor<1xf32>
+  return %0 : tensor<1xf32>
+}
+
+// -----
+
+func.func @gather_indices_batch_rank_too_large(
+    %source : tensor<10x10xf32>, %idx : tensor<1x2xi32>,
+    %output : tensor<10xf32>) -> tensor<10xf32> {
+  // expected-error @below {{'iree_linalg_ext.gather' op expected indices to be equal to batch rank or batch rank + 1}}
+  %0 = iree_linalg_ext.gather
+    dimension_map = [0]
+    ins(%source, %idx : tensor<10x10xf32>, tensor<1x2xi32>)
+    outs(%output : tensor<10xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32):
+      iree_linalg_ext.yield %arg0 : f32
+  } -> tensor<10xf32>
+  return %0 : tensor<10xf32>
+}
+
+// -----
+
+func.func @gather_dim_map_mismatch(
+    %source : tensor<2xf32>, %idx : tensor<1xi32>,
+    %output : tensor<1xf32>) -> tensor<1xf32> {
+  // expected-error @below {{'iree_linalg_ext.gather' op expected output to be at least the rank of non indexed source dims}}
+  %0 = iree_linalg_ext.gather
+    dimension_map = [0, 1]
+    ins(%source, %idx : tensor<2xf32>, tensor<1xi32>)
+    outs(%output : tensor<1xf32>) {
+    ^bb0(%arg0: f32, %arg1: f32):
+      iree_linalg_ext.yield %arg0 : f32
+  } -> tensor<1xf32>
+  return %0 : tensor<1xf32>
+}
+
+// -----
+
 func.func @topk_invalid(%input_values: tensor<2x10xf32>, %input_indices: tensor<2x10xi32>, %out_values : tensor<2x3xf32>, %out_indices: tensor<2x3xi32>) -> (tensor<2x3xf32>, tensor<2x3xi32>) {
   // expected-error@+1 {{expected one or two input operands}}
   %0:2 = iree_linalg_ext.topk