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

Fix how the output tensor is created in CUDA SpaceDepth ops (#11302)

Browse source
This commit is contained in:
Hariharan Seshadri 2022-04-21 19:58:57 -07:00 committed by GitHub
parent 86cdabbcfd
commit 23b01258b5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 27 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
onnxruntime/core/providers/cuda/tensor/space_depth_ops.cc
View file

@ -79,9 +79,11 @@ static Status SpaceDepthOpCudaImpl(const cudaDeviceProp& prop,
const std::vector<size_t>& permutation,
const int64_t batch_size,
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 in_dim4, const int64_t in_dim5,
const TensorShape& virtual_output_shape) {
TensorShape virtual_input_shape{batch_size, in_dim1, in_dim2, in_dim3, in_dim4, in_dim5};
return Transpose::DoTranspose(prop, stream, cublas_handle, permutation, input, output, &virtual_input_shape);
return Transpose::DoTranspose(prop, stream, cublas_handle, permutation, input, output,
&virtual_input_shape, &virtual_output_shape);
}
Status SpaceToDepth::ComputeInternal(OpKernelContext* context) const {
@ -105,17 +107,18 @@ Status SpaceToDepth::ComputeInternal(OpKernelContext* context) const {
output_depth, output_height, output_width,
true));
// We use the "virtual" output shape to construct the output tensor
Tensor& output = *context->Output(0,
{batch, blocksize_, blocksize_, input_depth, input_height / blocksize_, input_width / blocksize_});
// We use the "actual" output shape to construct the output tensor
Tensor& output = *context->Output(0, {batch, output_depth, output_height, output_width});
// We will pass in the "virtual" output shape to be used by DoTranspose() in SpaceDepthOpCudaImpl(...)
TensorShape virtual_output_shape{batch, blocksize_, blocksize_, input_depth,
input_height / blocksize_, input_width / blocksize_};
std::vector<size_t> permutation = {0, 3, 5, 1, 2, 4};
ORT_RETURN_IF_ERROR(SpaceDepthOpCudaImpl(GetDeviceProp(), Stream(), CublasHandle(), input, output, permutation, batch,
input_depth, input_height / blocksize_, blocksize_, input_width / blocksize_, blocksize_));
// Reshape to "actual" output shape
output.Reshape({batch, output_depth, output_height, output_width});
input_depth, input_height / blocksize_, blocksize_, input_width / blocksize_, blocksize_,
virtual_output_shape));
return Status::OK();
}
@ -141,9 +144,12 @@ Status DepthToSpace::ComputeInternal(OpKernelContext* context) const {
output_depth, output_height, output_width,
false));
// We use the "virtual" output shape to construct the output tensor
Tensor& output = *context->Output(0,
{batch, input_depth / blocksize_ / blocksize_, input_height, blocksize_, input_width, blocksize_});
// We use the "actual" output shape to construct the output tensor
Tensor& output = *context->Output(0, {batch, output_depth, output_height, output_width});
// We will pass in the "virtual" output shape to be used by DoTranspose() in SpaceDepthOpCudaImpl(...)
TensorShape virtual_output_shape{batch, input_depth / blocksize_ / blocksize_,
input_height, blocksize_, input_width, blocksize_};
std::vector<size_t> permutation;
permutation.reserve(6);
@ -170,10 +176,8 @@ Status DepthToSpace::ComputeInternal(OpKernelContext* context) const {
ORT_RETURN_IF_ERROR(SpaceDepthOpCudaImpl(GetDeviceProp(), Stream(), CublasHandle(), input, output,
permutation,
batch,
dim1, blocksize_, dim3, input_height, input_width));
// Reshape to "actual" output shape
output.Reshape({batch, output_depth, output_height, output_width});
dim1, blocksize_, dim3, input_height, input_width,
virtual_output_shape));
return Status::OK();
}

5
onnxruntime/core/providers/cuda/tensor/transpose.cc
View file

@ -86,13 +86,14 @@ Status Transpose::DoTranspose(const cudaDeviceProp& prop,
cudaStream_t stream,
const cublasHandle_t cublas_handle,
const gsl::span<const size_t>& permutations, const Tensor& input, Tensor& output,
const TensorShape* input_shape_override) {
const TensorShape* input_shape_override,
const TensorShape* output_shape_override) {
// special case when there is a dim value of 0 in the shape.
if (output.Shape().Size() == 0)
return Status::OK();
const auto input_dims = input_shape_override ? input_shape_override->GetDims() : input.Shape().GetDims();
const auto output_dims = output.Shape().GetDims();
const auto output_dims = output_shape_override ? output_shape_override->GetDims() : output.Shape().GetDims();
auto rank = static_cast<int32_t>(input_dims.size());
// flatten the adjacent dimensions which are contiguous

5
onnxruntime/core/providers/cuda/tensor/transpose.h
View file

@ -22,11 +22,14 @@ class Transpose final : public CudaKernel, public TransposeBase {
const gsl::span<const size_t>& permutations, const Tensor& input, Tensor& output);
// `input_shape_override` (if provided) overrides the shape of `input` for compute purposes
// `output_shape_override` (if provided) overrides the shape of `output` for compute purposes
static Status DoTranspose(const cudaDeviceProp& prop,
cudaStream_t stream,
const cublasHandle_t cublas_handle,
const gsl::span<const size_t>& permutations,
const Tensor& input, Tensor& output, const TensorShape* input_shape_override = nullptr);
const Tensor& input, Tensor& output,
const TensorShape* input_shape_override = nullptr,
const TensorShape* output_shape_override = nullptr);
};
} // namespace cuda