From 655f490c9564c27543d53121c20a157cb2aba76b Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Wed, 9 Feb 2022 07:39:57 +0800 Subject: [PATCH] Remove BFloat16 Specialized Code for ReduceSum (#10476) --- .../cuda/math/unary_elementwise_ops_impl.cu | 8 +- .../providers/cuda/reduction/reduction_ops.cc | 107 +----------------- .../core/providers/rocm/miopen_common.cc | 14 +++ .../core/providers/rocm/miopen_common.h | 12 ++ .../providers/rocm/reduction/reduction_ops.cc | 107 +----------------- 5 files changed, 34 insertions(+), 214 deletions(-) diff --git a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu index 9bb5a2d855..bb17158e29 100644 --- a/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu +++ b/onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu @@ -55,6 +55,10 @@ UNARY_OPS() SPECIALIZED_UNARY_ELEMENTWISE_IMPL(name, float) \ SPECIALIZED_UNARY_ELEMENTWISE_IMPL(name, double) +#define SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFDB(name) \ + SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(name) \ + SPECIALIZED_UNARY_ELEMENTWISE_IMPL(name, BFloat16) + #define SPECIALIZED_UNARY_ELEMENTWISE_IMPL_CSILHFD(name) \ SPECIALIZED_UNARY_ELEMENTWISE_IMPL(name, int8_t) \ SPECIALIZED_UNARY_ELEMENTWISE_IMPL(name, int16_t) \ @@ -75,8 +79,8 @@ SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Floor) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Ceil) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Reciprocal) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Sqrt) -SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Log) -SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Exp) +SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFDB(Log) +SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFDB(Exp) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Erf) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Round) SPECIALIZED_UNARY_ELEMENTWISE_IMPL_HFD(Sin) diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 5b05548728..7726c63573 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -521,7 +521,7 @@ Status ReduceComputeCore(CUDAExecutionProvider& cuda_ep, const Tensor& input, Pr } CudnnReduceDescriptor reduce_desc; - if (std::is_same::value) { + ORT_IF_CONSTEXPR (std::is_same::value || std::is_same::value) { ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, CudnnTensor::GetDataType(), ReduceTensorIndices)); } else { ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, ReduceTensorIndices)); @@ -835,111 +835,6 @@ SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int64_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int8_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(uint8_t) -template <> -template <> -Status ReduceKernel::ComputeImpl( - OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { - typedef typename ToCudaType::MappedType CudaT; - const Tensor* X = ctx->Input(0); - TensorShapeVector axes; - size_t num_inputs = ctx->InputCount(); - if (num_inputs == 2) { - const Tensor* axes_tensor = ctx->Input(1); - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); - ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); - auto nDims = static_cast(axes_tensor->Shape()[0]); - const auto* data = axes_tensor->template Data(); - axes.assign(data, data + nDims); - } else { - axes.assign(axes_.begin(), axes_.end()); - } - - if (axes.empty() && noop_with_empty_axes_) { - auto* Y = ctx->Output(0, X->Shape()); - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), - X->SizeInBytes(), cudaMemcpyDeviceToDevice, Stream())); - return Status::OK(); - } - - PrepareReduceMetadata prepare_reduce_metadata; - ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, axes, prepare_reduce_metadata)); - - Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); - - int64_t input_count = prepare_reduce_metadata.input_count; - int64_t output_count = prepare_reduce_metadata.output_count; - auto& input_dims_cudnn = prepare_reduce_metadata.input_dims_cudnn; - auto& output_dims_cudnn = prepare_reduce_metadata.output_dims_cudnn; - - if (input_count == 0) { - assert(Y->Shape().Size() == 0); - return Status::OK(); - } - - if (input_count == output_count) { - if (Y->template MutableData() != X->template Data()) { - CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), - input_count * sizeof(BFloat16), cudaMemcpyDeviceToDevice, Stream())); - } - return Status::OK(); - } - - if (fast_reduction_ && !ctx->GetUseDeterministicCompute()) { - int m{}, n{}; - const auto applicable_matrix_reduction = - get_applicable_matrix_reduction(cudnn_reduce_op, X->Shape().GetDims(), axes, m, n); - switch (applicable_matrix_reduction) { - case ApplicableMatrixReduction::Rows: { - return reduce_matrix_rows(Stream(), reinterpret_cast(X->template Data()), - reinterpret_cast(Y->template MutableData()), m, n); - } - case ApplicableMatrixReduction::Columns: { - const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size(m, n); - auto buffer = cuda_ep_->GetScratchBuffer(buffer_size_bytes); - return reduce_matrix_columns(Stream(), reinterpret_cast(X->template Data()), - reinterpret_cast(Y->template MutableData()), m, n, buffer.get(), - buffer_size_bytes); - } - default: - break; - } - } - - CUDA_RETURN_IF_ERROR(cudaMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream())); - - size_t indices_bytes = 0; - size_t workspace_bytes = 0; - CudnnTensor input_tensor; - CudnnTensor output_tensor; - CudnnReduceDescriptor reduce_desc; - - cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT; - IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); - Impl_Cast(Stream(), reinterpret_cast(X->template Data()), temp_X.get(), - X->Shape().Size()); - - ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, cudnn_type_X, CUDNN_REDUCE_TENSOR_NO_INDICES)); - ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_cudnn, cudnn_type_X)); - ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_cudnn, cudnn_type_X)); - CUDNN_RETURN_IF_ERROR( - cudnnGetReductionIndicesSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); - CUDNN_RETURN_IF_ERROR( - cudnnGetReductionWorkspaceSize(CudnnHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); - IAllocatorUniquePtr indices_cuda = GetScratchBuffer(indices_bytes); - IAllocatorUniquePtr workspace_cuda = GetScratchBuffer(workspace_bytes); - - const auto one = Consts::One; - const auto zero = Consts::Zero; - auto temp_Y = GetScratchBuffer(output_count); - CUDNN_RETURN_IF_ERROR(cudnnReduceTensor(CudnnHandle(), reduce_desc, indices_cuda.get(), indices_bytes, - workspace_cuda.get(), workspace_bytes, &one, input_tensor, temp_X.get(), - &zero, output_tensor, temp_Y.get())); - - Impl_Cast(Stream(), temp_Y.get(), reinterpret_cast(Y->template MutableData()), output_count); - - return Status::OK(); -} - namespace ReductionOps { template diff --git a/onnxruntime/core/providers/rocm/miopen_common.cc b/onnxruntime/core/providers/rocm/miopen_common.cc index 88c92df144..24f23853f1 100644 --- a/onnxruntime/core/providers/rocm/miopen_common.cc +++ b/onnxruntime/core/providers/rocm/miopen_common.cc @@ -122,10 +122,18 @@ const float Consts::Zero = 0; const float Consts::One = 1; +const float Consts::Zero = 0; + +const float Consts::One = 1; + #if ROCM_VERSION >= 40300 const float ReduceConsts::One = 1; const float ReduceConsts::Zero = 0; + +const float ReduceConsts::One = 1; + +const float ReduceConsts::Zero = 0; #else // Up until ROCm 4.2, miopenReduceTensor() required alpha/beta to be the same data // type as the input type. This differs from cudnnReduceTensor() and other @@ -135,6 +143,12 @@ const half ReduceConsts::One = 1.f; template <> const half ReduceConsts::Zero = 0.f; + +template <> +const BFloat16 ReduceConsts::One = 1.f; + +template <> +const BFloat16 ReduceConsts::Zero = 0.f; #endif template <> diff --git a/onnxruntime/core/providers/rocm/miopen_common.h b/onnxruntime/core/providers/rocm/miopen_common.h index 8140cad54b..b2da1ae990 100644 --- a/onnxruntime/core/providers/rocm/miopen_common.h +++ b/onnxruntime/core/providers/rocm/miopen_common.h @@ -64,6 +64,12 @@ struct Consts { static const float One; }; +template <> +struct Consts { + static const float Zero; + static const float One; +}; + template struct ReduceConsts { static const ElemType Zero; @@ -79,6 +85,12 @@ struct ReduceConsts { static const float Zero; static const float One; }; + +template <> +struct ReduceConsts { + static const float Zero; + static const float One; +}; #endif inline double ClampMiopenBatchNormEpsilon(double epsilon) { diff --git a/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc b/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc index 34537897b9..14f129214d 100644 --- a/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/rocm/reduction/reduction_ops.cc @@ -519,7 +519,7 @@ Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, Pr } MiopenReduceDescriptor reduce_desc; - if (std::is_same::value) { + ORT_IF_CONSTEXPR (std::is_same::value || std::is_same::value) { ORT_RETURN_IF_ERROR(reduce_desc.Set(miopen_reduce_op, MiopenTensor::GetDataType(), ReduceTensorIndices)); } else { ORT_RETURN_IF_ERROR(reduce_desc.Set(miopen_reduce_op, miopen_type_X, ReduceTensorIndices)); @@ -823,111 +823,6 @@ SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int64_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(int8_t) SPECIALIZED_REDUCEKERNEL_COMPUTEIMPL(uint8_t) -template <> -template <> -Status ReduceKernel::ComputeImpl( - OpKernelContext* ctx, miopenReduceTensorOp_t miopen_reduce_op) const { - typedef typename ToHipType::MappedType HipT; - const Tensor* X = ctx->Input(0); - TensorShapeVector axes; - size_t num_inputs = ctx->InputCount(); - if (num_inputs == 2) { - const Tensor* axes_tensor = ctx->Input(1); - ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); - ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); - auto nDims = static_cast(axes_tensor->Shape()[0]); - const auto* data = axes_tensor->template Data(); - axes.assign(data, data + nDims); - } else { - axes.assign(axes_.begin(), axes_.end()); - } - - if (axes.empty() && noop_with_empty_axes_) { - auto* Y = ctx->Output(0, X->Shape()); - HIP_RETURN_IF_ERROR(hipMemcpyAsync(Y->template MutableData(), X->template Data(), - X->SizeInBytes(), hipMemcpyDeviceToDevice, Stream())); - return Status::OK(); - } - - PrepareReduceMetadata prepare_reduce_metadata; - ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, axes, prepare_reduce_metadata)); - - Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); - - int64_t input_count = prepare_reduce_metadata.input_count; - int64_t output_count = prepare_reduce_metadata.output_count; - auto& input_dims_miopen = prepare_reduce_metadata.input_dims_miopen; - auto& output_dims_miopen = prepare_reduce_metadata.output_dims_miopen; - - if (input_count == 0) { - assert(Y->Shape().Size() == 0); - return Status::OK(); - } - - if (input_count == output_count) { - if (Y->template MutableData() != X->template Data()) { - HIP_RETURN_IF_ERROR(hipMemcpyAsync(Y->template MutableData(), X->template Data(), - input_count * sizeof(BFloat16), hipMemcpyDeviceToDevice, Stream())); - } - return Status::OK(); - } - - if (fast_reduction_ && !ctx->GetUseDeterministicCompute()) { - int m{}, n{}; - const auto applicable_matrix_reduction = - get_applicable_matrix_reduction(miopen_reduce_op, X->Shape().GetDims(), axes, m, n); - switch (applicable_matrix_reduction) { - case ApplicableMatrixReduction::Rows: { - return reduce_matrix_rows(Stream(), reinterpret_cast(X->template Data()), - reinterpret_cast(Y->template MutableData()), m, n); - } - case ApplicableMatrixReduction::Columns: { - const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size(m, n); - auto buffer = rocm_ep_->GetScratchBuffer(buffer_size_bytes); - return reduce_matrix_columns(Stream(), reinterpret_cast(X->template Data()), - reinterpret_cast(Y->template MutableData()), m, n, buffer.get(), - buffer_size_bytes); - } - default: - break; - } - } - - HIP_RETURN_IF_ERROR(hipMemsetAsync(Y->MutableDataRaw(), 0, Y->SizeInBytes(), Stream())); - - size_t indices_bytes = 0; - size_t workspace_bytes = 0; - MiopenTensor input_tensor; - MiopenTensor output_tensor; - MiopenReduceDescriptor reduce_desc; - - miopenDataType_t miopen_type_X = miopenFloat; - IAllocatorUniquePtr temp_X = GetScratchBuffer(input_count); - Impl_Cast(Stream(), reinterpret_cast(X->template Data()), temp_X.get(), - X->Shape().Size()); - - ORT_RETURN_IF_ERROR(reduce_desc.Set(miopen_reduce_op, miopen_type_X, MIOPEN_REDUCE_TENSOR_NO_INDICES)); - ORT_RETURN_IF_ERROR(input_tensor.Set(input_dims_miopen, miopen_type_X)); - ORT_RETURN_IF_ERROR(output_tensor.Set(output_dims_miopen, miopen_type_X)); - MIOPEN_RETURN_IF_ERROR( - miopenGetReductionIndicesSize(MiopenHandle(), reduce_desc, input_tensor, output_tensor, &indices_bytes)); - MIOPEN_RETURN_IF_ERROR( - miopenGetReductionIndicesSize(MiopenHandle(), reduce_desc, input_tensor, output_tensor, &workspace_bytes)); - IAllocatorUniquePtr indices_rocm = GetScratchBuffer(indices_bytes); - IAllocatorUniquePtr workspace_rocm = GetScratchBuffer(workspace_bytes); - - const auto one = Consts::One; - const auto zero = Consts::Zero; - auto temp_Y = GetScratchBuffer(output_count); - MIOPEN_RETURN_IF_ERROR(miopenReduceTensor(MiopenHandle(), reduce_desc, indices_rocm.get(), indices_bytes, - workspace_rocm.get(), workspace_bytes, &one, input_tensor, temp_X.get(), - &zero, output_tensor, temp_Y.get())); - - Impl_Cast(Stream(), temp_Y.get(), reinterpret_cast(Y->template MutableData()), output_count); - - return Status::OK(); -} - namespace ReductionOps { template