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

Add helper to create output to minimize binary size. (#2365)

Browse source 
Add ConstEigenTensorMap typedef so we don't unnecessarily const_cast the const input Tensor.
This commit is contained in:
Scott McKay 2019-11-12 09:08:04 +10:00 committed by GitHub
parent aa37e2de8f
commit 53ed36a3da
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 32 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

60
onnxruntime/core/providers/cpu/tensor/space_depth_ops.cc
View file

@ -25,7 +25,7 @@ ONNX_CPU_OPERATOR_KERNEL(
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
DepthToSpace<float>);
// intemediate tensor shapes are:
// intermediate tensor shapes are:
// (batch, blocksize, blocksize, input_depth / (blocksize * blocksize), input_height, input_width) for DepthToSpace
// (batch, input_depth, input_height / blocksize, blocksize, input_width / blocksize, blocksize) for SpaceToDepth
const int IntermediateTensorRank = 6;
@ -33,6 +33,24 @@ typedef Eigen::TensorMap<Eigen::Tensor<float, IntermediateTensorRank, Eigen::Row
Eigen::Aligned>
EigenTensorMap;
typedef Eigen::TensorMap<Eigen::Tensor<const float, IntermediateTensorRank, Eigen::RowMajor, int64_t>,
Eigen::Aligned>
ConstEigenTensorMap;
// helper to create output to minimize binary size
static void CreateOutput(const Tensor& input, Tensor& output,
const std::array<int64_t, IntermediateTensorRank>& permutation,
const int64_t batch_size, // dim0 in both input and output
const int64_t in_dim1, const int64_t in_dim2, const int64_t in_dim3,
const int64_t in_dim4, const int64_t in_dim5,
const int64_t out_dim1, const int64_t out_dim2, const int64_t out_dim3,
const int64_t out_dim4, const int64_t out_dim5) {
EigenTensorMap(output.template MutableData<float>(), batch_size, out_dim1, out_dim2, out_dim3, out_dim4, out_dim5) =
ConstEigenTensorMap(input.template Data<float>(), batch_size,
in_dim1, in_dim2, in_dim3, in_dim4, in_dim5)
.shuffle(permutation);
}
template <>
Status SpaceToDepth<float>::Compute(OpKernelContext* context) const {
const auto* tensor_pointer = context->Input<Tensor>(0);
@ -52,12 +70,9 @@ Status SpaceToDepth<float>::Compute(OpKernelContext* context) const {
Tensor& output = *context->Output(0, {batch, output_depth, output_height, output_width});
std::array<int64_t, IntermediateTensorRank> permutation{{0, 3, 5, 1, 2, 4}};
EigenTensorMap(output.template MutableData<float>(), batch, blocksize_, blocksize_,
input_depth, input_height / blocksize_, input_width / blocksize_) =
EigenTensorMap(const_cast<float*>(input.template Data<float>()), batch,
input_depth, input_height / blocksize_, blocksize_,
input_width / blocksize_, blocksize_)
.shuffle(permutation);
CreateOutput(input, output, permutation, batch,
input_depth, input_height / blocksize_, blocksize_, input_width / blocksize_, blocksize_,
blocksize_, blocksize_, input_depth, input_height / blocksize_, input_width / blocksize_);
return Status::OK();
}
@ -80,28 +95,17 @@ Status DepthToSpace<float>::Compute(OpKernelContext* context) const {
const int64_t output_width = input_width * blocksize_;
Tensor& output = *context->Output(0, {batch, output_depth, output_height, output_width});
// Process "DCR" mode
if (is_dcr_) {
std::array<int64_t, IntermediateTensorRank> permutation{{0, 3, 4, 1, 5, 2}};
EigenTensorMap(output.template MutableData<float>(), batch, input_depth / blocksize_ / blocksize_,
input_height, blocksize_, input_width, blocksize_) =
EigenTensorMap(const_cast<float*>(input.template Data<float>()), batch,
blocksize_, blocksize_, input_depth / blocksize_ / blocksize_,
input_height, input_width)
.shuffle(permutation);
}
// Process "CRD" mode
else {
std::array<int64_t, IntermediateTensorRank> permutation{{0, 1, 4, 2, 5, 3}};
EigenTensorMap(output.template MutableData<float>(), batch, input_depth / blocksize_ / blocksize_,
input_height, blocksize_, input_width, blocksize_) =
EigenTensorMap(const_cast<float*>(input.template Data<float>()), batch,
input_depth / blocksize_ / blocksize_, blocksize_, blocksize_,
input_height, input_width)
.shuffle(permutation);
}
// handle DCR and CRD format
auto dim1 = is_dcr_ ? blocksize_ : input_depth / blocksize_ / blocksize_;
auto dim3 = is_dcr_ ? input_depth / blocksize_ / blocksize_ : blocksize_;
auto permutation = is_dcr_ ? std::array<int64_t, IntermediateTensorRank>{{0, 3, 4, 1, 5, 2}}
: std::array<int64_t, IntermediateTensorRank>{{0, 1, 4, 2, 5, 3}};
CreateOutput(input, output, permutation, batch,
dim1, blocksize_, dim3, input_height, input_width,
input_depth / blocksize_ / blocksize_, input_height, blocksize_, input_width, blocksize_);
return Status::OK();
}