Files
UnrealEngine/Engine/Plugins/Experimental/OptiXDenoise/Source/ThirdParty/OptiXDenoiseBase/Private/OptiXDenoiseBase.cpp
2025-05-18 13:04:45 +08:00

483 lines
16 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
#pragma once
#include "OptiXDenoiseBase.h"
// avoid win api pollution
#define NOMINMAX
#include <optix.h>
#include <optix_function_table_definition.h>
#include <optix_stubs.h>
#include <optix_denoiser_tiling.h>
#include "optixOpticalFlow/optix_denoiser_opticalflow.h"
EOptiXDenoiseResult OptiXDenoiseResult(OptixResult Result)
{
return static_cast<EOptiXDenoiseResult>(static_cast<uint32_t>(Result));
}
OptixImage2D GetOptiXImage2D(const FOptiXImageData& OptiXImageData)
{
OptixImage2D Image;
Image.data = OptiXImageData.Data;
Image.width = OptiXImageData.Width;
Image.height = OptiXImageData.Height;
Image.rowStrideInBytes = OptiXImageData.RowStrideInBytes;
Image.pixelStrideInBytes = OptiXImageData.PixelStrideInBytes;
Image.format = static_cast<OptixPixelFormat>(OptiXImageData.Format);
return Image;
}
class FOpticalFlowContextImpl
{
public:
FOpticalFlowContextImpl() { OpticalFlow = new OptixUtilOpticalFlow();}
~FOpticalFlowContextImpl() { if (OpticalFlow) { delete OpticalFlow; OpticalFlow = nullptr;}}
EOptiXDenoiseResult Init(CUcontext Context, CUstream Stream, unsigned int Width, unsigned int Height);
EOptiXDenoiseResult ComputeFlow(FOptiXImageData& Flow, const FOptiXImageData* Input);
void Destroy();
private:
OptixUtilOpticalFlow* OpticalFlow = nullptr;
};
EOptiXDenoiseResult FOpticalFlowContextImpl::Init(CUcontext Context, CUstream Stream, unsigned int Width, unsigned int Height)
{
OptixResult Result = OpticalFlow->init(Context, Stream, Width, Height);
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOpticalFlowContextImpl::ComputeFlow(FOptiXImageData& Flow, const FOptiXImageData* Input)
{
OptixImage2D FlowRaw = GetOptiXImage2D(Flow);
OptixImage2D InputsRaw[2] = { GetOptiXImage2D(Input[0]), GetOptiXImage2D(Input[1]) };
return OpticalFlow->computeFlow(FlowRaw, (const OptixImage2D*)&InputsRaw[0]);
}
void FOpticalFlowContextImpl::Destroy()
{
OpticalFlow->destroy();
}
FOpticalFlowContext::FOpticalFlowContext()
{
OpticalFlowContextImpl = new FOpticalFlowContextImpl();
}
FOpticalFlowContext::~FOpticalFlowContext()
{
if (OpticalFlowContextImpl)
{
delete OpticalFlowContextImpl;
OpticalFlowContextImpl = nullptr;
}
}
EOptiXDenoiseResult FOpticalFlowContext::Init(CUcontext Context, CUstream Stream, unsigned int Width, unsigned int Height)
{
return OpticalFlowContextImpl->Init(Context, Stream, Width, Height);
}
EOptiXDenoiseResult FOpticalFlowContext::ComputeFlow(FOptiXImageData& Flow, const FOptiXImageData* Input)
{
return OpticalFlowContextImpl->ComputeFlow(Flow, Input);
}
void FOpticalFlowContext::Destroy()
{
OpticalFlowContextImpl->Destroy();
}
class FOptiXDenoiseContextImpl
{
public:
EOptiXDenoiseResult InitializeDenoiser(bool GuildNormal, bool GuildAlbedo, bool TemporalMode);
EOptiXDenoiseResult ComputeMemoryResource(uint32_t OutputWidth, uint32_t OutputHeight);
void SetPreviousOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData);
void SetOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData);
EOptiXDenoiseResult SetupDenoiser(CUstream CudaStream, uint32_t InputWidth, uint32_t InputHeight, CUdeviceptr State, size_t StateSizeInBytes, CUdeviceptr Scratch, size_t ScratchSizeInBytes);
void SetDenoiseAlpha(EOptiXDenoiseAlphaMode inMode);
void SetHdrIntensity(CUdeviceptr Intensity) { Params.hdrIntensity = Intensity; }
void SetHdrAverageColor(CUdeviceptr AverageColor) { Params.hdrAverageColor = AverageColor; }
void SetBlendFactor(float BlendFactor) { Params.blendFactor = BlendFactor; }
void SetTemporalModeUsePreviousLayers(bool Use) { Params.temporalModeUsePreviousLayers = Use; }
void SetLayerInput(const FOptiXImageData& OptiXImageData) { Layer.input = GetOptiXImage2D(OptiXImageData); }
void SetLayerOutput(const FOptiXImageData& OptiXImageData) { Layer.output = GetOptiXImage2D(OptiXImageData); }
void SetLayerPreviousOutput(const FOptiXImageData& OptiXImageData) { Layer.previousOutput = GetOptiXImage2D(OptiXImageData); }
void SetGuideLayerNormal(const FOptiXImageData& OptiXImageData) { GuideLayer.normal = GetOptiXImage2D(OptiXImageData); }
void SetGuideLayerAlbedo(const FOptiXImageData& OptiXImageData) { GuideLayer.albedo = GetOptiXImage2D(OptiXImageData); }
void SetGuideLayerFlow(const FOptiXImageData& OptiXImageData) { GuideLayer.flow = GetOptiXImage2D(OptiXImageData); }
EOptiXDenoiseResult ComputeIntensity(CUstream CudaStream, CUdeviceptr Intensity, CUdeviceptr CudaScratch, size_t CudaScratchSize);
EOptiXDenoiseResult ComputeAverageColor(CUstream CudaStream, CUdeviceptr AverageColor, CUdeviceptr CudaScratch, size_t CudaScratchSize);
EOptiXDenoiseResult InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, uint32_t OffsetX, uint32_t OffsetY, CUdeviceptr CudaScratch, size_t CudaScratchSize);
EOptiXDenoiseResult InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, CUdeviceptr CudaScratch, size_t CudaScratchSize, uint32_t CudaOverlap, uint32_t TileWidth, uint32_t TileHeight);
void Destroy();
public:
OptixDeviceContext Context = nullptr;
OptixDenoiser Denoiser = nullptr;
OptixDenoiserParams Params = {};
OptixDenoiserOptions Options;
OptixDenoiserSizes Sizes;
OptixDenoiserLayer Layer = {};
OptixDenoiserGuideLayer GuideLayer = {};
OptixDenoiserModelKind ModelKind;
};
EOptiXDenoiseResult FOptiXDenoiseContextImpl::InitializeDenoiser(bool GuildNormal, bool GuildAlbedo, bool TemporalMode)
{
OptixResult Result;
{
Options = {};
Options.guideAlbedo = GuildAlbedo ? 1 : 0;
Options.guideNormal = GuildNormal ? 1 : 0;
// @todo add AOVs denoising, Add KPMode supported
ModelKind = TemporalMode ? OPTIX_DENOISER_MODEL_KIND_TEMPORAL : OPTIX_DENOISER_MODEL_KIND_HDR;
Result = optixDenoiserCreate(Context, ModelKind, &Options, &Denoiser);
}
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::ComputeMemoryResource(uint32_t OutputWidth, uint32_t OutputHeight)
{
OptixResult Result;
Result = optixDenoiserComputeMemoryResources(Denoiser, OutputWidth, OutputHeight, &Sizes);
return OptiXDenoiseResult(Result);
}
void FOptiXDenoiseContextImpl::SetPreviousOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData)
{
GuideLayer.previousOutputInternalGuideLayer.data = OptiXImageData.Data;
GuideLayer.previousOutputInternalGuideLayer.width = OptiXImageData.Width;
GuideLayer.previousOutputInternalGuideLayer.height = OptiXImageData.Height;
GuideLayer.previousOutputInternalGuideLayer.pixelStrideInBytes = unsigned(Sizes.internalGuideLayerPixelSizeInBytes);
GuideLayer.previousOutputInternalGuideLayer.rowStrideInBytes =
GuideLayer.previousOutputInternalGuideLayer.width * GuideLayer.previousOutputInternalGuideLayer.pixelStrideInBytes;
GuideLayer.previousOutputInternalGuideLayer.format = OPTIX_PIXEL_FORMAT_INTERNAL_GUIDE_LAYER;
}
void FOptiXDenoiseContextImpl::SetOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData)
{
GuideLayer.outputInternalGuideLayer.data = OptiXImageData.Data;
GuideLayer.outputInternalGuideLayer.width = OptiXImageData.Width;
GuideLayer.outputInternalGuideLayer.height = OptiXImageData.Height;
GuideLayer.outputInternalGuideLayer.pixelStrideInBytes = unsigned(Sizes.internalGuideLayerPixelSizeInBytes);
GuideLayer.outputInternalGuideLayer.rowStrideInBytes =
GuideLayer.outputInternalGuideLayer.width * GuideLayer.outputInternalGuideLayer.pixelStrideInBytes;
GuideLayer.outputInternalGuideLayer.format = OPTIX_PIXEL_FORMAT_INTERNAL_GUIDE_LAYER;
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::SetupDenoiser(CUstream CudaStream, uint32_t InputWidth, uint32_t InputHeight, CUdeviceptr State, size_t StateSizeInBytes, CUdeviceptr Scratch, size_t ScratchSizeInBytes)
{
OptixResult Result;
Result = optixDenoiserSetup(
Denoiser,
CudaStream,
//TileWidth + 2 * CudaOverlap,
//TileHeight + 2 * CudaOverlap,
InputWidth,
InputHeight,
State,
StateSizeInBytes,
Scratch,
ScratchSizeInBytes);
return OptiXDenoiseResult(Result);
}
void FOptiXDenoiseContextImpl::SetDenoiseAlpha(EOptiXDenoiseAlphaMode InMode)
{
OptixDenoiserAlphaMode Mode;
switch (InMode)
{
case EOptiXDenoiseAlphaMode::COPY:
Mode = OptixDenoiserAlphaMode::OPTIX_DENOISER_ALPHA_MODE_COPY;
break;
case EOptiXDenoiseAlphaMode::ALPHA_AS_AOV:
Mode = OptixDenoiserAlphaMode::OPTIX_DENOISER_ALPHA_MODE_ALPHA_AS_AOV;
break;
case EOptiXDenoiseAlphaMode::FULL_DENOISE_PASS:
Mode = OptixDenoiserAlphaMode::OPTIX_DENOISER_ALPHA_MODE_FULL_DENOISE_PASS;
break;
default:
break;
}
Params.denoiseAlpha = Mode;
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::ComputeIntensity(CUstream CudaStream, CUdeviceptr Intensity, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
OptixResult Result = optixDenoiserComputeIntensity(
Denoiser,
CudaStream,
&Layer.input,
Intensity,
CudaScratch,
CudaScratchSize);
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::ComputeAverageColor(CUstream CudaStream, CUdeviceptr AverageColor, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
OptixResult Result = optixDenoiserComputeAverageColor(
Denoiser,
CudaStream,
&Layer.input,
AverageColor,
CudaScratch,
CudaScratchSize);
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, uint32_t OffsetX, uint32_t OffsetY, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
OptixResult Result = optixDenoiserInvoke(
Denoiser,
CudaStream,
&Params,
CudaState,
CudaStateSize,
&GuideLayer,
&Layer,
1u,
OffsetX, // input offset x
OffsetY, // input offset y
CudaScratch,
CudaScratchSize);
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContextImpl::InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, CUdeviceptr CudaScratch, size_t CudaScratchSize, uint32_t CudaOverlap, uint32_t TileWidth, uint32_t TileHeight)
{
OptixResult Result = optixUtilDenoiserInvokeTiled(
Denoiser,
CudaStream,
&Params,
CudaState,
CudaStateSize,
&GuideLayer,
&Layer,
1u,
CudaScratch,
CudaScratchSize,
CudaOverlap,
TileWidth,
TileHeight);
return OptiXDenoiseResult(Result);
}
void FOptiXDenoiseContextImpl::Destroy()
{
if (Denoiser)
{
optixDenoiserDestroy(Denoiser);
Denoiser = nullptr;
}
if (Context)
{
optixDeviceContextDestroy(Context);
Context = nullptr;
}
}
FOptiXDenoiseContext::FOptiXDenoiseContext()
{
OptiXDenoiseContextImpl = new FOptiXDenoiseContextImpl();
}
FOptiXDenoiseContext::~FOptiXDenoiseContext()
{
if (OptiXDenoiseContextImpl)
{
delete OptiXDenoiseContextImpl;
OptiXDenoiseContextImpl = nullptr;
}
}
EOptiXDenoiseResult FOptiXDenoiseContext::InitOptiX()
{
OptixResult Result;
Result = optixInit();
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContext::CreateContext(CUcontext CudaContext, OptiXDenoiseLogCallback LogCallBackFunction, int logCallbackLevel)
{
OptixResult Result;
OptixDeviceContextOptions DeviceContexOptions = {};
DeviceContexOptions.logCallbackFunction = (OptixLogCallback)LogCallBackFunction;
DeviceContexOptions.logCallbackLevel = logCallbackLevel;
Result = optixDeviceContextCreate(CudaContext, &DeviceContexOptions, &OptiXDenoiseContextImpl->Context);
return OptiXDenoiseResult(Result);
}
EOptiXDenoiseResult FOptiXDenoiseContext::InitializeDenoiser(bool GuildNormal, bool GuildAlbedo, bool TemporalMode)
{
return OptiXDenoiseContextImpl->InitializeDenoiser(GuildNormal, GuildAlbedo, TemporalMode);
}
EOptiXDenoiseResult FOptiXDenoiseContext::ComputeMemoryResource(uint32_t OutputWidth, uint32_t OutputHeight)
{
return OptiXDenoiseContextImpl->ComputeMemoryResource(OutputWidth, OutputHeight);
}
size_t FOptiXDenoiseContext::GetWithoutOverlapScratchSizeInBytes() const
{
OptixDenoiserSizes& Sizes = OptiXDenoiseContextImpl->Sizes;
return Sizes.withoutOverlapScratchSizeInBytes;
}
size_t FOptiXDenoiseContext::GetWithOverlapScratchSizeInBytes() const
{
OptixDenoiserSizes& Sizes = OptiXDenoiseContextImpl->Sizes;
return Sizes.withOverlapScratchSizeInBytes;
}
size_t FOptiXDenoiseContext::GetOverlapWindowSizeInPixels() const
{
OptixDenoiserSizes& Sizes = OptiXDenoiseContextImpl->Sizes;
return Sizes.overlapWindowSizeInPixels;
}
size_t FOptiXDenoiseContext::GetStateSizeInBytes() const
{
OptixDenoiserSizes& Sizes = OptiXDenoiseContextImpl->Sizes;
return Sizes.stateSizeInBytes;
}
size_t FOptiXDenoiseContext::GetInternalGuideLayerPixelSizeInBytes() const
{
OptixDenoiserSizes& Sizes = OptiXDenoiseContextImpl->Sizes;
return Sizes.internalGuideLayerPixelSizeInBytes;
}
void FOptiXDenoiseContext::SetPreviousOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetPreviousOutputInternalGuideLayer(OptiXImageData);
}
void FOptiXDenoiseContext::SetOutputInternalGuideLayer(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetOutputInternalGuideLayer(OptiXImageData);
}
EOptiXDenoiseResult FOptiXDenoiseContext::SetupDenoiser(CUstream CudaStream, uint32_t InputWidth, uint32_t InputHeight, CUdeviceptr State, size_t StateSizeInBytes, CUdeviceptr Scratch, size_t ScratchSizeInBytes)
{
return OptiXDenoiseContextImpl->SetupDenoiser(CudaStream, InputWidth, InputHeight, State, StateSizeInBytes, Scratch, ScratchSizeInBytes);
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetDenoiseAlpha(EOptiXDenoiseAlphaMode Mode)
{
OptiXDenoiseContextImpl->SetDenoiseAlpha(Mode);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetHdrIntensity(CUdeviceptr Intensity)
{
OptiXDenoiseContextImpl->SetHdrIntensity(Intensity);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetHdrAverageColor(CUdeviceptr AverageColor)
{
OptiXDenoiseContextImpl->SetHdrAverageColor(AverageColor);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetBlendFactor(float BlendFactor)
{
OptiXDenoiseContextImpl->SetBlendFactor(BlendFactor);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetTemporalModeUsePreviousLayers(bool Use)
{
OptiXDenoiseContextImpl->SetTemporalModeUsePreviousLayers(Use);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetLayerInput(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetLayerInput(OptiXImageData);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetLayerOutput(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetLayerOutput(OptiXImageData);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetLayerPreviousOutput(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetLayerPreviousOutput(OptiXImageData);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetGuideLayerNormal(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetGuideLayerNormal(OptiXImageData);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetGuideLayerAlbedo(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetGuideLayerAlbedo(OptiXImageData);
return *this;
}
FOptiXDenoiseContext& FOptiXDenoiseContext::SetGuideLayerFlow(const FOptiXImageData& OptiXImageData)
{
OptiXDenoiseContextImpl->SetGuideLayerFlow(OptiXImageData);
return *this;
}
EOptiXDenoiseResult FOptiXDenoiseContext::ComputeIntensity(CUstream CudaStream, CUdeviceptr Intensity, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
return OptiXDenoiseContextImpl->ComputeIntensity(CudaStream,Intensity, CudaScratch,CudaScratchSize);
}
EOptiXDenoiseResult FOptiXDenoiseContext::ComputeAverageColor(CUstream CudaStream, CUdeviceptr AverageColor, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
return OptiXDenoiseContextImpl->ComputeAverageColor(CudaStream, AverageColor, CudaScratch, CudaScratchSize);
}
EOptiXDenoiseResult FOptiXDenoiseContext::InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, uint32_t OffsetX, uint32_t OffsetY, CUdeviceptr CudaScratch, size_t CudaScratchSize)
{
return OptiXDenoiseContextImpl->InvokeOptiXDenoise(CudaStream, CudaState, CudaStateSize, OffsetX, OffsetY, CudaScratch, CudaScratchSize);
}
EOptiXDenoiseResult FOptiXDenoiseContext::InvokeOptiXDenoise(CUstream CudaStream, CUdeviceptr CudaState, size_t CudaStateSize, CUdeviceptr CudaScratch, size_t CudaScratchSize, uint32_t CudaOverlap, uint32_t TileWidth, uint32_t TileHeight)
{
return OptiXDenoiseContextImpl->InvokeOptiXDenoise(CudaStream, CudaState, CudaStateSize, CudaScratch, CudaScratchSize, CudaOverlap, TileWidth, TileHeight);
}
void FOptiXDenoiseContext::Destroy()
{
OptiXDenoiseContextImpl->Destroy();
}