chore: import upstream snapshot with attribution

2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
@@ -0,0 +1,354 @@
+/**
+ * Copyright 2017-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h> // @manual=//caffe2/aten:ATen-cu
+#include <cuda_runtime.h>
+#include <algorithm> // std::min/max
+#include <cub/cub.cuh>
+
+#include "alignment_train_cuda.h"
+#include "utils.h"
+
+namespace {
+
+// The thread block length in threads along the X dimension
+constexpr int BLOCK_DIM_X = 128;
+// The thread block length in threads along the Y dimension
+constexpr int BLOCK_DIM_Y = 8;
+// The thread block length in threads for scan operation
+constexpr int SCAN_BLOCK = 512;
+
+#define gpuErrchk(ans) \
+  { gpuAssert((ans), __FILE__, __LINE__); }
+
+inline void
+gpuAssert(cudaError_t code, const char* file, int line, bool abort = true) {
+  if (code != cudaSuccess) {
+    fprintf(
+        stderr,
+        "\nGPUassert: %s %s %d\n",
+        cudaGetErrorString(code),
+        file,
+        line);
+    if (abort)
+      exit(code);
+  }
+}
+
+template <typename T>
+struct Prod {
+  /// prod operator, returns <tt>a * b</tt>
+  __host__ __device__ __forceinline__ T
+  operator()(const T& a, const T& b) const {
+    return a * b;
+  }
+};
+
+template <typename T>
+struct BlockPrefixProdCallbackOp {
+  // Running prefix
+  T running_total;
+
+  // Constructor
+  __device__ BlockPrefixProdCallbackOp(T running_total)
+      : running_total(running_total) {}
+
+  // Callback operator to be entered by the first warp of threads in the block.
+  // Thread-0 is responsible for returning a value for seeding the block-wide
+  // scan.
+  __device__ T operator()(const T block_aggregate) {
+    T old_prefix = running_total;
+    running_total *= block_aggregate;
+    return old_prefix;
+  }
+};
+
+template <typename T>
+struct BlockPrefixSumCallbackOp {
+  // Running prefix
+  T running_total;
+
+  // Constructor
+  __device__ BlockPrefixSumCallbackOp(T running_total)
+      : running_total(running_total) {}
+
+  // Callback operator to be entered by the first warp of threads in the block.
+  // Thread-0 is responsible for returning a value for seeding the block-wide
+  // scan.
+  __device__ T operator()(const T block_aggregate) {
+    T old_prefix = running_total;
+    running_total += block_aggregate;
+    return old_prefix;
+  }
+};
+
+template <typename T>
+__global__ void oneMinusPKernel(
+    const T* __restrict__ p_choose,
+    T* __restrict__ cumprod_1mp,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len) {
+  for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
+    for (uint32_t tgt = threadIdx.y; tgt < tgt_len; tgt += blockDim.y) {
+      for (uint32_t src = threadIdx.x; src < src_len; src += blockDim.x) {
+        uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
+        cumprod_1mp[idx] = 1 - p_choose[idx];
+      }
+    }
+  }
+}
+
+template <typename T, int TPB>
+__global__ void innermostScanKernel(
+    T* __restrict__ cumprod_1mp,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len) {
+  for (uint32_t b = blockIdx.y; b < bsz; b += gridDim.y) {
+    for (uint32_t tgt = blockIdx.x; tgt < tgt_len; tgt += gridDim.x) {
+      // Specialize BlockScan for a 1D block of TPB threads on type T
+      typedef cub::BlockScan<T, TPB> BlockScan;
+      // Allocate shared memory for BlockScan
+      __shared__ typename BlockScan::TempStorage temp_storage;
+      // Initialize running total
+      BlockPrefixProdCallbackOp<T> prefix_op(1);
+
+      const uint32_t tid = threadIdx.x;
+      for (uint32_t block_src = 0; block_src < src_len;
+           block_src += blockDim.x) {
+        uint32_t src = block_src + tid;
+        uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
+        T thread_data = (src < src_len) ? cumprod_1mp[idx] : (T)0;
+
+        // Collectively compute the block-wide inclusive prefix sum
+        BlockScan(temp_storage)
+            .ExclusiveScan(thread_data, thread_data, Prod<T>(), prefix_op);
+        __syncthreads();
+
+        // write the scanned value to output
+        if (src < src_len) {
+          cumprod_1mp[idx] = thread_data;
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+__global__ void clampKernel(
+    const T* __restrict__ cumprod_1mp,
+    T* __restrict__ cumprod_1mp_clamp,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len,
+    T min_val,
+    T max_val) {
+  for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
+    for (uint32_t tgt = threadIdx.y; tgt < tgt_len; tgt += blockDim.y) {
+      for (uint32_t src = threadIdx.x; src < src_len; src += blockDim.x) {
+        uint32_t idx = b * tgt_len * src_len + tgt * src_len + src;
+        if (cumprod_1mp[idx] < min_val) {
+          cumprod_1mp_clamp[idx] = min_val;
+        } else if (cumprod_1mp[idx] > max_val) {
+          cumprod_1mp_clamp[idx] = max_val;
+        } else {
+          cumprod_1mp_clamp[idx] = cumprod_1mp[idx];
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+__global__ void initAlphaCUDAKernel(
+    T* alpha,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len) {
+  // alpha[:, 0, 0] = 1.0
+  for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
+    alpha[b * tgt_len * src_len] = (T)1.0;
+  }
+}
+
+template <typename T, int TPB>
+__global__ void alignmentTrainCUDAKernel(
+    const T* __restrict__ p_choose,
+    const T* __restrict__ cumprod_1mp,
+    const T* __restrict__ cumprod_1mp_clamp,
+    T* __restrict__ alpha,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len,
+    uint32_t tgt) {
+  for (uint32_t b = blockIdx.x; b < bsz; b += gridDim.x) {
+    // Specialize BlockScan for a 1D block of TPB threads on type T
+    typedef cub::BlockScan<T, TPB> BlockScan;
+
+    // Allocate shared memory for BlockScan
+    __shared__ typename BlockScan::TempStorage temp_storage;
+    // Initialize running total
+    BlockPrefixSumCallbackOp<T> prefix_op(0);
+
+    uint32_t b_offset = b * tgt_len * src_len;
+    const uint32_t tid = threadIdx.x;
+    for (uint32_t block_src = 0; block_src < src_len; block_src += blockDim.x) {
+      uint32_t src = block_src + tid;
+      // Obtain a segment of consecutive items that are blocked across threads
+      uint32_t inout_idx, alpha_idx;
+      if (tgt == 0) {
+        // both alpha and other input index is [b][0][src]
+        alpha_idx = b_offset + src;
+      } else {
+        // alpha index is [b][tgt-1][src]
+        alpha_idx = b_offset + (tgt - 1) * src_len + src;
+      }
+      inout_idx = b_offset + tgt * src_len + src;
+      T thread_data = (T)0;
+      if (src < src_len) {
+        thread_data = alpha[alpha_idx] / cumprod_1mp_clamp[inout_idx];
+      }
+
+      // Collectively compute the block-wide inclusive prefix sum
+      BlockScan(temp_storage).InclusiveSum(thread_data, thread_data, prefix_op);
+      __syncthreads();
+
+      if (src < src_len) {
+        T out = thread_data * p_choose[inout_idx] * cumprod_1mp[inout_idx];
+        // Clamps all elements into the range [ 0, 1.0 ]
+        alpha[inout_idx] = std::min<T>(std::max<T>(out, 0), (T)1.0);
+      }
+    }
+  }
+}
+
+template <typename T>
+void exclusiveCumprod(
+    const T* p_choose,
+    T* cumprod_1mp,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len,
+    uint32_t max_grid_x,
+    uint32_t max_grid_y,
+    cudaStream_t& stream) {
+  // cumprod_1mp = 1 - p_choose
+  dim3 grid(std::min<T>(max_grid_x, bsz), 1, 1);
+  dim3 block(BLOCK_DIM_X, BLOCK_DIM_Y, 1);
+  oneMinusPKernel<T><<<grid, block, 0, stream>>>(
+      p_choose, cumprod_1mp, bsz, tgt_len, src_len);
+  gpuErrchk(cudaGetLastError());
+
+  // scan on the innermost dimension of cumprod_1mp
+  // cumprod_1mp = cumprod(cumprod_1mp)
+  dim3 grid_scan(
+      std::min<T>(max_grid_x, tgt_len), std::min<T>(max_grid_y, bsz), 1);
+  innermostScanKernel<T, SCAN_BLOCK><<<grid_scan, SCAN_BLOCK, 0, stream>>>(
+      cumprod_1mp, bsz, tgt_len, src_len);
+  gpuErrchk(cudaGetLastError());
+}
+
+template <typename T>
+void alignmentTrainCUDAImpl(
+    const T* p_choose,
+    T* alpha,
+    uint32_t bsz,
+    uint32_t tgt_len,
+    uint32_t src_len,
+    float eps) {
+  // p_choose: bsz , tgt_len, src_len
+  // cumprod_1mp: bsz , tgt_len, src_len
+  // cumprod_1mp_clamp : bsz, tgt_len, src_len
+  // alpha: bsz, tgt_len, src_len
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  uint32_t max_grid_x = at::cuda::getCurrentDeviceProperties()->maxGridSize[0];
+  uint32_t max_grid_y = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
+
+  // Implementing exclusive cumprod.
+  // cumprod_1mp = cumprod(1 - p_choose)
+  // There is cumprod in pytorch, however there is no exclusive mode.
+  // cumprod(x) = [x1, x1x2, x2x3x4, ..., prod_{i=1}^n x_i]
+  // exclusive means
+  // cumprod(x) = [1, x1, x1x2, x1x2x3, ..., prod_{i=1}^{n-1} x_i]
+  uint32_t elements = bsz * tgt_len * src_len;
+  T* cumprod_1mp;
+  gpuErrchk(cudaMalloc(&cumprod_1mp, elements * sizeof(T)));
+  exclusiveCumprod<T>(
+      p_choose,
+      cumprod_1mp,
+      bsz,
+      tgt_len,
+      src_len,
+      max_grid_x,
+      max_grid_y,
+      stream);
+
+  // clamp cumprod_1mp to the range [eps, 1.0]
+  T* cumprod_1mp_clamp;
+  gpuErrchk(cudaMalloc(&cumprod_1mp_clamp, elements * sizeof(T)));
+  dim3 grid_clamp(std::min<T>(max_grid_x, bsz), 1, 1);
+  dim3 block_clamp(BLOCK_DIM_X, BLOCK_DIM_Y, 1);
+  clampKernel<T><<<grid_clamp, block_clamp, 0, stream>>>(
+      cumprod_1mp, cumprod_1mp_clamp, bsz, tgt_len, src_len, (T)eps, (T)1.0);
+  gpuErrchk(cudaGetLastError());
+
+  // ai = p_i * cumprod(1 − pi) * cumsum(a_i / cumprod(1 − pi))
+  dim3 grid_init(std::min<int>(max_grid_x, bsz), 1, 1);
+  initAlphaCUDAKernel<T>
+      <<<grid_init, 1, 0, stream>>>(alpha, bsz, tgt_len, src_len);
+  gpuErrchk(cudaGetLastError());
+
+  const int grid = std::min(bsz, max_grid_x);
+
+  for (uint32_t i = 0; i < tgt_len; i++) {
+    alignmentTrainCUDAKernel<T, SCAN_BLOCK><<<grid, SCAN_BLOCK, 0, stream>>>(
+        p_choose,
+        cumprod_1mp,
+        cumprod_1mp_clamp,
+        alpha,
+        bsz,
+        tgt_len,
+        src_len,
+        i);
+    gpuErrchk(cudaGetLastError());
+  }
+
+  gpuErrchk(cudaFree(cumprod_1mp));
+  gpuErrchk(cudaFree(cumprod_1mp_clamp));
+}
+
+} // namespace
+
+void alignmentTrainCUDAWrapper(
+    const torch::Tensor& p_choose,
+    torch::Tensor& alpha,
+    float eps) {
+  // p_choose dimension: bsz, tgt_len, src_len
+  uint32_t bsz = p_choose.size(0);
+  uint32_t tgt_len = p_choose.size(1);
+  uint32_t src_len = p_choose.size(2);
+
+  cudaSetDevice(p_choose.get_device());
+
+  AT_DISPATCH_FLOATING_TYPES_AND2(
+      torch::ScalarType::Half,
+      torch::ScalarType::BFloat16,
+      p_choose.scalar_type(),
+      "alignmentTrainCUDAImpl",
+      [&]() {
+        alignmentTrainCUDAImpl<scalar_t>(
+            p_choose.data_ptr<scalar_t>(),
+            alpha.data_ptr<scalar_t>(),
+            bsz,
+            tgt_len,
+            src_len,
+            eps);
+      });
+}