wehub-resource-sync
149 lines
4.9 KiB
Plaintext
149 lines
4.9 KiB
Plaintext
// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#if !defined(WITH_NV_JETSON) && !defined(PADDLE_WITH_HIP)
|
|
|
|
#include "paddle/phi/kernels/decode_jpeg_kernel.h"
|
|
|
|
#include "paddle/phi/backends/dynload/nvjpeg.h"
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/stream.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
|
|
namespace phi {
|
|
|
|
static cudaStream_t nvjpeg_stream = nullptr;
|
|
static nvjpegHandle_t nvjpeg_handle = nullptr;
|
|
|
|
void InitNvjpegImage(nvjpegImage_t* img) {
|
|
for (int c = 0; c < NVJPEG_MAX_COMPONENT; c++) {
|
|
img->channel[c] = nullptr;
|
|
img->pitch[c] = 0;
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void DecodeJpegKernel(const Context& dev_ctx,
|
|
const DenseTensor& x,
|
|
const std::string& mode,
|
|
DenseTensor* out) {
|
|
// Create nvJPEG handle
|
|
if (nvjpeg_handle == nullptr) {
|
|
nvjpegStatus_t create_status = dynload::nvjpegCreateSimple(&nvjpeg_handle);
|
|
|
|
PADDLE_ENFORCE_EQ(create_status,
|
|
NVJPEG_STATUS_SUCCESS,
|
|
errors::Fatal("nvjpegCreateSimple failed: %d.",
|
|
static_cast<int>(create_status)));
|
|
}
|
|
|
|
nvjpegJpegState_t nvjpeg_state;
|
|
nvjpegStatus_t state_status =
|
|
dynload::nvjpegJpegStateCreate(nvjpeg_handle, &nvjpeg_state);
|
|
|
|
PADDLE_ENFORCE_EQ(state_status,
|
|
NVJPEG_STATUS_SUCCESS,
|
|
errors::Fatal("nvjpegJpegStateCreate failed: %d",
|
|
static_cast<int>(state_status)));
|
|
|
|
int components;
|
|
nvjpegChromaSubsampling_t subsampling;
|
|
int widths[NVJPEG_MAX_COMPONENT];
|
|
int heights[NVJPEG_MAX_COMPONENT];
|
|
|
|
auto* x_data = x.data<T>();
|
|
|
|
nvjpegStatus_t info_status =
|
|
dynload::nvjpegGetImageInfo(nvjpeg_handle,
|
|
x_data,
|
|
(std::size_t)x.numel(),
|
|
&components,
|
|
&subsampling,
|
|
widths,
|
|
heights);
|
|
PADDLE_ENFORCE_EQ(info_status,
|
|
NVJPEG_STATUS_SUCCESS,
|
|
errors::Fatal("nvjpegGetImageInfo failed: %d",
|
|
static_cast<int>(info_status)));
|
|
|
|
int width = widths[0];
|
|
int height = heights[0];
|
|
|
|
nvjpegOutputFormat_t output_format;
|
|
int output_components;
|
|
|
|
if (mode == "unchanged") {
|
|
if (components == 1) {
|
|
output_format = NVJPEG_OUTPUT_Y;
|
|
output_components = 1;
|
|
} else if (components == 3) {
|
|
output_format = NVJPEG_OUTPUT_RGB;
|
|
output_components = 3;
|
|
} else {
|
|
dynload::nvjpegJpegStateDestroy(nvjpeg_state);
|
|
PADDLE_THROW(errors::Fatal(
|
|
"The provided mode is not supported for JPEG files on GPU"));
|
|
}
|
|
} else if (mode == "gray") {
|
|
output_format = NVJPEG_OUTPUT_Y;
|
|
output_components = 1;
|
|
} else if (mode == "rgb") {
|
|
output_format = NVJPEG_OUTPUT_RGB;
|
|
output_components = 3;
|
|
} else {
|
|
dynload::nvjpegJpegStateDestroy(nvjpeg_state);
|
|
PADDLE_THROW(errors::Fatal(
|
|
"The provided mode is not supported for JPEG files on GPU"));
|
|
}
|
|
|
|
nvjpegImage_t out_image;
|
|
InitNvjpegImage(&out_image);
|
|
|
|
// create nvjpeg stream
|
|
if (nvjpeg_stream == nullptr) {
|
|
cudaStreamCreateWithFlags(&nvjpeg_stream, cudaStreamNonBlocking);
|
|
}
|
|
|
|
int64_t sz = static_cast<int64_t>(widths[0]) * heights[0];
|
|
|
|
std::vector<int64_t> out_shape = {output_components, height, width};
|
|
out->Resize(out_shape);
|
|
|
|
T* data = dev_ctx.template Alloc<T>(out);
|
|
|
|
for (int c = 0; c < output_components; c++) {
|
|
out_image.channel[c] = data + c * sz;
|
|
out_image.pitch[c] = width;
|
|
}
|
|
|
|
nvjpegStatus_t decode_status = dynload::nvjpegDecode(nvjpeg_handle,
|
|
nvjpeg_state,
|
|
x_data,
|
|
x.numel(),
|
|
output_format,
|
|
&out_image,
|
|
nvjpeg_stream);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(decode_jpeg, // cuda_only
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::DecodeJpegKernel,
|
|
uint8_t) {
|
|
kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
|
|
}
|
|
|
|
#endif
|