From 6586afc8eb16abdbdb383d4752405721653f541b Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Wed, 11 Sep 2019 14:31:53 +1000 Subject: [PATCH] Refine the output shape calculation to avoid unnecessary re-allocations and vector insert operations. (#1781) --- .../core/providers/cpu/tensor/gather.cc | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/gather.cc b/onnxruntime/core/providers/cpu/tensor/gather.cc index 14eba04abe..817bf6794e 100644 --- a/onnxruntime/core/providers/cpu/tensor/gather.cc +++ b/onnxruntime/core/providers/cpu/tensor/gather.cc @@ -19,13 +19,23 @@ Status GatherBase::PrepareForCompute(OpKernelContext* context, Prepare& p) const p.indices_tensor = context->Input(1); const TensorShape& indices_shape = p.indices_tensor->Shape(); - p.axis = HandleNegativeAxis(axis_, input_data_shape.NumDimensions()); + const auto input_rank = input_data_shape.NumDimensions(); + p.axis = HandleNegativeAxis(axis_, input_rank); - std::vector shape(indices_shape.GetDims().begin(), indices_shape.GetDims().end()); - shape.insert(shape.begin(), input_data_shape.GetDims().begin(), input_data_shape.GetDims().begin() + p.axis); - shape.insert(shape.end(), input_data_shape.GetDims().begin() + p.axis + 1, input_data_shape.GetDims().end()); + std::vector shape; + shape.reserve(input_rank - 1 + indices_shape.NumDimensions()); - p.output_tensor = context->Output(0, TensorShape(shape)); + // replace the dimension for p.axis with the shape from the indices + for (int64_t i = 0; i < p.axis; ++i) + shape.push_back(input_data_shape[i]); + + for (const auto dim : indices_shape.GetDims()) + shape.push_back(dim); + + for (int64_t i = p.axis + 1; i < static_cast(input_rank); ++i) + shape.push_back(input_data_shape[i]); + + p.output_tensor = context->Output(0, TensorShape(std::move(shape))); return Status::OK(); }