如何在CUDA中为Transformer编写一个PyTorch自定义层

with torch.autograd.profiler.profile(use_cuda=True) as prof:
    # Execute ops here
print(prof)

template <typename scalar_t>
__global__ void __launch_bounds__(32) masked_softmax_cuda_forward_kernel(
    const scalar_t* __restrict__ input,
    const scalar_t* __restrict__ mask,
    scalar_t* __restrict__ output,
    unsigned int hidden_size,
    unsigned int m0, // size of mask dimension 0
    unsigned int m1, // size of mask dimension 1
    unsigned int m2, // size of mask dimension 2
    scalar_t scale) {

  // This threadIdx.x is a number between 0 and 31 because we only launched 32 threads.
  const int tid = threadIdx.x;
  // blockIdx.x, y, z are offsets of 0th, 1st, 2nd dimensions of input tensor.
  const unsigned int ibase = blockIdx.x * gridDim.y * gridDim.z * hidden_size +
                             blockIdx.y * gridDim.z * hidden_size +
                             blockIdx.z * hidden_size;
  const unsigned int mbase = blockIdx.x * (m0 > 1 ? m1 * m2 * hidden_size : 0) +
                             blockIdx.y * (m1 > 1 ? m2 * hidden_size : 0) +
                             blockIdx.z * (m2 > 1 ? hidden_size : 0);
  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < hidden_size);

  // Find a maximum input.
  scalar_t max_x = -FLT_MAX;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    max_x = fmaxf(max_x, m == 0 ? input[ibase + i] * scale : -FLT_MAX);
  }
  // Reduce values in threads to find a global maximum number.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  // Find a sum of exponential inputs.
  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    exp_sum += m == 0 ? std::exp(input[ibase + i] * scale - max_x) : 0;
  }
  // Reduce values in threads to find a global summation of exponential inputs.
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // Calculate outputs and save to global memory.
  for (unsigned int i = tid; i < hidden_size; i+=blockDim.x) {
    scalar_t m = mask[mbase + i];
    output[ibase + i] = m == 0 ? std::exp(input[ibase + i] * scale - max_x) / exp_sum : 0;
  }
}

for (unsigned int i = 16; i > 0; i >>= 1) {
  max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
}

mask = mask.size(2) - mask.sum(dim=2, dtype=torch.int32)

// Load a mask length.
  const unsigned int mask_offset = blockIdx.x * (m0 > 1 ? m1 : 0) +
                                   blockIdx.z * (m1 > 1 ? 1 : 0);
  unsigned int mask_size = min(static_cast<unsigned int>(mask[mask_offset]),
                               hidden_size);

  unsigned shfl_mask = __ballot_sync(0xffffffff, threadIdx.x < mask_size);

  scalar_t max_x = -FLT_MAX;
  // Iterate loop as much as the mask length.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    max_x = fmaxf(max_x, input[ibase + i] * scale);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    max_x = max(max_x, __shfl_xor_sync(shfl_mask, max_x, i));
  }

  scalar_t exp_sum = 0;
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    exp_sum += std::exp(input[ibase + i] * scale - max_x);
  }
  for (unsigned int i = 16; i > 0; i >>= 1) {
    exp_sum += __shfl_xor_sync(shfl_mask, exp_sum, i);
  }

  // We initialized "output" to zero, so remaining outputs will be zero.
  for (unsigned int i = tid; i < mask_size; i+=blockDim.x) {
    output[ibase + i] = std::exp(input[ibase + i] * scale - max_x) / exp_sum;
  }