saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] add hipblaslt into GemmFastGelu TunableOp (#15945)

Browse source 
add hipblaslt into GemmFastGelu TunableOp.
This commit is contained in:
PeixuanZuo 2023-05-23 11:07:09 +08:00 committed by GitHub
parent 684e900e96
commit 2fddc65c8c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 139 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
cmake/onnxruntime_kernel_explorer.cmake
View file

@ -19,6 +19,9 @@ elseif(onnxruntime_USE_ROCM)
if (onnxruntime_USE_COMPOSABLE_KERNEL)
include(composable_kernel)
endif()
if (onnxruntime_USE_HIPBLASLT)
find_package(hipblaslt REQUIRED)
endif()
set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/rocm/bert)
endif()
@ -69,6 +72,9 @@ elseif (onnxruntime_USE_ROCM)
if (onnxruntime_USE_TRITON_KERNEL)
target_compile_definitions(kernel_explorer PRIVATE USE_TRITON_KERNEL)
endif()
if (onnxruntime_USE_HIPBLASLT)
target_compile_definitions(kernel_explorer PRIVATE USE_HIPBLASLT)
endif()
endif()
add_dependencies(kernel_explorer onnxruntime_pybind11_state)

17
onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh
View file

@ -10,22 +10,25 @@
#include "contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh"
#include "contrib_ops/rocm/bert/gemm_fast_gelu_common.h"
#include "core/providers/rocm/tunable/gemm.h"
#include "core/providers/rocm/tunable/gemm_hipblaslt.h"
#include "core/providers/rocm/tunable/rocm_tunable.h"
namespace onnxruntime {
namespace contrib{
namespace contrib {
namespace rocm {
namespace blas {
namespace internal {
using namespace onnxruntime::rocm::tunable::blas::internal;
template <typename T>
Status GemmFastGeluUnfused(const GemmFastGeluParams<T>* params) {
namespace column_major = onnxruntime::rocm::tunable::blas::column_major;
ORT_RETURN_IF_ERROR(column_major::Gemm(params->tuning_ctx, params->stream, params->handle,
params->opb, params->opa,
params->n, params->m, params->k,
params->alpha, params->b, params->ldb, params->a, params->lda,
params->beta, params->c, params->ldc));
params->opb, params->opa,
params->n, params->m, params->k,
params->alpha, params->b, params->ldb, params->a, params->lda,
params->beta, params->c, params->ldc));
int64_t fast_gelu_input_length = params->m * params->n;
int64_t bias_length = (params->bias != nullptr) ? params->n : 0;
@ -65,6 +68,10 @@ class GemmFastGeluTunableOp : public TunableOp<GemmFastGeluParams<T>> {
this->RegisterOp(std::move(op));
}
#endif
#ifdef USE_HIPBLASLT
this->RegisterOp(HipBlasLtGemmFastGeluOp<T>);
#endif
}
};

15
onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h
View file

@ -7,6 +7,7 @@
#include <hipblaslt/hipblaslt.h>
#endif
#include "contrib_ops/rocm/bert/gemm_fast_gelu_common.h"
#include "core/common/common.h"
#include "core/providers/rocm/tunable/gemm_common.h"
#include "core/providers/rocm/tunable/rocm_tunable.h"
@ -17,9 +18,11 @@ namespace tunable {
namespace blas {
namespace internal {
using onnxruntime::contrib::rocm::blas::GemmFastGeluParams;
#ifdef USE_HIPBLASLT
// For large K and small M/N, K dim will be splited to multiple workgroups and buffers,
// For large K and small M/N, K dim will be split to multiple workgroups and buffers,
// which will require additional workspace. Here we set the max workspace size to 32MB.
constexpr const size_t kHipBlasLtMaxWorkSpaceSizeInBytes = 32 * 1024 * 1024;
// We only keep one heuristic result here. Note that for tuned input sizes, the first result
@ -142,7 +145,7 @@ Status HipBlasLtMatMul(const ParamsT* params, int64_t batch, ActivationType acti
HIPBLASLT_RETURN_IF_ERROR(hipblasLtMatmulPreferenceSetAttribute(
pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &max_workspace_size, sizeof(max_workspace_size)));
hipblasLtMatmulHeuristicResult_t heuristic_result[kHeuristicResultCount] = {0};
hipblasLtMatmulHeuristicResult_t heuristic_result[kHeuristicResultCount] = {};
int ret_algo_count = 0;
HIPBLASLT_RETURN_IF_ERROR(hipblasLtMatmulAlgoGetHeuristic(handle,
matmul,
@ -204,6 +207,14 @@ Status HipBlasLtStridedBatchedGemmOp(const StridedBatchedGemmParams<T>* params)
return HipBlasLtMatMul<T, StridedBatchedGemmParams<T>>(params, params->batch);
};
template <typename T>
Status HipBlasLtGemmFastGeluOp(const GemmFastGeluParams<T>* params) {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF((std::is_same_v<T, double>), "hipBLASLt does not support double inputs");
bool enable_bias = nullptr != params->bias;
return HipBlasLtMatMul<T, GemmFastGeluParams<T>>(params, /*batch=*/1, ActivationType::GELU,
enable_bias, params->bias);
};
#endif // USE_HIPBLASLT
} // namespace internal

8
onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc
View file

@ -39,6 +39,14 @@ KE_REGISTER(m) {
return true;
#else
return false;
#endif
});
m.def("is_hipblaslt_available", []() {
#ifdef USE_HIPBLASLT
return true;
#else
return false;
#endif
});
}

9
onnxruntime/python/tools/kernel_explorer/kernels/gemm_fast_gelu_test.py
View file

@ -94,6 +94,14 @@ def test_gemmfastgelu_ck_bert_cases(dtype, size, transab):
_test_gemmfastgelu(getattr(ke, wrapper_name), dtype, *size, *transab)
@pytest.mark.skipif(not ke.is_hipblaslt_available(), reason="hipblaslt is not available")
@pytest.mark.parametrize("dtype", dtypes)
@pytest.mark.parametrize("size", get_gemm_basic_sizes(full=False) + get_gemm_bert_sizes(full=False))
@pytest.mark.parametrize("transab", all_transabs)
def test_gemmfastgelu_hipblaslt_bert_cases(dtype, size, transab):
_test_gemmfastgelu(getattr(ke, "GemmFastGeluHipBlasLt_" + dtype_to_suffix(dtype)), dtype, *size, *transab)
@dataclass
class GemmFastGeluMetric(ke.ComputeMetric):
transa: bool
@ -156,6 +164,7 @@ def profile_with_args(transa, transb, dtype, m, n, k, sort):
profile_gemmfastgelu_func(
getattr(ke, "GemmFastGeluTunable" + dtype_suffix + transab_suffix), dtype, m, n, k, transa, transb
)
profile_gemmfastgelu_func(getattr(ke, "GemmFastGeluHipBlasLt" + dtype_suffix), dtype, m, n, k, transa, transb)
def profile():

91
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_hipblaslt.cu Normal file
View file

@ -0,0 +1,91 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <pybind11/stl.h>
#include <string>
#include <vector>
#ifdef USE_HIPBLASLT
#include "core/providers/rocm/tunable/gemm_hipblaslt.h"
#endif
#include "contrib_ops/rocm/bert/gemm_fast_gelu_common.h"
#include "core/providers/rocm/rocm_common.h"
#include "python/tools/kernel_explorer/device_array.h"
#include "python/tools/kernel_explorer/kernel_explorer_interface.h"
namespace py = pybind11;
namespace onnxruntime {
#ifdef USE_HIPBLASLT
template <typename T>
class GemmFastGeluHipBlasLt : public IKernelExplorer {
public:
GemmFastGeluHipBlasLt(BlasOp opa, BlasOp opb,
int64_t m, int64_t n, int64_t k,
double alpha,
DeviceArray& a, int64_t lda,
DeviceArray& b, int64_t ldb,
DeviceArray& bias,
double beta,
DeviceArray& c, int64_t ldc) : params_{} {
params_.tuning_ctx = TuningContext();
params_.stream = Stream();
params_.opa = opa;
params_.opb = opb;
params_.m = m;
params_.n = n;
params_.k = k;
params_.alpha = alpha;
params_.a = static_cast<T*>(a.ptr());
params_.lda = lda;
params_.b = static_cast<T*>(b.ptr());
params_.ldb = ldb;
params_.bias = static_cast<T*>(bias.ptr());
params_.beta = beta;
params_.c = static_cast<T*>(c.ptr());
params_.ldc = ldc;
}
void Run() override {
ORT_THROW_IF_ERROR((rocm::tunable::blas::internal::HipBlasLtGemmFastGeluOp<T>(&params_)));
}
std::vector<std::string> ListOps() const {
return {"GemmFastGeluHipBlasLt"};
}
bool SelectOp(const std::string& name) {
Status status = rocm::tunable::blas::internal::HipBlasLtGemmFastGeluOp<T>(&params_);
return status.IsOK() && name == "GemmFastGeluHipBlasLt";
}
private:
using ParamsT = contrib::rocm::blas::GemmFastGeluParams<T>;
ParamsT params_{};
};
#define REGISTER_OP(type) \
py::class_<GemmFastGeluHipBlasLt<type>>(m, "GemmFastGeluHipBlasLt_" #type) \
.def(py::init<BlasOp, BlasOp, int64_t, int64_t, int64_t, \
double, \
DeviceArray&, int64_t, \
DeviceArray&, int64_t, \
DeviceArray&, \
double, \
DeviceArray&, int64_t>()) \
.def("SetRepeats", &GemmFastGeluHipBlasLt<type>::SetRepeats) \
.def("Run", &GemmFastGeluHipBlasLt<type>::Run) \
.def("Profile", &GemmFastGeluHipBlasLt<type>::Profile) \
.def("ListOps", &GemmFastGeluHipBlasLt<type>::ListOps) \
.def("SelectOp", &GemmFastGeluHipBlasLt<type>::SelectOp);
KE_REGISTER(m) {
REGISTER_OP(float)
REGISTER_OP(half)
}
#endif // USE_HIPBLASLT
} // namespace onnxruntime