[CUDA] Optimize Slice Kernel (#13641)

The PR optimizes Slice CUDA kernel by two ways:
- Coalesce dimensions so less divmod during the kernel compute
- Split data load and write for better memory throughput

Below shows some perf results (cycles number from Nsight Compute) in
V100 using real cases from Huggingface's XLNet model:

  | Old | New
-- | -- | --
[8,12,2048,1024], axis=2, start=1, end=2048 | 1838687| 1539846
[8,12,1024,2047], axis=3, start=0, end=1024 | 951383| 722203

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

@@ -70,75 +70,91 @@ ONNX_CPU_OPERATOR_KERNEL(
         .TypeConstraint("Tind", BuildKernelDefConstraintsFromTypeList<EnabledIndicesTypes>()),
     Slice10);
 
-// Check if it's possible to combine innermost dimensions so we copy larger blocks.
-// Sets flattened_output_dims to nullptr if it is not.
-// Updates starts and steps to match flattened_output_dims if it is.
-// e.g. if input shape is { 2, 2, 2 }, output shape is { 1, 2, 2 }, and the 'steps' value for the last two dims is 1,
-// we are keeping all the data of the inner most two dimensions so can combine those into dims of { 1, 4 }
-static void FlattenOutputDims(gsl::span<const int64_t> input_dimensions,
-                              gsl::span<const int64_t> output_dims,
-                              TensorShapeVector& starts,
-                              TensorShapeVector& ends,
-                              TensorShapeVector& steps,
-                              TensorShapeVector*& flattened_output_dims) {
-  int num_to_combine = 0;
-  for (int64_t i = static_cast<int64_t>(starts.size()) - 1; i >= 0; --i) {
-    // if we're keeping all the data for the dimension and not reversing the direction we can potentially combine it
-    if (steps[narrow<size_t>(i)] == 1 && input_dimensions[narrow<size_t>(i)] == output_dims[narrow<size_t>(i)])
-      ++num_to_combine;
-    else
+// Coalesce contiguous non-slice dimensions into a single dimension.
+// Set p_flattened_input_dims_ and p_flattened_output_dims_ to nullptr if nothing coalesced.
+// Updates starts and steps to match the new dimensions.
+// e.g. if input shape is { 2, 2, 2, 2, 2 }, output shape is { 2, 2, 1, 2, 2 },
+// and the 'steps' value for all dims is 1 except dim-2, then the input shape is coalesced to { 4, 2, 4 }
+// and the output shape is coalesced to { 4, 1, 4 }.
+static void FlattenOutputDims(gsl::span<const int64_t> input_dimensions, gsl::span<const int64_t> output_dims,
+                              TensorShapeVector& starts, TensorShapeVector& ends, TensorShapeVector& steps,
+                              TensorShapeVector*& p_flattened_input_dims, TensorShapeVector*& p_flattened_output_dims) {
+  size_t cur = 0;
+  size_t nxt = 0;
+  while (true) {
+    // Skip all leading slicing dims.
+    while (nxt < starts.size() && (steps[nxt] != 1 || input_dimensions[nxt] != output_dims[nxt])) {
+      p_flattened_input_dims->emplace_back(input_dimensions[nxt]);
+      p_flattened_output_dims->emplace_back(output_dims[nxt]);
+      starts[cur] = starts[nxt];
+      ends[cur] = ends[nxt];
+      steps[cur] = steps[nxt];
+      ++cur;
+      ++nxt;
+    }
+    if (nxt == starts.size()) {
       break;
+    }
+    // Coalesce contiguous non-slicing dims.
+    int64_t running_size = 1;
+    while (nxt < starts.size() && steps[nxt] == 1 && input_dimensions[nxt] == output_dims[nxt]) {
+      running_size *= input_dimensions[nxt];
+      ++nxt;
+    }
+    if (running_size > 1) {
+      p_flattened_input_dims->emplace_back(running_size);
+      p_flattened_output_dims->emplace_back(running_size);
+      starts[cur] = 0LL;
+      ends[cur] = running_size;
+      steps[cur] = 1LL;
+      ++cur;
+    }
   }
 
-  if (num_to_combine > 1) {
-    auto num_dims = output_dims.size() - num_to_combine + 1;
-    flattened_output_dims->assign(output_dims.begin(), output_dims.end());
-    flattened_output_dims->resize(num_dims);
+  // No actual slice dim, and all dims are size 1.
+  if (cur == 0) {
+    p_flattened_input_dims->emplace_back(1LL);
+    p_flattened_output_dims->emplace_back(1LL);
+    starts[cur] = 0LL;
+    ends[cur] = 1LL;
+    steps[cur] = 1LL;
+    ++cur;
+  }
 
-    int64_t dim_value = 1;
-    for (size_t k = num_dims - 1, end = output_dims.size(); k < end; ++k) {
-      dim_value *= output_dims[k];
-    }
-
-    flattened_output_dims->back() = dim_value;
-
-    // the value of starts and steps for all the dims being combined are 0 and 1 respectively,
-    // so we can just shrink via resize so the number of entries matches flattened_output_dims
-    starts.resize(num_dims);
-    steps.resize(num_dims);
-
-    // update ends as well
-    ends.resize(num_dims);
-    ends.back() = dim_value;
+  if (p_flattened_output_dims->size() == output_dims.size()) {
+    p_flattened_input_dims->clear();
+    p_flattened_output_dims->clear();
+    p_flattened_input_dims = nullptr;
+    p_flattened_output_dims = nullptr;
   } else {
-    flattened_output_dims = nullptr;
+    starts.resize(cur);
+    ends.resize(cur);
+    steps.resize(cur);
   }
 }
 
 // Slice V1-9 & DynamicSlice
-Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
-                                    gsl::span<const int64_t> raw_ends,
+Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts, gsl::span<const int64_t> raw_ends,
                                     gsl::span<const int64_t> raw_axes,
                                     SliceOp::PrepareForComputeMetadata& compute_metadata) {
   ORT_RETURN_IF_ERROR(SliceOp::PrepareForComputeHelper(raw_starts, raw_ends, raw_axes, compute_metadata));
   FlattenOutputDims(compute_metadata.input_dimensions_, compute_metadata.output_dims_, compute_metadata.starts_,
-                    compute_metadata.ends_, compute_metadata.steps_, compute_metadata.p_flattened_output_dims_);
+                    compute_metadata.ends_, compute_metadata.steps_, compute_metadata.p_flattened_input_dims_,
+                    compute_metadata.p_flattened_output_dims_);
   return Status::OK();
 }
 
 // DynamicSlice & Slice V10
-Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts,
-                                    gsl::span<const int64_t> raw_ends,
-                                    gsl::span<const int64_t> raw_axes,
-                                    gsl::span<const int64_t> raw_steps,
+Status SliceBase::PrepareForCompute(gsl::span<const int64_t> raw_starts, gsl::span<const int64_t> raw_ends,
+                                    gsl::span<const int64_t> raw_axes, gsl::span<const int64_t> raw_steps,
                                     SliceOp::PrepareForComputeMetadata& compute_metadata) {
   ORT_RETURN_IF_ERROR(SliceOp::PrepareForComputeHelper(raw_starts, raw_ends, raw_axes, raw_steps, compute_metadata));
   FlattenOutputDims(compute_metadata.input_dimensions_, compute_metadata.output_dims_, compute_metadata.starts_,
-                    compute_metadata.ends_, compute_metadata.steps_, compute_metadata.p_flattened_output_dims_);
+                    compute_metadata.ends_, compute_metadata.steps_, compute_metadata.p_flattened_input_dims_,
+                    compute_metadata.p_flattened_output_dims_);
 
   return Status::OK();
 }
-
 namespace {
 template <typename T>
 void CopyData(const Tensor& start_tensor,
@@ -231,16 +247,11 @@ static Status SliceImpl(OpKernelContext* ctx,
      ORT_ENFORCE(output == output_end);
   };
 
-  if (compute_metadata.p_flattened_output_dims_) {
-    // if we have flattened output dims we need to also flatten the input dims.
-    // as we're combining the innermost dims and keeping all values we can just copy the size of the last dim
-    auto flattened_input_dims = input_tensor.Shape().AsShapeVector();
-    flattened_input_dims.resize(compute_metadata.p_flattened_output_dims_->size());
-    flattened_input_dims.back() = compute_metadata.p_flattened_output_dims_->back();
-    TensorShape input_shape(flattened_input_dims);
-
-    auto input_iterator2 = SliceIterator<T>(input_tensor, input_shape, compute_metadata.starts_,
-                                            *compute_metadata.p_flattened_output_dims_, compute_metadata.steps_);
+  if (compute_metadata.p_flattened_input_dims_) {
+    // If we were able to coalesce the input and output shapes, use the new shapes.
+    auto input_iterator2 =
+        SliceIterator<T>(input_tensor, TensorShape(compute_metadata.flattened_input_dims_), compute_metadata.starts_,
+                         compute_metadata.flattened_output_dims_, compute_metadata.steps_);
     create_output(input_iterator2);
   } else {
     auto input_iterator2 = SliceIterator<T>(input_tensor, compute_metadata.starts_, compute_metadata.output_dims_,

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

@@ -27,6 +27,8 @@ struct PrepareForComputeMetadata {
   TensorShapeVector ends_;
   TensorShapeVector steps_;
   TensorShapeVector output_dims_;
+  TensorShapeVector flattened_input_dims_;
+  TensorShapeVector* p_flattened_input_dims_ = &flattened_input_dims_;
   TensorShapeVector flattened_output_dims_;
   TensorShapeVector* p_flattened_output_dims_ = &flattened_output_dims_;
 };

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

@@ -103,40 +103,24 @@ static Status SliceImpCore(cudaStream_t stream,
 
 namespace SliceCuda {
 
-static Status ComputeSliceStrides(const TensorShape& input_shape,
-                                  TArray<int64_t>& input_strides,
+static Status ComputeSliceStrides(const TensorShape& input_shape, TArray<int64_t>& input_strides,
                                   TArray<fast_divmod>& output_strides,
                                   SliceOp::PrepareForComputeMetadata& compute_metadata) {
+  // If we were able to coalesce the input and output shapes, use the new shapes to compute the strides.
   const auto input_dimensions = input_shape.GetDims();
-  size_t dimension_count = input_dimensions.size();
-  // if we are able to flatten the output dims we updated 'starts' and 'steps' to match the smaller number of dims.
-  // update dimension_count to match.
-  if (compute_metadata.p_flattened_output_dims_) {
-    dimension_count = compute_metadata.p_flattened_output_dims_->size();
-  }
-
-  input_strides.SetSize(gsl::narrow_cast<int32_t>(dimension_count));
+  size_t rank = compute_metadata.p_flattened_input_dims_ ? compute_metadata.p_flattened_input_dims_->size()
+                                                         : input_dimensions.size();
+  input_strides.SetSize(gsl::narrow_cast<int32_t>(rank));
   const gsl::span<int64_t> input_strides_span = gsl::make_span(input_strides.Data(), input_strides.Size());
-  if (compute_metadata.p_flattened_output_dims_ != nullptr) {
-    // we were able to flatten the innermost dimensions as they're being copied in full to the output.
-    // do the same flattening to the innermost input dimensions in order to calculate pitches that match
-    // the flattened output dimensions.
-    int64_t aggregated_last_dim = 1;
-    for (size_t i = dimension_count - 1, end = input_dimensions.size(); i < end; ++i) {
-      aggregated_last_dim *= input_dimensions[i];
-    }
-
-    TensorShapeVector flattened_input_dims(input_dimensions.begin(), input_dimensions.end());
-    flattened_input_dims.resize(dimension_count);
-    flattened_input_dims.back() = aggregated_last_dim;
-    ORT_ENFORCE(TensorPitches::Calculate(input_strides_span, flattened_input_dims));
+  if (compute_metadata.p_flattened_input_dims_) {
+    ORT_ENFORCE(TensorPitches::Calculate(input_strides_span, compute_metadata.flattened_input_dims_));
   } else {
     ORT_ENFORCE(TensorPitches::Calculate(input_strides_span, input_dimensions));
   }
 
-  const auto output_dims = gsl::make_span(compute_metadata.p_flattened_output_dims_ != nullptr
-                                              ? compute_metadata.flattened_output_dims_
-                                              : compute_metadata.output_dims_);
+  const auto output_dims =
+      gsl::make_span(compute_metadata.p_flattened_output_dims_ != nullptr ? compute_metadata.flattened_output_dims_
+                                                                          : compute_metadata.output_dims_);
   TensorPitches original_output_strides(output_dims);
   output_strides.SetSize(gsl::narrow_cast<int32_t>(original_output_strides.size()));
   for (int32_t i = 0, limit = static_cast<int32_t>(original_output_strides.size()); i < limit; ++i) {

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

@@ -8,97 +8,83 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <bool is_grad, int DIMS, int NumThreadsPerBlock, int NumElementsPerThread, typename T>
-__global__ void _SliceKernel(const TArray<int64_t> starts,
-                             const TArray<int64_t> steps,
-                             const TArray<int64_t> input_strides,
-                             const TArray<fast_divmod> output_strides,
-                             const T* input_data,
-                             T* output_data,
-                             const CUDA_LONG N) {
-  CUDA_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
-  CUDA_LONG input_indices[NumElementsPerThread];
+namespace {
+#ifdef USE_ROCM
+constexpr int kNumElementsPerThread = 2;
+constexpr int kNumThreadsPerBlock = 512;
+#else
+constexpr int kNumElementsPerThread = GridDim::maxElementsPerThread;
+constexpr int kNumThreadsPerBlock = GridDim::maxThreadsPerBlock;
+#endif
+}  // namespace
 
-  CUDA_LONG id = start;
+template <bool is_grad, int DIMS, typename T>
+__global__ void _SliceKernel(const TArray<int64_t> starts, const TArray<int64_t> steps,
+                             const TArray<int64_t> input_strides, const TArray<fast_divmod> output_strides,
+                             const T* input_data, T* output_data, const CUDA_LONG N) {
+  CUDA_LONG start = kNumElementsPerThread * kNumThreadsPerBlock * blockIdx.x + threadIdx.x;
+  T values[kNumElementsPerThread];
+  CUDA_LONG id;
+  if (is_grad) {
+    id = start;
 #pragma unroll
-  for (int i = 0; i < NumElementsPerThread; i++) {
+    for (int i = 0; i < kNumElementsPerThread; ++i) {
+      if (id < N) {
+        values[i] = input_data[id];
+        id += kNumThreadsPerBlock;
+      }
+    }
+  }
+
+  id = start;
+#pragma unroll
+  for (int i = 0; i < kNumElementsPerThread; ++i) {
     if (id < N) {
       CUDA_LONG input_index = 0;
       int div;
       int mod = id;
-      int value = id;
       int dim = 0;
 #pragma unroll
       for (; dim < DIMS - 1; ++dim) {
-        output_strides[dim].divmod(value, div, mod);
+        output_strides[dim].divmod(mod, div, mod);
         input_index += (starts[dim] + div * steps[dim]) * input_strides[dim];
-        value = mod;
       }
       input_index += starts[dim] + mod * steps[dim];
-      input_indices[i] = input_index;
-      id += NumThreadsPerBlock;
+      if (is_grad) {
+        output_data[input_index] = values[i];
+      } else {
+        values[i] = input_data[input_index];
+      }
+      id += kNumThreadsPerBlock;
     }
   }
 
-  if (is_grad) {
+  if (!is_grad) {
     id = start;
 #pragma unroll
-    for (int i = 0; i < NumElementsPerThread; i++) {
+    for (int i = 0; i < kNumElementsPerThread; ++i) {
       if (id < N) {
-        output_data[input_indices[i]] = input_data[id];
-        id += NumThreadsPerBlock;
-      }
-    }
-  } else {
-    id = start;
-#pragma unroll
-    for (int i = 0; i < NumElementsPerThread; i++) {
-      if (id < N) {
-        output_data[id] = input_data[input_indices[i]];
-        id += NumThreadsPerBlock;
+        output_data[id] = values[i];
+        id += kNumThreadsPerBlock;
       }
     }
   }
 }
 
-Status SliceImpl(cudaStream_t stream,
-                 const size_t element_size,
-                 const int32_t dimension_count,
-                 const TArray<int64_t>& starts,
-                 const TArray<int64_t>& steps,
-                 const TArray<int64_t>& input_strides,
-                 const TArray<fast_divmod>& output_strides,
-                 const void* input_data,
-                 void* output_data,
-                 const size_t N) {
-  return SliceImplEx<false>(stream, element_size, dimension_count, starts, steps, input_strides, output_strides, input_data,
-                            output_data, N);
-}
-
-Status SliceImplGrad(cudaStream_t stream,
-                     const size_t element_size,
-                     const int32_t dimension_count,
-                     const TArray<int64_t>& starts,
-                     const TArray<int64_t>& steps,
-                     const TArray<int64_t>& input_strides,
-                     const TArray<fast_divmod>& output_strides,
-                     const void* input_data,
-                     void* output_data,
-                     const size_t N) {
-  return SliceImplEx<true>(stream, element_size, dimension_count, starts, steps, input_strides, output_strides, input_data,
-                           output_data, N);
-}
-
-#define HANDLE_DIMS(ELEMENT_TYPE, DIMS)                                                     \
-  case DIMS: {                                                                              \
-    _SliceKernel<is_grad, DIMS, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread> \
-        <<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0, stream>>>(                                \
-            starts, steps, input_strides, output_strides,                                   \
-            reinterpret_cast<const ToCudaType<ELEMENT_TYPE>::MappedType*>(input_data),      \
-            reinterpret_cast<ToCudaType<ELEMENT_TYPE>::MappedType*>(output_data),           \
-            (CUDA_LONG)N);                                                                  \
+template <bool is_grad>
+Status SliceImplEx(cudaStream_t stream, const size_t element_size, const int32_t dimension_count,
+                   const TArray<int64_t>& starts, const TArray<int64_t>& steps, const TArray<int64_t>& input_strides,
+                   const TArray<fast_divmod>& output_strides, const void* input_data, void* output_data,
+                   const size_t N) {
+  int blocksPerGrid = static_cast<int>(CeilDiv(N, kNumThreadsPerBlock * kNumElementsPerThread));
+  switch (element_size) {
+#define HANDLE_DIMS(ELEMENT_TYPE, DIMS)                                                           \
+  case DIMS: {                                                                                    \
+    _SliceKernel<is_grad, DIMS, ELEMENT_TYPE><<<blocksPerGrid, kNumThreadsPerBlock, 0, stream>>>( \
+        starts, steps, input_strides, output_strides,                                             \
+        reinterpret_cast<const ToCudaType<ELEMENT_TYPE>::MappedType*>(input_data),                \
+        reinterpret_cast<ToCudaType<ELEMENT_TYPE>::MappedType*>(output_data), (CUDA_LONG)N);      \
   } break
-
 #define HANDLE_ELEMENT_TYPE(ELEMENT_TYPE) \
   case sizeof(ELEMENT_TYPE): {            \
     switch (dimension_count) {            \
@@ -112,30 +98,35 @@ Status SliceImplGrad(cudaStream_t stream,
       HANDLE_DIMS(ELEMENT_TYPE, 8);       \
     }                                     \
   } break
-
-template <bool is_grad>
-Status SliceImplEx(cudaStream_t stream,
-                   const size_t element_size,
-                   const int32_t dimension_count,
-                   const TArray<int64_t>& starts,
-                   const TArray<int64_t>& steps,
-                   const TArray<int64_t>& input_strides,
-                   const TArray<fast_divmod>& output_strides,
-                   const void* input_data,
-                   void* output_data,
-                   const size_t N) {
-  int blocksPerGrid = static_cast<int>(CeilDiv(N, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
-  switch (element_size) {
     HANDLE_ELEMENT_TYPE(int8_t);
     HANDLE_ELEMENT_TYPE(int16_t);
     HANDLE_ELEMENT_TYPE(int32_t);
     HANDLE_ELEMENT_TYPE(int64_t);
     default:
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Type not supported for Slice operator");
+#undef HANDLE_ELEMENT_TYPE
+#undef HANDLE_DIMS
   }
 
   return Status::OK();
 }
 
+Status SliceImpl(cudaStream_t stream, const size_t element_size, const int32_t dimension_count,
+                 const TArray<int64_t>& starts, const TArray<int64_t>& steps, const TArray<int64_t>& input_strides,
+                 const TArray<fast_divmod>& output_strides, const void* input_data, void* output_data, const size_t N) {
+  return SliceImplEx<false>(stream, element_size, dimension_count, starts, steps, input_strides, output_strides,
+                            input_data, output_data, N);
+}
+
+#ifdef ENABLE_TRAINING
+Status SliceImplGrad(cudaStream_t stream, const size_t element_size, const int32_t dimension_count,
+                     const TArray<int64_t>& starts, const TArray<int64_t>& steps, const TArray<int64_t>& input_strides,
+                     const TArray<fast_divmod>& output_strides, const void* input_data, void* output_data,
+                     const size_t N) {
+  return SliceImplEx<true>(stream, element_size, dimension_count, starts, steps, input_strides, output_strides,
+                           input_data, output_data, N);
+}
+#endif  // ENABLE_TRAINING
+
 }  // namespace cuda
 }  // namespace onnxruntime

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

@@ -8,18 +8,6 @@
 namespace onnxruntime {
 namespace cuda {
 
-template <bool is_grad>
-Status SliceImplEx(cudaStream_t stream,
-                   const size_t element_size,
-                   const int32_t dimension_count,
-                   const TArray<int64_t>& starts,
-                   const TArray<int64_t>& steps,
-                   const TArray<int64_t>& input_strides,
-                   const TArray<fast_divmod>& output_strides,
-                   const void* input_data,
-                   void* output_data,
-                   const size_t N);
-
 Status SliceImpl(cudaStream_t stream,
                  const size_t element_size,
                  const int32_t dimension_count,
@@ -31,6 +19,7 @@ Status SliceImpl(cudaStream_t stream,
                  void* output_data,
                  const size_t N);
 
+#ifdef ENABLE_TRAINING
 Status SliceImplGrad(cudaStream_t stream,
                      const size_t element_size,
                      const int32_t dimension_count,
@@ -41,5 +30,7 @@ Status SliceImplGrad(cudaStream_t stream,
                      const void* input_data,
                      void* output_data,
                      const size_t N);
+#endif  // ENABLE_TRAINING
+
 }  // namespace cuda
 }  // namespace onnxruntime

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

@@ -731,5 +731,21 @@ TEST(SliceTest, EmptyDim) {
                       {0, 6},
                       {});
 }
+
+TEST(SliceTest, CoalesceDims) {
+  RunSliceTest<float>({2, 2, 2, 2},
+                      {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, -1.f, -2.f, -3.f, -4.f, -5.f, -6.f, -7.f, -8.f}, {1, 1},
+                      {0, 2}, {0, 1}, {-1, 1}, {1, 1, 2, 2}, {-5.f, -6.f, -7.f, -8.f}, true);
+  RunSliceTest<float>({1, 2, 2, 2, 2},
+                      {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, -1.f, -2.f, -3.f, -4.f, -5.f, -6.f, -7.f, -8.f}, {1},
+                      {std::numeric_limits<int64_t>::max()}, {2}, {}, {1, 2, 1, 2, 2},
+                      {5.f, 6.f, 7.f, 8.f, -5.f, -6.f, -7.f, -8.f}, true);
+  RunSliceTest<float>({1, 2, 2, 2, 2},
+                      {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, -1.f, -2.f, -3.f, -4.f, -5.f, -6.f, -7.f, -8.f}, {1, 1},
+                      {std::numeric_limits<int64_t>::max(), std::numeric_limits<int64_t>::max()}, {1, 3}, {},
+                      {1, 1, 2, 1, 2}, {-3.f, -4.f, -7.f, -8.f}, true);
+  RunSliceTest<float>({1, 1, 1}, {1.f}, {0}, {std::numeric_limits<int64_t>::max()}, {1}, {}, {1, 1, 1}, {1.f}, true);
+}
+
 }  // namespace test
 }  // namespace onnxruntime

orttraining/orttraining/test/training_ops/cpu/tensor/slice_grad_op_test.cc

@@ -56,5 +56,55 @@ TEST(SliceGradOpTest, SliceGrad_blockcopy_double) {
   test.Run();
 }
 
+TEST(SliceGradOpTest, CoalesceDims) {
+  {
+    OpTester test("SliceGrad", 1, onnxruntime::kMSDomain);
+    test.AddInput<float>("grad", {1, 1, 2, 2}, {1, 2, 3, 4});
+    test.AddInput<int64_t>("shape", {4}, {2, 2, 2, 2});
+    test.AddInput<int64_t>("starts", {2}, {1LL, 1LL});
+    test.AddInput<int64_t>("ends", {2}, {0LL, 2LL});
+    test.AddInput<int64_t>("axes", {2}, {0LL, 1LL});
+    test.AddInput<int64_t>("steps", {2}, {-1LL, 1LL});
+    test.AddOutput<float>("output", {2, 2, 2, 2}, {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 3, 4});
+    test.Run();
+  }
+
+  {
+    OpTester test("SliceGrad", 1, onnxruntime::kMSDomain);
+    test.AddInput<float>("grad", {1, 2, 1, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8});
+    test.AddInput<int64_t>("shape", {5}, {1, 2, 2, 2, 2});
+    test.AddInput<int64_t>("starts", {1}, {1LL});
+    test.AddInput<int64_t>("ends", {1}, {std::numeric_limits<int64_t>::max()});
+    test.AddInput<int64_t>("axes", {1}, {2LL});
+    test.AddInput<int64_t>("steps", {1}, {1LL});
+    test.AddOutput<float>("output", {1, 2, 2, 2, 2}, {0, 0, 0, 0, 1, 2, 3, 4, 0, 0, 0, 0, 5, 6, 7, 8});
+    test.Run();
+  }
+
+  {
+    OpTester test("SliceGrad", 1, onnxruntime::kMSDomain);
+    test.AddInput<float>("grad", {1, 1, 2, 1, 2}, {1, 2, 3, 4});
+    test.AddInput<int64_t>("shape", {5}, {1, 2, 2, 2, 2});
+    test.AddInput<int64_t>("starts", {2}, {1LL, 1LL});
+    test.AddInput<int64_t>("ends", {2}, {std::numeric_limits<int64_t>::max(), 2});
+    test.AddInput<int64_t>("axes", {2}, {1LL, 3LL});
+    test.AddInput<int64_t>("steps", {2}, {1LL, 1LL});
+    test.AddOutput<float>("output", {1, 2, 2, 2, 2}, {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 0, 0, 3, 4});
+    test.Run();
+  }
+
+  {
+    OpTester test("SliceGrad", 1, onnxruntime::kMSDomain);
+    test.AddInput<float>("grad", {1, 1, 1}, {1});
+    test.AddInput<int64_t>("shape", {3}, {1, 1, 1});
+    test.AddInput<int64_t>("starts", {1}, {0LL});
+    test.AddInput<int64_t>("ends", {1}, {std::numeric_limits<int64_t>::max()});
+    test.AddInput<int64_t>("axes", {1}, {1LL});
+    test.AddInput<int64_t>("steps", {1}, {1LL});
+    test.AddOutput<float>("output", {1, 1, 1}, {1});
+    test.Run();
+  }
+}
+
 }  // namespace test
 }  // namespace onnxruntime

orttraining/orttraining/training_ops/cpu/tensor/slice_grad.cc

@@ -41,13 +41,15 @@ Status SliceGrad::Compute(OpKernelContext* context) const {
 
   MLDataType T_type = grad.DataType();
   if (T_type == DataTypeImpl::GetType<float>()) {
-    return ComputeImpl<float>(context, output, compute_metadata.output_dims_, compute_metadata.p_flattened_output_dims_,
-                              compute_metadata.starts_, compute_metadata.steps_);
+    return ComputeImpl<float>(context, output, compute_metadata.output_dims_, compute_metadata.p_flattened_input_dims_,
+                              compute_metadata.p_flattened_output_dims_, compute_metadata.starts_,
+                              compute_metadata.steps_);
   }
 
   if (T_type == DataTypeImpl::GetType<double>()) {
-    return ComputeImpl<double>(context, output, compute_metadata.output_dims_, compute_metadata.p_flattened_output_dims_,
-                               compute_metadata.starts_, compute_metadata.steps_);
+    return ComputeImpl<double>(context, output, compute_metadata.output_dims_, compute_metadata.p_flattened_input_dims_,
+                               compute_metadata.p_flattened_output_dims_, compute_metadata.starts_,
+                               compute_metadata.steps_);
   }
 
   return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "Type for T or Tind not supported yet in SliceGrad.");
@@ -57,7 +59,8 @@ template <typename T>
 Status SliceGrad::ComputeImpl(OpKernelContext* ctx,
                               Tensor& output_grad_tensor,
                               const gsl::span<const int64_t>& output_dims,
-                              TensorShapeVector* flattened_output_dims,
+                              TensorShapeVector* p_flattened_input_dims,
+                              TensorShapeVector* p_flattened_output_dims,
                               const gsl::span<const int64_t>& starts,
                               const gsl::span<const int64_t>& steps) const {
   TensorShape output_shape(output_dims);
@@ -83,15 +86,10 @@ Status SliceGrad::ComputeImpl(OpKernelContext* ctx,
     ORT_ENFORCE(grad_data == grad_data_end);
   };
 
-  if (flattened_output_dims) {
-    // if we have flattened output dims we need to also flatten the input dims.
-    // as we're combining the innermost dims and keeping all values we can just copy the size of the last dim
-    TensorShapeVector flattened_input_dims(output_grad_tensor.Shape().AsShapeVector());
-    flattened_input_dims.resize(flattened_output_dims->size());
-    flattened_input_dims.back() = flattened_output_dims->back();
-    TensorShape input_shape(std::move(flattened_input_dims));
-
-    WritableSliceIterator<T> input_iterator(output_grad_tensor, input_shape, starts, *flattened_output_dims, steps);
+  if (p_flattened_input_dims) {
+    // If we were able to coalesce the input and output shapes, use the new shapes.
+    WritableSliceIterator<T> input_iterator(output_grad_tensor, TensorShape(*p_flattened_input_dims), starts,
+                                            *p_flattened_output_dims, steps);
     create_output(input_iterator);
   } else {
     WritableSliceIterator<T> input_iterator(output_grad_tensor, starts, output_dims, steps);

orttraining/orttraining/training_ops/cpu/tensor/slice_grad.h

@@ -22,7 +22,8 @@ class SliceGrad final : public OpKernel, public SliceBase {
   Status ComputeImpl(OpKernelContext* ctx,
                      Tensor& output_grad_tensor,
                      const gsl::span<const int64_t>& output_dims,
-                     TensorShapeVector* flattened_output_dims,
+                     TensorShapeVector* p_flattened_input_dims,
+                     TensorShapeVector* p_flattened_output_dims,
                      const gsl::span<const int64_t>& starts,
                      const gsl::span<const int64_t>& steps) const;
 };