Enable op Aten::_jagged_to_padded_dense_forward (#1517)

Enable op `Aten::_jagged_to_padded_dense_forward`

---------

Co-authored-by: Yutao Xu <[email protected]>

Author: chunhuanMeng

src/ATen/native/nested/xpu/NestedTensorTransformerFunctions.cpp

@@ -4,6 +4,8 @@
 #include <ATen/native/nested/NestedTensorUtils.h>
 #include <ATen/native/nested/xpu/sycl/NestedTensorTransformerFunctionKernels.h>
 
+#include <comm/XPUGuard.h>
+
 namespace at::native {
 
 namespace {
@@ -205,4 +207,24 @@ Tensor NestedTensor_to_padded_tensor_xpu(
   return NestedTensor_to_padded_tensor_generic(t, padding, output_size);
 }
 
+at::Tensor _fbgemm_jagged_to_padded_dense_forward(
+    const Tensor& values,
+    TensorList offsets,
+    c10::IntArrayRef max_lengths,
+    const double padding_value) {
+  const size_t num_jagged_dim = offsets.size();
+
+  TORCH_CHECK(
+      max_lengths.size() == num_jagged_dim,
+      "max_lengths.size(), ",
+      max_lengths.size(),
+      " != num_jagged_dim, ",
+      num_jagged_dim);
+  at::xpu::OptionalXPUGuard device_guard;
+  device_guard.set_index(values.get_device());
+
+  return at::native::xpu::_fbgemm_jagged_to_padded_dense_forward_kernel(
+      values, offsets, max_lengths, padding_value);
+}
+
 } // namespace at::native

src/ATen/native/nested/xpu/sycl/NestedTensorTransformerFunctionKernels.cpp

@@ -1,4 +1,7 @@
 #include <ATen/ATen.h>
+#include <ATen/Dispatch.h>
+#include <ATen/core/TensorAccessor.h>
+#include <ATen/native/StridedRandomAccessor.h>
 #include <ATen/native/nested/xpu/sycl/NestedTensorTransformerFunctionKernels.h>
 #include <comm/SYCLContext.h>
 
@@ -613,4 +616,343 @@ void add_padding_kernel(
       });
 }
 
+#define JAGGED_TENSOR_DISPATCH_DIMS()                                         \
+  AT_DISPATCH_INDEX_TYPES(x_offsets[0].scalar_type(), "jagged_indices", [=] { \
+    switch (num_jagged_dim) {                                                 \
+      case 1:                                                                 \
+        INVOKE_KERNEL_WITH_DIM(1);                                            \
+        break;                                                                \
+      case 2:                                                                 \
+        INVOKE_KERNEL_WITH_DIM(2);                                            \
+        break;                                                                \
+      case 3:                                                                 \
+        INVOKE_KERNEL_WITH_DIM(3);                                            \
+        break;                                                                \
+      case 4:                                                                 \
+        INVOKE_KERNEL_WITH_DIM(4);                                            \
+        break;                                                                \
+      case 5:                                                                 \
+        INVOKE_KERNEL_WITH_DIM(5);                                            \
+        break;                                                                \
+      default:                                                                \
+        TORCH_CHECK(                                                          \
+            false, "unsupported number of jagged dim ", num_jagged_dim);      \
+    }                                                                         \
+  });
+
+inline std::string torch_tensor_device_name(const at::Tensor& ten) {
+  return c10::DeviceTypeName(ten.device().type());
+}
+
+inline std::string torch_tensor_device_name(
+    const std::optional<at::Tensor>& ten) {
+  if (ten.has_value()) {
+    return torch_tensor_device_name(ten.value());
+  } else {
+    return "N/A";
+  }
+}
+
+inline bool torch_tensor_on_xpu_gpu_check(const at::Tensor& ten) {
+  return ten.is_xpu();
+}
+
+inline bool torch_tensor_on_xpu_gpu_check(
+    const std::optional<at::Tensor>& ten) {
+  return !ten.has_value() || torch_tensor_on_xpu_gpu_check(ten.value());
+}
+
+#define TENSOR_ON_XPU_GPU(x)                                  \
+  TORCH_CHECK(                                                \
+      torch_tensor_on_xpu_gpu_check(x),                       \
+      #x " must be a XPU tensor; it is currently on device ", \
+      torch_tensor_device_name(x))
+
+// A wrapper class for passing dynamically sized dimension information (e.g.
+// tensor.dims()) from the host to device.
+constexpr size_t kStackArrayMaxDims = 5;
+
+template <typename T>
+struct StackArray {
+  T vals[kStackArrayMaxDims];
+  size_t ndim;
+};
+
+template <typename scalar_t>
+struct PaddingValueFuncutor {
+  scalar_t operator()(scalar_t x, scalar_t /*unused*/) const {
+    return x;
+  }
+};
+
+// Subgroup size
+static constexpr int32_t kSubgroupSize = 32;
+// Max thread num in one thread workgroup
+static constexpr int32_t kMaxThreads = 1024;
+
+inline int32_t div_round_up(int32_t a, int32_t b) {
+  return (a + b - 1) / b;
+}
+
+inline int32_t round_down(int32_t a, int32_t b) {
+  return a / b * b;
+}
+
+inline std::tuple<sycl::range<2>, sycl::range<2>, StackArray<int64_t>>
+check_shape_and_partition_(
+    const Tensor& values,
+    const std::vector<Tensor>& offsets,
+    const Tensor& dense_tensor) {
+  const int outer_dense_size = dense_tensor.size(0);
+  TORCH_CHECK(
+      outer_dense_size == offsets[0].numel() - 1,
+      "outer_dense_size, ",
+      outer_dense_size,
+      " != offsets[0].numel() - 1, ",
+      offsets[0].numel() - 1);
+  const int inner_dense_size = dense_tensor.size(-1);
+  TORCH_CHECK(
+      inner_dense_size == values.size(-1),
+      "inner_dense_size, ",
+      inner_dense_size,
+      " != values.size(-1), ",
+      values.size(-1));
+  const int jagged_folded_size =
+      dense_tensor.numel() / (outer_dense_size * inner_dense_size);
+
+  const int wg_size_x =
+      inner_dense_size >= kSubgroupSize / 2 ? kSubgroupSize : inner_dense_size;
+  const int wg_size_y = kMaxThreads / kSubgroupSize;
+  const int num_group =
+      div_round_up(outer_dense_size * jagged_folded_size, wg_size_y);
+
+  StackArray<int64_t> jagged_dims_tensor{};
+  const int num_jagged_dim = dense_tensor.dim() - 2;
+  TORCH_CHECK(num_jagged_dim <= static_cast<int>(kStackArrayMaxDims));
+  jagged_dims_tensor.ndim = num_jagged_dim;
+  std::memcpy(
+      &(jagged_dims_tensor.vals[0]),
+      dense_tensor.sizes().data() + 1,
+      num_jagged_dim * sizeof(int64_t));
+  return {
+      sycl::range<2>(wg_size_x, wg_size_y),
+      sycl::range<2>(num_group * wg_size_x, wg_size_y),
+      jagged_dims_tensor};
+}
+
+template <int NUM_JAGGED_DIM, typename index_t>
+inline bool walk_down_tensor_storage_tree_(
+    int& offset,
+    const int flattened_jagged_idx,
+    const StackArray<int64_t>& jagged_dims,
+    const StackArray<index_t*>& x_offsets) {
+  // compute coorindates
+  int jagged_coords[NUM_JAGGED_DIM];
+  int j_temp = flattened_jagged_idx;
+#pragma unroll
+  for (int d = NUM_JAGGED_DIM - 1; d >= 0; --d) {
+    const int jagged_size = jagged_dims.vals[d];
+    jagged_coords[d] = j_temp % jagged_size;
+    j_temp /= jagged_size;
+  }
+
+  // walk down the tree
+  bool is_zero = false;
+#pragma unroll
+  for (int d = 0; d < NUM_JAGGED_DIM; ++d) {
+    const int begin = x_offsets.vals[d][offset];
+    const int end = x_offsets.vals[d][offset + 1];
+    if (jagged_coords[d] >= end - begin) {
+      is_zero = true;
+      break;
+    }
+    offset = begin + jagged_coords[d];
+  }
+  return is_zero;
+}
+
+template <int NUM_JAGGED_DIM, typename index_t, typename scalar_t, typename F>
+struct JaggedDenseElementwiseDenseFunctor {
+  void operator()(sycl::nd_item<2> item) const {
+    const int outer_dense_size = y_.size(0);
+    const int jagged_folded_size = y_.size(1);
+    const int inner_dense_size = y_.size(2);
+    auto output = output_;
+    const int outer_begin =
+        item.get_group(0) * item.get_local_range(1) + item.get_local_id(1);
+    const int outer_stride = item.get_group_range(0) * item.get_local_range(1);
+    for (int outer = outer_begin; outer < outer_dense_size * jagged_folded_size;
+         outer += outer_stride) {
+      const int oidx = outer / jagged_folded_size;
+      const int jidx = outer % jagged_folded_size;
+
+      int offset = oidx;
+      const bool is_zero = walk_down_tensor_storage_tree_<NUM_JAGGED_DIM>(
+          offset, jidx, jagged_dims_, x_offsets_);
+
+      if (is_zero) {
+        int iidx;
+        for (iidx = item.get_local_id(0); iidx * 2 + 1 < inner_dense_size;
+             iidx += item.get_local_range(0)) {
+          output[oidx][jidx][2 * iidx] =
+              f_(padding_value_, y_[oidx][jidx][2 * iidx]);
+          output[oidx][jidx][2 * iidx + 1] =
+              f_(padding_value_, y_[oidx][jidx][2 * iidx + 1]);
+        }
+        if (iidx * 2 + 1 == inner_dense_size) {
+          output[oidx][jidx][2 * iidx] =
+              f_(padding_value_, y_[oidx][jidx][2 * iidx]);
+        }
+      } else {
+        int iidx;
+        for (iidx = item.get_local_id(0); iidx * 2 + 1 < inner_dense_size;
+             iidx += item.get_local_range(0)) {
+          output[oidx][jidx][2 * iidx] =
+              f_(x_values_[offset][2 * iidx], y_[oidx][jidx][2 * iidx]);
+          output[oidx][jidx][2 * iidx + 1] =
+              f_(x_values_[offset][2 * iidx + 1], y_[oidx][jidx][2 * iidx + 1]);
+        }
+        if (iidx * 2 + 1 == inner_dense_size) {
+          output[oidx][jidx][2 * iidx] =
+              f_(x_values_[offset][2 * iidx], y_[oidx][jidx][2 * iidx]);
+        }
+      }
+    }
+  }
+  JaggedDenseElementwiseDenseFunctor(
+      const at::PackedTensorAccessor32<scalar_t, 2, RestrictPtrTraits> x_values,
+      StackArray<index_t*> x_offsets,
+      const at::PackedTensorAccessor32<scalar_t, 3, RestrictPtrTraits> y,
+      at::PackedTensorAccessor32<scalar_t, 3, RestrictPtrTraits> output,
+      StackArray<int64_t> jagged_dims,
+      F f,
+      const scalar_t padding_value)
+      : x_values_(x_values),
+        x_offsets_(x_offsets),
+        y_(y),
+        output_(output),
+        jagged_dims_(jagged_dims),
+        f_(f),
+        padding_value_(padding_value) {}
+
+ private:
+  const at::PackedTensorAccessor32<scalar_t, 2, RestrictPtrTraits> x_values_;
+  StackArray<index_t*> x_offsets_;
+  const at::PackedTensorAccessor32<scalar_t, 3, RestrictPtrTraits> y_;
+  at::PackedTensorAccessor32<scalar_t, 3, RestrictPtrTraits> output_;
+  StackArray<int64_t> jagged_dims_;
+  F f_;
+  const scalar_t padding_value_;
+};
+
+template <typename scalar_t, typename F>
+void jagged_dense_elementwise_dense_template(
+    const Tensor& x_values,
+    const std::vector<Tensor>& x_offsets,
+    const Tensor& y,
+    const Tensor& output,
+    F f,
+    const scalar_t padding_value = static_cast<scalar_t>(0)) {
+  TENSOR_ON_XPU_GPU(x_values);
+  for (auto& x_offset : x_offsets) {
+    TENSOR_ON_XPU_GPU(x_offset);
+  }
+
+  const int num_jagged_dim = y.dim() - 2;
+  TORCH_CHECK(
+      x_offsets.size() == static_cast<size_t>(num_jagged_dim),
+      "x_offsets.size(), ",
+      x_offsets.size(),
+      " != num_jagged_dim ",
+      num_jagged_dim);
+
+  if (y.numel() == 0) {
+    return;
+  }
+
+  sycl::range<2> global_range, local_range;
+  StackArray<int64_t> jagged_dims_tensor;
+  std::tie(local_range, global_range, jagged_dims_tensor) =
+      check_shape_and_partition_(x_values, x_offsets, y);
+
+  // Canonicalize y and output to 3D, collapsing jagged dimensions.
+  const Tensor y_reshaped = y.view({y.size(0), -1, y.size(-1)});
+  Tensor output_reshaped = output.view(y_reshaped.sizes());
+
+#define INVOKE_KERNEL_WITH_DIM(NUM_JAGGED_DIM)                                 \
+  {                                                                            \
+    std::vector<Tensor> x_offsets_contig;                                      \
+    x_offsets_contig.resize(num_jagged_dim);                                   \
+    StackArray<index_t*> x_offset_ptrs;                                        \
+    x_offset_ptrs.ndim = num_jagged_dim;                                       \
+    for (int d = 0; d < num_jagged_dim; ++d) {                                 \
+      x_offsets_contig[d] = x_offsets[d].contiguous();                         \
+      x_offset_ptrs.vals[d] =                                                  \
+          x_offsets_contig[d].template data_ptr<index_t>();                    \
+    }                                                                          \
+    auto kfn = JaggedDenseElementwiseDenseFunctor<                             \
+        NUM_JAGGED_DIM,                                                        \
+        index_t,                                                               \
+        scalar_t,                                                              \
+        F>(                                                                    \
+        x_values.packed_accessor32<scalar_t, 2, RestrictPtrTraits>(),          \
+        x_offset_ptrs,                                                         \
+        y_reshaped.packed_accessor32<scalar_t, 3, RestrictPtrTraits>(),        \
+        output_reshaped.packed_accessor32<scalar_t, 3, RestrictPtrTraits>(),   \
+        jagged_dims_tensor,                                                    \
+        f,                                                                     \
+        padding_value);                                                        \
+    sycl_kernel_submit(global_range, local_range, getCurrentSYCLQueue(), kfn); \
+  }
+
+  JAGGED_TENSOR_DISPATCH_DIMS();
+
+#undef INVOKE_KERNEL_WITH_DIM
+}
+
+at::Tensor _fbgemm_jagged_to_padded_dense_forward_kernel(
+    const Tensor& values,
+    TensorList offsets,
+    c10::IntArrayRef max_lengths,
+    const double padding_value) {
+  const Tensor values_canonicalized = values.view(
+      {values.size(0),
+       std::accumulate(
+           values.sizes().begin() + 1,
+           values.sizes().end(),
+           1,
+           std::multiplies<size_t>())});
+  at::SymDimVector padded_values_shape({at::SymInt(offsets[0].size(0) - 1)});
+  padded_values_shape.insert(
+      padded_values_shape.end(), max_lengths.begin(), max_lengths.end());
+
+  // Canonicalize padded_values by unsqueeze the last dim if the inner dense
+  // dimension is 1 and folded.
+  const bool D_folded = values.dim() == 1;
+  if (!D_folded) {
+    padded_values_shape.push_back(values.size(-1));
+  }
+  Tensor padded_values =
+      at::empty_symint(padded_values_shape, values.options());
+  Tensor padded_values_view =
+      D_folded ? padded_values.unsqueeze(-1) : padded_values;
+
+  AT_DISPATCH_ALL_TYPES_AND2(
+      at::ScalarType::Half,
+      at::ScalarType::BFloat16,
+      values.scalar_type(),
+      "jagged_to_padded_dense_xpu",
+      [&] {
+        jagged_dense_elementwise_dense_template<scalar_t>(
+            values_canonicalized,
+            offsets.vec(),
+            padded_values_view, // dummy not used in the lambda function
+            padded_values_view,
+            PaddingValueFuncutor<scalar_t>(),
+            static_cast<scalar_t>(padding_value));
+      });
+
+  return padded_values;
+}
+
 } // namespace at::native::xpu

src/ATen/native/nested/xpu/sycl/NestedTensorTransformerFunctionKernels.h

@@ -51,4 +51,10 @@ TORCH_XPU_API void add_padding_kernel(
     const int batch_size,
     const int output_batch_size);
 
+TORCH_XPU_API at::Tensor _fbgemm_jagged_to_padded_dense_forward_kernel(
+    const Tensor& values,
+    TensorList offsets,
+    c10::IntArrayRef max_lengths,
+    const double padding_value);
+
 } // namespace at::native::xpu