diff --git a/cmake/onnxruntime_kernel_explorer.cmake b/cmake/onnxruntime_kernel_explorer.cmake index e05321748e..09b153e998 100644 --- a/cmake/onnxruntime_kernel_explorer.cmake +++ b/cmake/onnxruntime_kernel_explorer.cmake @@ -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) diff --git a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh index 24058f80db..dcd2d45df5 100644 --- a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh @@ -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 Status GemmFastGeluUnfused(const GemmFastGeluParams* 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> { this->RegisterOp(std::move(op)); } #endif + +#ifdef USE_HIPBLASLT + this->RegisterOp(HipBlasLtGemmFastGeluOp); +#endif } }; diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h b/onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h index aa5630ebea..c3bd17ddf5 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h +++ b/onnxruntime/core/providers/rocm/tunable/gemm_hipblaslt.h @@ -7,6 +7,7 @@ #include #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* params) return HipBlasLtMatMul>(params, params->batch); }; +template +Status HipBlasLtGemmFastGeluOp(const GemmFastGeluParams* params) { + TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF((std::is_same_v), "hipBLASLt does not support double inputs"); + bool enable_bias = nullptr != params->bias; + return HipBlasLtMatMul>(params, /*batch=*/1, ActivationType::GELU, + enable_bias, params->bias); +}; + #endif // USE_HIPBLASLT } // namespace internal diff --git a/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc b/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc index fbcb55c781..553b237b5a 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc @@ -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 }); } diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_fast_gelu_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_fast_gelu_test.py index 697da9b9e8..b2154952f5 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_fast_gelu_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_fast_gelu_test.py @@ -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(): diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_hipblaslt.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_hipblaslt.cu new file mode 100644 index 0000000000..e927d78a8b --- /dev/null +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_hipblaslt.cu @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include + +#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 +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(a.ptr()); + params_.lda = lda; + params_.b = static_cast(b.ptr()); + params_.ldb = ldb; + params_.bias = static_cast(bias.ptr()); + params_.beta = beta; + params_.c = static_cast(c.ptr()); + params_.ldc = ldc; + } + + void Run() override { + ORT_THROW_IF_ERROR((rocm::tunable::blas::internal::HipBlasLtGemmFastGeluOp(¶ms_))); + } + + std::vector ListOps() const { + return {"GemmFastGeluHipBlasLt"}; + } + + bool SelectOp(const std::string& name) { + Status status = rocm::tunable::blas::internal::HipBlasLtGemmFastGeluOp(¶ms_); + return status.IsOK() && name == "GemmFastGeluHipBlasLt"; + } + + private: + using ParamsT = contrib::rocm::blas::GemmFastGeluParams; + ParamsT params_{}; +}; + +#define REGISTER_OP(type) \ + py::class_>(m, "GemmFastGeluHipBlasLt_" #type) \ + .def(py::init()) \ + .def("SetRepeats", &GemmFastGeluHipBlasLt::SetRepeats) \ + .def("Run", &GemmFastGeluHipBlasLt::Run) \ + .def("Profile", &GemmFastGeluHipBlasLt::Profile) \ + .def("ListOps", &GemmFastGeluHipBlasLt::ListOps) \ + .def("SelectOp", &GemmFastGeluHipBlasLt::SelectOp); + +KE_REGISTER(m) { + REGISTER_OP(float) + REGISTER_OP(half) +} +#endif // USE_HIPBLASLT + +} // namespace onnxruntime