From 00e24ae4fe4dde51d0ab1adfd3070b18222c2eaa Mon Sep 17 00:00:00 2001 From: Hector Li Date: Mon, 7 Oct 2019 13:17:49 -0700 Subject: [PATCH] refactor Cuda Ops Sum, Max, Min, remove dup code (#1946) refactor Cuda Ops Sum, Max, Min, remove dup code --- .../cuda/math/binary_elementwise_ops.cc | 164 +++++------------- .../cuda/math/binary_elementwise_ops.h | 74 +++++--- 2 files changed, 89 insertions(+), 149 deletions(-) diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc index 8d08e9047b..147762550d 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.cc @@ -223,33 +223,37 @@ BINARY_LOGICALOP_TYPED(Or, 7, bool) BINARY_LOGICALOP_TYPED(Xor, 7, bool) BINARY_OP_HFD(PRelu, 7) -template -Status Sum::ComputeInternal(OpKernelContext* context) const { - typedef typename ToCudaType::MappedType CudaT; +template +Status VariadicInputBase::ComputeMethod(OpKernelContext* context, ImplCompute Impl_Compute) 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"); + auto lhs_tensor = context->Input(0); if (input_count == 1) { - auto input_tensor = context->Input(0); - const auto& input_shape = input_tensor->Shape(); + const auto& input_shape = lhs_tensor->Shape(); auto output_tensor = context->Output(0, input_shape); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice)); + if (lhs_tensor->DataRaw() != output_tensor->DataRaw()) { + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), lhs_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice)); + } } else { // compute output shape first, using broadcast rule TensorShape output_shape; - ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input(0)->Shape(), context->Input(1)->Shape(), output_shape)); - for (int index = 2; index < input_count; index++) { - TensorShape previous_output_shape = output_shape; + TensorShape previous_output_shape = lhs_tensor->Shape(); + for (int index = 1; index < input_count; index++) { ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input(index)->Shape(), output_shape)); + previous_output_shape = output_shape; } Tensor* output_tensor = context->Output(0, output_shape); BinaryElementwisePreparation prepare(this); + + auto rhs_tensor = context->Input(1); if (input_count == 2) { // special case for 2 tensors to avoid memset zero - ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(context->Input(0), context->Input(1), output_tensor, &prepare)); - Impl_Add( + ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(lhs_tensor, rhs_tensor, output_tensor, &prepare)); + ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); + Impl_Compute( prepare.output_rank_or_simple_broadcast, prepare.lhs_padded_strides.GpuPtr(), reinterpret_cast(prepare.lhs_tensor->template Data()), @@ -263,9 +267,24 @@ Status Sum::ComputeInternal(OpKernelContext* context) const { } else { // for more than 2 inputs, we need to accumulate into output tensor, as the shape from input0 + input1 might be different from output shape CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT))); - for (int index = 0; index < input_count; index++) { + + ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, lhs_tensor, output_tensor, &prepare)); + ORT_RETURN_IF_ERROR(prepare.CopyToGpu()); + Impl_Add( + prepare.output_rank_or_simple_broadcast, + prepare.lhs_padded_strides.GpuPtr(), + reinterpret_cast(prepare.lhs_tensor->template Data()), + prepare.rhs_padded_strides.GpuPtr(), + reinterpret_cast(prepare.rhs_tensor->template Data()), + prepare.fdm_output_strides.GpuPtr(), + prepare.fdm_H, + prepare.fdm_C, + reinterpret_cast(prepare.output_tensor->template MutableData()), + prepare.output_tensor->Shape().Size()); + + for (int index = 1; index < input_count; index++) { ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input(index), output_tensor, &prepare)); - Impl_Add( + Impl_Compute( prepare.output_rank_or_simple_broadcast, prepare.lhs_padded_strides.GpuPtr(), reinterpret_cast(prepare.lhs_tensor->template Data()), @@ -281,133 +300,34 @@ Status Sum::ComputeInternal(OpKernelContext* context) const { } return Status::OK(); } +template +Status Sum::ComputeInternal(OpKernelContext* context) const { + + this->ComputeMethod(context, &Impl_Add); + + return Status::OK(); +} template Status Max::ComputeInternal(OpKernelContext* context) const { - typedef typename ToCudaType::MappedType CudaT; - 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"); - if (input_count == 1) { - auto input_tensor = context->Input(0); - const auto& input_shape = input_tensor->Shape(); - auto output_tensor = context->Output(0, input_shape); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice)); - } else { - // compute output shape first, using broadcast rule - TensorShape output_shape; - ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input(0)->Shape(), context->Input(1)->Shape(), output_shape)); - for (int index = 2; index < input_count; index++) { - TensorShape previous_output_shape = output_shape; - ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input(index)->Shape(), output_shape)); - } - Tensor* output_tensor = context->Output(0, output_shape); - BinaryElementwisePreparation prepare(this); + this->ComputeMethod(context, &Impl_Max); - // More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly - CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT))); - ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input(0), output_tensor, &prepare)); - Impl_Add( - prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), - prepare.fdm_H, - prepare.fdm_C, - reinterpret_cast(prepare.output_tensor->template MutableData()), - prepare.output_tensor->Shape().Size()); - for (int index = 1; index < input_count; index++) { - ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input(index), output_tensor, &prepare)); - Impl_Max( - prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), - prepare.fdm_H, - prepare.fdm_C, - reinterpret_cast(prepare.output_tensor->template MutableData()), - prepare.output_tensor->Shape().Size()); - } - } return Status::OK(); } template Status Min::ComputeInternal(OpKernelContext* context) const { - typedef typename ToCudaType::MappedType CudaT; - 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"); - if (input_count == 1) { - auto input_tensor = context->Input(0); - const auto& input_shape = input_tensor->Shape(); - auto output_tensor = context->Output(0, input_shape); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(output_tensor->MutableDataRaw(), input_tensor->DataRaw(), sizeof(CudaT) * input_shape.Size(), cudaMemcpyDeviceToDevice)); - } else { - // compute output shape first, using broadcast rule - TensorShape output_shape; - ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, context->Input(0)->Shape(), context->Input(1)->Shape(), output_shape)); - for (int index = 2; index < input_count; index++) { - TensorShape previous_output_shape = output_shape; - ORT_RETURN_IF_ERROR(ComputeOutputShape(node_name, previous_output_shape, context->Input(index)->Shape(), output_shape)); - } - Tensor* output_tensor = context->Output(0, output_shape); - BinaryElementwisePreparation prepare(this); + this->ComputeMethod(context, &Impl_Min); - // More than 2 inputs, set output to 0, add input0 to output, so that input0 can be broadcast with output shape correctly - CUDA_RETURN_IF_ERROR(cudaMemset(output_tensor->MutableDataRaw(), 0, output_shape.Size() * sizeof(CudaT))); - ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input(0), output_tensor, &prepare)); - Impl_Add( - prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), - prepare.fdm_H, - prepare.fdm_C, - reinterpret_cast(prepare.output_tensor->template MutableData()), - prepare.output_tensor->Shape().Size()); - for (int index = 1; index < input_count; index++) { - ORT_RETURN_IF_ERROR(BinaryElementwiseBroadcastPrepare(output_tensor, context->Input(index), output_tensor, &prepare)); - Impl_Min( - prepare.output_rank_or_simple_broadcast, - prepare.lhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.lhs_tensor->template Data()), - prepare.rhs_padded_strides.GpuPtr(), - reinterpret_cast(prepare.rhs_tensor->template Data()), - prepare.fdm_output_strides.GpuPtr(), - prepare.fdm_H, - prepare.fdm_C, - reinterpret_cast(prepare.output_tensor->template MutableData()), - prepare.output_tensor->Shape().Size()); - } - } return Status::OK(); } //Greater op output tensor type is bool, so it cannot directly fit in the macros //for other elementwise ops template -Status CompareFunction::CompareMethod(OpKernelContext* context, void (*Impl_Compare)( - size_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, - const CudaT* lhs_data, - const int64_t* rhs_padded_strides, - const CudaT* rhs_data, - const fast_divmod* fdm_output_strides, - const fast_divmod& fdm_H, - const fast_divmod& fdm_C, - CudaT* output_data, - size_t count)) const { +Status CompareFunction::CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const { BinaryElementwisePreparation prepare(this); Prepare(context, &prepare); size_t output_size = prepare.output_tensor->Shape().Size(); diff --git a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h index e662f99ebc..958fa5a196 100644 --- a/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h +++ b/onnxruntime/core/providers/cuda/math/binary_elementwise_ops.h @@ -186,31 +186,50 @@ class PRelu final : public BinaryElementwise { Status ComputeInternal(OpKernelContext* context) const override; }; +template +class VariadicInputBase : public CudaKernel { + public: + VariadicInputBase(const OpKernelInfo& info) : CudaKernel(info) {} + + Status ComputeInternal(OpKernelContext*) const override { + return Status(common::ONNXRUNTIME, common::FAIL); // should not reach here + } + + typedef void (*ImplCompute)(size_t output_rank_or_simple_broadcast, + const int64_t* lhs_padded_strides, + const CudaT* lhs_data, + const int64_t* rhs_padded_strides, + const CudaT* rhs_data, + const fast_divmod* fdm_output_strides, + const fast_divmod& fdm_H, + const fast_divmod& fdm_C, + CudaT* output_data, + size_t count); + + Status ComputeMethod(OpKernelContext* context, ImplCompute Impl_Compute) const; +}; + // Sum allows varadic inputs, and it uses binary elementwise Add in implementation template -class Sum final : public CudaKernel { +class Sum final : public VariadicInputBase::MappedType> { public: - Sum(const OpKernelInfo& info) : CudaKernel(info) { - } - - Status ComputeInternal(OpKernelContext* context) const override; -}; - - -template -class Max final : public CudaKernel { - public: - Max(const OpKernelInfo& info) : CudaKernel(info) { - } + Sum(const OpKernelInfo& info) : VariadicInputBase::MappedType>(info) {} Status ComputeInternal(OpKernelContext* context) const override; }; template -class Min final : public CudaKernel { +class Max final : public VariadicInputBase::MappedType> { public: - Min(const OpKernelInfo& info) : CudaKernel(info) { - } + Max(const OpKernelInfo& info) : VariadicInputBase::MappedType>(info) {} + + Status ComputeInternal(OpKernelContext* context) const override; +}; + +template +class Min final : public VariadicInputBase::MappedType> { + public: + Min(const OpKernelInfo& info) : VariadicInputBase::MappedType>(info) {} Status ComputeInternal(OpKernelContext* context) const override; }; @@ -220,17 +239,18 @@ class CompareFunction : public BinaryElementwise { public: CompareFunction(const OpKernelInfo& info) : BinaryElementwise(info) {} - Status CompareMethod(OpKernelContext* context, void (*Impl_Compare)( - size_t output_rank_or_simple_broadcast, - const int64_t* lhs_padded_strides, - const CudaT* lhs_data, - const int64_t* rhs_padded_strides, - const CudaT* rhs_data, - const fast_divmod* fdm_output_strides, - const fast_divmod& fdm_H, - const fast_divmod& fdm_C, - CudaT* output_data, - size_t count)) const; + typedef void (*ImplCompare)(size_t output_rank_or_simple_broadcast, + const int64_t* lhs_padded_strides, + const CudaT* lhs_data, + const int64_t* rhs_padded_strides, + const CudaT* rhs_data, + const fast_divmod* fdm_output_strides, + const fast_divmod& fdm_H, + const fast_divmod& fdm_C, + CudaT* output_data, + size_t count); + + Status CompareMethod(OpKernelContext* context, ImplCompare Impl_Compare) const; }; template