saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ROCm] simplify ck data type Adaptor (#15734)

Browse source 
DataTypeAdaptor is defined many times in every file that integrates CK.
This PR refactor the code to put DataTypeAdaptor in a header file.
This commit is contained in:
PeixuanZuo 2023-05-06 17:39:52 +08:00 committed by GitHub
parent 42d62b8f2b
commit 2735e0d031
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 55 additions and 80 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh
View file

@ -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<GemmSoftmaxGem
};
#ifdef USE_COMPOSABLE_KERNEL
namespace {
template <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct DataTypeAdaptor<BFloat16> {
using type = ck::bhalf16_t;
};
} // namespace
template <typename T, bool USE_BIAS, bool USE_MASK>
auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() {
@ -475,7 +460,7 @@ auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() {
using Nop = ck::tensor_operation::element_wise::PassThrough;
using Acc0ElementOp = internal::PreSoftmaxAttentionScoreOp;
using CKDataType = typename DataTypeAdaptor<T>::type;
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using D0DataType = typename ck::detail::tuple_concat<
std::conditional_t<USE_BIAS, ck::Tuple<CKDataType>, ck::Tuple<>>,
std::conditional_t<USE_MASK, ck::Tuple<CKDataType>, ck::Tuple<>>>::type;

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

@ -8,6 +8,8 @@
#include <vector>
#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 <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct DataTypeAdaptor<BFloat16> {
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 <typename T, typename ALayout, typename BLayout>
auto GetCKGemmAddFastGeluTypeStringAndOps() {
using CKDataType = typename DataTypeAdaptor<T>::type;
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using DeviceGemmAddFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD<
ALayout, BLayout, ck::Tuple<Row>, Row,
CKDataType, CKDataType, ck::Tuple<CKDataType>, CKDataType,
@ -89,7 +78,7 @@ auto GetCKGemmAddFastGeluTypeStringAndOps() {
template <typename T, typename ALayout, typename BLayout>
auto GetCKGemmFastGeluTypeStringAndOps() {
using CKDataType = typename DataTypeAdaptor<T>::type;
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using DeviceGemmFastGelu = ck::tensor_operation::device::DeviceGemmMultipleD<
ALayout, BLayout, ck::Tuple<>, Row,
CKDataType, CKDataType, ck::Tuple<>, CKDataType,

18
onnxruntime/contrib_ops/rocm/diffusion/group_norm_ck.cuh
View file

@ -8,6 +8,8 @@
#include <vector>
#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 <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
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 <typename T, typename AccT, bool WithSwish>
auto GetCKGroupNormNHWCTypeStringAndOps() {
using InDataType = typename DataTypeAdaptor<T>::type;
using OutDataType = typename DataTypeAdaptor<T>::type;
using AccDataType = typename DataTypeAdaptor<AccT>::type;
using InDataType = typename CKDataTypeAdaptor<T>::type;
using OutDataType = typename CKDataTypeAdaptor<T>::type;
using AccDataType = typename CKDataTypeAdaptor<AccT>::type;
using GammaDataType = float;
using BetaDataType = float;

33
onnxruntime/core/providers/rocm/composable_kernel_common.h Normal file
View file

@ -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 <typename T>
struct CKDataTypeAdaptor {
using type = T;
};
template <>
struct CKDataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct CKDataTypeAdaptor<BFloat16> {
using type = ck::bhalf16_t;
};
#endif
} // namespace rocm
} // namespace onnxruntime

23
onnxruntime/core/providers/rocm/math/softmax_ck.cuh
View file

@ -8,6 +8,8 @@
#include <vector>
#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 <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct DataTypeAdaptor<BFloat16> {
using type = ck::bhalf16_t;
};
using Nop = ck::tensor_operation::element_wise::PassThrough;
constexpr int Rank = 4;
constexpr int NumReduceDim = 1;
template <typename InputT, typename OutputT, typename AccT>
auto GetCKSoftmaxTypeStringAndOps() {
using InDataType = typename DataTypeAdaptor<InputT>::type;
using OutDataType = typename DataTypeAdaptor<OutputT>::type;
using AccDataType = typename DataTypeAdaptor<AccT>::type;
using InDataType = typename CKDataTypeAdaptor<InputT>::type;
using OutDataType = typename CKDataTypeAdaptor<OutputT>::type;
using AccDataType = typename CKDataTypeAdaptor<AccT>::type;
using DeviceSoftmax = ck::tensor_operation::device::
DeviceSoftmax<InDataType, AccDataType, OutDataType, Nop, Nop, Rank>;
using InstanceFactory = ck::tensor_operation::device::instance::DeviceOperationInstanceFactory<DeviceSoftmax>;

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

@ -8,6 +8,8 @@
#include <vector>
#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 <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct DataTypeAdaptor<BFloat16> {
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 <typename T, typename ALayout, typename BLayout>
auto GetCKGemmTypeStringAndOps() {
using CKDataType = typename DataTypeAdaptor<T>::type;
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using DeviceGemm = ck::tensor_operation::device::DeviceGemm<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,
@ -95,7 +82,7 @@ auto GetCKGemmTypeStringAndOps() {
template <typename T, typename ALayout, typename BLayout>
auto GetCKStridedBatchedGemmTypeStringAndOps() {
using CKDataType = typename DataTypeAdaptor<T>::type;
using CKDataType = typename CKDataTypeAdaptor<T>::type;
using DeviceStridedBatchedGemm = ck::tensor_operation::device::DeviceBatchedGemm<
ALayout, BLayout, Row,
CKDataType, CKDataType, CKDataType,