saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-17 18:40:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make CK an optional dependencies and only built with ck if ROCm >= 5.3 (#14232)

Browse source 
Recently, ck dropped ROCm 5.2 support, which is causing packaging
pipeline failures. This PR workaround it.
This commit is contained in:
cloudhan 2023-01-12 17:09:40 +08:00 committed by GitHub
parent b9ecd428c1
commit 712f781702
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 74 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
cmake/CMakeLists.txt
View file

@ -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)

9
cmake/onnxruntime_kernel_explorer.cmake
View file

@ -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)

25
cmake/onnxruntime_providers.cmake
View file

@ -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")

11
onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh
View file

@ -7,12 +7,14 @@
#include <utility>
#include <vector>
#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 <typename T>
struct DataTypeAdaptor {
using type = T;
@ -84,7 +88,6 @@ auto GetCKGemmAddFastGeluTypeStringAndOps() {
return ret;
}
template <typename T, typename ALayout, typename BLayout>
auto GetCKGemmFastGeluTypeStringAndOps() {
using CKDataType = typename DataTypeAdaptor<T>::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

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

@ -55,6 +55,7 @@ class GemmFastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOp<GemmF
public:
GemmFastGeluTunableOp() {
this->RegisterOp(GemmFastGeluUnfused<T>);
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKGemmAddFastGeluTypeStringAndOps<T, ALayout, BLayout>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
@ -63,6 +64,7 @@ class GemmFastGeluTunableOp : public onnxruntime::rocm::tunable::TunableOp<GemmF
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
#endif
}
};

8
onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh
View file

@ -7,6 +7,7 @@
#include <utility>
#include <vector>
#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 <typename T>
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

4
onnxruntime/core/providers/rocm/tunable/gemm_tunable.cuh
View file

@ -42,10 +42,12 @@ class GemmTunableOp : public tunable::TunableOp<GemmParams<T>> {
this->RegisterNestedTunableOp(&rocblas_gemm_tunable_op_);
#endif /* #ifdef USE_ROCBLAS_EXTENSION_API */
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKGemmTypeStringAndOps<T, ALayout, BLayout>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
#endif
}
const GemmParams<T>* PreTuning(const GemmParams<T>* params) override {
@ -127,10 +129,12 @@ class StridedBatchedGemmTunableOp : public tunable::TunableOp<StridedBatchedGemm
public:
StridedBatchedGemmTunableOp() {
this->RegisterOp(RocBlasStridedBatchedGemmOp<T>);
#ifdef USE_COMPOSABLE_KERNEL
for (auto&& [_, op] : GetCKStridedBatchedGemmTypeStringAndOps<T, ALayout, BLayout>()) {
ORT_UNUSED_PARAMETER(_);
this->RegisterOp(std::move(op));
}
#endif
}
const StridedBatchedGemmParams<T>* PreTuning(const StridedBatchedGemmParams<T>* params) override {

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

@ -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

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

@ -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)

2
onnxruntime/python/tools/kernel_explorer/kernels/gemm_test.py
View file

@ -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)

2
onnxruntime/python/tools/kernel_explorer/kernels/kernel_explorer.pyi
View file

@ -5,3 +5,5 @@ class DeviceArray:
class blas_op:
T: int
N: int
def is_composable_kernel_available(*args, **kwargs): ...

4
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_ck.cu
View file

@ -24,6 +24,7 @@ namespace py = pybind11;
namespace onnxruntime {
#ifdef USE_COMPOSABLE_KERNEL
template <typename T, typename ALayout, typename BLayout>
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

4
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu_ck.cu
View file

@ -24,6 +24,7 @@ namespace py = pybind11;
namespace onnxruntime {
#ifdef USE_COMPOSABLE_KERNEL
template <typename T, typename ALayout, typename BLayout>
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

2
onnxruntime/python/tools/kernel_explorer/kernels/strided_batched_gemm_test.py
View file

@ -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)