Skip to content

Commit 35bb9b8

Browse files
mmpmmp
authored
Merge pull request #513 from jflatt-gia/scale2x
Add Optix 2X scaler to imgtool
2 parents 70f8b9e + d8d18ff commit 35bb9b8

File tree

4 files changed

+311
-1
lines changed

4 files changed

+311
-1
lines changed

‎CMakeLists.txt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -707,13 +707,15 @@ if (PBRT_CUDA_ENABLED)
707707
set (PBRT_GPU_SOURCE
708708
src/pbrt/gpu/optix/aggregate.cpp
709709
src/pbrt/gpu/optix/denoiser.cpp
710+
src/pbrt/gpu/optix/scaler.cpp
710711
src/pbrt/gpu/memory.cpp
711712
src/pbrt/gpu/util.cpp
712713
)
713714
set (PBRT_GPU_SOURCE_HEADERS
714715
src/pbrt/gpu/optix/aggregate.h
715716
src/pbrt/gpu/cudagl.h
716717
src/pbrt/gpu/optix/denoiser.h
718+
src/pbrt/gpu/optix/scaler.h
717719
src/pbrt/gpu/memory.h
718720
src/pbrt/gpu/optix/optix.h
719721
src/pbrt/gpu/util.h

‎src/pbrt/cmd/imgtool.cpp‎

Lines changed: 131 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#ifdef PBRT_BUILD_GPU_RENDERER
1010
#ifndef __HIP_PLATFORM_AMD__
1111
#include <pbrt/gpu/optix/denoiser.h>
12+
#include <pbrt/gpu/optix/scaler.h>
1213
#endif // __HIP_PLATFORM_AMD__
1314
#include <pbrt/gpu/util.h>
1415
#endif // PBRT_BUILD_GPU_RENDERER
@@ -149,6 +150,14 @@ static std::map<std::string, CommandUsage> commandUsage = {
149150
" be a multi-channel EXR as generated by pbrt's \"gbuffer\" film.",
150151
std::string(R"( options:
151152
--outfile <name> Filename to use for the denoised image.
153+
)")}},
154+
{"scale-optix",
155+
{"scale-optix [options] <filename>",
156+
"2X Scales the image using NVIDIA's OptiX denoiser which is\n"
157+
" based on a deep neural network. The provided image should\n"
158+
" be a multi-channel EXR as generated by pbrt's \"gbuffer\" film.",
159+
std::string(R"( options:
160+
--outfile <name> Filename to use for the scaled image.
152161
)")}},
153162
#endif // PBRT_BUILD_GPU_RENDERER
154163
{"error",
@@ -222,7 +231,7 @@ static std::map<std::string, CommandUsage> commandUsage = {
222231
{"scalenormalmap",
223232
{"scalenormalmap [options] <filename>",
224233
"Scale the provided normal map by applying the given factor for x and y\n"
225-
" and output the resulting normal map.\n",
234+
" and output the resulting normal map.",
226235
std::string(R"(
227236
--scale <s> Scale factor. Default: 1
228237
--outfile <name> Filename to store final image in.
@@ -2335,6 +2344,125 @@ int denoise_optix(std::vector<std::string> args) {
23352344

23362345
return 0;
23372346
}
2347+
2348+
int scale_optix(std::vector<std::string> args) {
2349+
std::string inFilename, outFilename;
2350+
2351+
auto onError = [](const std::string &err) {
2352+
usage("scale-optix", "%s", err.c_str());
2353+
exit(1);
2354+
};
2355+
for (auto iter = args.begin(); iter != args.end(); ++iter) {
2356+
if (ParseArg(&iter, args.end(), "outfile", &outFilename, onError)) {
2357+
// success
2358+
} else if ((*iter)[0] == '-')
2359+
usage("scale-optix", "%s: unknown command flag", iter->c_str());
2360+
else if (inFilename.empty()) {
2361+
inFilename = *iter;
2362+
} else
2363+
usage("scale-optix", "multiple input filenames provided.");
2364+
}
2365+
if (inFilename.empty())
2366+
usage("scale-optix", "input image filename must be provided.");
2367+
if (outFilename.empty())
2368+
usage("scale-optix", "output image filename must be provided.");
2369+
2370+
ImageAndMetadata im = Image::Read(inFilename);
2371+
Image &image = im.image;
2372+
2373+
CUDA_CHECK(cudaFree(nullptr));
2374+
2375+
int nLayers = 3;
2376+
bool oldNormalNaming = false;
2377+
ImageChannelDesc desc[3] = {
2378+
image.GetChannelDesc({"R", "G", "B"}),
2379+
image.GetChannelDesc({"Albedo.R", "Albedo.G", "Albedo.B"}),
2380+
image.GetChannelDesc({"Ns.X", "Ns.Y", "Ns.Z"})};
2381+
if (!desc[0]) {
2382+
Error("%s: image doesn't have R, G, B channels.", inFilename);
2383+
return 1;
2384+
}
2385+
if (!desc[1]) {
2386+
Warning("%s: image doesn't have Albedo.{R,G,B} channels. "
2387+
"Denoising quality may suffer.",
2388+
inFilename);
2389+
nLayers = 1;
2390+
}
2391+
if (!desc[2]) {
2392+
// Try the old naming scheme
2393+
desc[2] = image.GetChannelDesc({"Nsx", "Nsy", "Nsz"});
2394+
if (desc[2])
2395+
oldNormalNaming = true;
2396+
else {
2397+
Warning("%s: image doesn't have Ns.X, Ns.Y, Ns.Z channels. "
2398+
"Denoising quality may suffer.",
2399+
inFilename);
2400+
nLayers = 1;
2401+
}
2402+
}
2403+
2404+
Scaler scaler((Vector2i)image.Resolution(), nLayers == 3);
2405+
2406+
size_t imageBytes = 3 * image.Resolution().x * image.Resolution().y * sizeof(float);
2407+
2408+
auto copyChannelsToGPU = [&](std::array<std::string, 3> ch, bool flipZ = false) {
2409+
void *bufGPU;
2410+
CUDA_CHECK(cudaMalloc(&bufGPU, imageBytes));
2411+
std::vector<float> hostStaging(imageBytes / sizeof(float));
2412+
2413+
ImageChannelDesc desc = image.GetChannelDesc(ch);
2414+
CHECK(desc);
2415+
int offset = 0;
2416+
for (int y = 0; y < image.Resolution().y; ++y)
2417+
for (int x = 0; x < image.Resolution().x; ++x) {
2418+
ImageChannelValues v = image.GetChannels({x, y}, desc);
2419+
if (flipZ)
2420+
v[2] *= -1; // flip normal's z--right handed...
2421+
for (int c = 0; c < 3; ++c)
2422+
hostStaging[offset++] = v[c];
2423+
}
2424+
CUDA_CHECK(
2425+
cudaMemcpy(bufGPU, hostStaging.data(), imageBytes, cudaMemcpyHostToDevice));
2426+
return bufGPU;
2427+
};
2428+
RGB *rgbGPU = (RGB *)copyChannelsToGPU({"R", "G", "B"});
2429+
2430+
RGB *albedoGPU = nullptr;
2431+
Normal3f *normalGPU = nullptr;
2432+
if (nLayers == 3) {
2433+
albedoGPU = (RGB *)copyChannelsToGPU({"Albedo.R", "Albedo.G", "Albedo.B"});
2434+
if (oldNormalNaming)
2435+
normalGPU = (Normal3f *)copyChannelsToGPU({"Nsx", "Nsy", "Nsz"}, true);
2436+
else
2437+
normalGPU = (Normal3f *)copyChannelsToGPU({"Ns.X", "Ns.Y", "Ns.Z"}, true);
2438+
}
2439+
2440+
RGB *rgbResultGPU;
2441+
Point2i destRes(image.Resolution().x * 2, image.Resolution().y * 2);
2442+
size_t destBytes = 3 * destRes.x * destRes.y * sizeof(float);
2443+
CUDA_CHECK(cudaMalloc(&rgbResultGPU, destBytes));
2444+
2445+
scaler.Scale(rgbGPU, normalGPU, albedoGPU, rgbResultGPU);
2446+
2447+
CUDA_CHECK(cudaDeviceSynchronize());
2448+
2449+
Image result(PixelFormat::Float, destRes, {"R", "G", "B"});
2450+
CUDA_CHECK(cudaMemcpy(result.RawPointer({0, 0}), (const void *)rgbResultGPU,
2451+
destBytes, cudaMemcpyDeviceToHost));
2452+
2453+
ImageMetadata outMetadata;
2454+
outMetadata.cameraFromWorld = im.metadata.cameraFromWorld;
2455+
outMetadata.NDCFromWorld = im.metadata.NDCFromWorld;
2456+
outMetadata.pixelBounds = Bounds2i(
2457+
{im.metadata.pixelBounds->pMin.x * 2, im.metadata.pixelBounds->pMin.y * 2},
2458+
{im.metadata.pixelBounds->pMax.x * 2, im.metadata.pixelBounds->pMax.y * 2});
2459+
outMetadata.fullResolution = destRes;
2460+
outMetadata.colorSpace = im.metadata.colorSpace;
2461+
2462+
CHECK(result.Write(outFilename, outMetadata));
2463+
2464+
return 0;
2465+
}
23382466
#endif // PBRT_BUILD_GPU_RENDERER
23392467

23402468
int main(int argc, char *argv[]) {
@@ -2365,6 +2493,8 @@ int main(int argc, char *argv[]) {
23652493
#ifdef PBRT_BUILD_GPU_RENDERER
23662494
else if (cmd == "denoise-optix")
23672495
return denoise_optix(args);
2496+
else if (cmd == "scale-optix")
2497+
return scale_optix(args);
23682498
#endif // PBRT_BUILD_GPU_RENDERER
23692499
else if (cmd == "error")
23702500
return error(args);

‎src/pbrt/gpu/optix/scaler.cpp‎

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
1+
// pbrt is Copyright(c) 1998-2020 Matt Pharr, Wenzel Jakob, and Greg Humphreys.
2+
// The pbrt source code is licensed under the Apache License, Version 2.0.
3+
// SPDX: Apache-2.0
4+
5+
#include <pbrt/gpu/optix/scaler.h>
6+
7+
#include <pbrt/gpu/memory.h>
8+
#include <pbrt/gpu/util.h>
9+
10+
#include <cuda.h>
11+
#include <cuda_runtime.h>
12+
#include <array>
13+
14+
#include <optix.h>
15+
#include <optix_stubs.h>
16+
17+
#define OPTIX_CHECK(EXPR) \
18+
do { \
19+
OptixResult res = EXPR; \
20+
if (res != OPTIX_SUCCESS) \
21+
LOG_FATAL("OptiX call " #EXPR " failed with code %d: \"%s\"", int(res), \
22+
optixGetErrorString(res)); \
23+
} while (false) /* eat semicolon */
24+
25+
// Stop that, Windows.
26+
#ifdef RGB
27+
#undef RGB
28+
#endif
29+
30+
namespace pbrt {
31+
32+
Scaler::Scaler(Vector2i resolution, bool haveAlbedoAndNormal)
33+
: resolution(resolution), haveAlbedoAndNormal(haveAlbedoAndNormal) {
34+
CUcontext cudaContext;
35+
CU_CHECK(cuCtxGetCurrent(&cudaContext));
36+
CHECK(cudaContext != nullptr);
37+
38+
OPTIX_CHECK(optixInit());
39+
OptixDeviceContext optixContext;
40+
OPTIX_CHECK(optixDeviceContextCreate(cudaContext, 0, &optixContext));
41+
42+
OptixDenoiserOptions options = {};
43+
#if (OPTIX_VERSION >= 80000)
44+
options.denoiseAlpha = OPTIX_DENOISER_ALPHA_MODE_COPY;
45+
#endif
46+
47+
#if (OPTIX_VERSION >= 70300)
48+
if (haveAlbedoAndNormal)
49+
options.guideAlbedo = options.guideNormal = 1;
50+
51+
OPTIX_CHECK(optixDenoiserCreate(optixContext, OPTIX_DENOISER_MODEL_KIND_UPSCALE2X, &options,
52+
&denoiserHandle));
53+
#else
54+
options.inputKind = haveAlbedoAndNormal ? OPTIX_DENOISER_INPUT_RGB_ALBEDO_NORMAL
55+
: OPTIX_DENOISER_INPUT_RGB;
56+
57+
OPTIX_CHECK(optixDenoiserCreate(optixContext, &options, &denoiserHandle));
58+
59+
OPTIX_CHECK(
60+
optixDenoiserSetModel(denoiserHandle, OPTIX_DENOISER_MODEL_KIND_UPSCALE2X, nullptr, 0));
61+
#endif
62+
63+
OPTIX_CHECK(optixDenoiserComputeMemoryResources(denoiserHandle, resolution.x,
64+
resolution.y, &memorySizes));
65+
66+
CUDA_CHECK(cudaMalloc(&denoiserState, memorySizes.stateSizeInBytes));
67+
CUDA_CHECK(cudaMalloc(&scratchBuffer, memorySizes.withoutOverlapScratchSizeInBytes));
68+
69+
OPTIX_CHECK(optixDenoiserSetup(
70+
denoiserHandle, 0 /* stream */, resolution.x, resolution.y,
71+
CUdeviceptr(denoiserState), memorySizes.stateSizeInBytes,
72+
CUdeviceptr(scratchBuffer), memorySizes.withoutOverlapScratchSizeInBytes));
73+
74+
CUDA_CHECK(cudaMalloc(&intensity, sizeof(float)));
75+
}
76+
77+
void Scaler::Scale(RGB *rgb, Normal3f *n, RGB *albedo, RGB *result) {
78+
std::array<OptixImage2D, 3> inputLayers;
79+
int nLayers = haveAlbedoAndNormal ? 3 : 1;
80+
for (int i = 0; i < nLayers; ++i) {
81+
inputLayers[i].width = resolution.x;
82+
inputLayers[i].height = resolution.y;
83+
inputLayers[i].rowStrideInBytes = resolution.x * 3 * sizeof(float);
84+
inputLayers[i].pixelStrideInBytes = 0;
85+
inputLayers[i].format = OPTIX_PIXEL_FORMAT_FLOAT3;
86+
}
87+
inputLayers[0].data = CUdeviceptr(rgb);
88+
if (haveAlbedoAndNormal) {
89+
CHECK(n != nullptr && albedo != nullptr);
90+
inputLayers[1].data = CUdeviceptr(albedo);
91+
inputLayers[2].data = CUdeviceptr(n);
92+
} else
93+
CHECK(n == nullptr && albedo == nullptr);
94+
95+
OptixImage2D outputImage;
96+
outputImage.width = resolution.x * 2;
97+
outputImage.height = resolution.y * 2;
98+
outputImage.rowStrideInBytes = resolution.x * 2 * 3 * sizeof(float);
99+
outputImage.pixelStrideInBytes = 0;
100+
outputImage.format = OPTIX_PIXEL_FORMAT_FLOAT3;
101+
outputImage.data = CUdeviceptr(result);
102+
103+
OPTIX_CHECK(optixDenoiserComputeIntensity(
104+
denoiserHandle, 0 /* stream */, &inputLayers[0], CUdeviceptr(intensity),
105+
CUdeviceptr(scratchBuffer), memorySizes.withoutOverlapScratchSizeInBytes));
106+
107+
OptixDenoiserParams params = {};
108+
#if (OPTIX_VERSION >= 80000)
109+
// denoiseAlpha is moved to OptixDenoiserOptions in OptiX 8.0
110+
#elif (OPTIX_VERSION >= 70500)
111+
params.denoiseAlpha = OPTIX_DENOISER_ALPHA_MODE_COPY;
112+
#else
113+
params.denoiseAlpha = 0;
114+
#endif
115+
params.hdrIntensity = CUdeviceptr(intensity);
116+
params.blendFactor = 0; // TODO what should this be??
117+
118+
#if (OPTIX_VERSION >= 70300)
119+
OptixDenoiserGuideLayer guideLayer;
120+
if (haveAlbedoAndNormal) {
121+
guideLayer.albedo = inputLayers[1];
122+
guideLayer.normal = inputLayers[2];
123+
}
124+
125+
OptixDenoiserLayer layers;
126+
layers.input = inputLayers[0];
127+
layers.output = outputImage;
128+
129+
OPTIX_CHECK(optixDenoiserInvoke(
130+
denoiserHandle, 0 /* stream */, &params, CUdeviceptr(denoiserState),
131+
memorySizes.stateSizeInBytes, &guideLayer, &layers, 1 /* # layers to denoise */,
132+
0 /* offset x */, 0 /* offset y */, CUdeviceptr(scratchBuffer),
133+
memorySizes.withoutOverlapScratchSizeInBytes));
134+
#else
135+
OPTIX_CHECK(optixDenoiserInvoke(
136+
denoiserHandle, 0 /* stream */, &params, CUdeviceptr(denoiserState),
137+
memorySizes.stateSizeInBytes, inputLayers.data(), nLayers, 0 /* offset x */,
138+
0 /* offset y */, &outputImage, CUdeviceptr(scratchBuffer),
139+
memorySizes.withoutOverlapScratchSizeInBytes));
140+
#endif
141+
}
142+
143+
} // namespace pbrt

‎src/pbrt/gpu/optix/scaler.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
// pbrt is Copyright(c) 1998-2020 Matt Pharr, Wenzel Jakob, and Greg Humphreys.
2+
// The pbrt source code is licensed under the Apache License, Version 2.0.
3+
// SPDX: Apache-2.0
4+
5+
#ifndef PBRT_GPU_SCALER_H
6+
#define PBRT_GPU_SCALER_H
7+
8+
#include <pbrt/pbrt.h>
9+
10+
#include <pbrt/util/color.h>
11+
#include <pbrt/util/vecmath.h>
12+
13+
#include <optix.h>
14+
15+
namespace pbrt {
16+
17+
class Scaler {
18+
public:
19+
Scaler(Vector2i resolution, bool haveAlbedoAndNormal);
20+
21+
// All pointers should be to GPU memory.
22+
// |n| and |albedo| should be nullptr iff \haveAlbedoAndNormal| is false.
23+
void Scale(RGB *rgb, Normal3f *n, RGB *albedo, RGB *result);
24+
25+
private:
26+
Vector2i resolution;
27+
bool haveAlbedoAndNormal;
28+
OptixDenoiser denoiserHandle;
29+
OptixDenoiserSizes memorySizes;
30+
void *denoiserState, *scratchBuffer, *intensity;
31+
};
32+
33+
} // namespace pbrt
34+
35+
#endif // PBRT_GPU_SCALER_H

0 commit comments

Comments
 (0)