chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/row_conv_kernel.cu

@@ -0,0 +1,162 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/kernels/row_conv_kernel.h"
+#include "paddle/phi/backends/gpu/gpu_device_function.h"
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/core/mixed_vector.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace phi {
+
+namespace {
+
+static inline int DivUp(int x, int y) { return (x + y - 1) / y; }
+
+// Forward prop (shared memory version, for small future_context)
+template <typename T>
+__global__ void RowConvForwardSharedMemory(const T *in,
+                                           const T *wt,
+                                           int num_sequence,
+                                           int input_dim,
+                                           int future_context,
+                                           const size_t *batch_indices,
+                                           T *out) {
+  int blx = blockDim.x;
+  int bly = blockDim.y;
+  int thx = threadIdx.x;
+  int thy = threadIdx.y;
+  int d = blockIdx.x * blx + thx;  // index along input dim
+
+  extern __shared__ T mem[];
+  T *sw = mem;
+
+  if (thy < future_context) {
+    sw[thy * blx + thx] =
+        (d < input_dim) ? wt[thy * input_dim + d] : static_cast<T>(0);
+  }
+  __syncthreads();
+  for (size_t i = 0; i < num_sequence; i++) {
+    size_t start = batch_indices[i];
+    size_t end = batch_indices[i + 1];
+    size_t current_timesteps = end - start;
+
+    for (size_t k = thy; k < current_timesteps; k += bly) {
+      T sum = 0;
+      for (int w = 0; (w < future_context) && ((k + w) < current_timesteps);
+           w++) {
+        sum += (d < input_dim)
+                   ? sw[w * blx + thx] * in[(start + k + w) * input_dim + d]
+                   : static_cast<T>(0);
+      }
+      if (d < input_dim) {
+        out[(start + k) * input_dim + d] = sum;
+      }
+    }
+  }
+}
+
+// Forward prop (naive version)
+template <typename T>
+__global__ void RowConvForward(const T *in,
+                               const T *wt,
+                               int num_sequence,
+                               int input_dim,
+                               int future_context,
+                               const size_t *batch_indices,
+                               T *out) {
+  int64_t d =
+      static_cast<int64_t>(blockIdx.x) * static_cast<int64_t>(blockDim.x) +
+      static_cast<int64_t>(threadIdx.x);  // index along input_dim
+  int bly = blockDim.y;
+  int thy = threadIdx.y;
+
+  if (d >= input_dim) return;
+  for (size_t i = 0; i < num_sequence; i++) {
+    size_t start = batch_indices[i];
+    size_t end = batch_indices[i + 1];
+    size_t current_timesteps = end - start;
+
+    for (size_t k = thy; k < current_timesteps; k += bly) {
+      T sum = 0;
+      for (size_t w = 0; (w < future_context) && ((k + w) < current_timesteps);
+           w++) {
+        sum += (wt[w * input_dim + d] * in[(start + k + w) * input_dim + d]);
+      }
+      out[(start + k) * input_dim + d] = sum;
+    }
+  }
+}
+}  // namespace
+
+template <typename T, typename Context>
+void RowConvKernel(const Context &dev_ctx,
+                   const DenseTensor &x_in,
+                   const DenseTensor &filter_in,
+                   DenseTensor *Out) {
+  auto *X = &x_in;
+  auto *Filter = &filter_in;
+
+  const T *in = X->data<T>();
+  const T *weight = Filter->data<T>();
+  T *out = dev_ctx.template Alloc<T>(Out);
+  bool is_tensor = X->lod().empty();
+  int batch_size = 0;
+  if (is_tensor) {
+    batch_size = X->dims()[0];
+  } else {
+    batch_size = X->lod()[0].size() - 1;
+  }
+  int input_dim = 0;
+  Vector<size_t> batch_indices(batch_size + 1);
+  int64_t timesteps = X->dims()[1];
+
+  if (is_tensor) {
+    for (int i = 0; i < batch_size + 1; i++) {
+      batch_indices[i] = i * timesteps;
+    }
+    input_dim = X->dims()[2];
+  } else {
+    batch_indices = X->lod()[0];
+    input_dim = X->dims()[1];
+  }
+
+  int num_sequence = batch_indices.size() - 1;
+  int64_t future_context = Filter->dims()[0];
+  // TODO(large-tensor): CUDA kernel future_context not support int64
+  PADDLE_ENFORCE_LE_INT_MAX(future_context, "future_context");
+  int future_context_int = static_cast<int>(future_context);
+
+  MixVector<size_t> mix_vector(&batch_indices);
+  size_t *idx = mix_vector.CUDAMutableData(dev_ctx.GetPlace());
+  auto stream = dev_ctx.stream();
+
+  if (future_context_int <= 32) {
+    dim3 block_dim = dim3(32, 32);
+    dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
+    int mem_per_block = (future_context_int * block_dim.x) * sizeof(T);
+    RowConvForwardSharedMemory<T>
+        <<<grid_dim, block_dim, mem_per_block, stream>>>(
+            in, weight, num_sequence, input_dim, future_context_int, idx, out);
+  } else {
+    dim3 block_dim = dim3(32, 32);
+    dim3 grid_dim = dim3(DivUp(input_dim, block_dim.x), 1);
+    RowConvForward<T><<<grid_dim, block_dim, 0, stream>>>(
+        in, weight, num_sequence, input_dim, future_context_int, idx, out);
+  }
+  mix_vector.CopyToCPU();
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(row_conv, GPU, ALL_LAYOUT, phi::RowConvKernel, float) {}