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

Use TArray in AMD element-wise kernels, rather than manually copying memory to device.

Browse source
This commit is contained in:
Jesse Benson 2020-12-21 16:46:41 -08:00 committed by Jesse Benson
parent 1fc7f92f25
commit 52228a703c
13 changed files with 172 additions and 202 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -19,24 +19,23 @@ namespace rocm {
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(this); \
Prepare(context, &prepare); \
ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); \
#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.GpuPtr(), \
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.GpuPtr(), \
&prepare.rhs_padded_strides, \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()), \
prepare.fdm_output_strides.GpuPtr(), \
prepare.fdm_H, \
prepare.fdm_C, \
&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(); \
return Status::OK(); \
}
#define CONTRIB_BINARY_OP_TYPED(name, ver, T) \

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

@ -33,11 +33,16 @@ namespace rocm {
count); \
}
#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \
template void Impl_##x<T>(int32_t output_rank, \
const int64_t* lhs_padded_strides, const T* lhs_data, \
const int64_t* rhs_padded_strides, const T* rhs_data, \
const fast_divmod* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_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) \

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

@ -18,19 +18,18 @@ namespace rocm {
// 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 int64_t* lhs_padded_strides, \
const T* lhs_data, \
const int64_t* rhs_padded_strides, \
const T* rhs_data, \
const fast_divmod* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
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

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

@ -14,11 +14,11 @@ 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 int64_t* lhs_padded_strides,
const TArray<int64_t> lhs_padded_strides,
const T1* lhs_data,
const int64_t* rhs_padded_strides,
const TArray<int64_t> rhs_padded_strides,
const T2* rhs_data,
const fast_divmod* fdm_output_strides,
const TArray<fast_divmod> fdm_output_strides,
T* output_data,
const FuncT& functor,
HIP_LONG N) {
@ -34,7 +34,11 @@ __global__ void _BinaryElementWise(
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;
for (auto dim = 0; dim < output_rank; dim++) {
#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) {
@ -197,11 +201,11 @@ void BinaryElementWiseNoBroadcastImpl(
template <typename T, typename T1, typename T2, typename FuncT>
void BinaryElementWiseImpl(
int32_t output_rank_or_simple_broadcast,
const int64_t* lhs_padded_strides,
const TArray<int64_t>* lhs_padded_strides,
const T1* lhs_data,
const int64_t* rhs_padded_strides,
const TArray<int64_t>* rhs_padded_strides,
const T2* rhs_data,
const fast_divmod* fdm_output_strides,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,
@ -252,36 +256,36 @@ void BinaryElementWiseImpl(
func,
N);
} else {
if (lhs_padded_strides && rhs_padded_strides)
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_padded_strides,
lhs_data,
rhs_padded_strides,
*rhs_padded_strides,
rhs_data,
fdm_output_strides,
*fdm_output_strides,
output_data,
func,
N);
else if (lhs_padded_strides)
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_padded_strides,
lhs_data,
rhs_padded_strides,
*rhs_padded_strides,
rhs_data,
fdm_output_strides,
*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_padded_strides,
lhs_data,
rhs_padded_strides,
*rhs_padded_strides,
rhs_data,
fdm_output_strides,
*fdm_output_strides,
output_data,
func,
N);

70
onnxruntime/core/providers/rocm/math/binary_elementwise_ops.cc
View file

@ -142,16 +142,15 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, Bin
#define BINARY_ELEMENTWISE_COMPUTE(x, T) \
template <> \
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
BinaryElementwisePreparation prepare(this); \
Prepare(context, &prepare); \
ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); \
Impl_##x<typename ToHipType<T>::MappedType>( \
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.GpuPtr(), \
&prepare.lhs_padded_strides, \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()), \
prepare.rhs_padded_strides.GpuPtr(), \
&prepare.rhs_padded_strides, \
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.rhs_tensor->template Data<T>()), \
prepare.fdm_output_strides.GpuPtr(), \
&prepare.fdm_output_strides, \
prepare.fdm_H, \
prepare.fdm_C, \
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()), \
@ -159,7 +158,7 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, Bin
return Status::OK(); \
}
#define BINARY_OP_VERSIONED_TYPED(name, startver, endver, T) \
#define BINARY_OP_VERSIONED_TYPED(name, startver, endver, T) \
BINARY_ELEMENTWISE_REGISTER_KERNEL_VERSIONED_TYPED(name, startver, endver, T)
#define BINARY_OP_TYPED(name, ver, T) \
@ -189,16 +188,16 @@ Status BinaryElementwise<ShouldBroadcast>::Prepare(OpKernelContext* context, Bin
// D: double
// O: bool
#define BINARY_OP_VERSIONED_HFD(name, startver, endver) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, float) \
#define BINARY_OP_VERSIONED_HFD(name, startver, endver) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, MLFloat16) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, float) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, double)
#define BINARY_OP_VERSIONED_UZILHFD(name, startver, endver) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, uint32_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, uint64_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, int32_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, int64_t) \
#define BINARY_OP_VERSIONED_UZILHFD(name, startver, endver) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, uint32_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, uint64_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, int32_t) \
BINARY_OP_VERSIONED_TYPED(name, startver, endver, int64_t) \
BINARY_OP_VERSIONED_HFD(name, startver, endver)
#define BINARY_OP_HFD(name, ver) \
@ -309,11 +308,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) {
case on::TensorProto_DataType_INT32:
ImplT1_Pow<typename ToHipType<T>::MappedType, typename ToHipType<int32_t>::MappedType>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const typename ToHipType<int32_t>::MappedType*>(prepare.rhs_tensor->template Data<int32_t>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
@ -322,11 +321,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) {
case on::TensorProto_DataType_INT64:
ImplT1_Pow<typename ToHipType<T>::MappedType, typename ToHipType<int64_t>::MappedType>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const typename ToHipType<int64_t>::MappedType*>(prepare.rhs_tensor->template Data<int64_t>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
@ -335,11 +334,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) {
case on::TensorProto_DataType_FLOAT:
ImplT1_Pow<typename ToHipType<T>::MappedType, typename ToHipType<float>::MappedType>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const typename ToHipType<float>::MappedType*>(prepare.rhs_tensor->template Data<float>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
@ -348,11 +347,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) {
case on::TensorProto_DataType_DOUBLE:
ImplT1_Pow<typename ToHipType<T>::MappedType, typename ToHipType<double>::MappedType>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const typename ToHipType<T>::MappedType*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const typename ToHipType<double>::MappedType*>(prepare.rhs_tensor->template Data<double>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<typename ToHipType<T>::MappedType*>(prepare.output_tensor->template MutableData<T>()),
@ -367,8 +366,8 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) {
} // namespace pow12_internal
Status Pow::ComputeInternal(OpKernelContext* context) const {
BinaryElementwisePreparation prepare(this);
Prepare(context, &prepare);
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(Prepare(context, &prepare));
namespace on = ONNX_NAMESPACE;
using namespace pow12_internal;
@ -398,17 +397,16 @@ Status Pow::ComputeInternal(OpKernelContext* context) const {
//for other elementwise ops
template <typename T, typename HipT>
Status CompareFunction<T, HipT>::CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const {
BinaryElementwisePreparation prepare(this);
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(Prepare(context, &prepare));
size_t output_size = prepare.output_tensor->Shape().Size();
IAllocatorUniquePtr<T> output_buffer = GetScratchBuffer<T>(output_size);
Impl_Compare(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<ToHipType<bool>::MappedType*>(prepare.output_tensor->template MutableData<bool>()),

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

@ -14,28 +14,17 @@ 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
int32_t output_rank_or_simple_broadcast = 0; // for no_broadcast|left_scalar|right_scalar cases, output_rank uses SimpleBroadcast enums
// TODO: Unlike to CUDA EP, "pass by pointer" has been used in ROCM EP since ROCm compiler has a perf issue with "pass by value".
// The fix for ROCm compiler will be released in ROCm 3.10
RocmKernel::RocmAsyncBuffer<int64_t> lhs_padded_strides; // for lhs shape == output shape, this is nullptr
RocmKernel::RocmAsyncBuffer<int64_t> rhs_padded_strides; // for rhs shape == output shape, this is nullptr
RocmKernel::RocmAsyncBuffer<fast_divmod> fdm_output_strides;
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(const RocmKernel* op_kernel) : lhs_padded_strides(op_kernel),
rhs_padded_strides(op_kernel),
fdm_output_strides(op_kernel) {}
Status CopyToGpu() {
ORT_RETURN_IF_ERROR(lhs_padded_strides.CopyToGpu());
ORT_RETURN_IF_ERROR(rhs_padded_strides.CopyToGpu());
ORT_RETURN_IF_ERROR(fdm_output_strides.CopyToGpu());
return Status::OK();
}
BinaryElementwisePreparation() {}
Status BinaryElementwiseBroadcastPrepareHelper(const TensorShape& lhs_shape,
const TensorShape& rhs_shape,
@ -53,8 +42,8 @@ struct BinaryElementwisePreparation {
// 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);
? SimpleBroadcast::LeftScalar
: SimpleBroadcast::RightScalar);
return Status::OK();
}
@ -85,66 +74,35 @@ struct BinaryElementwisePreparation {
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.size_ = gsl::narrow_cast<int32_t>(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.size_ = gsl::narrow_cast<int32_t>(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.size_ = gsl::narrow_cast<int32_t>(out_rank);
// for (auto i = 0; i < out_rank; ++i) {
// fdm_output_strides[i] = fast_divmod(gsl::narrow_cast<int>(original_output_strides[i]));
// }
if (lhs_shape != output_shape) {
TensorPitches original_lhs_padded_strides(lhs_shape.GetDims(), out_rank);
lhs_padded_strides.AllocCpuPtr(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.CpuPtr()[i] = original_lhs_padded_strides[i];
} else {
lhs_padded_strides.CpuPtr()[i] = 0;
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.AllocCpuPtr(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.CpuPtr()[i] = original_rhs_padded_strides[i];
} else {
rhs_padded_strides.CpuPtr()[i] = 0;
rhs_padded_strides[i] = original_rhs_padded_strides[i];
}
}
}
fdm_output_strides.AllocCpuPtr(out_rank);
ORT_RETURN_IF_NOT(CalculateFdmStrides(fdm_output_strides.CpuSpan(), output_shape.GetDims()));
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();
}
@ -262,11 +220,11 @@ class CompareFunction : public BinaryElementwise<ShouldBroadcast> {
CompareFunction(const OpKernelInfo& info) : BinaryElementwise(info) {}
typedef void (*ImplCompare)(int32_t output_rank_or_simple_broadcast,
const int64_t* lhs_padded_strides,
const TArray<int64_t>* lhs_padded_strides,
const HipT* lhs_data,
const int64_t* rhs_padded_strides,
const TArray<int64_t>* rhs_padded_strides,
const HipT* rhs_data,
const fast_divmod* fdm_output_strides,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
bool* output_data,

25
onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.cu
View file

@ -57,21 +57,21 @@ namespace rocm {
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \
template void Impl_##x<T>(int32_t output_rank, \
const int64_t* lhs_padded_strides, const T* lhs_data, \
const int64_t* rhs_padded_strides, const T* rhs_data, \
const fast_divmod* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count);
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 SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, T1) \
template void ImplT1_##x<T, T1>(int32_t output_rank, \
const int64_t* lhs_padded_strides, const T* lhs_data, \
const int64_t* rhs_padded_strides, const T1* rhs_data, \
const fast_divmod* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count);
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);
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T2(x, T, T1, T2) \
template void ImplT2_##x<T, T1, T2>(int32_t output_rank, \
const int64_t* lhs_padded_strides, const T1* lhs_data, \
const int64_t* rhs_padded_strides, const T2* rhs_data, \
const fast_divmod* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, size_t count);
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 SPECIALIZED_BINARY_ELEMENTWISE_IMPL_UZILHFD(x) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, uint32_t) \
@ -82,6 +82,13 @@ namespace rocm {
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, float) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, double)
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1_ILHFD(x, T) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, int32_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, int64_t) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, half) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, float) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL_T1(x, T, double)
#define SPECIALIZED_BINARY_ELEMENTWISE_IMPL_OIL(x) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, bool) \
SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, int32_t) \

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

@ -35,11 +35,11 @@ namespace rocm {
template <typename T> \
void Impl_##name( \
int32_t output_rank_or_simple_broadcast, \
const int64_t* lhs_padded_strides, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const int64_t* rhs_padded_strides, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const fast_divmod* fdm_output_strides, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
@ -53,11 +53,11 @@ BINARY_OPS()
template <typename T, typename T1> \
void ImplT1_##name( \
int32_t output_rank_or_simple_broadcast, \
const int64_t* lhs_padded_strides, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const int64_t* rhs_padded_strides, \
const TArray<int64_t>* rhs_padded_strides, \
const T1* rhs_data, \
const fast_divmod* fdm_output_strides, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
@ -69,11 +69,11 @@ BINARY_ELEMENTWISE_IMPL_DECLARATION_T1(Pow);
template <typename T, typename T1, typename T2> \
void ImplT2_##name( \
int32_t output_rank_or_simple_broadcast, \
const int64_t* lhs_padded_strides, \
const TArray<int64_t>* lhs_padded_strides, \
const T1* lhs_data, \
const int64_t* rhs_padded_strides, \
const TArray<int64_t>* rhs_padded_strides, \
const T2* rhs_data, \
const fast_divmod* fdm_output_strides, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \

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

@ -39,19 +39,19 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
template <typename T>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::
BinaryImplDispatchTarget<T>::operator()(const RocmKernel* kernel, const Tensor& lhs, const Tensor& rhs, Tensor& output) const {
BinaryImplDispatchTarget<T>::operator()(const Tensor& lhs, const Tensor& rhs, Tensor& output) const {
using HipT = typename ToHipType<T>::MappedType;
BinaryElementwisePreparation prepare(kernel);
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.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
@ -64,23 +64,23 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
template <typename VariadicElementwiseOpTag, typename... SupportedElementTypes>
template <typename T>
Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>::
GeneralImplDispatchTarget<T>::operator()(const RocmKernel* kernel, const InputTensorVector& inputs, Tensor& output) const {
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(kernel);
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.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
@ -91,11 +91,11 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
Impl_General<HipT, VariadicElementwiseOpTag>(
prepare.output_rank_or_simple_broadcast,
prepare.lhs_padded_strides.GpuPtr(),
&prepare.lhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.lhs_tensor->template Data<T>()),
prepare.rhs_padded_strides.GpuPtr(),
&prepare.rhs_padded_strides,
reinterpret_cast<const HipT*>(prepare.rhs_tensor->template Data<T>()),
prepare.fdm_output_strides.GpuPtr(),
&prepare.fdm_output_strides,
prepare.fdm_H,
prepare.fdm_C,
reinterpret_cast<HipT*>(prepare.output_tensor->template MutableData<T>()),
@ -152,7 +152,7 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
// 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(this, input_tensors[0], input_tensors[1], output_tensor));
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors[0], input_tensors[1], output_tensor));
return Status::OK();
}
@ -177,7 +177,7 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
// special case for 2 inputs
if (input_count == 2) {
utils::MLTypeCallDispatcherRet<Status, BinaryImplDispatchTarget, SupportedElementTypes...> dispatcher(element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(this, input_tensors[0], input_tensors[1], output_tensor));
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors[0], input_tensors[1], output_tensor));
return Status::OK();
}
@ -186,7 +186,7 @@ Status VariadicElementwiseOp<VariadicElementwiseOpTag, SupportedElementTypes...>
{
utils::MLTypeCallDispatcherRet<Status, GeneralImplDispatchTarget, SupportedElementTypes...> dispatcher(
element_type);
ORT_RETURN_IF_ERROR(dispatcher.Invoke(this, input_tensors, output_tensor));
ORT_RETURN_IF_ERROR(dispatcher.Invoke(input_tensors, output_tensor));
}
return Status::OK();

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

@ -29,12 +29,12 @@ class VariadicElementwiseOp : public RocmKernel {
template <typename T>
struct BinaryImplDispatchTarget {
Status operator()(const RocmKernel* kernel, const Tensor& lhs, const Tensor& rhs, Tensor& output) const;
Status operator()(const Tensor& lhs, const Tensor& rhs, Tensor& output) const;
};
template <typename T>
struct GeneralImplDispatchTarget {
Status operator()(const RocmKernel* kernel, const InputTensorVector& inputs, Tensor& output) const;
Status operator()(const InputTensorVector& inputs, Tensor& output) const;
};
};

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

@ -21,11 +21,11 @@ struct VariadicElementwiseOpTraits;
\
static void ComputeFn( \
int32_t output_rank_or_simple_broadcast, \
const int64_t* lhs_padded_strides, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const int64_t* rhs_padded_strides, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const fast_divmod* fdm_output_strides, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \
@ -53,11 +53,11 @@ DEFINE_TRAITS(variadic_elementwise_ops::Max, Max)
template <typename T, typename VariadicElementwiseOpTag>
void Impl_General(
int32_t output_rank_or_simple_broadcast,
const int64_t* lhs_padded_strides,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const int64_t* rhs_padded_strides,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const fast_divmod* fdm_output_strides,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,
@ -92,11 +92,11 @@ void Impl_NoBroadcastInputBatch(
#define SPECIALIZE_IMPL(T, VariadicElementwiseOpTag) \
template void Impl_General<T, VariadicElementwiseOpTag>( \
int32_t output_rank_or_simple_broadcast, \
const int64_t* lhs_padded_strides, \
const TArray<int64_t>* lhs_padded_strides, \
const T* lhs_data, \
const int64_t* rhs_padded_strides, \
const TArray<int64_t>* rhs_padded_strides, \
const T* rhs_data, \
const fast_divmod* fdm_output_strides, \
const TArray<fast_divmod>* fdm_output_strides, \
const fast_divmod& fdm_H, \
const fast_divmod& fdm_C, \
T* output_data, \

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

@ -13,11 +13,11 @@ namespace rocm {
template <typename T, typename VariadicElementwiseOpTag>
void Impl_General(
int32_t output_rank_or_simple_broadcast,
const int64_t* lhs_padded_strides,
const TArray<int64_t>* lhs_padded_strides,
const T* lhs_data,
const int64_t* rhs_padded_strides,
const TArray<int64_t>* rhs_padded_strides,
const T* rhs_data,
const fast_divmod* fdm_output_strides,
const TArray<fast_divmod>* fdm_output_strides,
const fast_divmod& fdm_H,
const fast_divmod& fdm_C,
T* output_data,

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

@ -19,19 +19,19 @@ namespace rocm {
.MayInplace(0, 0), \
x<T>);
#define BINARY_ELEMENTWISE_COMPUTE(x, T) \
template <> \
Status x<T>::ComputeInternal(OpKernelContext* context) const { \
BinaryElementwisePreparation prepare(this); \
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()); \
#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(); \
func_ctx.GpuPtr(), prepare.output_tensor->Shape().Size()); \
return Status::OK(); \
}
#define ACTIVATION_GRAD_OP_TYPED(name, ver, domain, T) \