From 98ab4a62d647a95011667a3704ff46e7bc2d2920 Mon Sep 17 00:00:00 2001 From: cloudhan Date: Fri, 17 Mar 2023 15:59:57 +0800 Subject: [PATCH] Fix ROCm 5.2.3 pipeline (#15073) Make CK optional again. --- .../impl.cuh | 6 +-- .../impl_fp16.cu | 2 + .../impl_fp16_biased.cu | 2 + .../impl_fp16_biased_biased.cu | 2 + ...ed_gemm_softmax_gemm_permute_pipelines.cuh | 41 +++++++++++-------- .../kernels/gemm_softmax_gemm_permute_test.py | 6 ++- .../kernels/rocm/gemm_softmax_gemm_permute.cu | 4 ++ 7 files changed, 41 insertions(+), 22 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh index c127a1089a..ae4aef5f99 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh @@ -4,6 +4,7 @@ #pragma once +#ifdef USE_COMPOSABLE_KERNEL #include #include "ck/ck.hpp" @@ -75,7 +76,6 @@ using device_batched_gemm_softmax_gemm_permute_instances = // clang-format on >; - struct PreSoftmaxAttentionScoreOp { PreSoftmaxAttentionScoreOp(float scale) : scale_(scale) {} @@ -97,7 +97,6 @@ struct PreSoftmaxAttentionScoreOp { float scale_; }; - // Use this function to gat implementation template std::vector, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>(); -} // namespace ck +} // namespace internal } // namespace rocm } // namespace contrib } // namespace onnxruntime +#endif // USE_COMPOSABLE_KERNEL diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu index 9499eb84c8..181e47f012 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#ifdef USE_COMPOSABLE_KERNEL #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh" #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp" @@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances< } // namespace rocm } // namespace contrib } // namespace onnxruntime +#endif // USE_COMPOSABLE_KERNEL diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu index fe1256b920..1577bdf397 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#ifdef USE_COMPOSABLE_KERNEL #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh" #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp" @@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances< } // namespace rocm } // namespace contrib } // namespace onnxruntime +#endif // USE_COMPOSABLE_KERNEL diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu index a0d99d2e84..14de592343 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#ifdef USE_COMPOSABLE_KERNEL #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh" #include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp" #include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp" @@ -35,3 +36,4 @@ GetDeviceBatchedGemmSoftmaxGemmPermuteInstances< } // namespace rocm } // namespace contrib } // namespace onnxruntime +#endif // USE_COMPOSABLE_KERNEL diff --git a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh index f501f20692..86b761986a 100644 --- a/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh @@ -63,11 +63,13 @@ are in composable kernels. The scale and add logic is performed via Acc0ElementO #include "contrib_ops/cpu/bert/attention_base.h" #include "contrib_ops/rocm/bert/attention_impl.h" #include "contrib_ops/rocm/bert/attention_softmax.h" +#ifdef USE_COMPOSABLE_KERNEL #include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh" #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/tensor_layout.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" +#endif // USE_COMPOSABLE_KERNEL #include #include @@ -284,24 +286,6 @@ struct GemmSoftmaxGemmPermuteGenericPipeline { } }; -namespace { -template -struct DataTypeAdaptor { - using type = T; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::half_t; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::bhalf16_t; -}; - -} // namespace - template class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOp> { public: @@ -391,6 +375,24 @@ class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOp +struct DataTypeAdaptor { + using type = T; +}; + +template <> +struct DataTypeAdaptor { + using type = ck::half_t; +}; + +template <> +struct DataTypeAdaptor { + using type = ck::bhalf16_t; +}; +} // namespace + template auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() { constexpr const int kNumBiasBuffer = static_cast(USE_BIAS) + static_cast(USE_MASK); @@ -501,6 +503,7 @@ auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() { } return ret; } +#endif // USE_COMPOSABLE_KERNEL template GemmSoftmaxGemmPermuteTunableOp::GemmSoftmaxGemmPermuteTunableOp() { @@ -508,6 +511,7 @@ GemmSoftmaxGemmPermuteTunableOp::GemmSoftmaxGemmPermuteTunableOp() { return GemmSoftmaxGemmPermuteGenericPipeline::Run(params, false); }); +#ifdef USE_COMPOSABLE_KERNEL for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps()) { this->RegisterOp(std::move(op)); } @@ -523,6 +527,7 @@ GemmSoftmaxGemmPermuteTunableOp::GemmSoftmaxGemmPermuteTunableOp() { for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps()) { this->RegisterOp(std::move(op)); } +#endif } } // namespace rocm diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py index 2f7df713c8..e7b59baed3 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py +++ b/onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py @@ -184,6 +184,7 @@ def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, n _test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale) +@pytest.mark.skipif(not ke.is_composable_kernel_available(), reason="ck is not enabled") @pytest.mark.parametrize("mask_dim", mask_dims, ids=get_mask_dim_id) @pytest.mark.parametrize("biased", biaseds, ids=get_biased_id) @pytest.mark.parametrize("head_size", head_sizes) @@ -329,7 +330,10 @@ def profile_with_args(dtype, batch, seqlen, total_seqlen, num_heads, head_size, profile_gemm_softmax_gemm_permute_func( getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)), *args ) - profile_gemm_softmax_gemm_permute_func(getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args) + if ke.is_composable_kernel_available(): + profile_gemm_softmax_gemm_permute_func( + getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args + ) profile_gemm_softmax_gemm_permute_func( getattr(ke, "GemmSoftmaxGemmPermuteTunable_" + dtype_to_suffix(dtype)), *args ) diff --git a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu index 85c3ce1ad7..e39bf534bd 100644 --- a/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu +++ b/onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu @@ -158,6 +158,7 @@ class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplor } }; +#ifdef USE_COMPOSABLE_KERNEL template class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer { public: @@ -215,6 +216,7 @@ class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer std::vector type_strings_; size_t selected_op_{}; }; +#endif // USE_COMPOSABLE_KERNEL // The pipeline composed from rocblas api calls and kernel launches. template @@ -287,10 +289,12 @@ class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplor void InitGemmSoftmaxGemmPermute(py::module m) { REGISTER_GENERIC(half); +#ifdef USE_COMPOSABLE_KERNEL REGISTER_CK(half, false, false, ""); REGISTER_CK(half, true, false, "Biased"); REGISTER_CK(half, false, true, "Masked"); REGISTER_CK(half, true, true, "BiasedMasked"); +#endif REGISTER_TUNABLE(half); }