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

Fix scratch buffer early free.

Browse source
This commit is contained in:
Jesse Benson 2021-01-21 11:02:34 -08:00 committed by Jesse Benson
parent 86ac11af1a
commit 76fcebd0a4
2 changed files with 14 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -202,9 +202,11 @@ Status ReduceKernel<allow_multi_axes>::ReduceKernelShared(
auto output_count = output_shape.Size();
if (ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES) {
IAllocatorUniquePtr<T> input_data_buffer(nullptr, [](T*) {});
CudaT* input_data = nullptr;
if (calculate_sqt_) {
input_data = reinterpret_cast<CudaT*>(GetScratchBuffer<T>(input_count).get());
input_data_buffer = GetScratchBuffer<T>(input_count);
input_data = reinterpret_cast<CudaT*>(input_data_buffer.get());
fast_divmod tmp_div;
Impl_Mul<CudaT>(static_cast<int32_t>(SimpleBroadcast::NoBroadcast), nullptr,
reinterpret_cast<const CudaT*>(X), nullptr,
@ -225,8 +227,10 @@ Status ReduceKernel<allow_multi_axes>::ReduceKernelShared(
// Exp(X-ReduceMax)
const TensorShape rhs_shape(output_dims);
auto exp_result = GetScratchBuffer<T>(input_count).get();
auto log_sum_result = GetScratchBuffer<T>(output_count).get();
auto exp_result_buffer = GetScratchBuffer<T>(input_count);
auto exp_result = exp_result_buffer.get();
auto log_sum_result_buffer = GetScratchBuffer<T>(output_count);
auto log_sum_result = log_sum_result_buffer.get();
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(prepare.BinaryElementwiseBroadcastPrepareHelper(input_shape, rhs_shape, input_shape));
Impl_Sub<CudaT>(prepare.output_rank_or_simple_broadcast,

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

@ -202,9 +202,11 @@ Status ReduceKernel<allow_multi_axes>::ReduceKernelShared(
auto output_count = output_shape.Size();
if (ReduceTensorIndices == MIOPEN_REDUCE_TENSOR_NO_INDICES) {
IAllocatorUniquePtr<T> input_data_buffer(nullptr, [](T*) {});
HipT* input_data = nullptr;
if (calculate_sqt_) {
input_data = reinterpret_cast<HipT*>(GetScratchBuffer<T>(input_count).get());
input_data_buffer = GetScratchBuffer<T>(input_count);
input_data = reinterpret_cast<HipT*>(input_data_buffer.get());
fast_divmod tmp_div;
Impl_Mul<HipT>(static_cast<int32_t>(SimpleBroadcast::NoBroadcast), nullptr,
reinterpret_cast<const HipT*>(X), nullptr,
@ -225,8 +227,10 @@ Status ReduceKernel<allow_multi_axes>::ReduceKernelShared(
// Exp(X-ReduceMax)
const TensorShape rhs_shape(output_dims);
auto exp_result = GetScratchBuffer<T>(input_count).get();
auto log_sum_result = GetScratchBuffer<T>(output_count).get();
auto exp_result_buffer = GetScratchBuffer<T>(input_count);
auto exp_result = exp_result_buffer.get();
auto log_sum_result_buffer = GetScratchBuffer<T>(output_count);
auto log_sum_result = log_sum_result_buffer.get();
BinaryElementwisePreparation prepare;
ORT_RETURN_IF_ERROR(prepare.BinaryElementwiseBroadcastPrepareHelper(input_shape, rhs_shape, input_shape));
Impl_Sub<HipT>(prepare.output_rank_or_simple_broadcast,