Support multi-dimensional lengths in segment_reduce to support pytorch_scatter.segment_* functionalities (CUDA)

Pull Request resolved: pytorch#77061

Approved by: https://github.com/cpuhrsch

Author: mikaylagawarecki

aten/src/ATen/native/cuda/SegmentReduce.cu

@@ -13,6 +13,8 @@
 #else
 #include <ATen/ops/empty.h>
 #include <ATen/ops/zeros.h>
+#include <ATen/ops/cat.h>
+#include <ATen/ops/cumsum.h>
 #endif
 
 namespace at {
@@ -68,7 +70,7 @@ Tensor _get_complete_sum(const Tensor& lengths) {
   offsets[0].zero_();
 
   AT_DISPATCH_INDEX_TYPES(
-      lengths.type(), "_segment_reduce_cuda_backward_kernel1", ([&] {
+      lengths.scalar_type(), "_segment_reduce_cuda_backward_kernel1", ([&] {
         auto* lengths_data_ptr = lengths.data_ptr<index_t>();
         auto* offsets_data_ptr = offsets.data_ptr<index_t>();
         at::cuda::cub::inclusive_sum(
@@ -108,22 +110,34 @@ __global__ void segment_reduce_forward_kernel(
     const index_t* lengths_data,
     const index_t* lengths_cumsum_data,
     const int64_t segment_count,
-    const int64_t stride_count,
+    const int64_t lengths_stride_axis,
     bool is_initial_set,
-    scalar_t initial_value) {
+    scalar_t initial_value,
+    const int64_t outer_offset,
+    const int64_t inner_offset,
+    const int64_t data_stride_axis,
+    const int64_t data_size_axis,
+    const int64_t output_stride_axis,
+    const int64_t output_size_axis,
+    const int64_t lengths_cumsum_stride_axis) {
   int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-  int64_t row_id = idx / stride_count;
-  int64_t lane_id = idx % stride_count;
-  if (idx >= (segment_count * stride_count)) {
+  if (idx >= (outer_offset * segment_count * inner_offset)) {
     return;
   }
-  int64_t offset_start = lengths_cumsum_data[row_id];
-  int64_t offset_end = lengths_cumsum_data[row_id + 1];
+  int64_t row_id = idx / inner_offset;
+  int64_t lane_id = idx % inner_offset;   // lane_id is the inner_idx
+  int64_t outer_idx = row_id / segment_count;
+  int64_t dim_idx = row_id % segment_count;
+
+  int64_t offset_idx = outer_idx * lengths_cumsum_stride_axis * (segment_count + 1) + dim_idx;
+  index_t offset_start = lengths_cumsum_data[offset_idx];
+  index_t offset_end = lengths_cumsum_data[offset_idx + 1];
 
   // ===== step2: apply reduction
-  for (int64_t j = offset_start; j < offset_end; ++j) {
-    int64_t starting_index = (j * stride_count) + lane_id;
-    const auto data = values_data[starting_index];
+  for (index_t j = offset_start; j < offset_end; ++j) {
+    int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                         + j * data_stride_axis + lane_id;
+    const auto data = values_data[data_index];
     // TODO: There is no need to branch with every element
     if (reduction == SegmentReductionType::MAX) {
       initial_value =
@@ -142,19 +156,22 @@ __global__ void segment_reduce_forward_kernel(
   }
 
   // ===== step3: finalize reduction
-  CUDA_KERNEL_ASSERT(lengths_data[row_id] >= 0);
-  if (lengths_data[row_id] == 0 && !is_initial_set &&
+  int64_t lengths_idx = outer_idx * lengths_stride_axis * segment_count + dim_idx;
+  CUDA_KERNEL_ASSERT(lengths_data[lengths_idx] >= 0);
+  if (lengths_data[lengths_idx] == 0 && !is_initial_set &&
       reduction == SegmentReductionType::MEAN) {
     initial_value = static_cast<scalar_t>(NAN);
   } else if (
-      reduction == SegmentReductionType::MEAN && lengths_data[row_id] > 0 &&
+      reduction == SegmentReductionType::MEAN && lengths_data[lengths_idx] > 0 &&
       !at::_isnan(initial_value)) {
-    initial_value = initial_value / lengths_data[row_id];
+    initial_value = initial_value / lengths_data[lengths_idx];
   }
-  int64_t output_index = (row_id * stride_count) + lane_id;
+  int64_t output_index = outer_idx * output_stride_axis * output_size_axis
+                         + dim_idx * output_stride_axis + lane_id;
   output_data[output_index] = initial_value;
 }
 
+
 template <typename scalar_t, typename index_t>
 __global__ void segment_reduce_backward_kernel(
     SegmentReductionType reduction,
@@ -165,32 +182,46 @@ __global__ void segment_reduce_backward_kernel(
     const index_t* lengths_data,
     const index_t* lengths_cumsum_data,
     const int64_t segment_count,
-    const int64_t stride_count,
-    scalar_t initial_prod_value) {
+    const int64_t lengths_stride_axis,
+    scalar_t initial_prod_value,
+    const int64_t outer_offset,
+    const int64_t inner_offset,
+    const int64_t data_stride_axis,
+    const int64_t data_size_axis,
+    const int64_t output_stride_axis,
+    const int64_t output_size_axis,
+    const int64_t lengths_cumsum_stride_axis) {
   int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
-  int64_t row_id = idx / stride_count;
-  int64_t lane_id = idx % stride_count;
-
-  if (idx >= (segment_count * stride_count)) {
+  if (idx >= (outer_offset * segment_count * inner_offset)) {
     return;
   }
-  if (lengths_data[row_id] == 0) {
+  int64_t row_id = idx / inner_offset;
+  int64_t lane_id = idx % inner_offset;  // lane_id is the inner_idx
+  int64_t outer_idx = row_id / segment_count;
+  int64_t dim_idx = row_id % segment_count;
+
+  int64_t lengths_idx = outer_idx * lengths_stride_axis * segment_count + dim_idx;
+  auto segment_length = lengths_data[lengths_idx];
+  if (segment_length == 0) {
     return;
   }
 
-  int64_t offset_start = lengths_cumsum_data[row_id];
-  int64_t offset_end = lengths_cumsum_data[row_id + 1];
+  int64_t offset_idx = outer_idx * lengths_cumsum_stride_axis * (segment_count + 1) + dim_idx;
+  index_t offset_start = lengths_cumsum_data[offset_idx];
+  index_t offset_end = lengths_cumsum_data[offset_idx + 1];
 
-  int64_t output_index = (row_id * stride_count) + lane_id;
+  int64_t output_index = outer_idx * output_stride_axis * output_size_axis
+                         + dim_idx * output_stride_axis + lane_id;
 
   if (reduction == SegmentReductionType::MAX ||
       reduction == SegmentReductionType::MIN) {
     int64_t counter = 0;
     for (int64_t j = offset_start; j < offset_end; ++j) {
-      int64_t starting_index = (j * stride_count) + lane_id;
-      if (at::_isnan(values_data[starting_index]) ||
-          values_data[starting_index] == output_data[output_index]) {
-        grad_input_data[starting_index] = grad_data[output_index];
+      int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                           + j * data_stride_axis + lane_id;
+      if (at::_isnan(values_data[data_index]) ||
+          values_data[data_index] == output_data[output_index]) {
+        grad_input_data[data_index] = grad_data[output_index];
         counter++;
       }
     }
@@ -200,47 +231,51 @@ __global__ void segment_reduce_backward_kernel(
       return;
     }
     for (int64_t j = offset_start; j < offset_end; ++j) {
-      int64_t starting_index = (j * stride_count) + lane_id;
-      if (grad_input_data[starting_index] > 0) {
-        grad_input_data[starting_index] =
-            grad_input_data[starting_index] / counter;
+      int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                           + j * data_stride_axis + lane_id;
+      if (grad_input_data[data_index] > 0) {
+        grad_input_data[data_index] =
+            grad_input_data[data_index] / counter;
       }
     }
   } else if (reduction == SegmentReductionType::MEAN) {
-    auto grad_val = grad_data[output_index] / lengths_data[row_id];
+    auto grad_val = grad_data[output_index] / segment_length;
     for (int64_t j = offset_start; j < offset_end; ++j) {
-      int64_t starting_index = (j * stride_count) + lane_id;
-      grad_input_data[starting_index] = grad_val;
+      int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                           + j * data_stride_axis + lane_id;
+      grad_input_data[data_index] = grad_val;
     }
   } else if (reduction == SegmentReductionType::SUM) {
     const auto& grad_val = grad_data[output_index];
     for (int64_t j = offset_start; j < offset_end; ++j) {
-      int64_t starting_index = (j * stride_count) + lane_id;
-      grad_input_data[starting_index] = grad_val;
+      int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                           + j * data_stride_axis + lane_id;
+      grad_input_data[data_index] = grad_val;
     }
   } else if (reduction == SegmentReductionType::PROD) {
     const auto& grad_val = grad_data[output_index] * output_data[output_index];
     for (int64_t j = offset_start; j < offset_end; ++j) {
-      int64_t starting_index = (j * stride_count) + lane_id;
-      if (at::_isnan(values_data[starting_index]) ||
-          values_data[starting_index] == 0) {
+      int64_t data_index = outer_idx * data_stride_axis * data_size_axis
+                           + j * data_stride_axis + lane_id;
+      if (at::_isnan(values_data[data_index]) ||
+          values_data[data_index] == 0) {
         // explicitly compute exclusive prod
         scalar_t exclusive_prod = initial_prod_value;
-        int64_t idx;
+        int64_t prod_idx;
         for (int64_t k = offset_start; k < offset_end; ++k) {
           if (k != j) {
-            idx = (k * stride_count) + lane_id;
-            exclusive_prod *= values_data[idx];
+            prod_idx = outer_idx * data_stride_axis * data_size_axis
+                       + k * data_stride_axis + lane_id;
+            exclusive_prod *= values_data[prod_idx];
           }
         }
-        grad_input_data[starting_index] = grad_data[output_index] * exclusive_prod;
+        grad_input_data[data_index] = grad_data[output_index] * exclusive_prod;
       } else {
-        grad_input_data[starting_index] = grad_val / values_data[starting_index];
+        grad_input_data[data_index] = grad_val / values_data[data_index];
       }
     }
   }
 }
-
 } // namespace
 
 Tensor _segment_reduce_cuda_backward_kernel(
@@ -251,28 +286,43 @@ Tensor _segment_reduce_cuda_backward_kernel(
     const Tensor& lengths_contig,
     int64_t axis,
     const c10::optional<Scalar>& initial) {
-  int64_t segment_count = lengths_contig.numel();
-  auto output_shape = data_contig.sizes().vec();
-  output_shape[axis] = segment_count;
+  axis = lengths_contig.dim() - 1;
+  int64_t segment_count = lengths_contig.size(axis);
+  int64_t lengths_stride_axis = lengths_contig.stride(axis);
   auto grad_input = at::zeros({data_contig.sizes()}, grad_contig.options());
 
-  int64_t stride_count = data_contig.numel() / data_contig.size(axis);
+  auto zeros_shape = lengths_contig.sizes().vec();
+  zeros_shape[axis] = 1;
+  auto offsets = at::cat({at::zeros(zeros_shape, lengths_contig.options()), lengths_contig}, axis);
+  offsets.cumsum_(axis);
 
-  auto offsets = _get_complete_sum(lengths_contig);
+  // outer_offset is the size of the outer dimensions of output (before axis)
+  // inner_offset is the size of the inner dimensions of output (after axis)
+  int64_t outer_offset = 1, inner_offset = 1;
+  for (int64_t d = 0; d < axis; d++) {
+    outer_offset *= output_contig.size(d);
+  }
+  for (int64_t d = axis + 1; d < output_contig.dim(); d++) {
+    inner_offset *= output_contig.size(d);
+  }
 
   constexpr int threads_per_block = 256;
-  int64_t num_blocks =
-      ((segment_count * stride_count) + threads_per_block - 1) /
-      threads_per_block;
+  int64_t num_blocks = (outer_offset * inner_offset * segment_count + threads_per_block - 1) / threads_per_block;
 
   num_blocks = std::max(num_blocks, (int64_t)1);
 
+  auto data_stride_axis = data_contig.stride(axis);
+  auto data_size_axis = data_contig.size(axis);
+  auto output_stride_axis = output_contig.stride(axis);
+  auto output_size_axis = output_contig.size(axis);
+  auto offsets_stride_axis = offsets.stride(axis);
+
   AT_DISPATCH_INDEX_TYPES(
-      lengths_contig.type(), "_segment_reduce_cuda_backward_kernel1", ([&] {
+      lengths_contig.scalar_type(), "_segment_reduce_cuda_backward_kernel1", ([&] {
         const auto* lengths_data = lengths_contig.data_ptr<index_t>();
         auto* offsets_data = offsets.data_ptr<index_t>();
 
-        // TODO: Swtich to TensorIterator for better maintainablility and
+        // TODO: Switch to TensorIterator for better maintainablility and
         // readability
         AT_DISPATCH_FLOATING_TYPES_AND2(
             kBFloat16,
@@ -305,8 +355,16 @@ Tensor _segment_reduce_cuda_backward_kernel(
                       lengths_data,
                       offsets_data,
                       segment_count,
-                      stride_count,
-                      initial_prod_value);
+                      lengths_stride_axis,
+                      initial_prod_value,
+                      outer_offset,
+                      inner_offset,
+                      data_stride_axis,
+                      data_size_axis,
+                      output_stride_axis,
+                      output_size_axis,
+                      offsets_stride_axis
+                    );
               C10_CUDA_KERNEL_LAUNCH_CHECK();
             }));
       }));
@@ -319,24 +377,46 @@ Tensor _segment_reduce_cuda_kernel(
     const Tensor& lengths,
     int64_t axis,
     const c10::optional<Scalar>& initial) {
-  int64_t segment_count = lengths.numel();
+  // data and lengths should be contiguous from the call to .contiguous in segment_reduce_kernel
+  TORCH_CHECK(data.is_contiguous(), "Expected data to be contiguous.");
+  TORCH_CHECK(lengths.is_contiguous(), "Expected lengths to be contiguous.");
+  axis = lengths.dim() - 1;
+  int64_t segment_count = lengths.size(axis);
+  int64_t lengths_stride_axis = lengths.stride(axis);
   auto output_shape = data.sizes().vec();
   output_shape[axis] = segment_count;
   auto output = at::empty(output_shape, data.options());
 
-  int64_t stride_count = data.numel() / data.size(axis);
-
-  auto offsets = _get_complete_sum(lengths);
+  // _get_complete_sum only supports 1D?
+  auto zeros_shape = lengths.sizes().vec();
+  zeros_shape[axis] = 1;
+  auto offsets = at::cat({at::zeros(zeros_shape, lengths.options()), lengths}, axis);
+  offsets.cumsum_(axis);
+
+  // outer_offset is the size of the outer dimensions of output (before axis)
+  // inner_offset is the size of the inner dimensions of output (after axis)
+  int64_t outer_offset = 1, inner_offset = 1;
+  for (int64_t d = 0; d < axis; d++) {
+    outer_offset *= output.size(d);
+  }
+  for (int64_t d = axis + 1; d < output.dim(); d++) {
+    inner_offset *= output.size(d);
+  }
 
   constexpr int threads_per_block = 256;
-  int64_t num_blocks =
-      ((segment_count * stride_count) + threads_per_block - 1) /
-      threads_per_block;
+  // segment_count * stride_count is just output.numel() ?
+  int64_t num_blocks = (output.numel() + threads_per_block - 1) / threads_per_block;
 
   num_blocks = std::max(num_blocks, (int64_t)1);
 
+  auto data_stride_axis = data.stride(axis);
+  auto data_size_axis = data.size(axis);
+  auto output_stride_axis = output.stride(axis);
+  auto output_size_axis = output.size(axis);
+  auto offsets_stride_axis = offsets.stride(axis);
+
   AT_DISPATCH_INDEX_TYPES(
-      lengths.type(), "_segment_reduce_cuda_kernel1", ([&] {
+      lengths.scalar_type(), "_segment_reduce_cuda_kernel1", ([&] {
         auto* offsets_data_ptr = offsets.data_ptr<index_t>();
         auto* lengths_data_ptr = lengths.data_ptr<index_t>();
         AT_DISPATCH_FLOATING_TYPES_AND2(
@@ -376,9 +456,17 @@ Tensor _segment_reduce_cuda_kernel(
                         lengths_data_ptr,
                         offsets_data_ptr,
                         segment_count,
-                        stride_count,
+                        lengths_stride_axis,
                         initial.has_value(),
-                        initial_value);
+                        initial_value,
+                        outer_offset,
+                        inner_offset,
+                        data_stride_axis,
+                        data_size_axis,
+                        output_stride_axis,
+                        output_size_axis,
+                        offsets_stride_axis
+                      );
                 C10_CUDA_KERNEL_LAUNCH_CHECK();
               } else {
                 if (reduction == SegmentReductionType::MAX) {

test/test_ops.py

@@ -1547,7 +1547,7 @@ def test_refs_are_in_python_ref_db(self, op):
     "to_sparse",  # Could not run 'aten::to_sparse' with arguments from the 'Meta' backend
     "tensor_split",  # The tensor has a non-zero number of elements, but its data is not allocated yet
     "repeat_interleave",  # cannot repeat_interleave a meta tensor without output_size
-    "segment_reduce",  # Could not run 'aten::segment_reduce' with arguments from the 'Meta' backend.
+    "segment_reduce.lengths",  # Could not run 'aten::segment_reduce' with arguments from the 'Meta' backend.
     "sparse.sampled.addmm",  # sparsity not supported
     # Can not infer total number of classes from meta. no way at present to throw DynamicOutputShapeException
     "nn.functional.one_hot",