wehub-resource-sync
257 lines
9.4 KiB
Plaintext
257 lines
9.4 KiB
Plaintext
// 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) {}
|