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

Remove most ROCm-specific element-wise code and reuse CUDA element-wise code.

Browse source
This commit is contained in:
Jesse Benson 2020-12-22 12:19:40 -08:00 committed by Jesse Benson
parent 52228a703c
commit 7ccdfed1a6
13 changed files with 0 additions and 1418 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

69
onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.cc
View file

@ -1,69 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/rocm/math/binary_elementwise_ops.h"
#include "contrib_ops/rocm/math/binary_elementwise_ops_impl.h"
using namespace onnxruntime::common;
namespace onnxruntime {
namespace contrib {
namespace rocm {
#define CONTRIB_BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(x, ver, T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
x, \
kMSDomain, \
ver, \
T, \
kRocmExecutionProvider, \
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
x<T>);
#define CONTRIB_BINARY_ELEMENTWISE_COMPUTE(x, T) \
template <> \
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
BinaryElementwisePreparation prepare; \
ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \
Impl_##x<typename ToHipType<T>::MappedType>( \
prepare.output_rank_or_simple_broadcast, \
&prepare.lhs_padded_strides, \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()), \
&prepare.rhs_padded_strides, \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()), \
&prepare.fdm_output_strides, \
prepare.fdm_H, \
prepare.fdm_C, \
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()), \
prepare.output_tensor->Shape().Size()); \
return Status::OK(); \
}
#define CONTRIB_BINARY_OP_TYPED(name, ver, T) \
CONTRIB_BINARY_ELEMENTWISE_REGISTER_KERNEL_TYPED(name, ver, T) \
CONTRIB_BINARY_ELEMENTWISE_COMPUTE(name, T)
// since different ops has different types, we cannot use BINARY_OPS() directly
// the postfix of means the types supported by the op:
// B: uint8_t
// W: uint16_t
// U: uint32_t
// Z: uint64_t
// C: int8_t
// S: int16_t
// I: int32_t
// L: int64_t
// H: float16
// F: float
// D: double
// O: bool
#define CONTRIB_BINARY_OP_HFD(name, ver) \
CONTRIB_BINARY_OP_TYPED(name, ver, MLFloat16) \
CONTRIB_BINARY_OP_TYPED(name, ver, float) \
CONTRIB_BINARY_OP_TYPED(name, ver, double)
CONTRIB_BINARY_OP_HFD(BiasGelu, 1)
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

29
onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.h
View file

@ -1,29 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/rocm/math/binary_elementwise_ops.h"
#include "core/providers/rocm/rocm_common.h"
#include "core/providers/rocm/shared_inc/fast_divmod.h"
#include "core/providers/cpu/tensor/utils.h"
using namespace onnxruntime::rocm;
namespace onnxruntime {
namespace contrib {
namespace rocm {
// AddGelu fuse Add + Gelu
template <typename T>
class BiasGelu final : public BinaryElementwise<ShouldBroadcast> {
public:
BiasGelu(const OpKernelInfo& info) : BinaryElementwise(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
};
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

94
onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.cu
View file

@ -1,94 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <hip/hip_runtime.h>
#include "contrib_ops/rocm/math/binary_elementwise_ops_impl.h"
#include "core/providers/rocm/cu_inc/common.cuh"
#include "core/providers/rocm/cu_inc/binary_elementwise_impl.cuh"
namespace onnxruntime {
namespace contrib {
namespace rocm {
#define OP(name, expr) \
template <class T> \
struct OP_##name { \
__device__ __inline__ T operator()(T a, T b) const { \
return (expr); \
} \
};
#define CONTRIB_BINARY_ELEMENTWISE_IMPL(name) \
CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name) { \
BinaryElementWiseImpl(output_rank_or_simple_broadcast, \
lhs_padded_strides, \
lhs_data, \
rhs_padded_strides, \
rhs_data, \
fdm_output_strides, \
fdm_H, \
fdm_C, \
output_data, \
OP_##name<T>(), \
count); \
}
#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \
template void Impl_##x<T>(int32_t output_rank, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<onnxruntime::rocm::fast_divmod>* fdm_output_strides, \
const onnxruntime::rocm::fast_divmod& fdm_H, \
const onnxruntime::rocm::fast_divmod& fdm_C, \
T* output_data, size_t count);
#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint64_t) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double)
#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_OIL(x) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, bool) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int64_t)
#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(x) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, half) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double)
// create declarations for op and impl
#define CONTRIB_BINARY_OP_NAME_EXPR(name, expr) \
OP(name, expr) \
CONTRIB_BINARY_ELEMENTWISE_IMPL(name)
CONTRIB_BINARY_OPS()
#undef CONTRIB_BINARY_OP_NAME_EXPR
// create specialized impl
// the postfix of means the types supported by the op:
// B: uint8_t
// W: uint16_t
// U: uint32_t
// Z: uint64_t
// C: int8_t
// S: int16_t
// I: int32_t
// L: int64_t
// H: float16
// F: float
// D: double
// O: bool
CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL_HFD(BiasGelu)
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

38
onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.h
View file

@ -1,38 +0,0 @@
#include <stdint.h>
#include "core/providers/rocm/shared_inc/rocm_utils.h"
using namespace onnxruntime::rocm;
namespace onnxruntime {
namespace contrib {
namespace rocm {
// These macros simplifies coding. To add a new op with following steps:
// 1. Add a new entry in CONTRIB_BINARY_OPS() list
// 2. (optional) Define templated single element operator in binary_elementwise_ops_impl.cu
// 3. (optional) Implement specialized single element operator
// 4. Add op kernel class definition in binary_elementwise_ops.h
// 5. Add op kernel registration and compute specialization in binary_elementwise_ops.cc
#define CONTRIB_BINARY_OPS() \
CONTRIB_BINARY_OP_NAME_EXPR(BiasGelu, _Gelu(a + b))
// NOTE that cu files are compiled with nvcc and should not refer to any onnxruntime headers
// so struct BinaryElementwisePreparation cannot be used here
#define CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name) \
template <typename T> \
void Impl_##name( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count)
#define CONTRIB_BINARY_OP_NAME_EXPR(name, expr) CONTRIB_BINARY_ELEMENTWISE_IMPL_DECLARATION(name);
CONTRIB_BINARY_OPS()
#undef CONTRIB_BINARY_OP_NAME_EXPR
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

296
onnxruntime/core/providers/rocm/cu_inc/binary_elementwise_impl.cuh
View file

@ -1,296 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <stdint.h>
#include "core/providers/rocm/shared_inc/rocm_utils.h"
#include "core/providers/rocm/cu_inc/common.cuh"
namespace onnxruntime {
namespace rocm {
// broadcast by computing output coordinate from offset, using fast_divmod
template <typename T, typename T1, typename T2, typename FuncT,
bool lhs_need_compute, bool rhs_need_compute, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _BinaryElementWise(
int32_t output_rank,
const TArray<int64_t> lhs_padded_strides,
const T1* lhs_data,
const TArray<int64_t> rhs_padded_strides,
const T2* rhs_data,
const TArray<fast_divmod> fdm_output_strides,
T* output_data,
const FuncT& functor,
HIP_LONG N) {
HIP_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
T1 lvalue[NumElementsPerThread];
T2 rvalue[NumElementsPerThread];
HIP_LONG id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
HIP_LONG lhs_index = (lhs_need_compute ? 0 : id);
HIP_LONG rhs_index = (rhs_need_compute ? 0 : id);
// compute indexes with broadcasting rules: https://github.com/onnx/onnx/blob/master/docs/Broadcasting.md
HIP_LONG offset = id;
#pragma unroll
for (auto dim = 0; dim < fdm_output_strides.Capacity(); dim++) {
if (dim >= output_rank) {
break;
}
int q, r;
fdm_output_strides[dim].divmod(offset, q, r);
if (lhs_need_compute) {
lhs_index += static_cast<int>(lhs_padded_strides[dim]) * q;
}
if (rhs_need_compute) {
rhs_index += static_cast<int>(rhs_padded_strides[dim]) * q;
}
offset = r;
}
lvalue[i] = lhs_data[lhs_index];
rvalue[i] = rhs_data[rhs_index];
id += NumThreadsPerBlock;
}
}
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[id] = functor(lvalue[i], rvalue[i]);
id += NumThreadsPerBlock;
}
}
}
// for scalar broadcast or non-broadcast case
template <bool IncL, bool IncR, typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _BinaryElementWiseSimple(
const T1* lhs_data,
const T2* rhs_data,
T* output_data,
const FuncT& func,
HIP_LONG N) {
HIP_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
T1 lvalue[NumElementsPerThread];
T2 rvalue[NumElementsPerThread];
HIP_LONG id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
lvalue[i] = lhs_data[IncL ? id : 0];
rvalue[i] = rhs_data[IncR ? id : 0];
id += NumThreadsPerBlock;
}
}
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[id] = func(lvalue[i], rvalue[i]);
id += NumThreadsPerBlock;
}
}
}
// for rhs per-channel broadcast case
template <typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _BinaryElementWiseRhsPerChannelBatch1(
const T1* lhs_data,
const T2* rhs_data,
const fast_divmod fdm_H,
T* output_data,
FuncT func,
HIP_LONG N) {
HIP_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
T1 lvalue[NumElementsPerThread];
T2 rvalue[NumElementsPerThread];
HIP_LONG id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
HIP_LONG rhs_id = fdm_H.div(id);
lvalue[i] = lhs_data[id];
rvalue[i] = rhs_data[rhs_id];
id += NumThreadsPerBlock;
}
}
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[id] = func(lvalue[i], rvalue[i]);
id += NumThreadsPerBlock;
}
}
}
template <typename T, typename T1, typename T2, typename FuncT, int NumThreadsPerBlock, int NumElementsPerThread>
__global__ void _BinaryElementWiseRhsPerChannelBatchN(
const T1* lhs_data,
const T2* rhs_data,
const fast_divmod fdm_H,
const fast_divmod fdm_C,
T* output_data,
FuncT func,
HIP_LONG N) {
HIP_LONG start = NumElementsPerThread * NumThreadsPerBlock * blockIdx.x + threadIdx.x;
T1 lvalue[NumElementsPerThread];
T2 rvalue[NumElementsPerThread];
HIP_LONG id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
HIP_LONG rhs_id = fdm_H.div(id);
int q, r;
fdm_C.divmod(rhs_id, q, r);
rhs_id = r;
lvalue[i] = lhs_data[id];
rvalue[i] = rhs_data[rhs_id];
id += NumThreadsPerBlock;
}
}
id = start;
#pragma unroll
for (int i = 0; i < NumElementsPerThread; i++) {
if (id < N) {
output_data[id] = func(lvalue[i], rvalue[i]);
id += NumThreadsPerBlock;
}
}
}
template <typename T, typename T1, typename T2, typename FuncT>
void BinaryElementWiseNoBroadcastImpl(
const T1* lhs_data,
const T2* rhs_data,
T* output_data,
const FuncT& func,
size_t count) {
if (count == 0) // special case where there's a dim value of 0 in the output shape
return;
int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
HIP_LONG N = static_cast<HIP_LONG>(count);
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWiseSimple<true, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
lhs_data,
rhs_data,
output_data,
func,
N);
}
template <typename T, typename T1, typename T2, typename FuncT>
void BinaryElementWiseImpl(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T1* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T2* rhs_data,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,
const FuncT& func,
size_t count) {
if (count == 0) // special case where there's a dim value of 0 in the output shape
return;
int blocksPerGrid = static_cast<int>(CeilDiv(count, GridDim::maxThreadsPerBlock * GridDim::maxElementsPerThread));
HIP_LONG N = static_cast<HIP_LONG>(count);
if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::NoBroadcast)) {
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWiseSimple<true, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
lhs_data,
rhs_data,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::LeftScalar)) {
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWiseSimple<false, true, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
lhs_data,
rhs_data,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightScalar)) {
_BinaryElementWiseSimple<true, false, T, T1, T2, FuncT, GridDim::maxThreadsPerBlock,
GridDim::maxElementsPerThread><<<blocksPerGrid, GridDim::maxThreadsPerBlock, 0>>>(
lhs_data,
rhs_data,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatch1)) {
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWiseRhsPerChannelBatch1<T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
lhs_data,
rhs_data,
fdm_H,
output_data,
func,
N);
} else if (output_rank_or_simple_broadcast == static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatchN)) {
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWiseRhsPerChannelBatchN<T, T1, T2, FuncT, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
lhs_data,
rhs_data,
fdm_H,
fdm_C,
output_data,
func,
N);
} else {
if (lhs_padded_strides && rhs_padded_strides && lhs_padded_strides->Size() && rhs_padded_strides->Size())
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWise<T, T1, T2, FuncT, true, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
output_rank_or_simple_broadcast,
*lhs_padded_strides,
lhs_data,
*rhs_padded_strides,
rhs_data,
*fdm_output_strides,
output_data,
func,
N);
else if (lhs_padded_strides && lhs_padded_strides->Size())
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWise<T, T1, T2, FuncT, true, false, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
output_rank_or_simple_broadcast,
*lhs_padded_strides,
lhs_data,
*rhs_padded_strides,
rhs_data,
*fdm_output_strides,
output_data,
func,
N);
else
hipLaunchKernelGGL(HIP_KERNEL_NAME(_BinaryElementWise<T, T1, T2, FuncT, false, true, GridDim::maxThreadsPerBlock, GridDim::maxElementsPerThread>), dim3(blocksPerGrid), dim3(GridDim::maxThreadsPerBlock), 0, 0,
output_rank_or_simple_broadcast,
*lhs_padded_strides,
lhs_data,
*rhs_padded_strides,
rhs_data,
*fdm_output_strides,
output_data,
func,
N);
}
}
} // namespace rocm
} // namespace onnxruntime

260
onnxruntime/core/providers/rocm/math/binary_elementwise_ops.h
View file

@ -1,260 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/providers/rocm/rocm_kernel.h"
#include "core/providers/rocm/shared_inc/fast_divmod.h"
#include "core/providers/cpu/tensor/utils.h"
namespace onnxruntime {
namespace rocm {
struct BinaryElementwisePreparation {
const Tensor* lhs_tensor = nullptr;
const Tensor* rhs_tensor = nullptr;
Tensor* output_tensor = nullptr;
int32_t output_rank_or_simple_broadcast = 0; // for no_broadcast|left_scalar|right_scalar cases, output_rank uses SimpleBroadcast enums
TArray<int64_t> lhs_padded_strides;
TArray<int64_t> rhs_padded_strides;
TArray<fast_divmod> fdm_output_strides;
// these are for RightPerChannel case
fast_divmod fdm_H;
fast_divmod fdm_C;
BinaryElementwisePreparation() {}
Status BinaryElementwiseBroadcastPrepareHelper(const TensorShape& lhs_shape,
const TensorShape& rhs_shape,
const TensorShape& output_shape) {
int32_t lhs_rank = gsl::narrow_cast<int32_t>(lhs_shape.NumDimensions());
int32_t rhs_rank = gsl::narrow_cast<int32_t>(rhs_shape.NumDimensions());
int32_t out_rank = std::max(lhs_rank, rhs_rank);
// early return when shapes match
if (lhs_shape == rhs_shape) {
output_rank_or_simple_broadcast = static_cast<int32_t>(SimpleBroadcast::NoBroadcast);
return Status::OK();
}
// early return if one operand is scalar
if (lhs_shape.Size() == 1 || rhs_shape.Size() == 1) {
output_rank_or_simple_broadcast = static_cast<int32_t>(lhs_shape.Size() == 1
? SimpleBroadcast::LeftScalar
: SimpleBroadcast::RightScalar);
return Status::OK();
}
// special case for lhs(N,C,H) and rhs (C,1) which is used in conv bias
// when N == 1: out[id] = op(lhs[id], rhs[id / H])
// When N > 1: out[id] = op(lhs[id], rhs[id / H % C])
if (lhs_shape == output_shape) {
const auto& rhs_dims = rhs_shape.GetDims();
int64_t C = 0;
if (1 == std::count_if(rhs_dims.begin(), rhs_dims.end(),
[&C](int64_t dim) { if (dim != 1) C = dim; return (dim != 1); })) {
int32_t dim_C = gsl::narrow_cast<int32_t>(std::find(rhs_dims.begin(), rhs_dims.end(), C) - rhs_dims.begin() + output_shape.NumDimensions() - rhs_shape.NumDimensions());
int64_t N = output_shape.SizeToDimension(dim_C);
int64_t H = (dim_C < out_rank - 1 ? output_shape.SizeFromDimension(dim_C + 1) : 1);
std::vector<int64_t> new_output_dims;
if (N == 1) {
output_rank_or_simple_broadcast = static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatch1);
fdm_H = fast_divmod(gsl::narrow_cast<int>(H));
} else {
output_rank_or_simple_broadcast = static_cast<int32_t>(SimpleBroadcast::RightPerChannelBatchN);
fdm_H = fast_divmod(gsl::narrow_cast<int>(H));
fdm_C = fast_divmod(gsl::narrow_cast<int>(C));
}
return Status::OK();
}
}
output_rank_or_simple_broadcast = out_rank;
if (lhs_shape != output_shape) {
TensorPitches original_lhs_padded_strides(lhs_shape.GetDims(), out_rank);
lhs_padded_strides.SetSize(out_rank);
auto offset = out_rank - lhs_rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (lhs_shape.GetDims()[i - offset] != 1) {
lhs_padded_strides[i] = original_lhs_padded_strides[i];
}
}
}
if (rhs_shape != output_shape) {
TensorPitches original_rhs_padded_strides(rhs_shape.GetDims(), out_rank);
rhs_padded_strides.SetSize(out_rank);
auto offset = out_rank - rhs_rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (rhs_shape.GetDims()[i - offset] != 1) {
rhs_padded_strides[i] = original_rhs_padded_strides[i];
}
}
}
TensorPitches original_output_strides(output_shape.GetDims());
fdm_output_strides.SetSize(out_rank);
for (auto i = 0; i < out_rank; ++i) {
fdm_output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
}
return Status::OK();
}
};
Status ComputeOutputShape(
const std::string& node_name,
const TensorShape& lhs_shape,
const TensorShape& rhs_shape,
TensorShape& out_shape);
Status BinaryElementwiseBroadcastPrepare(
const Tensor* lhs_tensor,
const Tensor* rhs_tensor,
Tensor* output_tensor,
BinaryElementwisePreparation* p,
const TensorShape* override_lhs_shape = nullptr,
const TensorShape* override_rhs_shape = nullptr);
// trait classes to indicate if the kernel supports broadcast
class ShouldBroadcast {
};
class ShouldNotBroadcast {
};
template <typename BroadcastTrait>
class BinaryElementwise : public RocmKernel {
protected:
typedef BroadcastTrait broadcast_type;
BinaryElementwise(const OpKernelInfo& info) : RocmKernel(info) {}
Status ComputeInternal(OpKernelContext*) const override {
return Status(common::ONNXRUNTIME, common::FAIL); // should not reach here
}
Status Prepare(OpKernelContext* context, BinaryElementwisePreparation* p) const;
};
template <typename T>
class Add final : public BinaryElementwise<ShouldBroadcast> {
public:
Add(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Sub final : public BinaryElementwise<ShouldBroadcast> {
public:
Sub(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Mul final : public BinaryElementwise<ShouldBroadcast> {
public:
Mul(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Div final : public BinaryElementwise<ShouldBroadcast> {
public:
Div(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Pow_7 final : public BinaryElementwise<ShouldBroadcast> {
public:
Pow_7(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
// Since version 12
class Pow final : public BinaryElementwise<ShouldBroadcast> {
public:
Pow(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class And final : public BinaryElementwise<ShouldBroadcast> {
public:
And(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Or final : public BinaryElementwise<ShouldBroadcast> {
public:
Or(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Xor final : public BinaryElementwise<ShouldBroadcast> {
public:
Xor(const OpKernelInfo& info) : BinaryElementwise(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
// PRelu is activation function, but it's closer to binary elementwise ops in implementation
template <typename T>
class PRelu final : public BinaryElementwise<ShouldBroadcast> {
public:
PRelu(const OpKernelInfo& info) : BinaryElementwise(info) {
}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T, typename HipT>
class CompareFunction : public BinaryElementwise<ShouldBroadcast> {
public:
CompareFunction(const OpKernelInfo& info) : BinaryElementwise(info) {}
typedef void (*ImplCompare)(int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const HipT* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const HipT* rhs_data,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
bool* output_data,
size_t count);
Status CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const;
};
template <typename T>
class Greater final : public CompareFunction<T, typename ToHipType<T>::MappedType> {
public:
Greater(const OpKernelInfo& info) : CompareFunction<T, typename ToHipType<T>::MappedType>(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Equal final : public CompareFunction<T, typename ToHipType<T>::MappedType> {
public:
Equal(const OpKernelInfo& info) : CompareFunction<T, typename ToHipType<T>::MappedType>(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
template <typename T>
class Less final : public CompareFunction<T, typename ToHipType<T>::MappedType> {
public:
Less(const OpKernelInfo& info) : CompareFunction<T, typename ToHipType<T>::MappedType>(info) {}
Status ComputeInternal(OpKernelContext* context) const override;
};
} // namespace rocm
} // namespace onnxruntime

92
onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.h
View file

@ -1,92 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <stdint.h>
#include "core/providers/rocm/shared_inc/rocm_utils.h"
namespace onnxruntime {
namespace rocm {
// These macros simplifies coding. To add a new op with following steps:
// 1. Add a new entry in BINARY_OPS() list
// 2. (optional) Define templated single element operator in binary_elementwise_ops_impl.cu
// 3. (optional) Implement specialized single element operator
// 4. Add op kernel class definition in binary_elementwise_ops.h
// 5. Add op kernel registration and compute specialization in binary_elementwise_ops.cc
#define BINARY_OPS() \
BINARY_OP_NAME_EXPR(Add, (a + b)) \
BINARY_OP_NAME_EXPR(Sub, (a - b)) \
BINARY_OP_NAME_EXPR(Mul, (a * b)) \
BINARY_OP_NAME_EXPR(Div, (a / b)) \
BINARY_OP_NAME_EXPR(Pow_7, _Pow(a, b)) \
BINARY_OP_NAME_EXPR(And, (a & b)) \
BINARY_OP_NAME_EXPR(Or, (a | b)) \
BINARY_OP_NAME_EXPR(Xor, (a ^ b)) \
BINARY_OP_NAME_EXPR(PRelu, (a > (T)0 ? a : a * b)) \
BINARY_OP_NAME_EXPR(Max, _Max(a, b)) \
BINARY_OP_NAME_EXPR(Min, _Min(a, b))
// NOTE that cu files are compiled with nvcc and should not refer to any onnxruntime headers
// so struct BinaryElementwisePreparation cannot be used here
#define BINARY_ELEMENTWISE_IMPL_DECLARATION(name) \
template <typename T> \
void Impl_##name( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count)
#define BINARY_OP_NAME_EXPR(name, expr) BINARY_ELEMENTWISE_IMPL_DECLARATION(name);
BINARY_OPS()
#undef BINARY_OP_NAME_EXPR
#define BINARY_ELEMENTWISE_IMPL_DECLARATION_T1(name) \
template <typename T, typename T1> \
void ImplT1_##name( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T1* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count)
BINARY_ELEMENTWISE_IMPL_DECLARATION_T1(Pow);
#define BINARY_ELEMENTWISE_IMPL_DECLARATION_T2(name) \
template <typename T, typename T1, typename T2> \
void ImplT2_##name( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T1* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T2* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count)
#define BINARY_OPS2() \
BINARY_OP_NAME_EXPR2(Greater, (a > b)) \
BINARY_OP_NAME_EXPR2(Equal, (a == b)) \
BINARY_OP_NAME_EXPR2(Less, (a < b))
#define BINARY_OP_NAME_EXPR2(name, expr) BINARY_ELEMENTWISE_IMPL_DECLARATION_T2(name);
BINARY_OPS2()
#undef BINARY_OP_NAME_EXPR2
} // namespace rocm
} // namespace onnxruntime

252
onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.cc
View file

@ -1,252 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/rocm/math/variadic_elementwise_ops.h"
#include <cassert>
#include "core/framework/data_types_internal.h"
#include "core/providers/rocm/math/binary_elementwise_ops.h"
#include "core/providers/rocm/math/binary_elementwise_ops_impl.h"
#include "core/providers/rocm/math/variadic_elementwise_ops_impl.h"
#include "core/providers/rocm/math/variadic_elementwise_ops_tags.h"
namespace onnxruntime {
namespace rocm {
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
template <typename T>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::
NoBroadcastBatchImplDispatchTarget<T>::operator()(const InputTensorVector& inputs, Tensor& output) const {
assert(inputs.size() > 1);
using HipT = typename ToHipType<T>::MappedType;
InputBatchArray<HipT> input_data_batch{static_cast<int32_t>(inputs.size())};
for (size_t i = 0; i < inputs.size(); ++i) {
input_data_batch[static_cast<int32_t>(i)] = reinterpret_cast<const HipT*>(inputs[i].get().template Data<T>());
}
HipT* output_data = reinterpret_cast<HipT*>(output.template MutableData<T>());
Impl_NoBroadcastInputBatch<HipT, VariadicElementwiseOpTag>(
input_data_batch, output_data, output.Shape().Size());
return Status::OK();
}
// special case for 2 tensors to avoid memset zero
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
template <typename T>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::
BinaryImplDispatchTarget<T>::operator()(const Tensor& lhs, const Tensor& rhs, Tensor& output) const {
using HipT = typename ToHipType<T>::MappedType;
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(&lhs, &rhs, &output, &prepare));
Impl_General<HipT, VariadicElementwiseOpTag>(
prepare.output_rank_or_simple_broadcast,
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
prepare.output_tensor->Shape().Size());
return Status::OK();
}
// for more than 2 inputs, we need to accumulate into output tensor, as the shape from input0 + input1 might be different from output shape
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
template <typename T>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::
GeneralImplDispatchTarget<T>::operator()(const InputTensorVector& inputs, Tensor& output) const {
assert(inputs.size() > 1);
using HipT = typename ToHipType<T>::MappedType;
HIP_RETURN_IF_ERROR(hipMemsetAsync(output.MutableDataRaw(), 0, output.SizeInBytes()));
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(&output, &inputs[0].get(), &output, &prepare));
Impl_Add(
prepare.output_rank_or_simple_broadcast,
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
prepare.output_tensor->Shape().Size());
for (size_t index = 1; index < inputs.size(); index++) {
ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(&output, &inputs[index].get(), &output, &prepare));
Impl_General<HipT, VariadicElementwiseOpTag>(
prepare.output_rank_or_simple_broadcast,
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
prepare.output_tensor->Shape().Size());
}
return Status::OK();
}
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::ComputeInternal(
OpKernelContext* context) const {
const auto& node = Node();
const auto& node_name = node.Name();
auto input_count = node.InputArgCount().front();
ORT_RETURN_IF_NOT(input_count >= 1, "Must have 1 or more inputs");
const InputTensorVector input_tensors =
[&context, input_count]() {
InputTensorVector result{};
result.reserve(input_count);
for (int i = 0; i < input_count; ++i) {
const auto& tensor = context->RequiredInput<Tensor>(i);
result.push_back(std::cref(tensor));
}
return result;
}();
const auto& first_input_tensor = input_tensors[0].get();
// special case for 1 input
if (input_count == 1) {
auto& output_tensor = context->RequiredOutput(0, first_input_tensor.Shape());
if (first_input_tensor.DataRaw() != output_tensor.DataRaw()) {
HIP_RETURN_IF_ERROR(hipMemcpyAsync(
output_tensor.MutableDataRaw(), first_input_tensor.DataRaw(), first_input_tensor.SizeInBytes(),
hipMemcpyDeviceToDevice));
}
return Status::OK();
}
const auto element_type = first_input_tensor.GetElementType();
// special case for no broadcasting and few enough inputs
if (input_count <= k_max_input_batch_size &&
std::all_of(
input_tensors.begin() + 1, input_tensors.end(),
[&first_input_tensor](InputTensorVector::value_type t) {
return first_input_tensor.Shape() == t.get().Shape();
})) {
auto& output_tensor = context->RequiredOutput(0, first_input_tensor.Shape());
// special case for no broadcasting and 2 inputs
if (input_count == 2) {
utils::MLTypeCallDispatcherRet<Status, BinaryImplDispatchTarget, SupportedElementTypes...> dispatcher(element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors[0], input_tensors[1], output_tensor));
return Status::OK();
}
utils::MLTypeCallDispatcherRet<Status, NoBroadcastBatchImplDispatchTarget, SupportedElementTypes...> dispatcher(
element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors, output_tensor));
return Status::OK();
}
// compute output shape first, using broadcast rule
TensorShape output_shape;
TensorShape previous_output_shape = first_input_tensor.Shape();
for (int index = 1; index < input_count; index++) {
ORT_RETURN_IF_ERROR(ComputeOutputShape(
node_name, previous_output_shape, input_tensors[index].get().Shape(), output_shape));
previous_output_shape = output_shape;
}
Tensor& output_tensor = context->RequiredOutput(0, output_shape);
// special case for 2 inputs
if (input_count == 2) {
utils::MLTypeCallDispatcherRet<Status, BinaryImplDispatchTarget, SupportedElementTypes...> dispatcher(element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors[0], input_tensors[1], output_tensor));
return Status::OK();
}
// general case for more than 2 inputs
{
utils::MLTypeCallDispatcherRet<Status, GeneralImplDispatchTarget, SupportedElementTypes...> dispatcher(
element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors, output_tensor));
}
return Status::OK();
}
namespace {
using SumOp = VariadicElementwiseOp<
variadic_elementwise_ops::Sum,
MLFloat16, float, double>;
using MinOp = VariadicElementwiseOp<
variadic_elementwise_ops::Min,
uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>;
using MaxOp = VariadicElementwiseOp<
variadic_elementwise_ops::Max,
uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>;
const auto k_uzilhfd_datatypes =
BuildKernelDefConstraints<uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double>();
const auto k_hfd_datatypes =
BuildKernelDefConstraints<MLFloat16, float, double>();
} // namespace
// kernel registration
#define REGISTER_KERNEL(name, impl_class, version, datatypes) \
ONNX_OPERATOR_KERNEL_EX( \
name, \
kOnnxDomain, \
version, \
kRocmExecutionProvider, \
KernelDefBuilder().TypeConstraint("T", datatypes), \
impl_class)
#define REGISTER_VERSIONED_KERNEL(name, impl_class, start_version, end_version, datatypes) \
ONNX_OPERATOR_VERSIONED_KERNEL_EX( \
name, \
kOnnxDomain, \
start_version, end_version, \
kRocmExecutionProvider, \
KernelDefBuilder().TypeConstraint("T", datatypes), \
impl_class)
REGISTER_KERNEL(Sum, SumOp, 13, k_hfd_datatypes)
REGISTER_VERSIONED_KERNEL(Sum, SumOp, 8, 12, k_hfd_datatypes)
REGISTER_VERSIONED_KERNEL(Sum, SumOp, 6, 7, k_hfd_datatypes)
REGISTER_KERNEL(Min, MinOp, 13, k_uzilhfd_datatypes)
REGISTER_VERSIONED_KERNEL(Min, MinOp, 12, 12, k_uzilhfd_datatypes)
REGISTER_VERSIONED_KERNEL(Min, MinOp, 6, 11, k_hfd_datatypes)
REGISTER_KERNEL(Max, MaxOp, 13, k_uzilhfd_datatypes)
REGISTER_VERSIONED_KERNEL(Max, MaxOp, 12, 12, k_uzilhfd_datatypes)
REGISTER_VERSIONED_KERNEL(Max, MaxOp, 6, 11, k_hfd_datatypes)
#undef REGISTER_VERSIONED_KERNEL
#undef REGISTER_KERNEL
} // namespace rocm
} // namespace onnxruntime

42
onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.h
View file

@ -1,42 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <functional>
#include <vector>
#include "core/providers/rocm/rocm_kernel.h"
namespace onnxruntime {
namespace rocm {
using InputTensorVector = std::vector<std::reference_wrapper<const Tensor>>;
template <typename VariadicElementwiseOpTag,
typename... SupportedElementTypes>
class VariadicElementwiseOp : public RocmKernel {
public:
VariadicElementwiseOp(const OpKernelInfo& info) : RocmKernel(info) {}
private:
Status ComputeInternal(OpKernelContext* context) const override;
template <typename T>
struct NoBroadcastBatchImplDispatchTarget {
Status operator()(const InputTensorVector& inputs, Tensor& output) const;
};
template <typename T>
struct BinaryImplDispatchTarget {
Status operator()(const Tensor& lhs, const Tensor& rhs, Tensor& output) const;
};
template <typename T>
struct GeneralImplDispatchTarget {
Status operator()(const InputTensorVector& inputs, Tensor& output) const;
};
};
} // namespace rocm
} // namespace onnxruntime

145
onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.cu
View file

@ -1,145 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/rocm/math/variadic_elementwise_ops_impl.h"
#include "core/providers/rocm/cu_inc/variadic_elementwise_impl.cuh"
#include "core/providers/rocm/math/binary_elementwise_ops_impl.h"
#include "core/providers/rocm/math/binary_elementwise_ops_impl_functors.cuh"
#include "core/providers/rocm/math/variadic_elementwise_ops_tags.h"
namespace onnxruntime {
namespace rocm {
template <typename T, typename VariadicElementwiseOpTag>
struct VariadicElementwiseOpTraits;
#define DEFINE_TRAITS(VariadicElementwiseOpTag, ImplName) \
template <typename T> \
struct VariadicElementwiseOpTraits<T, VariadicElementwiseOpTag> { \
using ScalarComputeFunctor = OP_##ImplName<T, T, T>; \
\
static void ComputeFn( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count) { \
Impl_##ImplName( \
output_rank_or_simple_broadcast, \
lhs_padded_strides, \
lhs_data, \
rhs_padded_strides, \
rhs_data, \
fdm_output_strides, \
fdm_H, \
fdm_C, \
output_data, \
count); \
} \
};
DEFINE_TRAITS(variadic_elementwise_ops::Sum, Add)
DEFINE_TRAITS(variadic_elementwise_ops::Min, Min)
DEFINE_TRAITS(variadic_elementwise_ops::Max, Max)
#undef DEFINE_TRAITS
template <typename T, typename VariadicElementwiseOpTag>
void Impl_General(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,
size_t count) {
VariadicElementwiseOpTraits<T, VariadicElementwiseOpTag>::ComputeFn(
output_rank_or_simple_broadcast,
lhs_padded_strides,
lhs_data,
rhs_padded_strides,
rhs_data,
fdm_output_strides,
fdm_H,
fdm_C,
output_data,
count);
}
template <typename T, typename VariadicElementwiseOpTag>
void Impl_NoBroadcastInputBatch(
InputBatchArray<T> input_data_batch,
T* output_data,
size_t count) {
VariadicElementWiseNoBroadcastInputBatchImpl<
T, typename VariadicElementwiseOpTraits<T, VariadicElementwiseOpTag>::ScalarComputeFunctor,
k_max_input_batch_size>(
typename VariadicElementwiseOpTraits<T, VariadicElementwiseOpTag>::ScalarComputeFunctor{},
count,
input_data_batch,
output_data);
}
#define SPECIALIZE_IMPL(T, VariadicElementwiseOpTag) \
template void Impl_General<T, VariadicElementwiseOpTag>( \
int32_t output_rank_or_simple_broadcast, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
size_t count); \
\
template void Impl_NoBroadcastInputBatch<T, VariadicElementwiseOpTag>( \
InputBatchArray<T> input_data_batch, \
T * output_data, \
size_t count);
// the postfix means the types supported by the op:
// B: uint8_t
// W: uint16_t
// U: uint32_t
// Z: uint64_t
// C: int8_t
// S: int16_t
// I: int32_t
// L: int64_t
// H: float16
// F: float
// D: double
// O: bool
#define SPECIALIZE_IMPL_HFD(VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(half, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(float, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(double, VariadicElementwiseOpTag)
#define SPECIALIZE_IMPL_UZILHFD(VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(uint32_t, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(uint64_t, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(int32_t, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL(int64_t, VariadicElementwiseOpTag) \
SPECIALIZE_IMPL_HFD(VariadicElementwiseOpTag)
SPECIALIZE_IMPL_HFD(variadic_elementwise_ops::Sum)
SPECIALIZE_IMPL_UZILHFD(variadic_elementwise_ops::Min)
SPECIALIZE_IMPL_UZILHFD(variadic_elementwise_ops::Max)
#undef SPECIALIZE_IMPL_UZILHFD
#undef SPECIALIZE_IMPL_HFD
#undef SPECIALIZE_IMPL
} // namespace rocm
} // namespace onnxruntime

38
onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.h
View file

@ -1,38 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <cstdint>
#include "core/providers/rocm/shared_inc/rocm_utils.h"
namespace onnxruntime {
namespace rocm {
template <typename T, typename VariadicElementwiseOpTag>
void Impl_General(
int32_t output_rank_or_simple_broadcast,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,
size_t count);
constexpr int32_t k_max_input_batch_size = 8;
template <typename T>
using InputBatchArray = TArray<const T*, k_max_input_batch_size>;
template <typename T, typename VariadicElementwiseOpTag>
void Impl_NoBroadcastInputBatch(
InputBatchArray<T> input_data_batch,
T* output_data,
size_t count);
} // namespace rocm
} // namespace onnxruntime

51
orttraining/orttraining/training_ops/rocm/activation/activations_grad.cc
View file

@ -1,51 +0,0 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/rocm/activation/activations_grad.h"
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace rocm {
#define REGISTER_ACTIVATION_GRAD_KERNEL(x, ver, domain, T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
x, \
domain, \
ver, \
T, \
kRocmExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()) \
.MayInplace(0, 0), \
x<T>);
#define BINARY_ELEMENTWISE_COMPUTE(x, T) \
template <> \
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
BinaryElementwisePreparation prepare; \
ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \
RocmAsyncBuffer<Ctx##x> func_ctx(this, MakeFuncCtx(), 1); \
if (!std::is_same<CtxNull, Ctx##x>::value) ORT_RETURN_IF_ERROR(func_ctx.CopyToGpu()); \
Impl_##x<typename ToHipType<T>::MappedType>( \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()), \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()), \
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()), \
func_ctx.GpuPtr(), prepare.output_tensor->Shape().Size()); \
return Status::OK(); \
}
#define ACTIVATION_GRAD_OP_TYPED(name, ver, domain, T) \
REGISTER_ACTIVATION_GRAD_KERNEL(name, ver, domain, T) \
BINARY_ELEMENTWISE_COMPUTE(name, T)
#define ACTIVATION_GRAD_OP_HFD(name, ver, domain) \
ACTIVATION_GRAD_OP_TYPED(name, ver, domain, MLFloat16) \
ACTIVATION_GRAD_OP_TYPED(name, ver, domain, float) \
ACTIVATION_GRAD_OP_TYPED(name, ver, domain, double)
ACTIVATION_GRAD_OP_HFD(GeluGrad, 1, kMSDomain);
ACTIVATION_GRAD_OP_HFD(FastGeluGrad, 1, kMSDomain);
ACTIVATION_GRAD_OP_HFD(ReluGrad, 1, kMSDomain);
} //namespace rocm
} // namespace onnxruntime

12
tools/ci_build/amd_hipify.py
View file

@ -35,10 +35,6 @@ contrib_ops_excluded_files = [
'math/bias_softmax.cc',
'math/bias_softmax.h',
'math/bias_softmax_impl.cu',
'math/binary_elementwise_ops.cc',
'math/binary_elementwise_ops.h',
'math/binary_elementwise_ops_impl.cu',
'math/binary_elementwise_ops_impl.h',
'math/complex_mul.cc',
'math/complex_mul.h',
'math/complex_mul_impl.cu',
@ -77,7 +73,6 @@ provider_excluded_files = [
'controlflow/loop.h',
'controlflow/scan.cc',
'controlflow/scan.h',
'cu_inc/binary_elementwise_impl.cuh',
'cu_inc/common.cuh',
'generator/constant_of_shape.cc',
'generator/constant_of_shape.h',
@ -90,9 +85,7 @@ provider_excluded_files = [
'math/einsum_utils/einsum_auxiliary_ops_diagonal.cu',
'math/einsum_utils/einsum_auxiliary_ops_diagonal.h',
'math/binary_elementwise_ops.cc',
'math/binary_elementwise_ops.h',
'math/binary_elementwise_ops_impl.cu',
'math/binary_elementwise_ops_impl.h',
'math/cumsum.cc',
'math/cumsum.h',
'math/cumsum_impl.cu',
@ -111,10 +104,6 @@ provider_excluded_files = [
'math/topk.h',
'math/topk_impl.cu',
'math/topk_impl.h',
'math/variadic_elementwise_ops.cc',
'math/variadic_elementwise_ops.h',
'math/variadic_elementwise_ops_impl.cu',
'math/variadic_elementwise_ops_impl.h',
'nn/batch_norm.cc',
'nn/batch_norm.h',
'nn/conv.cc',
@ -235,7 +224,6 @@ provider_excluded_files = [
]
training_ops_excluded_files = [
'activation/activations_grad.cc',
'collective/adasum_kernels.cc',
'collective/adasum_kernels.h',
'collective/nccl_common.cc',