chore: import upstream snapshot with attribution

paddle/phi/kernels/gpu/cuda_gemm_kernel.cu

@@ -0,0 +1,199 @@
+// 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/gpu/cuda_gemm_kernel.h"
+#include <glog/logging.h>
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/enforce.h"
+#include "paddle/phi/core/kernel_registry.h"
+
+namespace phi {
+
+template <typename Type, int CtaM, int CtaN, int Threads>
+__global__ void int8_gemm(int8_t const* act,
+                          int8_t const* weight,
+                          Type* output,
+                          int m,
+                          int n,
+                          int k) {
+#if defined(PADDLE_WITH_CUDA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700
+  return;
+#else
+  using VecType = int4;
+  static constexpr int kStepK = 128 / (8 * sizeof(int8_t));
+  static constexpr int CtaK = kStepK * Threads;
+  int tile_id_m = blockIdx.x * CtaM;
+  int tile_id_n = blockIdx.y * CtaN;
+  int tid = threadIdx.x;
+  int8_t tile_a[kStepK], tile_w[CtaN * kStepK];
+  int acc[CtaM * CtaN];
+#pragma unroll
+  for (int i = 0; i < CtaM * CtaN; ++i) {
+    acc[i] = 0;
+  }
+  act += tile_id_m * k;
+  weight += tile_id_n * k;
+  output += tile_id_m * n + tile_id_n;
+  for (int idx_k = tid * kStepK; idx_k < k; idx_k += CtaK) {
+#pragma unroll
+    for (int i = 0; i < CtaN; ++i) {
+      reinterpret_cast<VecType*>(tile_w)[i] =
+          reinterpret_cast<VecType const*>(weight + i * k + idx_k)[0];
+    }
+#pragma unroll
+    for (int i = 0; i < CtaM; ++i) {
+      reinterpret_cast<VecType*>(tile_a)[0] =
+          reinterpret_cast<VecType const*>(act + i * k + idx_k)[0];
+#pragma unroll
+      for (int j = 0; j < CtaN; ++j) {
+#pragma unroll
+        for (int l = 0; l < kStepK; l += 4) {
+          acc[i * CtaN + j] =
+              __dp4a(reinterpret_cast<int*>(tile_a + l)[0],
+                     reinterpret_cast<int*>(tile_w + j * kStepK + l)[0],
+                     acc[i * CtaN + j]);
+        }
+      }
+    }
+  }
+
+  static constexpr int kWarpSize = 32;
+  static constexpr int kWarpNum = Threads / kWarpSize;
+  __shared__ int shmem[CtaM * CtaN * kWarpNum];
+  int warp_id = tid / kWarpSize, lane_id = tid % kWarpSize;
+#pragma unroll
+  for (int i = 0; i < CtaM; ++i) {
+#pragma unroll
+    for (int j = 0; j < CtaN; ++j) {
+      int val = acc[i * CtaN + j];
+      val += __shfl_xor_sync(~0, val, 16);
+      val += __shfl_xor_sync(~0, val, 8);
+      val += __shfl_xor_sync(~0, val, 4);
+      val += __shfl_xor_sync(~0, val, 2);
+      val += __shfl_xor_sync(~0, val, 1);
+      if (lane_id == 0) {
+        shmem[i * CtaN + j + warp_id * CtaM * CtaN] = val;
+      }
+    }
+  }
+  __syncthreads();
+#pragma unroll
+  for (int ii = tid; ii < CtaM * CtaN; ii += Threads) {
+    int mid = ii / CtaN, nid = ii % CtaN;
+    int val = 0;
+#pragma unroll
+    for (int jj = 0; jj < kWarpNum; ++jj) {
+      val += shmem[jj * CtaM * CtaN + ii];
+    }
+    output[mid * n + nid] = static_cast<Type>(static_cast<float>(val));
+  }
+#endif
+}
+
+template <typename InputType,
+          typename OutputType,
+          int32_t TILE_M,
+          int32_t TILE_N,
+          int32_t BLOCK_SIZE>
+void cudaCoreGemmKernel(GemmParams const& params) {
+  dim3 block(BLOCK_SIZE);
+  dim3 grid(params.m / TILE_M, params.n / TILE_N);
+  int8_gemm<OutputType, TILE_M, TILE_N, BLOCK_SIZE>
+      <<<grid, block, 0, params.stream>>>(
+          reinterpret_cast<InputType const*>(params.act),
+          reinterpret_cast<InputType const*>(params.weight),
+          reinterpret_cast<OutputType*>(params.output),
+          params.m,
+          params.n,
+          params.k);
+}
+
+template <typename InputType,
+          typename OutputType,
+          int TILE_M,
+          int TILE_N,
+          int BLOCK_SIZE>
+bool cudaCoreGemmTemplateCaller(GemmParams const& params) {
+  constexpr int cudaCoreGemmTemplateMaxM = 16;
+  if (params.m == TILE_M) {
+    cudaCoreGemmKernel<InputType, OutputType, TILE_M, TILE_N, BLOCK_SIZE>(
+        params);
+    return true;
+  }
+  if constexpr (TILE_M < cudaCoreGemmTemplateMaxM) {
+    return cudaCoreGemmTemplateCaller<InputType,
+                                      OutputType,
+                                      TILE_M + 1,
+                                      TILE_N,
+                                      BLOCK_SIZE>(params);
+  }
+  return false;
+}
+
+template <typename InputType, typename OutputType>
+bool cudaCoreGemmLauncher(GemmParams const& params) {
+  return cudaCoreGemmTemplateCaller<InputType, OutputType, 1, 2, 256>(params);
+}
+
+template <typename T, typename Context>
+void CudaGemm(const Context& dev_ctx,
+              const DenseTensor& input,
+              const DenseTensor& w,
+              DenseTensor* output) {
+  dev_ctx.template Alloc<int32_t>(output);
+  auto input_dims = input.dims();
+  PADDLE_ENFORCE_EQ(input_dims.size(),
+                    2UL,
+                    common::errors::InvalidArgument(
+                        "The input tensor dimensions should be 2, but got %d.",
+                        input_dims.size()));
+  auto weight_dims = w.dims();
+  PADDLE_ENFORCE_EQ(weight_dims.size(),
+                    2UL,
+                    common::errors::InvalidArgument(
+                        "The weight tensor dimensions should be 2, but got %d.",
+                        weight_dims.size()));
+
+  auto out_dims = output->dims();
+
+  const int m = input_dims[0];
+  const int n = weight_dims[0];
+
+  PADDLE_ENFORCE_EQ(
+      input_dims[1],
+      weight_dims[1],
+      common::errors::InvalidArgument(
+          "The input dims[1] %d should be equal to weight dims[1] %d.",
+          input_dims[1],
+          weight_dims[1]));
+  const int k = weight_dims[1];
+
+  GemmParams params = {
+      reinterpret_cast<const void*>(input.data<T>()),
+      reinterpret_cast<const void*>(w.data<T>()),
+      reinterpret_cast<void*>(output->data<int32_t>()),
+      m,
+      n,
+      k,
+      dev_ctx.stream(),
+  };
+
+  if (!cudaCoreGemmLauncher<int8_t, int32_t>(params)) {
+    PADDLE_THROW(common::errors::Fatal("cuda gemm kernel run error"));
+  }
+}
+}  // namespace phi
+
+PD_REGISTER_KERNEL(cuda_gemm, GPU, ALL_LAYOUT, phi::CudaGemm, int8_t) {}