diff --git a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.cc b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.cc index 9f3eead04a..c779db9064 100644 --- a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.cc +++ b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops.cc @@ -19,24 +19,23 @@ namespace rocm { KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ x); -#define CONTRIB_BINARY_ELEMENTWISE_COMPUTE(x, T) \ - template <> \ - Status x::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::ComputeInternal(OpKernelContext* context) const { \ + BinaryElementwisePreparation prepare; \ + ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \ Impl_##x::MappedType>( \ - prepare.output_rank_or_simple_broadcast, \ - prepare.lhs_padded_strides.GpuPtr(), \ + prepare.output_rank_or_simple_broadcast, \ + &prepare.lhs_padded_strides, \ reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), \ - prepare.rhs_padded_strides.GpuPtr(), \ + &prepare.rhs_padded_strides, \ reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), \ - prepare.fdm_output_strides.GpuPtr(), \ - prepare.fdm_H, \ - prepare.fdm_C, \ + &prepare.fdm_output_strides, \ + prepare.fdm_H, \ + prepare.fdm_C, \ reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), \ prepare.output_tensor->Shape().Size()); \ - return Status::OK(); \ + return Status::OK(); \ } #define CONTRIB_BINARY_OP_TYPED(name, ver, T) \ diff --git a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.cu b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.cu index fc325ac549..20d6c02ed8 100644 --- a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.cu @@ -33,11 +33,16 @@ namespace rocm { count); \ } -#define CONTRIB_SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \ -template void Impl_##x(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(int32_t output_rank, \ + const TArray* lhs_padded_strides, \ + const T* lhs_data, \ + const TArray* rhs_padded_strides, \ + const T* rhs_data, \ + const TArray* 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) \ diff --git a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.h b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.h index 111984e995..cb31e1f9dd 100644 --- a/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.h +++ b/onnxruntime/contrib_ops/rocm/math/binary_elementwise_ops_impl.h @@ -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 \ - 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 \ + void Impl_##name( \ + int32_t output_rank_or_simple_broadcast, \ + const TArray* lhs_padded_strides, \ + const T* lhs_data, \ + const TArray* rhs_padded_strides, \ + const T* rhs_data, \ + const TArray* 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 diff --git a/onnxruntime/core/providers/rocm/cu_inc/binary_elementwise_impl.cuh b/onnxruntime/core/providers/rocm/cu_inc/binary_elementwise_impl.cuh index 7d98a7ffe3..54424aabdd 100644 --- a/onnxruntime/core/providers/rocm/cu_inc/binary_elementwise_impl.cuh +++ b/onnxruntime/core/providers/rocm/cu_inc/binary_elementwise_impl.cuh @@ -14,11 +14,11 @@ template __global__ void _BinaryElementWise( int32_t output_rank, - const int64_t* lhs_padded_strides, + const TArray lhs_padded_strides, const T1* lhs_data, - const int64_t* rhs_padded_strides, + const TArray rhs_padded_strides, const T2* rhs_data, - const fast_divmod* fdm_output_strides, + const TArray 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 void BinaryElementWiseImpl( int32_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, + const TArray* lhs_padded_strides, const T1* lhs_data, - const int64_t* rhs_padded_strides, + const TArray* rhs_padded_strides, const T2* rhs_data, - const fast_divmod* fdm_output_strides, + const TArray* 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), 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), 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), 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); diff --git a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.cc index 0acd228134..8e7ebebfc1 100644 --- a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.cc @@ -142,16 +142,15 @@ Status BinaryElementwise::Prepare(OpKernelContext* context, Bin #define BINARY_ELEMENTWISE_COMPUTE(x, T) \ template <> \ Status x::ComputeInternal(OpKernelContext* context) const { \ - BinaryElementwisePreparation prepare(this); \ - Prepare(context, &prepare); \ - ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); \ - Impl_##x::MappedType>( \ + BinaryElementwisePreparation prepare; \ + ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \ + Impl_##x::MappedType>( \ prepare.output_rank_or_simple_broadcast, \ - prepare.lhs_padded_strides.GpuPtr(), \ + &prepare.lhs_padded_strides, \ reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), \ - prepare.rhs_padded_strides.GpuPtr(), \ + &prepare.rhs_padded_strides, \ reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), \ - prepare.fdm_output_strides.GpuPtr(), \ + &prepare.fdm_output_strides, \ prepare.fdm_H, \ prepare.fdm_C, \ reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), \ @@ -159,7 +158,7 @@ Status BinaryElementwise::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::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::MappedType, typename ToHipType::MappedType>( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), @@ -322,11 +321,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) { case on::TensorProto_DataType_INT64: ImplT1_Pow::MappedType, typename ToHipType::MappedType>( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), @@ -335,11 +334,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) { case on::TensorProto_DataType_FLOAT: ImplT1_Pow::MappedType, typename ToHipType::MappedType>( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), @@ -348,11 +347,11 @@ Status DispatchOnFirstArg(const BinaryElementwisePreparation& prepare) { case on::TensorProto_DataType_DOUBLE: ImplT1_Pow::MappedType, typename ToHipType::MappedType>( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), @@ -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 Status CompareFunction::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 output_buffer = GetScratchBuffer(output_size); + Impl_Compare( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), diff --git a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.h b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.h index 4f957bf62e..5687594030 100644 --- a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.h +++ b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops.h @@ -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 lhs_padded_strides; // for lhs shape == output shape, this is nullptr - RocmKernel::RocmAsyncBuffer rhs_padded_strides; // for rhs shape == output shape, this is nullptr - RocmKernel::RocmAsyncBuffer fdm_output_strides; + TArray lhs_padded_strides; + TArray rhs_padded_strides; + TArray 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(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(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(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(out_rank); - // for (auto i = 0; i < out_rank; ++i) { - // fdm_output_strides[i] = fast_divmod(gsl::narrow_cast(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(original_output_strides[i])); + } return Status::OK(); } @@ -262,11 +220,11 @@ class CompareFunction : public BinaryElementwise { CompareFunction(const OpKernelInfo& info) : BinaryElementwise(info) {} typedef void (*ImplCompare)(int32_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, + const TArray* lhs_padded_strides, const HipT* lhs_data, - const int64_t* rhs_padded_strides, + const TArray* rhs_padded_strides, const HipT* rhs_data, - const fast_divmod* fdm_output_strides, + const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, bool* output_data, diff --git a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.cu b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.cu index a469398a95..f4b3f90762 100644 --- a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.cu @@ -57,21 +57,21 @@ namespace rocm { #define SPECIALIZED_BINARY_ELEMENTWISE_IMPL(x, T) \ template void Impl_##x(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* lhs_padded_strides, const T* lhs_data, \ + const TArray* rhs_padded_strides, const T* rhs_data, \ + const TArray* 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(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* lhs_padded_strides, const T* lhs_data, \ + const TArray* rhs_padded_strides, const T1* rhs_data, \ + const TArray* 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(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* lhs_padded_strides, const T1* lhs_data, \ + const TArray* rhs_padded_strides, const T2* rhs_data, \ + const TArray* 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) \ diff --git a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.h b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.h index 2175303363..4ce74516d7 100644 --- a/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.h +++ b/onnxruntime/core/providers/rocm/math/binary_elementwise_ops_impl.h @@ -35,11 +35,11 @@ namespace rocm { template \ void Impl_##name( \ int32_t output_rank_or_simple_broadcast, \ - const int64_t* lhs_padded_strides, \ + const TArray* lhs_padded_strides, \ const T* lhs_data, \ - const int64_t* rhs_padded_strides, \ + const TArray* rhs_padded_strides, \ const T* rhs_data, \ - const fast_divmod* fdm_output_strides, \ + const TArray* fdm_output_strides, \ const fast_divmod& fdm_H, \ const fast_divmod& fdm_C, \ T* output_data, \ @@ -53,11 +53,11 @@ BINARY_OPS() template \ void ImplT1_##name( \ int32_t output_rank_or_simple_broadcast, \ - const int64_t* lhs_padded_strides, \ + const TArray* lhs_padded_strides, \ const T* lhs_data, \ - const int64_t* rhs_padded_strides, \ + const TArray* rhs_padded_strides, \ const T1* rhs_data, \ - const fast_divmod* fdm_output_strides, \ + const TArray* 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 \ void ImplT2_##name( \ int32_t output_rank_or_simple_broadcast, \ - const int64_t* lhs_padded_strides, \ + const TArray* lhs_padded_strides, \ const T1* lhs_data, \ - const int64_t* rhs_padded_strides, \ + const TArray* rhs_padded_strides, \ const T2* rhs_data, \ - const fast_divmod* fdm_output_strides, \ + const TArray* fdm_output_strides, \ const fast_divmod& fdm_H, \ const fast_divmod& fdm_C, \ T* output_data, \ diff --git a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.cc b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.cc index c5e86343fe..4e1a72b7a7 100644 --- a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.cc +++ b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.cc @@ -39,19 +39,19 @@ Status VariadicElementwiseOp template template Status VariadicElementwiseOp:: - BinaryImplDispatchTarget::operator()(const RocmKernel* kernel, const Tensor& lhs, const Tensor& rhs, Tensor& output) const { + BinaryImplDispatchTarget::operator()(const Tensor& lhs, const Tensor& rhs, Tensor& output) const { using HipT = typename ToHipType::MappedType; - BinaryElementwisePreparation prepare(kernel); + BinaryElementwisePreparation prepare; ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(&lhs, &rhs, &output, &prepare)); Impl_General( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast(prepare.output_tensor->template MutableData()), @@ -64,23 +64,23 @@ Status VariadicElementwiseOp template template Status VariadicElementwiseOp:: - GeneralImplDispatchTarget::operator()(const RocmKernel* kernel, const InputTensorVector& inputs, Tensor& output) const { + GeneralImplDispatchTarget::operator()(const InputTensorVector& inputs, Tensor& output) const { assert(inputs.size() > 1); using HipT = typename ToHipType::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(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast(prepare.output_tensor->template MutableData()), @@ -91,11 +91,11 @@ Status VariadicElementwiseOp Impl_General( prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), + &prepare.lhs_padded_strides, reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), + &prepare.rhs_padded_strides, reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), + &prepare.fdm_output_strides, prepare.fdm_H, prepare.fdm_C, reinterpret_cast(prepare.output_tensor->template MutableData()), @@ -152,7 +152,7 @@ Status VariadicElementwiseOp // special case for no broadcasting and 2 inputs if (input_count == 2) { utils::MLTypeCallDispatcherRet 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 // special case for 2 inputs if (input_count == 2) { utils::MLTypeCallDispatcherRet 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 { utils::MLTypeCallDispatcherRet 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(); diff --git a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.h b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.h index 81478003f5..0dd485c3e9 100644 --- a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.h +++ b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops.h @@ -29,12 +29,12 @@ class VariadicElementwiseOp : public RocmKernel { template 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 struct GeneralImplDispatchTarget { - Status operator()(const RocmKernel* kernel, const InputTensorVector& inputs, Tensor& output) const; + Status operator()(const InputTensorVector& inputs, Tensor& output) const; }; }; diff --git a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.cu b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.cu index 8c81eece51..189e8701c6 100644 --- a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.cu @@ -21,11 +21,11 @@ struct VariadicElementwiseOpTraits; \ static void ComputeFn( \ int32_t output_rank_or_simple_broadcast, \ - const int64_t* lhs_padded_strides, \ + const TArray* lhs_padded_strides, \ const T* lhs_data, \ - const int64_t* rhs_padded_strides, \ + const TArray* rhs_padded_strides, \ const T* rhs_data, \ - const fast_divmod* fdm_output_strides, \ + const TArray* 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 void Impl_General( int32_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, + const TArray* lhs_padded_strides, const T* lhs_data, - const int64_t* rhs_padded_strides, + const TArray* rhs_padded_strides, const T* rhs_data, - const fast_divmod* fdm_output_strides, + const TArray* 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( \ int32_t output_rank_or_simple_broadcast, \ - const int64_t* lhs_padded_strides, \ + const TArray* lhs_padded_strides, \ const T* lhs_data, \ - const int64_t* rhs_padded_strides, \ + const TArray* rhs_padded_strides, \ const T* rhs_data, \ - const fast_divmod* fdm_output_strides, \ + const TArray* fdm_output_strides, \ const fast_divmod& fdm_H, \ const fast_divmod& fdm_C, \ T* output_data, \ diff --git a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.h b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.h index ad1ff02f2f..adc5ca07dc 100644 --- a/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.h +++ b/onnxruntime/core/providers/rocm/math/variadic_elementwise_ops_impl.h @@ -13,11 +13,11 @@ namespace rocm { template void Impl_General( int32_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, + const TArray* lhs_padded_strides, const T* lhs_data, - const int64_t* rhs_padded_strides, + const TArray* rhs_padded_strides, const T* rhs_data, - const fast_divmod* fdm_output_strides, + const TArray* fdm_output_strides, const fast_divmod& fdm_H, const fast_divmod& fdm_C, T* output_data, diff --git a/orttraining/orttraining/training_ops/rocm/activation/activations_grad.cc b/orttraining/orttraining/training_ops/rocm/activation/activations_grad.cc index 9c1b68fe30..a0e0d708f0 100644 --- a/orttraining/orttraining/training_ops/rocm/activation/activations_grad.cc +++ b/orttraining/orttraining/training_ops/rocm/activation/activations_grad.cc @@ -19,19 +19,19 @@ namespace rocm { .MayInplace(0, 0), \ x); -#define BINARY_ELEMENTWISE_COMPUTE(x, T) \ - template <> \ - Status x::ComputeInternal(OpKernelContext* context) const { \ - BinaryElementwisePreparation prepare(this); \ - Prepare(context, &prepare); \ - RocmAsyncBuffer func_ctx(this, MakeFuncCtx(), 1); \ - if (!std::is_same::value) ORT_RETURN_IF_ERROR(func_ctx.CopyToGpu()); \ +#define BINARY_ELEMENTWISE_COMPUTE(x, T) \ + template <> \ + Status x::ComputeInternal(OpKernelContext* context) const { \ + BinaryElementwisePreparation prepare; \ + ORT_RETURN_IF_ERROR(Prepare(context, &prepare)); \ + RocmAsyncBuffer func_ctx(this, MakeFuncCtx(), 1); \ + if (!std::is_same::value) ORT_RETURN_IF_ERROR(func_ctx.CopyToGpu()); \ Impl_##x::MappedType>( \ reinterpret_cast::MappedType*>(prepare.lhs_tensor->template Data()), \ reinterpret_cast::MappedType*>(prepare.rhs_tensor->template Data()), \ reinterpret_cast::MappedType*>(prepare.output_tensor->template MutableData()), \ - 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) \