chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,256 @@
|
||||
// Copyright (c) 2023 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 <algorithm>
|
||||
#include <mutex>
|
||||
#include <unordered_map>
|
||||
#include "glog/logging.h"
|
||||
#include "paddle/phi/common/data_type.h"
|
||||
#include "paddle/phi/core/kernel_registry.h"
|
||||
#include "paddle/phi/core/tensor_utils.h"
|
||||
|
||||
#include "paddle/common/enforce.h"
|
||||
#include "paddle/phi/kernels/addmm_kernel.h"
|
||||
#include "paddle/phi/kernels/elementwise_add_kernel.h"
|
||||
#include "paddle/phi/kernels/impl/matmul_kernel_impl.h"
|
||||
#include "paddle/phi/kernels/linear_v2_kernel.h"
|
||||
#include "paddle/phi/kernels/reshape_kernel.h"
|
||||
#include "paddle/phi/kernels/tile_kernel.h"
|
||||
|
||||
#ifdef PADDLE_WITH_HIP
|
||||
#include <hip/hip_runtime.h>
|
||||
#include <hip/hip_runtime_api.h>
|
||||
#else
|
||||
#include <cuda_runtime_api.h> // NOLINT
|
||||
#include "cuda.h" // NOLINT
|
||||
#endif
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
|
||||
|
||||
#include "paddle/common/flags.h"
|
||||
#include "paddle/phi/backends/all_context.h"
|
||||
#include "paddle/phi/common/amp_type_traits.h"
|
||||
#include "paddle/phi/common/memory_utils.h"
|
||||
#include "paddle/phi/core/dense_tensor.h"
|
||||
#include "paddle/phi/core/enforce.h"
|
||||
#include "paddle/utils/optional.h"
|
||||
#if defined(PADDLE_WITH_CUDA)
|
||||
#include "paddle/phi/backends/dynload/cublasLt.h"
|
||||
#include "paddle/phi/backends/gpu/cuda/cuda_helper.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.cu.h"
|
||||
#elif defined(PADDLE_WITH_HIP)
|
||||
#include "paddle/phi/backends/dynload/hipblasLt.h"
|
||||
#include "paddle/phi/backends/gpu/rocm/rocm_helper.h"
|
||||
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.hip.h"
|
||||
#endif
|
||||
|
||||
#endif
|
||||
COMMON_DECLARE_bool(use_legacy_linear);
|
||||
|
||||
namespace phi {
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && \
|
||||
!defined(_WIN32) && CUDA_VERSION >= 11060
|
||||
|
||||
// Direct cublasLt matmul+bias, bypassing MatmulPlanner/DescriptorSetter/
|
||||
// CublasLtBase. Uses persistent workspace from GPUContext.
|
||||
template <typename T>
|
||||
static void CublasLtMatmulBias(const GPUContext& ctx,
|
||||
const T* x,
|
||||
const T* w,
|
||||
const T* bias,
|
||||
T* out,
|
||||
int64_t M,
|
||||
int64_t N,
|
||||
int64_t K,
|
||||
bool trans_w) {
|
||||
using MT = typename MPTypeTrait<T>::Type;
|
||||
constexpr auto compute =
|
||||
std::is_same<T, double>::value ? CUBLAS_COMPUTE_64F : CUBLAS_COMPUTE_32F;
|
||||
const auto dtype = backends::gpu::ToCudaDataType<T>();
|
||||
const auto stype = backends::gpu::ToCudaDataType<MT>();
|
||||
|
||||
MT alpha = static_cast<MT>(1), beta = static_cast<MT>(0);
|
||||
auto lt = ctx.cublaslt_handle();
|
||||
|
||||
// op desc
|
||||
cublasLtMatmulDesc_t op = nullptr;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cublasLtMatmulDescCreate(&op, compute, stype));
|
||||
|
||||
// col-major: C(N,M) = A(weight) * B(input)
|
||||
cublasOperation_t ta = trans_w ? CUBLAS_OP_T : CUBLAS_OP_N;
|
||||
cublasOperation_t tb = CUBLAS_OP_N;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
|
||||
op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
|
||||
op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));
|
||||
|
||||
cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_BIAS;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
|
||||
op, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
|
||||
op, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));
|
||||
|
||||
// matrix layouts (col-major)
|
||||
cublasLtMatrixLayout_t la = nullptr, lb = nullptr, lc = nullptr;
|
||||
if (trans_w) {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cublasLtMatrixLayoutCreate(&la, dtype, K, N, K));
|
||||
} else {
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cublasLtMatrixLayoutCreate(&la, dtype, N, K, N));
|
||||
}
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cublasLtMatrixLayoutCreate(&lb, dtype, K, M, K));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(
|
||||
dynload::cublasLtMatrixLayoutCreate(&lc, dtype, N, M, N));
|
||||
|
||||
// persistent workspace from context
|
||||
constexpr size_t kWsSize = 1024 * 1024;
|
||||
auto [ws, ws_sz] = ctx.cublaslt_workspace(kWsSize);
|
||||
|
||||
// heuristic
|
||||
cublasLtMatmulPreference_t pref = nullptr;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulPreferenceCreate(&pref));
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulPreferenceSetAttribute(
|
||||
pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_sz, sizeof(ws_sz)));
|
||||
|
||||
cublasLtMatmulHeuristicResult_t heur = {};
|
||||
int n_res = 0;
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulAlgoGetHeuristic(
|
||||
lt, op, la, lb, lc, lc, pref, 1, &heur, &n_res));
|
||||
PADDLE_ENFORCE_GT(
|
||||
n_res,
|
||||
0,
|
||||
common::errors::Unavailable("No cublasLt GEMM algorithm available."));
|
||||
dynload::cublasLtMatmulPreferenceDestroy(pref);
|
||||
|
||||
// execute
|
||||
PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmul(lt,
|
||||
op,
|
||||
&alpha,
|
||||
w,
|
||||
la,
|
||||
x,
|
||||
lb,
|
||||
&beta,
|
||||
out,
|
||||
lc,
|
||||
out,
|
||||
lc,
|
||||
&heur.algo,
|
||||
ws,
|
||||
ws_sz,
|
||||
ctx.stream()));
|
||||
|
||||
// cleanup
|
||||
dynload::cublasLtMatmulDescDestroy(op);
|
||||
dynload::cublasLtMatrixLayoutDestroy(la);
|
||||
dynload::cublasLtMatrixLayoutDestroy(lb);
|
||||
dynload::cublasLtMatrixLayoutDestroy(lc);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
||||
template <typename T, typename Context>
|
||||
void LinearV2Kernel(const Context& dev_ctx,
|
||||
const DenseTensor& input,
|
||||
const DenseTensor& weight,
|
||||
const DenseTensor& bias,
|
||||
const bool transpose_weight,
|
||||
DenseTensor* out) {
|
||||
dev_ctx.template Alloc<T>(out);
|
||||
if (out->numel() == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && \
|
||||
!defined(_WIN32) && CUDA_VERSION >= 11060
|
||||
if (!FLAGS_use_legacy_linear) {
|
||||
const auto out_dim_original = out->dims();
|
||||
const auto [M, N, K] = canonicalize_dims(input, weight, transpose_weight);
|
||||
|
||||
DenseTensor input_processed = input;
|
||||
DenseTensor weight_processed = weight;
|
||||
input_processed.Resize({M, K});
|
||||
if (transpose_weight) {
|
||||
weight_processed.Resize({N, K});
|
||||
} else {
|
||||
weight_processed.Resize({K, N});
|
||||
}
|
||||
out->Resize({M, N});
|
||||
|
||||
if (N > 1 && K > 1) {
|
||||
DenseTensor bias_processed;
|
||||
if (bias.numel() != N) {
|
||||
TileKernel<T, Context>(dev_ctx, bias, {N}, &bias_processed);
|
||||
} else {
|
||||
bias_processed = bias;
|
||||
}
|
||||
CublasLtMatmulBias<T>(dev_ctx,
|
||||
input_processed.data<T>(),
|
||||
weight_processed.data<T>(),
|
||||
bias_processed.data<T>(),
|
||||
out->data<T>(),
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
transpose_weight);
|
||||
} else {
|
||||
// When N=1 or K=1, {N,K} and {K,N} have identical memory layout,
|
||||
// so just reshape to {K,N} which is what AddmmKernel expects.
|
||||
weight_processed.Resize({K, N});
|
||||
DenseTensor bias_processed = bias;
|
||||
if (bias.numel() != (M * N)) {
|
||||
bias_processed.Resize({1, bias.numel()});
|
||||
TileKernel<T, Context>(
|
||||
dev_ctx, bias_processed, {M, 1}, &bias_processed);
|
||||
}
|
||||
AddmmKernel<T>(dev_ctx,
|
||||
bias_processed,
|
||||
input_processed,
|
||||
weight_processed,
|
||||
1.0f,
|
||||
1.0f,
|
||||
out);
|
||||
}
|
||||
out->Resize(out_dim_original);
|
||||
} else // NOLINT
|
||||
#endif
|
||||
{ // NOLINT
|
||||
std::vector<std::int64_t> input_dims_vec = vectorize(input.dims());
|
||||
std::vector<std::int64_t> weight_dims_vec = vectorize(weight.dims());
|
||||
|
||||
MatMulFunction<Context, T>(dev_ctx,
|
||||
input,
|
||||
weight,
|
||||
input_dims_vec,
|
||||
weight_dims_vec,
|
||||
out,
|
||||
false,
|
||||
transpose_weight);
|
||||
AddKernel<T, Context>(dev_ctx, *out, bias, out);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace phi
|
||||
|
||||
PD_REGISTER_KERNEL(linear_v2,
|
||||
GPU,
|
||||
ALL_LAYOUT,
|
||||
phi::LinearV2Kernel,
|
||||
float,
|
||||
double,
|
||||
phi::float16,
|
||||
phi::bfloat16) {}
|
||||
Reference in New Issue
Block a user