Skip to content

Commit

Permalink
Adds some support for shader gen'ing writable textures in Hgi. Also m…
Browse files Browse the repository at this point in the history
…akes some

progress towards getting the dome light computations to run on Metal.

See PixarAnimationStudios#1656

(Internal change: 2200704)
(Internal change: 2201198)
  • Loading branch information
clach authored and lkerley committed Jan 7, 2022
1 parent a02e484 commit 471ce60
Show file tree
Hide file tree
Showing 20 changed files with 582 additions and 121 deletions.
14 changes: 12 additions & 2 deletions pxr/imaging/hdSt/domeLightComputations.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ HdSt_DomeLightComputationGPU::Execute(
static_cast<HdStResourceRegistry*>(resourceRegistry);

constexpr int localSize = 8;
const bool hasUniforms = _roughness >= 0.0f;

HdStGLSLProgramSharedPtr const computeProgram =
HdStGLSLProgram::GetComputeProgram(
Expand All @@ -149,6 +150,17 @@ HdSt_DomeLightComputationGPU::Execute(
computeDesc.shaderStage = HgiShaderStageCompute;
computeDesc.computeDescriptor.localSize =
GfVec3i(localSize, localSize, 1);

HgiShaderFunctionAddTexture(&computeDesc, "inTexture");
HgiShaderFunctionAddWritableTexture(&computeDesc, "outTexture",
2, HgiFormatFloat16Vec4);
if (hasUniforms) {
HgiShaderFunctionAddConstantParam(
&computeDesc, "inRoughness", HdStTokens->_float);
}
HgiShaderFunctionAddStageInput(
&computeDesc, "hd_GlobalInvocationID", "uvec3",
HgiShaderKeywordTokens->hdGlobalInvocationID);
});
if (!TF_VERIFY(computeProgram)) {
return;
Expand Down Expand Up @@ -249,8 +261,6 @@ HdSt_DomeLightComputationGPU::Execute(

uniform.roughness = _roughness;

bool hasUniforms = uniform.roughness >= 0.0f;

HgiComputePipelineDesc desc;
desc.debugName = "DomeLightComputation";
desc.shaderProgram = computeProgram->GetProgram();
Expand Down
31 changes: 11 additions & 20 deletions pxr/imaging/hdSt/shaders/domeLight.glslfx
Original file line number Diff line number Diff line change
Expand Up @@ -51,13 +51,10 @@

const float PI = 3.1415926536;

layout(binding = 0) uniform sampler2D inTexture;
layout(rgba16f, binding = 1) uniform image2D outTexture;

// compute texture coords based on the size of the output image/texture
vec2 GetTexCoords(ivec2 outCoords)
{
vec2 outDims = imageSize(outTexture);
vec2 outDims = HdGetSize_outTexture();
// apply a (0.5, 0.5) offset to use pixel centers and not pixel corners
vec2 texCoords = (outCoords + vec2(0.5, 0.5)) / outDims;
return texCoords;
Expand All @@ -67,7 +64,7 @@ vec2 GetTexCoords(ivec2 outCoords)
vec3 SampleEnvMapLod(vec3 sampleVec, float sampleLod) {
vec2 coord = vec2((atan(sampleVec.z, sampleVec.x) + PI) / (2.0 * PI),
acos(sampleVec.y) / PI);
return textureLod(inTexture, coord, sampleLod).rgb;
return HdTextureLod_inTexture(coord, sampleLod).rgb;
}

// compute world position from texture coords
Expand Down Expand Up @@ -124,7 +121,7 @@ vec3 SampleEnvMap(vec3 sampleVec)
{
// sample from a mipmap level of the environment map determined by the
// size of the environment map and the number of samples we are taking
ivec2 inDims = textureSize(inTexture, 0);
ivec2 inDims = HdGetSize_inTexture();
float mipLevel = ceil(log2(inDims.x * deltaPhi/(2.0 * PI)) * 2.0f);

return SampleEnvMapLod(sampleVec, mipLevel);
Expand Down Expand Up @@ -155,22 +152,18 @@ vec3 ComputeIrradiance(vec3 inPos)

void main(void)
{
ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);
ivec2 outCoords = ivec2(hd_GlobalInvocationID.xy);

vec2 texCoords = GetTexCoords(outCoords);
vec3 pos3D = GetWorldPos(texCoords);
vec4 outColor = vec4(ComputeIrradiance(pos3D), 1.0);

imageStore(outTexture, outCoords, outColor);
HdSet_outTexture(outCoords, outColor);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.Prefilter

layout(std140, binding=0) uniform Uniforms {
float roughness;
} uniforms;

// Normal Distribution function
float Distribution_GGX(float dotNH, float roughness)
{
Expand All @@ -186,7 +179,7 @@ vec3 PrefilterEnvMap(vec3 R, float roughness)
vec3 V = R;
vec3 color = vec3(0.0);
float totalWeight = 0.0;
float envMapDim = float(textureSize(inTexture, 0).s);
float envMapDim = float(HdGetSize_inTexture().s);
const uint numSamples = 1024u;
for (uint i = 0u; i < numSamples; i++) {
vec2 Xi = Hammersley2d(i, numSamples);
Expand Down Expand Up @@ -218,21 +211,19 @@ vec3 PrefilterEnvMap(vec3 R, float roughness)

void main(void)
{
ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);
ivec2 outCoords = ivec2(hd_GlobalInvocationID.xy);

vec2 texCoords = GetTexCoords(outCoords);
vec3 pos3D = GetWorldPos(texCoords);
vec3 R = normalize(pos3D);
vec4 outColor = vec4(PrefilterEnvMap(R, uniforms.roughness), 1.0);
vec4 outColor = vec4(PrefilterEnvMap(R, inRoughness), 1.0);

imageStore(outTexture, outCoords, outColor);
HdSet_outTexture(outCoords, outColor);
}

--- --------------------------------------------------------------------------
-- glsl DomeLight.BRDF

uniform int sampleLevel = 0;

float Geometry_SchlicksmithGGX(float dotNL, float dotNV, float roughness)
{
float k = (roughness * roughness) / 2.0;
Expand Down Expand Up @@ -274,11 +265,11 @@ vec2 ComputeBRDF(float NoV, float roughness)

void main(void)
{
ivec2 outCoords = ivec2(gl_GlobalInvocationID.xy);
ivec2 outCoords = ivec2(hd_GlobalInvocationID.xy);

vec2 texCoords = GetTexCoords(outCoords);
// texCoords.x represents N dot E and texCoords.y represents roughness
vec4 outColor = vec4(ComputeBRDF(texCoords.x, texCoords.y), 0.0, 1.0);

imageStore(outTexture, outCoords, outColor);
HdSet_outTexture(outCoords, outColor);
}
23 changes: 21 additions & 2 deletions pxr/imaging/hgi/shaderFunctionDesc.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ PXR_NAMESPACE_OPEN_SCOPE
HgiShaderFunctionTextureDesc::HgiShaderFunctionTextureDesc()
: dimensions(2)
, format(HgiFormatInvalid)
, writable(false)
{
}

Expand Down Expand Up @@ -67,7 +68,8 @@ bool operator==(
{
return lhs.nameInShader == rhs.nameInShader &&
lhs.dimensions == rhs.dimensions &&
lhs.format == rhs.format;
lhs.format == rhs.format &&
lhs.writable == rhs.writable;
}

bool operator!=(
Expand Down Expand Up @@ -156,7 +158,7 @@ bool operator==(
lhs.constantParams == rhs.constantParams &&
lhs.stageInputs == rhs.stageInputs &&
lhs.stageOutputs == rhs.stageOutputs &&
lhs.computeDescriptor == rhs.computeDescriptor;
lhs.computeDescriptor == rhs.computeDescriptor &&
lhs.tessellationDescriptor == rhs.tessellationDescriptor;
}

Expand All @@ -178,6 +180,23 @@ HgiShaderFunctionAddTexture(
texDesc.nameInShader = nameInShader;
texDesc.dimensions = dimensions;
texDesc.format = format;
texDesc.writable = false;

desc->textures.push_back(std::move(texDesc));
}

void
HgiShaderFunctionAddWritableTexture(
HgiShaderFunctionDesc * const desc,
const std::string &nameInShader,
const uint32_t dimensions /* = 2 */,
const HgiFormat &format /* = HgiFormatFloat32Vec4*/)
{
HgiShaderFunctionTextureDesc texDesc;
texDesc.nameInShader = nameInShader;
texDesc.dimensions = dimensions;
texDesc.format = format;
texDesc.writable = true;

desc->textures.push_back(std::move(texDesc));
}
Expand Down
12 changes: 12 additions & 0 deletions pxr/imaging/hgi/shaderFunctionDesc.h
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ PXR_NAMESPACE_OPEN_SCOPE
/// <li>format
/// Format of the texture. This is required in APIs where sampler types depend
/// on the texture (e.g., GL) </li>
/// <li>writable
/// Whether the texture is writable.</li>
/// </ul>
///
struct HgiShaderFunctionTextureDesc
Expand All @@ -58,6 +60,7 @@ struct HgiShaderFunctionTextureDesc
std::string nameInShader;
uint32_t dimensions;
HgiFormat format;
bool writable;
};

using HgiShaderFunctionTextureDescVector =
Expand Down Expand Up @@ -293,6 +296,15 @@ HgiShaderFunctionAddTexture(
uint32_t dimensions = 2,
const HgiFormat &format = HgiFormatFloat32Vec4);

/// Adds writable texture descriptor to given shader function descriptor.
HGI_API
void
HgiShaderFunctionAddWritableTexture(
HgiShaderFunctionDesc *desc,
const std::string &nameInShader,
uint32_t dimensions = 2,
const HgiFormat &format = HgiFormatFloat32Vec4);

/// Adds buffer descriptor to given shader function descriptor.
HGI_API
void
Expand Down
58 changes: 58 additions & 0 deletions pxr/imaging/hgi/types.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,64 @@ HgiGetDataSize(
bpt;
}

HgiFormat
HgiGetComponentBaseFormat(
const HgiFormat f)
{
switch(f) {
case HgiFormatUNorm8:
case HgiFormatUNorm8Vec2:
case HgiFormatUNorm8Vec4:
case HgiFormatUNorm8Vec4srgb:
case HgiFormatBC7UNorm8Vec4:
case HgiFormatBC7UNorm8Vec4srgb:
case HgiFormatBC1UNorm8Vec4:
case HgiFormatBC3UNorm8Vec4:
return HgiFormatUNorm8;
case HgiFormatSNorm8:
case HgiFormatSNorm8Vec2:
case HgiFormatSNorm8Vec4:
return HgiFormatSNorm8;
case HgiFormatFloat16:
case HgiFormatFloat16Vec2:
case HgiFormatFloat16Vec3:
case HgiFormatFloat16Vec4:
return HgiFormatFloat16;
case HgiFormatInt16:
case HgiFormatInt16Vec2:
case HgiFormatInt16Vec3:
case HgiFormatInt16Vec4:
return HgiFormatInt16;
case HgiFormatUInt16:
case HgiFormatUInt16Vec2:
case HgiFormatUInt16Vec3:
case HgiFormatUInt16Vec4:
return HgiFormatUInt16;
case HgiFormatFloat32:
case HgiFormatFloat32Vec2:
case HgiFormatFloat32Vec3:
case HgiFormatFloat32Vec4:
return HgiFormatFloat32;
case HgiFormatInt32:
case HgiFormatInt32Vec2:
case HgiFormatInt32Vec3:
case HgiFormatInt32Vec4:
return HgiFormatInt32;
case HgiFormatFloat32UInt8:
return HgiFormatFloat32UInt8;
case HgiFormatBC6FloatVec3:
return HgiFormatBC6FloatVec3;
case HgiFormatBC6UFloatVec3:
return HgiFormatBC6UFloatVec3;
case HgiFormatCount:
case HgiFormatInvalid:
TF_CODING_ERROR("Invalid Format");
return HgiFormatInvalid;
}
TF_CODING_ERROR("Missing Format");
return HgiFormatInvalid;
}

uint16_t
_ComputeNumMipLevels(const GfVec3i &dimensions)
{
Expand Down
6 changes: 6 additions & 0 deletions pxr/imaging/hgi/types.h
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,12 @@ size_t HgiGetDataSize(
HgiFormat f,
const GfVec3i &dimensions);

/// Returns the scalar type of the format, in the form of an HgiFormat, if
/// possible.
HGI_API
HgiFormat HgiGetComponentBaseFormat(
HgiFormat f);

/// Returns mip infos.
///
/// If dataByteSize is specified, the levels stops when the total memory
Expand Down
51 changes: 51 additions & 0 deletions pxr/imaging/hgiGL/conversions.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,45 @@ _primitiveTypeTable[HgiPrimitiveTypeCount][2] =
{HgiPrimitiveTypePatchList, GL_PATCHES}
};

static const std::string
_imageLayoutFormatTable[HgiFormatCount][2] =
{
{"HgiFormatUNorm8", "r8"},
{"HgiFormatUNorm8Vec2", "rg8"},
{"HgiFormatUNorm8Vec4", "rgba8"},
{"HgiFormatSNorm8", "r8_snorm"},
{"HgiFormatSNorm8Vec2", "rg8_snorm"},
{"HgiFormatSNorm8Vec4", "rgba8_snorm"},
{"HgiFormatFloat16", "r16f"},
{"HgiFormatFloat16Vec2", "rg16f"},
{"HgiFormatFloat16Vec3", ""},
{"HgiFormatFloat16Vec4", "rgba16f"},
{"HgiFormatFloat32", "r32f"},
{"HgiFormatFloat32Vec2", "rg32f"},
{"HgiFormatFloat32Vec3", ""},
{"HgiFormatFloat32Vec4", "rgba32f" },
{"HgiFormatInt16", "r16i"},
{"HgiFormatInt16Vec2", "rg16i"},
{"HgiFormatInt16Vec3", ""},
{"HgiFormatInt16Vec4", "rgba16i"},
{"HgiFormatUInt16", "r16ui"},
{"HgiFormatUInt16Vec2", "rg16ui"},
{"HgiFormatUInt16Vec3", ""},
{"HgiFormatUInt16Vec4", "rgba16ui"},
{"HgiFormatInt32", "r32i"},
{"HgiFormatInt32Vec2", "rg32i"},
{"HgiFormatInt32Vec3", ""},
{"HgiFormatInt32Vec4", "rgba32i"},
{"HgiFormatUNorm8Vec4srgb", ""},
{"HgiFormatBC6FloatVec3", ""},
{"HgiFormatBC6UFloatVec3", ""},
{"HgiFormatBC7UNorm8Vec4", ""},
{"HgiFormatBC7UNorm8Vec4srgb", ""},
{"HgiFormatBC1UNorm8Vec4", ""},
{"HgiFormatBC3UNorm8Vec4", ""},
{"HgiFormatFloat32UInt8", ""},
};

void
HgiGLConversions::GetFormat(
HgiFormat inFormat,
Expand Down Expand Up @@ -385,4 +424,16 @@ HgiGLConversions::GetPrimitiveType(HgiPrimitiveType pt)
return _primitiveTypeTable[pt][1];
}

std::string
HgiGLConversions::GetImageLayoutFormatQualifier(HgiFormat inFormat)
{
const std::string layoutQualifier = _imageLayoutFormatTable[inFormat][1];
if (layoutQualifier.empty()) {
TF_WARN("Given HgiFormat is not a supported image unit format, "
"defaulting to rgba16f");
return _imageLayoutFormatTable[9][1];
}
return layoutQualifier;
}

PXR_NAMESPACE_CLOSE_SCOPE
3 changes: 3 additions & 0 deletions pxr/imaging/hgiGL/conversions.h
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,9 @@ class HgiGLConversions final

HGIGL_API
static GLenum GetPrimitiveType(HgiPrimitiveType pt);

HGIGL_API
static std::string GetImageLayoutFormatQualifier(HgiFormat inFormat);
};


Expand Down
Loading

0 comments on commit 471ce60

Please sign in to comment.