gather elements optimization (#8154)

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

@@ -61,8 +61,11 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const {
   if (indices_shape.Size() == 0)
     return Status::OK();
 
-  TensorPitches input_strides(input_dims);
-  TArray<int64_t> gpu_input_strides(input_strides);
+  // Set stride along axis to 0 so we don't need IF statment to check in kernel.
+  TensorPitches masked_input_strides(input_dims);
+  int64_t input_stride_along_axis = masked_input_strides[axis];
+  masked_input_strides[axis] = 0;
+  TArray<int64_t> gpu_masked_input_strides(masked_input_strides);
 
   TArray<fast_divmod> fdm_indices_strides(indices_rank);
   TensorPitches indices_strides(indices_dims);
@@ -77,14 +80,15 @@ Status GatherElements::ComputeInternal(OpKernelContext* context) const {
       indices_tensor->IsDataType<int64_t>()) {
     GatherElementsImpl(
         Stream(),
-        input_rank,
+        // Save one divmod in kernel if axis is the last dim.
+        input_rank == axis + 1 ? input_rank - 1 : input_rank,
         input_tensor->DataRaw(),
         input_dims[axis],
-        gpu_input_strides,
+        input_stride_along_axis,
+        gpu_masked_input_strides,
         indices_tensor->DataRaw(),
         indices_size,
         fdm_indices_strides,
-        axis,
         output_tensor->MutableDataRaw(),
         element_size,
         index_element_size);

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

@@ -33,12 +33,12 @@ __global__ void _GatherElementsKernel(
     const int64_t rank,
     const T* input_data,
     const int64_t input_dim_along_axis,
-    const TArray<int64_t> input_strides,
+    const int64_t input_stride_along_axis,
+    const TArray<int64_t> masked_input_strides,   // the stride along axis is 0
     const void* indices_data,
     const int64_t indices_size,
     const size_t index_element_size,
     const TArray<fast_divmod> indices_strides,
-    const int64_t axis,
     T* output_data) {
 
   CUDA_LONG indices_index = threads_per_block * thread_worksize * blockIdx.x + threadIdx.x;
@@ -46,19 +46,20 @@ __global__ void _GatherElementsKernel(
   #pragma unroll
   for (int work = 0; work < thread_worksize; ++work) {
     if (indices_index < indices_size) {
-
       int dim = 0;
       int remain = indices_index;
       int64_t data_idx = 0;
 
-      int i = 0;
-      for (; i < axis && remain > 0; ++i) {
+      #pragma unroll
+      for (auto i = 0; i < indices_strides.Capacity(); i++) {
+        if (i >= rank) {
+          break;
+        }
+
         indices_strides[i].divmod(remain, dim, remain);
-        data_idx += input_strides[i] * dim;
+        data_idx += masked_input_strides[i] * dim;
       }
 
-      i = axis;
-      indices_strides[i].divmod(remain, dim, remain);
       dim = GetIndexValue(indices_data, index_element_size, indices_index);
       if (dim < -input_dim_along_axis || dim >= input_dim_along_axis) {
         return;  // Invalid index
@@ -68,13 +69,7 @@ __global__ void _GatherElementsKernel(
         dim += input_dim_along_axis;
       }
 
-      data_idx += input_strides[i] * dim;
-
-      ++i;  // past axis
-      for (; i < rank && remain > 0; ++i) {
-        indices_strides[i].divmod(remain, dim, remain);
-        data_idx += input_strides[i] * dim;
-      }
+      data_idx += input_stride_along_axis * dim;
       output_data[indices_index] = input_data[data_idx];
 
       indices_index += threads_per_block;
@@ -87,11 +82,11 @@ void GatherElementsImpl(
     const int64_t rank,
     const void* input_data,
     const int64_t input_dim_along_axis,
-    const TArray<int64_t>& input_strides,
+    const int64_t input_stride_along_axis,
+    const TArray<int64_t>& masked_input_strides,
     const void* indices_data,
     const int64_t indices_size,
     const TArray<fast_divmod>& indices_strides,
-    const int64_t axis,
     void* output_data,
     size_t element_size,
     size_t index_element_size) {
@@ -105,33 +100,33 @@ void GatherElementsImpl(
       case sizeof(int8_t): {
         using CudaType = typename ToCudaType<int8_t>::MappedType;
         _GatherElementsKernel<<<blocksPerGrid, block, 0, stream>>>(
-            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
-            indices_data, indices_size, index_element_size, indices_strides,
-            axis, reinterpret_cast<CudaType*>(output_data));
+            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_stride_along_axis,
+            masked_input_strides, indices_data, indices_size, index_element_size, indices_strides,
+            reinterpret_cast<CudaType*>(output_data));
       } break;
 
       case sizeof(int16_t): {
         using CudaType = typename ToCudaType<int16_t>::MappedType;
         _GatherElementsKernel<<<blocksPerGrid, block, 0, stream>>>(
-            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
-            indices_data, indices_size, index_element_size, indices_strides,
-            axis, reinterpret_cast<CudaType*>(output_data));
+            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_stride_along_axis,
+            masked_input_strides, indices_data, indices_size, index_element_size, indices_strides,
+            reinterpret_cast<CudaType*>(output_data));
       } break;
 
       case sizeof(int32_t): {
         using CudaType = typename ToCudaType<int32_t>::MappedType;
         _GatherElementsKernel<<<blocksPerGrid, block, 0, stream>>>(
-            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
-            indices_data, indices_size, index_element_size, indices_strides,
-            axis, reinterpret_cast<CudaType*>(output_data));
+            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_stride_along_axis,
+            masked_input_strides, indices_data, indices_size, index_element_size, indices_strides,
+            reinterpret_cast<CudaType*>(output_data));
       } break;
 
       case sizeof(int64_t): {
         using CudaType = typename ToCudaType<int64_t>::MappedType;
         _GatherElementsKernel<<<blocksPerGrid, block, 0, stream>>>(
-            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_strides,
-            indices_data, indices_size, index_element_size, indices_strides,
-            axis, reinterpret_cast<CudaType*>(output_data));
+            rank, reinterpret_cast<const CudaType*>(input_data), input_dim_along_axis, input_stride_along_axis,
+            masked_input_strides, indices_data, indices_size, index_element_size, indices_strides,
+            reinterpret_cast<CudaType*>(output_data));
       } break;
 
       // should not reach here as we validate if the all relevant types are supported in the Compute method

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

@@ -14,11 +14,11 @@ 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_dim_along_axis,
-    const TArray<int64_t>& input_strides,
+    const int64_t input_stride_along_axis,
+    const TArray<int64_t>& masked_input_strides,
     const void* indices_data,
     const int64_t indices_size,
     const TArray<fast_divmod>& indices_strides,
-    const int64_t axis,
     void* output_data,
     size_t element_size,
     size_t index_element_size);