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 4b38ccafc5..40b8b54b7f 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 @@ -67,6 +67,7 @@ are in composable kernels. The scale and add logic is performed via Acc0ElementO #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 "core/providers/rocm/composable_kernel_common.h" #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/tensor_layout.hpp" @@ -451,22 +452,6 @@ 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() { @@ -475,7 +460,7 @@ auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() { using Nop = ck::tensor_operation::element_wise::PassThrough; using Acc0ElementOp = internal::PreSoftmaxAttentionScoreOp; - using CKDataType = typename DataTypeAdaptor::type; + using CKDataType = typename CKDataTypeAdaptor::type; using D0DataType = typename ck::detail::tuple_concat< std::conditional_t, ck::Tuple<>>, std::conditional_t, ck::Tuple<>>>::type; 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 f1fa3a7f5c..1c29bf4b3a 100644 --- a/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh +++ b/onnxruntime/contrib_ops/rocm/bert/gemm_fast_gelu_ck.cuh @@ -8,6 +8,8 @@ #include #ifdef USE_COMPOSABLE_KERNEL +#include "core/providers/rocm/composable_kernel_common.h" + #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" @@ -28,20 +30,7 @@ namespace internal { #ifdef USE_COMPOSABLE_KERNEL -template -struct DataTypeAdaptor { - using type = T; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::half_t; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::bhalf16_t; -}; +using onnxruntime::rocm::CKDataTypeAdaptor; using Row = ck::tensor_layout::gemm::RowMajor; using Col = ck::tensor_layout::gemm::ColumnMajor; @@ -52,7 +41,7 @@ using FastGelu = ck::tensor_operation::element_wise::FastGelu; template auto GetCKGemmAddFastGeluTypeStringAndOps() { - using CKDataType = typename DataTypeAdaptor::type; + using CKDataType = typename CKDataTypeAdaptor::type; using DeviceGemmAddFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD< ALayout, BLayout, ck::Tuple, Row, CKDataType, CKDataType, ck::Tuple, CKDataType, @@ -89,7 +78,7 @@ auto GetCKGemmAddFastGeluTypeStringAndOps() { template auto GetCKGemmFastGeluTypeStringAndOps() { - using CKDataType = typename DataTypeAdaptor::type; + using CKDataType = typename CKDataTypeAdaptor::type; using DeviceGemmFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD< ALayout, BLayout, ck::Tuple<>, Row, CKDataType, CKDataType, ck::Tuple<>, CKDataType, diff --git a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_ck.cuh b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_ck.cuh index 4ce7b1284d..054b310715 100644 --- a/onnxruntime/contrib_ops/rocm/diffusion/group_norm_ck.cuh +++ b/onnxruntime/contrib_ops/rocm/diffusion/group_norm_ck.cuh @@ -8,6 +8,8 @@ #include #ifdef USE_COMPOSABLE_KERNEL +#include "core/providers/rocm/composable_kernel_common.h" + #include "ck/ck.hpp" #include "ck/tensor_operation/gpu/device/tensor_layout.hpp" #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp" @@ -22,15 +24,7 @@ namespace rocm { #ifdef USE_COMPOSABLE_KERNEL -template -struct DataTypeAdaptor { - using type = T; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::half_t; -}; +using onnxruntime::rocm::CKDataTypeAdaptor; using Swish = ck::tensor_operation::element_wise::Swish; using Pass = ck::tensor_operation::element_wise::PassThrough; @@ -40,9 +34,9 @@ constexpr int NumReduceDim = 3; template auto GetCKGroupNormNHWCTypeStringAndOps() { - using InDataType = typename DataTypeAdaptor::type; - using OutDataType = typename DataTypeAdaptor::type; - using AccDataType = typename DataTypeAdaptor::type; + using InDataType = typename CKDataTypeAdaptor::type; + using OutDataType = typename CKDataTypeAdaptor::type; + using AccDataType = typename CKDataTypeAdaptor::type; using GammaDataType = float; using BetaDataType = float; diff --git a/onnxruntime/core/providers/rocm/composable_kernel_common.h b/onnxruntime/core/providers/rocm/composable_kernel_common.h new file mode 100644 index 0000000000..f2ef9c9dd0 --- /dev/null +++ b/onnxruntime/core/providers/rocm/composable_kernel_common.h @@ -0,0 +1,33 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#ifdef USE_COMPOSABLE_KERNEL +#include "ck/utility/data_type.hpp" +#endif + +#include "core/providers/rocm/rocm_common.h" + +namespace onnxruntime { +namespace rocm { + +#ifdef USE_COMPOSABLE_KERNEL +template +struct CKDataTypeAdaptor { + using type = T; +}; + +template <> +struct CKDataTypeAdaptor { + using type = ck::half_t; +}; + +template <> +struct CKDataTypeAdaptor { + using type = ck::bhalf16_t; +}; +#endif + +} // namespace rocm +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/rocm/math/softmax_ck.cuh b/onnxruntime/core/providers/rocm/math/softmax_ck.cuh index 710f6e540b..4e96b92132 100644 --- a/onnxruntime/core/providers/rocm/math/softmax_ck.cuh +++ b/onnxruntime/core/providers/rocm/math/softmax_ck.cuh @@ -8,6 +8,8 @@ #include #ifdef USE_COMPOSABLE_KERNEL +#include "core/providers/rocm/composable_kernel_common.h" + #include "ck/ck.hpp" #include "ck/library/tensor_operation_instance/gpu/softmax.hpp" #include "ck/tensor_operation/gpu/device/device_softmax.hpp" @@ -22,30 +24,15 @@ namespace rocm { #ifdef USE_COMPOSABLE_KERNEL -template -struct DataTypeAdaptor { - using type = T; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::half_t; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::bhalf16_t; -}; - using Nop = ck::tensor_operation::element_wise::PassThrough; constexpr int Rank = 4; constexpr int NumReduceDim = 1; template auto GetCKSoftmaxTypeStringAndOps() { - using InDataType = typename DataTypeAdaptor::type; - using OutDataType = typename DataTypeAdaptor::type; - using AccDataType = typename DataTypeAdaptor::type; + using InDataType = typename CKDataTypeAdaptor::type; + using OutDataType = typename CKDataTypeAdaptor::type; + using AccDataType = typename CKDataTypeAdaptor::type; using DeviceSoftmax = ck::tensor_operation::device:: DeviceSoftmax; using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory; diff --git a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh index 0b5ec71450..d0ce0915c6 100644 --- a/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh +++ b/onnxruntime/core/providers/rocm/tunable/gemm_ck.cuh @@ -8,6 +8,8 @@ #include #ifdef USE_COMPOSABLE_KERNEL +#include "core/providers/rocm/composable_kernel_common.h" + #include "ck/ck.hpp" #include "ck/library/tensor_operation_instance/gpu/batched_gemm.hpp" #include "ck/library/tensor_operation_instance/gpu/gemm.hpp" @@ -27,21 +29,6 @@ namespace internal { #ifdef USE_COMPOSABLE_KERNEL -template -struct DataTypeAdaptor { - using type = T; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::half_t; -}; - -template <> -struct DataTypeAdaptor { - using type = ck::bhalf16_t; -}; - using Row = ck::tensor_layout::gemm::RowMajor; using Col = ck::tensor_layout::gemm::ColumnMajor; @@ -49,7 +36,7 @@ using Nop = ck::tensor_operation::element_wise::PassThrough; template auto GetCKGemmTypeStringAndOps() { - using CKDataType = typename DataTypeAdaptor::type; + using CKDataType = typename CKDataTypeAdaptor::type; using DeviceGemm = ck::tensor_operation::device::DeviceGemm< ALayout, BLayout, Row, CKDataType, CKDataType, CKDataType, @@ -95,7 +82,7 @@ auto GetCKGemmTypeStringAndOps() { template auto GetCKStridedBatchedGemmTypeStringAndOps() { - using CKDataType = typename DataTypeAdaptor::type; + using CKDataType = typename CKDataTypeAdaptor::type; using DeviceStridedBatchedGemm = ck::tensor_operation::device::DeviceBatchedGemm< ALayout, BLayout, Row, CKDataType, CKDataType, CKDataType,