Skip to content

[Contrib] Workspace for cuBLAS backend #16413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
masahi merged 1 commit into from
Jan 25, 2024
Merged

[Contrib] Workspace for cuBLAS backend #16413

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions src/runtime/contrib/cublas/cublas.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,9 +135,10 @@ int roundoff(int v, int d) { return (v + d - 1) / d * d; }

#if CUDART_VERSION >= 10010

void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream, const DLTensor* A, const DLTensor* B,
void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream,
cublasLtMatmulPreference_t matmul_pref_desc, const DLTensor* A, const DLTensor* B,
const DLTensor* bias, const DLTensor* C, bool transa, bool transb,
cublasLtEpilogue_t epilogue) {
void* workspace_ptr, size_t workspace_size, cublasLtEpilogue_t epilogue) {
ICHECK(TypeEqual(A->dtype, B->dtype));
// Reversed strides indicates an in-place transpose operation.
transa = IsInPlaceTransposed(A) ? !transa : transa;
Expand Down Expand Up @@ -265,8 +266,21 @@ void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream, const DLTensor* A,
auto B_data = static_cast<char*>(B->data) + B->byte_offset;
auto C_data = static_cast<char*>(C->data) + C->byte_offset;

cublasLtMatmulPreferenceSetAttribute(matmul_pref_desc, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&workspace_size, sizeof(size_t));

cublasLtMatmulHeuristicResult_t heuristic_result = {};
int returned_result = 0;
CHECK_CUBLAS_ERROR(cublasLtMatmulAlgoGetHeuristic(hdl, op_desc, A_desc, B_desc, C_desc, C_desc,
matmul_pref_desc, 1, &heuristic_result,
&returned_result));
if (returned_result == 0) {
CHECK_CUBLAS_ERROR(CUBLAS_STATUS_NOT_SUPPORTED);
}

CHECK_CUBLAS_ERROR(cublasLtMatmul(hdl, op_desc, alpha, B_data, A_desc, A_data, B_desc, beta,
C_data, C_desc, C_data, C_desc, nullptr, nullptr, 0, stream));
C_data, C_desc, C_data, C_desc, &heuristic_result.algo,
workspace_ptr, workspace_size, stream));

cublasLtMatmulDescDestroy(op_desc);
cublasLtMatrixLayoutDestroy(A_desc);
Expand Down
5 changes: 3 additions & 2 deletions src/runtime/contrib/cublas/cublas_json_runtime.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,8 +129,9 @@ class CublasJSONRuntime : public JSONRuntimeBase {

auto [a_ptr, b_ptr, bias_ptr] = get_inputs(node, epilogue != CUBLASLT_EPILOGUE_DEFAULT);

tvm::contrib::CallCublasLt(entry_ptr->handle, stream, a_ptr, b_ptr, bias_ptr, out_ptr,
transa, transb, epilogue);
tvm::contrib::CallCublasLt(entry_ptr->handle, stream, entry_ptr->matmul_pref_desc, a_ptr,
b_ptr, bias_ptr, out_ptr, transa, transb,
entry_ptr->workspace_ptr, entry_ptr->workspace_size, epilogue);
}
}
}
Expand Down
14 changes: 13 additions & 1 deletion src/runtime/contrib/cublas/cublas_utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,25 @@ CuBlasThreadEntry* CuBlasThreadEntry::ThreadLocal() {
return retval;
}

CuBlasLtThreadEntry::CuBlasLtThreadEntry() { CHECK_CUBLAS_ERROR(cublasLtCreate(&handle)); }
CuBlasLtThreadEntry::CuBlasLtThreadEntry() {
CHECK_CUBLAS_ERROR(cublasLtCreate(&handle));
CHECK_CUBLAS_ERROR(cublasLtMatmulPreferenceCreate(&matmul_pref_desc));
CUDA_CALL(cudaMalloc(&workspace_ptr, workspace_size));
}

CuBlasLtThreadEntry::~CuBlasLtThreadEntry() {
if (handle) {
cublasLtDestroy(handle);
handle = nullptr;
}
if (matmul_pref_desc) {
cublasLtMatmulPreferenceDestroy(matmul_pref_desc);
matmul_pref_desc = nullptr;
}
if (workspace_ptr != nullptr) {
cudaFree(workspace_ptr);
workspace_ptr = nullptr;
}
}

typedef dmlc::ThreadLocalStore<CuBlasLtThreadEntry> CuBlasLtThreadStore;
Expand Down
11 changes: 10 additions & 1 deletion src/runtime/contrib/cublas/cublas_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,14 @@ struct CuBlasThreadEntry {
struct CuBlasLtThreadEntry {
CuBlasLtThreadEntry();
~CuBlasLtThreadEntry();

cublasLtHandle_t handle{nullptr};
cublasLtMatmulPreference_t matmul_pref_desc{nullptr};
void* workspace_ptr{nullptr};
// 32MB workspace as suggested by NVIDIA
// https://docs.nvidia.com/cuda/cublas/index.html#cublassetworkspace.
static constexpr const size_t workspace_size = 33554432;

@masahi masahi Jan 25, 2024

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm assuming that 32MB is also good for pre-Hopper since it is bigger than the recommended size, 4MB. @vinx13


static CuBlasLtThreadEntry* ThreadLocal();
}; // CuBlasLtThreadEntry

Expand Down Expand Up @@ -113,8 +120,10 @@ inline cudaDataType_t GetCudaDataType(DLDataType type) {
}

/*! \brief Execute matrix multiply followed by the specified epilogue, using cuBLASLt. */
void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream, const DLTensor* A, const DLTensor* B,
void CallCublasLt(cublasLtHandle_t hdl, cudaStream_t stream,
cublasLtMatmulPreference_t matmul_pref_desc, const DLTensor* A, const DLTensor* B,
const DLTensor* bias, const DLTensor* C, bool transa, bool transb,
void* workspace_ptr, size_t workspace_size,
cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT);

} // namespace contrib
Expand Down