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

Remove BFloat16 Specialized Code for ReduceSum (#10476)

Browse source
This commit is contained in:
Vincent Wang 2022-02-09 07:39:57 +08:00 committed by GitHub
parent 7e5d68eea6
commit 655f490c95
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 34 additions and 214 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
onnxruntime/core/providers/cuda/math/unary_elementwise_ops_impl.cu
View file

@ -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)

107
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
View file

@ -521,7 +521,7 @@ Status ReduceComputeCore(CUDAExecutionProvider& cuda_ep, const Tensor& input, Pr
}
CudnnReduceDescriptor reduce_desc;
if (std::is_same<T, MLFloat16>::value) {
ORT_IF_CONSTEXPR (std::is_same<T, MLFloat16>::value || std::is_same<T, BFloat16>::value) {
ORT_RETURN_IF_ERROR(reduce_desc.Set(cudnn_reduce_op, CudnnTensor::GetDataType<float>(), 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<true>::ComputeImpl<BFloat16, CUDNN_REDUCE_TENSOR_NO_INDICES>(
OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const {
typedef typename ToCudaType<BFloat16>::MappedType CudaT;
const Tensor* X = ctx->Input<Tensor>(0);
TensorShapeVector axes;
size_t num_inputs = ctx->InputCount();
if (num_inputs == 2) {
const Tensor* axes_tensor = ctx->Input<Tensor>(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<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->template Data<int64_t>();
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<BFloat16>(), X->template Data<BFloat16>(),
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<BFloat16>() != X->template Data<BFloat16>()) {
CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData<BFloat16>(), X->template Data<BFloat16>(),
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<const CudaT*>(X->template Data<BFloat16>()),
reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), m, n);
}
case ApplicableMatrixReduction::Columns: {
const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size<CudaT>(m, n);
auto buffer = cuda_ep_->GetScratchBuffer<void>(buffer_size_bytes);
return reduce_matrix_columns(Stream(), reinterpret_cast<const CudaT*>(X->template Data<BFloat16>()),
reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), 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<float> temp_X = GetScratchBuffer<float>(input_count);
Impl_Cast<CudaT, float>(Stream(), reinterpret_cast<const CudaT*>(X->template Data<BFloat16>()), 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<uint32_t> indices_cuda = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<CudaT> workspace_cuda = GetScratchBuffer<CudaT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<float, CudaT>(Stream(), temp_Y.get(), reinterpret_cast<CudaT*>(Y->template MutableData<BFloat16>()), output_count);
return Status::OK();
}
namespace ReductionOps {
template <typename T, cudnnReduceTensorIndices_t ReduceTensorIndices>

14
onnxruntime/core/providers/rocm/miopen_common.cc
View file

@ -122,10 +122,18 @@ const float Consts<half>::Zero = 0;
const float Consts<half>::One = 1;
const float Consts<BFloat16>::Zero = 0;
const float Consts<BFloat16>::One = 1;
#if ROCM_VERSION >= 40300
const float ReduceConsts<half>::One = 1;
const float ReduceConsts<half>::Zero = 0;
const float ReduceConsts<BFloat16>::One = 1;
const float ReduceConsts<BFloat16>::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<half>::One = 1.f;
template <>
const half ReduceConsts<half>::Zero = 0.f;
template <>
const BFloat16 ReduceConsts<BFloat16>::One = 1.f;
template <>
const BFloat16 ReduceConsts<BFloat16>::Zero = 0.f;
#endif
template <>

12
onnxruntime/core/providers/rocm/miopen_common.h
View file

@ -64,6 +64,12 @@ struct Consts<half> {
static const float One;
};
template <>
struct Consts<BFloat16> {
static const float Zero;
static const float One;
};
template <typename ElemType>
struct ReduceConsts {
static const ElemType Zero;
@ -79,6 +85,12 @@ struct ReduceConsts<half> {
static const float Zero;
static const float One;
};
template <>
struct ReduceConsts<BFloat16> {
static const float Zero;
static const float One;
};
#endif
inline double ClampMiopenBatchNormEpsilon(double epsilon) {

107
onnxruntime/core/providers/rocm/reduction/reduction_ops.cc
View file

@ -519,7 +519,7 @@ Status ReduceComputeCore(ROCMExecutionProvider& rocm_ep, const Tensor& input, Pr
}
MiopenReduceDescriptor reduce_desc;
if (std::is_same<T, MLFloat16>::value) {
ORT_IF_CONSTEXPR (std::is_same<T, MLFloat16>::value || std::is_same<T, BFloat16>::value) {
ORT_RETURN_IF_ERROR(reduce_desc.Set(miopen_reduce_op, MiopenTensor::GetDataType<float>(), 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<true>::ComputeImpl<BFloat16, MIOPEN_REDUCE_TENSOR_NO_INDICES>(
OpKernelContext* ctx, miopenReduceTensorOp_t miopen_reduce_op) const {
typedef typename ToHipType<BFloat16>::MappedType HipT;
const Tensor* X = ctx->Input<Tensor>(0);
TensorShapeVector axes;
size_t num_inputs = ctx->InputCount();
if (num_inputs == 2) {
const Tensor* axes_tensor = ctx->Input<Tensor>(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<size_t>(axes_tensor->Shape()[0]);
const auto* data = axes_tensor->template Data<int64_t>();
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<BFloat16>(), X->template Data<BFloat16>(),
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<BFloat16>() != X->template Data<BFloat16>()) {
HIP_RETURN_IF_ERROR(hipMemcpyAsync(Y->template MutableData<BFloat16>(), X->template Data<BFloat16>(),
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<const HipT*>(X->template Data<BFloat16>()),
reinterpret_cast<HipT*>(Y->template MutableData<BFloat16>()), m, n);
}
case ApplicableMatrixReduction::Columns: {
const auto buffer_size_bytes = compute_reduce_matrix_columns_buffer_size<HipT>(m, n);
auto buffer = rocm_ep_->GetScratchBuffer<void>(buffer_size_bytes);
return reduce_matrix_columns(Stream(), reinterpret_cast<const HipT*>(X->template Data<BFloat16>()),
reinterpret_cast<HipT*>(Y->template MutableData<BFloat16>()), 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<float> temp_X = GetScratchBuffer<float>(input_count);
Impl_Cast<HipT, float>(Stream(), reinterpret_cast<const HipT*>(X->template Data<BFloat16>()), 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<uint32_t> indices_rocm = GetScratchBuffer<uint32_t>(indices_bytes);
IAllocatorUniquePtr<HipT> workspace_rocm = GetScratchBuffer<HipT>(workspace_bytes);
const auto one = Consts<float>::One;
const auto zero = Consts<float>::Zero;
auto temp_Y = GetScratchBuffer<float>(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<float, HipT>(Stream(), temp_Y.get(), reinterpret_cast<HipT*>(Y->template MutableData<BFloat16>()), output_count);
return Status::OK();
}
namespace ReductionOps {
template <typename T, miopenReduceTensorIndices_t ReduceTensorIndices>