diff --git a/cmake/CMakeLists.txt b/cmake/CMakeLists.txt index 89d827b5d8..9d3c94f91d 100644 --- a/cmake/CMakeLists.txt +++ b/cmake/CMakeLists.txt @@ -193,6 +193,8 @@ option(onnxruntime_ENABLE_CPUINFO "Enable cpuinfo" ON) # ATen fallback support option(onnxruntime_ENABLE_ATEN "Enable ATen fallback" OFF) +# composable kernel is managed automatically, unless user want to explicitly disable it, it should not be manually set +option(onnxruntime_USE_COMPOSABLE_KERNEL "Enable composable kernel for ROCm EP" ON) option(onnxruntime_BUILD_KERNEL_EXPLORER "Build Kernel Explorer for testing and profiling GPU kernels" OFF) option(onnxruntime_BUILD_CACHE "onnxruntime build with cache" OFF) @@ -305,6 +307,11 @@ if (onnxruntime_USE_ROCM) string(TOUPPER ${CMAKE_BUILD_TYPE} BUILD_TYPE) message("CMAKE_HIP_FLAGS_${BUILD_TYPE}: ${CMAKE_HIP_FLAGS_${BUILD_TYPE}}") add_definitions(-DROCM_VERSION=${ROCM_VERSION_DEV_INT}) + + if (onnxruntime_USE_COMPOSABLE_KERNEL AND ROCM_VERSION_DEV VERSION_LESS "5.3") + message(WARNING "composable kernel is only supported on ROCm >= 5.3") + set(onnxruntime_USE_COMPOSABLE_KERNEL OFF) + endif() endif() if (APPLE) diff --git a/cmake/onnxruntime_kernel_explorer.cmake b/cmake/onnxruntime_kernel_explorer.cmake index e525e7005b..f30b518bf9 100644 --- a/cmake/onnxruntime_kernel_explorer.cmake +++ b/cmake/onnxruntime_kernel_explorer.cmake @@ -16,7 +16,9 @@ if (onnxruntime_USE_CUDA) elseif(onnxruntime_USE_ROCM) check_language(HIP) set(LANGUAGE HIP) - include(composable_kernel) + if (onnxruntime_USE_COMPOSABLE_KERNEL) + include(composable_kernel) + endif() set(BERT_DIR ${ONNXRUNTIME_ROOT}/contrib_ops/rocm/bert) endif() @@ -59,7 +61,10 @@ elseif (onnxruntime_USE_ROCM) auto_set_source_files_hip_language(${kernel_explorer_kernel_srcs} ${kernel_explorer_rocm_kernel_srcs}) target_sources(kernel_explorer PRIVATE ${kernel_explorer_rocm_kernel_srcs}) target_compile_definitions(kernel_explorer PRIVATE __HIP_PLATFORM_AMD__=1 __HIP_PLATFORM_HCC__=1) - target_link_libraries(kernel_explorer PRIVATE onnxruntime_composable_kernel_includes) + if (onnxruntime_USE_COMPOSABLE_KERNEL) + target_compile_definitions(kernel_explorer PRIVATE USE_COMPOSABLE_KERNEL) + target_link_libraries(kernel_explorer PRIVATE onnxruntime_composable_kernel_includes) + endif() endif() add_dependencies(kernel_explorer onnxruntime_pybind11_state) diff --git a/cmake/onnxruntime_providers.cmake b/cmake/onnxruntime_providers.cmake index 88014a1e53..e850a3703b 100644 --- a/cmake/onnxruntime_providers.cmake +++ b/cmake/onnxruntime_providers.cmake @@ -1437,17 +1437,20 @@ if (onnxruntime_USE_ROCM) #endif() endif() - include(composable_kernel) - target_link_libraries(onnxruntime_providers_rocm PRIVATE - onnxruntime_composable_kernel_includes - # Currently we shall not use composablekernels::device_operations, the target includes all conv dependencies, which - # are extremely slow to compile. Instead, we only link all gemm related objects. See the following link on updating. - # https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/85978e0201/library/src/tensor_operation_instance/gpu/CMakeLists.txt#L33-L54 - device_gemm_instance - device_gemm_add_fastgelu_instance - device_gemm_fastgelu_instance - device_batched_gemm_instance - ) + if (onnxruntime_USE_COMPOSABLE_KERNEL) + include(composable_kernel) + target_link_libraries(onnxruntime_providers_rocm PRIVATE + onnxruntime_composable_kernel_includes + # Currently we shall not use composablekernels::device_operations, the target includes all conv dependencies, which + # are extremely slow to compile. Instead, we only link all gemm related objects. See the following link on updating. + # https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/85978e0201/library/src/tensor_operation_instance/gpu/CMakeLists.txt#L33-L54 + device_gemm_instance + device_gemm_add_fastgelu_instance + device_gemm_fastgelu_instance + device_batched_gemm_instance + ) + target_compile_definitions(onnxruntime_providers_rocm PRIVATE USE_COMPOSABLE_KERNEL) + endif() if(UNIX) set_property(TARGET onnxruntime_providers_rocm APPEND_STRING PROPERTY LINK_FLAGS "-Xlinker --version-script=${ONNXRUNTIME_ROOT}/core/providers/rocm/version_script.lds -Xlinker --gc-sections") diff --git a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh index 77c138a30b..3cafa7d0df 100644 --- a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh @@ -7,12 +7,14 @@ #include #include +#ifdef USE_COMPOSABLE_KERNEL #include "ck/ck.hpp" #include "ck/library/tensor_operation_instance/gpu/gemm_add_fastgelu.hpp" #include "ck/library/tensor_operation_instance/gpu/gemm_fastgelu.hpp" #include "ck/tensor_operation/gpu/device/tensor_layout.hpp" #include "ck/tensor_operation/gpu/device/device_gemm_multiple_d.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#endif #include "contrib_ops/rocm/bert/gemm_fast_gelu_common.h" @@ -25,6 +27,8 @@ namespace rocm { namespace blas { namespace internal { +#ifdef USE_COMPOSABLE_KERNEL + template struct DataTypeAdaptor { using type = T; @@ -84,7 +88,6 @@ auto GetCKGemmAddFastGeluTypeStringAndOps() { return ret; } - template auto GetCKGemmFastGeluTypeStringAndOps() { using CKDataType = typename DataTypeAdaptor::type; @@ -117,7 +120,7 @@ auto GetCKGemmFastGeluTypeStringAndOps() { params->ldc, nop, nop, fastgelu); TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!impl->IsSupportedArgument(arg.get()), - impl->GetTypeString(), " does not support ", params->Signature()); + impl->GetTypeString(), " does not support ", params->Signature()); invoker->Run(arg.get(), StreamConfig{params->stream}); return Status::OK(); }; @@ -125,6 +128,10 @@ auto GetCKGemmFastGeluTypeStringAndOps() { } return ret; } +#else +struct Row {}; +struct Col {}; +#endif // USE_COMPOSABLE_KERNEL } // namespace internal } // namespace blas 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 de0e11b52a..230951dc45 100644 --- a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_tunable.cuh @@ -55,6 +55,7 @@ class GemmFastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOpRegisterOp(GemmFastGeluUnfused); +#ifdef USE_COMPOSABLE_KERNEL for (auto&& [_, op] : GetCKGemmAddFastGeluTypeStringAndOps()) { ORT_UNUSED_PARAMETER(_); this->RegisterOp(std::move(op)); @@ -63,6 +64,7 @@ class GemmFastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOpRegisterOp(std::move(op)); } +#endif } }; diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh index 84c0e68d55..0b5ec71450 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh +++ b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh @@ -7,6 +7,7 @@ #include #include +#ifdef USE_COMPOSABLE_KERNEL #include "ck/ck.hpp" #include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp" #include "ck/library/tensor_operation_instance/gpu/gemm.hpp" @@ -14,6 +15,7 @@ #include "ck/tensor_operation/gpu/device/device_batched_gemm.hpp" #include "ck/tensor_operation/gpu/device/device_gemm.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#endif #include "core/providers/rocm/tunable/gemm_common.h" @@ -23,6 +25,8 @@ namespace tunable { namespace blas { namespace internal { +#ifdef USE_COMPOSABLE_KERNEL + template struct DataTypeAdaptor { using type = T; @@ -127,6 +131,10 @@ auto GetCKStridedBatchedGemmTypeStringAndOps() { } return ret; } +#else +struct Row {}; +struct Col {}; +#endif // USE_COMPOSABLE_KERNEL } // namespace internal } // namespace blas diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh index 6b4716c01e..ccaf0e06d8 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh +++ b/onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh @@ -42,10 +42,12 @@ class GemmTunableOp : public tunable::TunableOp> { this->RegisterNestedTunableOp(&rocblas_gemm_tunable_op_); #endif /* #ifdef USE_ROCBLAS_EXTENSION_API */ +#ifdef USE_COMPOSABLE_KERNEL for (auto&& [_, op] : GetCKGemmTypeStringAndOps()) { ORT_UNUSED_PARAMETER(_); this->RegisterOp(std::move(op)); } +#endif } const GemmParams* PreTuning(const GemmParams* params) override { @@ -127,10 +129,12 @@ class StridedBatchedGemmTunableOp : public tunable::TunableOpRegisterOp(RocBlasStridedBatchedGemmOp); +#ifdef USE_COMPOSABLE_KERNEL for (auto&& [_, op] : GetCKStridedBatchedGemmTypeStringAndOps()) { ORT_UNUSED_PARAMETER(_); this->RegisterOp(std::move(op)); } +#endif } const StridedBatchedGemmParams* PreTuning(const StridedBatchedGemmParams* params) override { diff --git a/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc b/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc index ac638a0229..733c3ee5c5 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc +++ b/onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc @@ -25,6 +25,14 @@ PYBIND11_MODULE(_kernel_explorer, m) { InitSkipLayerNorm(m); InitGemmFastGelu(m); #endif + + m.def("is_composable_kernel_available", []() { +#ifdef USE_COMPOSABLE_KERNEL + return true; +#else + return false; +#endif + }); } } // namespace onnxruntime 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 7f6d9368c6..96348d3803 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 @@ -85,6 +85,7 @@ def test_gemmfastgelu_tunable_bert_cases(dtype, size, transab): _test_gemmfastgelu(getattr(ke, wrapper_name), dtype, *size, *transab) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @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) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py index 6c0c5e0ede..c2e0e38fc0 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py @@ -78,6 +78,7 @@ def test_rocblas_gemm_all_cases(dtype, transa, transb, m, n, k): _test_gemm(getattr(ke, "RocBlasGemm_" + dtype_to_suffix(dtype)), dtype, transa, transb, m, n, k) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("m, n, k", get_gemm_basic_sizes(full=False) + get_gemm_bert_sizes(full=False)) @pytest.mark.parametrize("transa, transb", all_transabs) @pytest.mark.parametrize("dtype", dtypes) @@ -103,6 +104,7 @@ def test_rocblas_gemm_alpha_beta(dtype, transa, transb, alpha, beta): _test_gemm(getattr(ke, wrapper_name), dtype, transa, transb, 128, 256, 768, alpha=alpha, beta=beta) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("alpha, beta", [(0.5, 0.5)]) @pytest.mark.parametrize("transa, transb", all_transabs) @pytest.mark.parametrize("dtype", dtypes) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/kernel_explorer.pyi b/onnxruntime/python/tools/kernel_explorer/kernels/kernel_explorer.pyi index f5bc7c1fe8..79da3c11b9 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/kernel_explorer.pyi +++ b/onnxruntime/python/tools/kernel_explorer/kernels/kernel_explorer.pyi @@ -5,3 +5,5 @@ class DeviceArray: class blas_op: T: int N: int + +def is_composable_kernel_available(*args, **kwargs): ... diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu index ed134554ef..e7ba0e01e9 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu @@ -24,6 +24,7 @@ namespace py = pybind11; namespace onnxruntime { +#ifdef USE_COMPOSABLE_KERNEL template class CKGemm : public IKernelExplorer { public: @@ -212,5 +213,8 @@ void InitComposableKernelGemm(py::module m) { REGISTER_CKSTRIDEDBATCHEDGEMM_FOR_ALL_TRANSAB(float); REGISTER_CKSTRIDEDBATCHEDGEMM_FOR_ALL_TRANSAB(half); } +#else +void InitComposableKernelGemm(py::module) {} +#endif // USE_COMPOSABLE_KERNEL } // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu index d9e9d5cdea..5ae4e626c1 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu @@ -24,6 +24,7 @@ namespace py = pybind11; namespace onnxruntime { +#ifdef USE_COMPOSABLE_KERNEL template class CKGemmFastGelu : public IKernelExplorer { public: @@ -122,5 +123,8 @@ void InitComposableKernelGemmFastGelu(py::module m) { REGISTER_OP_FOR_ALL_TRANSAB(float); REGISTER_OP_FOR_ALL_TRANSAB(half); } +#else +void InitComposableKernelGemmFastGelu(py::module) {} +#endif // USE_COMPOSABLE_KERNEL } // namespace onnxruntime diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/strided_batched_gemm_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/strided_batched_gemm_test.py index 84af45f4da..bcd9f8db7a 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/strided_batched_gemm_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/strided_batched_gemm_test.py @@ -98,6 +98,7 @@ def test_rocblas_gemm_all_cases(dtype, transa, transb, m, n, k, batch): _test_strided_batched_gemm(getattr(ke, wrapper_name), dtype, transa, transb, m, n, k, batch) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("batch", [1, 64]) @pytest.mark.parametrize("m, n, k", get_gemm_basic_sizes(full=False) + get_gemm_bert_sizes(full=False)) @pytest.mark.parametrize("transa, transb", all_transabs) @@ -127,6 +128,7 @@ def test_rocblas_gemm_alpha_beta(dtype, transa, transb, alpha, beta): ) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("alpha, beta", [(0.5, 0.5)]) @pytest.mark.parametrize("transa, transb", all_transabs) @pytest.mark.parametrize("dtype", dtypes)