Add CUDA GatherElements kernel (#2310)

* Updates

* Update test

* Update

* Updates

* nits

* PR feedback

* Update

* Update

* PR feedback

* PR comments

* Update

* Fix build

* Fix build

* Nits

* Fix

onnxruntime/core/providers/cpu/tensor/gather_elements.cc

@@ -41,7 +41,7 @@ static inline int64_t compute_base_offset(const std::vector<int64_t>& shape, con
 // Example: input = [2, 3]     output = 2
 //          input = [3, 2, 4]  output = 3 * 2 = 6
 //          input  = [2]       output = 1
-int64_t calculate_num_inner_dim(const TensorShape& dims) {
+static int64_t calculate_num_inner_dim(const TensorShape& dims) {
   // in this context, rank can never be < 1, so saving checking overhead
   return dims.SizeToDimension(dims.NumDimensions() - 1);
 }
@@ -79,21 +79,6 @@ static inline void increment_over_inner_dim(std::vector<int64_t>& current_dims,
 // parse indices_tensor and along the way validate its shape and contents
 static std::vector<int64_t> parse_and_validate_indices_tensor(const Tensor* indices_tensor,
                                                               int64_t axis, const TensorShape& input_shape) {
-  const auto& indices_shape = indices_tensor->Shape().GetDims();
-  int64_t indices_rank = static_cast<int64_t>(indices_shape.size());
-  for (int64_t i = 0; i < indices_rank; ++i) {
-    // for all axes except the axis of interest,
-    // make sure that the corresponding 'indices' shape
-    // value if within bounds of the corresponding 'data' shape
-    if (i != axis) {
-      if (indices_shape[i] < 0 || indices_shape[i] > input_shape[i])
-        ORT_THROW(
-            "GatherElements op: 'indices' shape should have values within bounds of 'data' shape. "
-            "Invalid value in indices shape is: ",
-            indices_shape[i]);
-    }
-  }
-
   // first parse 'indices' data
   auto num_elements = indices_tensor->Shape().Size();
   std::vector<int64_t> indices_data;
@@ -128,6 +113,7 @@ static std::vector<int64_t> parse_and_validate_indices_tensor(const Tensor* indi
 
   return indices_data;
 }
+
 #ifdef __GNUC__
 #pragma GCC diagnostic push
 #ifdef HAS_CLASS_MEMACCESS
@@ -215,23 +201,54 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor,
 #ifdef __GNUC__
 #pragma GCC diagnostic pop
 #endif
-Status GatherElements::Compute(OpKernelContext* context) const {
-  const Tensor* input_tensor = context->Input<Tensor>(0);
-  const TensorShape& input_data_shape = input_tensor->Shape();
-  int64_t input_data_rank = static_cast<int64_t>(input_data_shape.NumDimensions());
 
-  const Tensor* indices_tensor = context->Input<Tensor>(1);
-  const TensorShape& indices_shape = indices_tensor->Shape();
+Status GatherElements::ValidateInputShapes(const TensorShape& input_data_shape,
+                                           const TensorShape& indices_shape,
+                                           int64_t axis) {
+  int64_t input_data_rank = static_cast<int64_t>(input_data_shape.NumDimensions());
   int64_t indices_rank = static_cast<int64_t>(indices_shape.NumDimensions());
 
+  // GatherElements cannot operate on scalars
   if (input_data_rank < 1)
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "GatherElements op: Cannot operate on scalar input");
 
+  // The ranks of the inputs must be the same
   if (input_data_rank != indices_rank)
     return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                            "GatherElements op: Rank of input 'data' needs to be equal to rank of input 'indices'");
 
+  // Except for the axis of interest all other dim values of the 'indices' input must be within bounds
+  // of the corresponding 'data' input dim value
+  for (int64_t i = 0; i < indices_rank; ++i) {
+    // for all axes except the axis of interest,
+    // make sure that the corresponding 'indices' shape
+    // value if within bounds of the corresponding 'data' shape
+    if (i != axis) {
+      if (indices_shape[i] < 0 || indices_shape[i] > input_data_shape[i])
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "GatherElements op: 'indices' shape should have values within bounds of 'data' shape. "
+                               "Invalid value in indices shape is: ",
+                               indices_shape[i]);
+    }
+  }
+
+  return Status::OK();
+}
+
+Status GatherElements::Compute(OpKernelContext* context) const {
+  const Tensor* input_tensor = context->Input<Tensor>(0);
+  const TensorShape& input_data_shape = input_tensor->Shape();
+
+  const Tensor* indices_tensor = context->Input<Tensor>(1);
+  const TensorShape& indices_shape = indices_tensor->Shape();
+
+  int64_t axis = HandleNegativeAxis(axis_, input_data_shape.NumDimensions());
+
+  auto status = ValidateInputShapes(input_data_shape, indices_shape, axis);
+  if (!status.IsOK())
+    return status;
+
   Tensor* output_tensor = context->Output(0, TensorShape(indices_shape));
 
   const auto& input_data_type = input_tensor->DataType();
@@ -243,8 +260,6 @@ Status GatherElements::Compute(OpKernelContext* context) const {
   if (indices_shape.Size() == 0)
     return Status::OK();
 
-  int64_t axis = HandleNegativeAxis(axis_, input_data_shape.NumDimensions());
-
   if (input_data_type == DataTypeImpl::GetType<std::string>())
     core_impl<true, std::string>(input_tensor, indices_tensor, output_tensor, axis);
 

onnxruntime/core/providers/cpu/tensor/gather_elements.h

@@ -18,6 +18,11 @@ class GatherElements final : public OpKernel {
 
   Status Compute(OpKernelContext* context) const override;
 
+  // holds common checks for the CPU and CUDA GatherElements kernel
+  static Status ValidateInputShapes(const TensorShape& input_data_shape,
+                                    const TensorShape& indices_shape,
+                                    int64_t axis);  // axis might be different from the member axis_ based on the input being processed
+
  private:
   Status CoreImplString(const Tensor* input_tensor, const Tensor* indices_tensor,
                         Tensor* output_tensor, int64_t axis) const;

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

@@ -580,6 +580,7 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain,
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, NonMaxSuppression);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Range);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, ScatterElements);
+class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, GatherElements);
 
 class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 9, TopK);
 class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, TopK);
@@ -948,11 +949,10 @@ static void RegisterCudaKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, NonMaxSuppression)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, Range)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, ScatterElements)>,
-
+      BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, GatherElements)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 1, 9, TopK)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 10, 10, TopK)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 11, TopK)>,
-
   };
 
   for (auto& function_table_entry : function_table) {

onnxruntime/core/providers/cuda/tensor/gather_elements.cc

@@ -0,0 +1,104 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gather_elements.h"
+#include "gather_elements_impl.h"
+#include "core/providers/cpu/tensor/utils.h"
+#include "core/providers/common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+ONNX_OPERATOR_KERNEL_EX(
+    GatherElements,
+    kOnnxDomain,
+    11,
+    kCudaExecutionProvider,
+    KernelDefBuilder()
+        .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes())
+        .TypeConstraint("Tind", std::vector<MLDataType>{
+                                    DataTypeImpl::GetTensorType<int32_t>(),
+                                    DataTypeImpl::GetTensorType<int64_t>()}),
+    GatherElements);
+
+Status GatherElements::ComputeInternal(OpKernelContext* context) const {
+  // Process input data tensor
+  const auto* input_tensor = context->Input<Tensor>(0);
+  const auto& input_shape = input_tensor->Shape();
+  const auto& input_dims = input_shape.GetDims();
+  const int64_t input_rank = static_cast<int64_t>(input_dims.size());
+  const int64_t input_size = input_shape.Size();
+
+  // Process indices tensor
+  const auto* indices_tensor = context->Input<Tensor>(1);
+  const auto& indices_shape = indices_tensor->Shape();
+  const auto& indices_dims = indices_shape.GetDims();
+  const int64_t indices_rank = static_cast<int64_t>(indices_dims.size());
+  const int64_t indices_size = indices_shape.Size();
+
+  // Handle negative axis if any
+  const int64_t axis = static_cast<int64_t>(HandleNegativeAxis(axis_, input_rank));
+
+  // Validate input shapes and ranks (invoke the static method in the CPU GatherElements kernel that hosts the shared checks)
+  auto status = onnxruntime::GatherElements::ValidateInputShapes(input_shape, indices_shape, axis);
+  if (!status.IsOK())
+    return status;
+
+  // create output tensor
+  auto* output_tensor = context->Output(0, TensorShape(indices_shape));
+
+  // if there are no elements in 'indices' - nothing to process
+  if (indices_shape.Size() == 0)
+    return Status::OK();
+
+  TensorPitches input_strides(input_dims);
+  CudaAsyncBuffer<int64_t> gpu_input_strides(this, input_strides);
+
+  CudaAsyncBuffer<fast_divmod> fdm_indices_strides(this, indices_rank);
+  ORT_ENFORCE(CalculateFdmStrides(fdm_indices_strides.CpuSpan(), indices_dims));
+
+  ORT_RETURN_IF_ERROR(gpu_input_strides.CopyToGpu());
+  ORT_RETURN_IF_ERROR(fdm_indices_strides.CopyToGpu());
+
+  size_t element_size = input_tensor->DataType()->Size();
+  MLDataType Tin_type = indices_tensor->DataType();
+
+  if (Tin_type == DataTypeImpl::GetType<int32_t>()) {
+    const int32_t* indices_data = indices_tensor->template Data<int32_t>();
+    GatherElementsImpl<int32_t>(
+        input_rank,
+        input_tensor->DataRaw(),
+        input_size,
+        input_dims[axis],
+        gpu_input_strides.GpuPtr(),
+        indices_data,
+        indices_size,
+        fdm_indices_strides.GpuPtr(),
+        axis,
+        output_tensor->MutableDataRaw(),
+        element_size);
+    return Status::OK();
+  } else if (Tin_type == DataTypeImpl::GetType<int64_t>()) {
+    const int64_t* indices_data = indices_tensor->template Data<int64_t>();
+    GatherElementsImpl<int64_t>(
+        input_rank,
+        input_tensor->DataRaw(),
+        input_size,
+        input_dims[axis],
+        gpu_input_strides.GpuPtr(),
+        indices_data,
+        indices_size,
+        fdm_indices_strides.GpuPtr(),
+        axis,
+        output_tensor->MutableDataRaw(),
+        element_size);
+    return Status::OK();
+  } else {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "GatherElements op: Type of 'indices' must be int32 or int64");
+  }
+
+  return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "String type is not supported yet for the GatherElements op");
+}
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tensor/gather_elements.h

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+#pragma once
+
+#include "core/common/common.h"
+#include "core/framework/op_kernel.h"
+#include "core/providers/cpu/tensor/gather_elements.h"
+#include "core/providers/cuda/cuda_common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+class GatherElements final : public CudaKernel {
+ public:
+  GatherElements(const OpKernelInfo& info) : CudaKernel(info) {
+    ORT_ENFORCE(info.GetAttr<int64_t>("axis", &axis_).IsOK(),
+                "Missing/Invalid 'axis' attribute value");
+  }
+  ~GatherElements() = default;
+  Status ComputeInternal(OpKernelContext* context) const override;
+
+ private:
+  int64_t axis_;
+};
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/tensor/gather_elements_impl.cu

@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/cu_inc/common.cuh"
+#include "gather_elements_impl.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename T, typename Tin>
+__global__ void _GatherElementsKernel(
+    const int64_t rank,
+    const T* input_data,
+    const int64_t input_dim_along_axis,
+    const int64_t* input_strides,
+    const Tin* indices_data,
+    const int64_t indices_size,
+    const fast_divmod* indices_strides,
+    const int64_t axis,
+    T* output_data) {
+  CALCULATE_ELEMENTWISE_INDEX_OR_EXIT(indices_index, indices_size);
+  int dim = 0; 
+  int remain = indices_index;
+  size_t data_idx = 0;
+  for (int i = 0; i < rank; ++i) {
+    indices_strides[i].divmod(remain, dim, remain);
+    if (i == axis) {
+      dim = static_cast<int>(indices_data[indices_index]);
+      if (dim < -input_dim_along_axis || dim >= input_dim_along_axis) {
+        return; // Invalid index
+      }
+      if (dim < 0) {
+        dim += input_dim_along_axis;
+      }
+    }
+    data_idx += input_strides[i] * dim;
+  }
+  output_data[indices_index] = input_data[data_idx];
+}
+
+template <typename Tin>
+void GatherElementsImpl(
+    const int64_t rank,
+    const void* input_data,
+    const int64_t input_size,
+    const int64_t input_dim_along_axis,
+    const int64_t* input_strides,
+    const Tin* indices_data,
+    const int64_t indices_size,
+    const fast_divmod* indices_strides,
+    const int64_t axis,
+    void* output_data,
+    size_t element_size) {
+
+  if (indices_size > 0) {
+
+    int blocksPerGrid = static_cast<int>((indices_size + GridDim::maxThreadsPerBlock - 1) / GridDim::maxThreadsPerBlock);
+
+    switch (element_size) {
+      case sizeof(int8_t):
+         _GatherElementsKernel<int8_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            rank, reinterpret_cast<const ToCudaType<int8_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
+            indices_data, indices_size, indices_strides,
+            axis, reinterpret_cast<ToCudaType<int8_t>::MappedType*>(output_data));
+        break;
+
+      case sizeof(int16_t):
+        _GatherElementsKernel<int16_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            rank, reinterpret_cast<const ToCudaType<int16_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
+            indices_data, indices_size, indices_strides,
+            axis, reinterpret_cast<ToCudaType<int16_t>::MappedType*>(output_data));
+        break;
+
+      case sizeof(int32_t):
+        _GatherElementsKernel<int32_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            rank, reinterpret_cast<const ToCudaType<int32_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
+            indices_data, indices_size, indices_strides,
+            axis, reinterpret_cast<ToCudaType<int32_t>::MappedType*>(output_data));
+        break;
+
+      case sizeof(int64_t):
+        _GatherElementsKernel<int64_t, Tin><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
+            rank, reinterpret_cast<const ToCudaType<int64_t>::MappedType*>(input_data), input_dim_along_axis, input_strides,
+            indices_data, indices_size, indices_strides,
+            axis, reinterpret_cast<ToCudaType<int64_t>::MappedType*>(output_data));
+        break;
+
+      // should not reach here as we validate if the all relevant types are supported in the Compute method 
+      default:
+        ORT_THROW("Unsupported element size by the GatherElements CUDA kernel");
+    }
+  }
+}
+
+template void GatherElementsImpl<int32_t>(
+    const int64_t rank,
+    const void* input_data,
+    const int64_t input_size,
+    const int64_t input_dim_along_axis,
+    const int64_t* input_strides,
+    const int32_t* indices_data,
+    const int64_t indices_size,
+    const fast_divmod* indices_strides,
+    const int64_t axis,
+    void* output_data,
+    size_t element_size);
+
+template void GatherElementsImpl<int64_t>(
+    const int64_t rank,
+    const void* input_data,
+    const int64_t input_size,
+    const int64_t input_dim_along_axis,
+    const int64_t* input_strides,
+    const int64_t* indices_data,
+    const int64_t indices_size,
+    const fast_divmod* indices_strides,
+    const int64_t axis,
+    void* output_data,
+    size_t element_size);
+
+}  // namespace cuda
+}  // namespace onnxruntime
+

onnxruntime/core/providers/cuda/tensor/gather_elements_impl.h

@@ -0,0 +1,27 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <stdint.h>
+#include "core/providers/cuda/shared_inc/cuda_utils.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+template <typename Tin>
+void GatherElementsImpl(
+    const int64_t rank,  // both inputs have same rank and this is validated in the main Compute
+    const void* input_data,
+    const int64_t input_size,
+    const int64_t input_dim_along_axis,
+    const int64_t* input_strides,
+    const Tin* indices_data,
+    const int64_t indices_size,
+    const fast_divmod* indices_strides,
+    const int64_t axis,
+    void* output_data,
+    size_t element_size);
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc

@@ -75,9 +75,10 @@ void RunTypedTest() {
                      {1, 1,
                       4, 4});
   // skip nuphar, which will not throw error message but will ensure no out-of-bound access
+  // skip cuda as the cuda kernel won't throw the error message
   test5.Run(OpTester::ExpectResult::kExpectFailure,
             "GatherElements op: Value in indices must be within bounds [-2 , 1]. Actual value is 2",
-            {kNupharExecutionProvider});
+            {kNupharExecutionProvider, kCudaExecutionProvider});
 
   // 3D input - axis 1
   OpTester test6("GatherElements", 11);
@@ -140,6 +141,22 @@ void RunTypedTest() {
   test10.Run();
 }
 
+template <>
+void RunTypedTest<bool>() {
+  // 3D input - axis 2
+  OpTester test1("GatherElements", 11);
+  test1.AddAttribute<int64_t>("axis", 2LL);
+  test1.AddInput<bool>("data", {2, 2, 2},
+                       {true, false,
+                        true, false,
+                        true, false,
+                        true, false});
+  test1.AddInput<int64_t>("indices", {1, 2, 1},
+                          {0, 1});
+  test1.AddOutput<bool>("output", {1, 2, 1}, {true, false});
+  test1.Run();
+}
+
 template <>
 void RunTypedTest<std::string>() {
   // int32_t indices - axis 0
@@ -282,6 +299,18 @@ TEST(GatherElementsOpTest, uint64_t) {
   RunTypedTest<uint64_t>();
 }
 
+TEST(GatherElementsOpTest, float) {
+  RunTypedTest<float>();
+}
+
+TEST(GatherElementsOpTest, double) {
+  RunTypedTest<double>();
+}
+
+TEST(GatherElementsOpTest, bool) {
+  RunTypedTest<bool>();
+}
+
 TEST(GatherElementsOpTest, string) {
   RunTypedTest<std::string>();
 }