Tensor: Handle oversized GPU convolution kernels

libeigen/eigen!3438

Closes #2762

contrib/Eigen/src/Tensor/TensorConvolution.h

diff --git a/contrib/Eigen/src/Tensor/TensorConvolution.h b/contrib/Eigen/src/Tensor/TensorConvolution.h
index 70357aa..8dab047 100644
--- a/contrib/Eigen/src/Tensor/TensorConvolution.h
+++ b/contrib/Eigen/src/Tensor/TensorConvolution.h
@@ -653,6 +653,40 @@
   }
 }
 
+// Oversized kernels cannot stage even one output's input footprint in shared memory.
+template <typename InputEvaluator, typename Index, typename InputDims>
+__global__ void EigenConvolutionKernel3DGlobal(
+    InputEvaluator eval, const internal::IndexMapper<Index, InputDims, 3, InputEvaluator::Layout> indexMapper,
+    const float* kernel, Index total, Index numX, Index numY, Index numZ, Index kernelSizeX, Index kernelSizeY,
+    Index kernelSizeZ, float* buffer) {
+  const Index step = Index(blockDim.x) * gridDim.x;
+  for (Index output = Index(blockIdx.x) * blockDim.x + threadIdx.x; output < total;) {
+    Index remaining = output;
+    const Index x = remaining % numX;
+    remaining /= numX;
+    const Index y = remaining % numY;
+    remaining /= numY;
+    const Index z = remaining % numZ;
+    const Index plane = remaining / numZ;
+    const Index inputOffset = indexMapper.mapGpuInputPlaneToTensorInputOffset(plane);
+    float result = 0;
+    Index kernelIndex = 0;
+    for (Index k = 0; k < kernelSizeZ; ++k) {
+      for (Index j = 0; j < kernelSizeY; ++j) {
+        for (Index i = 0; i < kernelSizeX; ++i) {
+          result += eval.coeff(inputOffset + indexMapper.mapGpuInputKernelToTensorInputOffset(x + i, y + j, z + k)) *
+                    kernel[kernelIndex++];
+        }
+      }
+    }
+    buffer[indexMapper.mapGpuOutputPlaneToTensorOutputOffset(plane) +
+           indexMapper.mapGpuOutputKernelToTensorOutputOffset(x, y, z)] = result;
+    // Test the remaining distance before adding: even the final increment must not overflow Index.
+    if (step >= total - output) break;
+    output += step;
+  }
+}
+
 template <typename InputEvaluator, typename Index, typename InputDims>
 __global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void EigenConvolutionKernel3D(
     InputEvaluator eval, const internal::IndexMapper<Index, InputDims, 3, InputEvaluator::Layout> indexMapper,
@@ -1014,8 +1048,28 @@
         const int numX = dimensions()[m_indices[idxX]];
         const int numY = dimensions()[m_indices[idxY]];
         const int numZ = dimensions()[m_indices[idxZ]];
-        const int numP = dimensions().TotalSize() / (numX * numY * numZ);
+        const Index total = dimensions().TotalSize();
+        if (total == 0) return;
 
+        const array<Index, 3> indices{m_indices[idxX], m_indices[idxY], m_indices[idxZ]};
+        const array<Index, 3> kernel_dims{m_kernelImpl.dimensions()[idxX], m_kernelImpl.dimensions()[idxY],
+                                          m_kernelImpl.dimensions()[idxZ]};
+        internal::IndexMapper<Index, InputDims, 3, Layout> indexMapper(m_inputImpl.dimensions(), kernel_dims, indices);
+
+        if (m_kernelImpl.dimensions().TotalSize() > Index(maxSharedMem / sizeof(Scalar))) {
+          constexpr int preferredBlockSize = 256;
+          const int threads = numext::mini(preferredBlockSize, maxThreadsPerBlock);
+          const int maxBlocks = static_cast<int>(
+              numext::mini<int64_t>(int64_t(numMultiProcessors) * 32, NumTraits<Index>::highest() / threads));
+          const int blocks = static_cast<int>(numext::mini<Index>((total - 1) / threads + 1, Index(maxBlocks)));
+          LAUNCH_GPU_KERNEL(
+              (EigenConvolutionKernel3DGlobal<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims>), blocks,
+              threads, 0, m_device, m_inputImpl, indexMapper, m_kernel, total, Index(numX), Index(numY), Index(numZ),
+              Index(kernel_size_x), Index(kernel_size_y), Index(kernel_size_z), data);
+          break;
+        }
+
+        const int numP = total / (Index(numX) * numY * numZ);
         const int maxX = numext::mini<int>(
             128, numext::mini<int>(maxSharedMem / (sizeof(Scalar) * kernel_size_y * kernel_size_z) - kernel_size_x + 1,
                                    numX));
@@ -1039,11 +1093,6 @@
             (maxX + kernel_size_x - 1) * (maxY + kernel_size_y - 1) * (maxZ + kernel_size_z - 1) * sizeof(Scalar);
         gpu_assert(shared_mem <= maxSharedMem);
 
-        const array<Index, 3> indices{m_indices[idxX], m_indices[idxY], m_indices[idxZ]};
-        const array<Index, 3> kernel_dims{m_kernelImpl.dimensions()[idxX], m_kernelImpl.dimensions()[idxY],
-                                          m_kernelImpl.dimensions()[idxZ]};
-        internal::IndexMapper<Index, InputDims, 3, Layout> indexMapper(m_inputImpl.dimensions(), kernel_dims, indices);
-
         LAUNCH_GPU_KERNEL((EigenConvolutionKernel3D<TensorEvaluator<InputArgType, GpuDevice>, Index, InputDims>),
                           num_blocks, block_size, shared_mem, m_device, m_inputImpl, indexMapper, m_kernel, numP, numX,
                           maxX, numY, maxY, numZ, maxZ, kernel_size_x, kernel_size_y, kernel_size_z, data);

contrib/test/CMakeLists.txt

diff --git a/contrib/test/CMakeLists.txt b/contrib/test/CMakeLists.txt
index 3469f3a..c7ba51a 100644
--- a/contrib/test/CMakeLists.txt
+++ b/contrib/test/CMakeLists.txt
@@ -268,6 +268,7 @@
   ei_add_test(tensor_device)
   ei_add_test(tensor_gpu_device_attributes)
   ei_add_test(tensor_gpu)
+  ei_add_test(tensor_convolution_large_gpu)
   ei_add_test(tensor_gpu_runtime_check)
   ei_add_test(tensor_gpu_runtime_check_ndebug)
   ei_add_test(tensor_contract_gpu)

contrib/test/tensor_convolution_large_gpu.cu

diff --git a/contrib/test/tensor_convolution_large_gpu.cu b/contrib/test/tensor_convolution_large_gpu.cu
new file mode 100644
index 0000000..a70d7ff
--- /dev/null
+++ b/contrib/test/tensor_convolution_large_gpu.cu
@@ -0,0 +1,108 @@
+// SPDX-FileCopyrightText: The Eigen Authors
+// SPDX-License-Identifier: MPL-2.0
+
+#define EIGEN_USE_GPU
+#define EIGEN_TEST_NO_LONGDOUBLE
+#define EIGEN_TEST_NO_COMPLEX
+#include "main.h"
+#include <contrib/Eigen/Tensor>
+
+template <int Layout, int NumDims, typename TensorIndex>
+void verify_gpu_convolution(const Tensor<float, NumDims, Layout, TensorIndex>& input,
+                            const Tensor<float, 3, Layout, TensorIndex>& kernel, const array<TensorIndex, 3>& axes,
+                            const Tensor<float, NumDims, Layout, TensorIndex>& expected) {
+  GpuStreamDevice stream;
+  GpuDevice device(&stream);
+  bool nonzero = false, varying = false;
+  for (Index i = 0; i < expected.size(); ++i) {
+    nonzero |= expected.data()[i] != 0;
+    varying |= expected.data()[i] != expected.data()[0];
+  }
+  VERIFY(nonzero);
+  VERIFY(varying);
+  Tensor<float, NumDims, Layout, TensorIndex> result;
+  result.resize(expected.dimensions());
+  float* a = static_cast<float*>(device.allocate(input.size() * sizeof(float)));
+  float* k = static_cast<float*>(device.allocate(kernel.size() * sizeof(float)));
+  float* out = static_cast<float*>(device.allocate(result.size() * sizeof(float)));
+  device.memcpyHostToDevice(a, input.data(), input.size() * sizeof(float));
+  device.memcpyHostToDevice(k, kernel.data(), kernel.size() * sizeof(float));
+  // NaN sentinels make missing output writes fail even when the reference is zero.
+  device.memset(out, 0xff, result.size() * sizeof(float));
+  TensorMap<Tensor<float, NumDims, Layout, TensorIndex>> gpu_input(a, input.dimensions());
+  TensorMap<Tensor<float, 3, Layout, TensorIndex>> gpu_kernel(k, kernel.dimensions());
+  TensorMap<Tensor<float, NumDims, Layout, TensorIndex>> gpu_result(out, result.dimensions());
+  gpu_result.device(device) = gpu_input.convolve(gpu_kernel, axes);
+  device.memcpyDeviceToHost(result.data(), out, result.size() * sizeof(float));
+  device.synchronize();
+  for (Index i = 0; i < result.size(); ++i) VERIFY_IS_EQUAL(result.data()[i], expected.data()[i]);
+  device.deallocate(a);
+  device.deallocate(k);
+  device.deallocate(out);
+}
+
+template <int Layout>
+void check_large_convolution(const array<Index, 3>& kernel_dims, const array<Index, 3>& axes) {
+  array<Index, 5> input_dims{{2, 3, 2, 3, 2}};
+  for (int i = 0; i < 3; ++i) input_dims[axes[i]] = kernel_dims[i] + i + 1;
+  Tensor<float, 5, Layout> input(input_dims);
+  Tensor<float, 3, Layout> kernel(kernel_dims);
+  // Small integers keep every product and partial sum exactly representable in float.
+  for (Index i = 0; i < input.size(); ++i) input.data()[i] = float((i * 17 + i / 7) % 19 - 9);
+  for (Index i = 0; i < kernel.size(); ++i) kernel.data()[i] = float((i * 7 + i / 5) % 7 - 2);
+  Tensor<float, 5, Layout> expected = input.convolve(kernel, axes);
+  verify_gpu_convolution(input, kernel, axes, expected);
+}
+
+template <int Layout, typename TensorIndex = Index>
+void check_convolution_grid_stride(Index output_size = 0) {
+  GpuStreamDevice stream;
+  GpuDevice device(&stream);
+  const Index kernel_size = device.sharedMemPerBlock() / sizeof(float) + 1;
+  const Index threads = numext::mini(256, device.maxGpuThreadsPerBlock());
+  const Index grid_stride = Index(device.getNumGpuMultiProcessors()) * 32 * threads;
+  if (output_size == 0) output_size = 2 * grid_stride + 17;
+  const Index max_input =
+      static_cast<Index>(numext::mini<int64_t>(NumTraits<Index>::highest(), NumTraits<TensorIndex>::highest()));
+  const Index max_output = max_input - kernel_size + 1;
+  // Some devices cannot fit an oversized convolution's input in this narrow index type.
+  if (max_output < 2) return;
+  output_size = numext::mini(output_size, max_output);
+  const array<TensorIndex, 3> axes{{0, 1, 2}};
+  Tensor<float, 3, Layout, TensorIndex> input(1, 1, TensorIndex(kernel_size + output_size - 1));
+  Tensor<float, 3, Layout, TensorIndex> kernel(1, 1, TensorIndex(kernel_size));
+  Tensor<float, 3, Layout, TensorIndex> expected(1, 1, TensorIndex(output_size));
+  input.setConstant(1);
+  input.data()[kernel_size + output_size / 2 - 1] = 2;
+  kernel.setConstant(1);
+  // An independent sliding sum avoids a CPU convolution of O(kernel_size * output_size).
+  Index sum = kernel_size;
+  for (Index i = 0; i < output_size; ++i) {
+    expected.data()[i] = float(sum);
+    if (i + 1 < output_size) sum += Index(input.data()[i + kernel_size]) - Index(input.data()[i]);
+  }
+  if (std::is_same<TensorIndex, Index>::value) VERIFY(output_size > grid_stride);
+  verify_gpu_convolution(input, kernel, axes, expected);
+}
+
+EIGEN_DECLARE_TEST(tensor_convolution_large_gpu) {
+  GpuStreamDevice stream;
+  GpuDevice device(&stream);
+  // Always exceed this device's shared-memory limit, including newer GPUs.
+  const Index oversized = device.sharedMemPerBlock() / (sizeof(float) * 3 * 50) + 1;
+  const Index capacity = device.sharedMemPerBlock() / sizeof(float);
+  const array<Index, 3> shapes[] = {{{3, 4, 2}},           {{3, 50, 120}},         {{3, 50, 145}},
+                                    {{3, 50, oversized}},  {{1, 1, capacity - 1}}, {{1, 1, capacity}},
+                                    {{1, 1, capacity + 1}}};
+  const array<Index, 3> axes[] = {{{1, 2, 3}}, {{0, 2, 4}}, {{4, 0, 2}}};
+  for (const auto& shape : shapes) {
+    for (const auto& axis : axes) {
+      CALL_SUBTEST(check_large_convolution<ColMajor>(shape, axis));
+      CALL_SUBTEST(check_large_convolution<RowMajor>(shape, axis));
+    }
+  }
+  CALL_SUBTEST(check_convolution_grid_stride<ColMajor>());
+  CALL_SUBTEST(check_convolution_grid_stride<RowMajor>());
+  CALL_SUBTEST((check_convolution_grid_stride<ColMajor, int16_t>(18000)));
+  CALL_SUBTEST((check_convolution_grid_stride<RowMajor, int16_t>(18000)));
+}