Add an environment pre-convolution shader

libraries/pbrlib/genglsl/lib/mx_environment_prefilter.glsl

@@ -17,10 +17,18 @@ vec3 mx_environment_radiance(vec3 N, vec3 V, vec3 X, vec2 alpha, int distributio
     float avgAlpha = mx_average_alpha(alpha);
     vec3 F = mx_compute_fresnel(NdotV, fd);
     float G = mx_ggx_smith_G2(NdotV, NdotV, avgAlpha);
-    vec3 FG = fd.refraction ? vec3(1.0) - (F * G) : F * G;
+    // Convert n and k to f0. f90 will be equal to 1.
+    // Source: Section 2.2: https://jcgt.org/published/0003/04/03/paper.pdf
+    vec3 k2 = fd.extinction * fd.extinction;
+    vec3 F0 = ((fd.ior - 1)*(fd.ior - 1) + k2) / ((fd.ior + 1)*(fd.ior + 1) + k2);
+    vec3 FGD = mx_ggx_dir_albedo(NdotV, avgAlpha, F0, 1.0);
+    if (fd.refraction) {
+        FGD = 1.0 - FGD;
+    }
+
 
     vec3 Li = mx_latlong_map_lookup(L, $envMatrix, mx_latlong_compute_lod(avgAlpha), $envRadiance);
-    return Li * FG;
+    return Li * FGD;
 }
 
 vec3 mx_environment_irradiance(vec3 N)

libraries/pbrlib/genglsl/lib/mx_pre_convolve_environment.glsl

@@ -0,0 +1,9 @@
+#include "mx_microfacet_specular.glsl"
+
+vec3 mx_pre_convolve_environment()
+{
+    vec2 uv = gl_FragCoord.xy / vec2(2048.0, 1024.0);
+    // vec2 ggxDirAlbedo = mx_ggx_dir_albedo(uv.x, uv.y, vec3(1, 0, 0), vec3(0, 1, 0)).xy;
+    return textureLod($envRadiance, uv, 0).rgb * vec3(uv, 0.0);
+    // return vec3(uv, 0.0);
+}

libraries/pbrlib/genglsl/mx_conductor_bsdf.glsl

@@ -66,5 +66,6 @@ void mx_conductor_bsdf_indirect(vec3 V, float weight, vec3 ior_n, vec3 ior_k, ve
 
     vec3 Li = mx_environment_radiance(N, V, X, safeAlpha, distribution, fd);
 
-    bsdf.response = Li * comp * weight;
+    // bsdf.response = Li * comp * weight;
+    bsdf.response = Li;
 }

source/JsMaterialX/JsMaterialXGenShader/JsGenOptions.cpp

@@ -42,5 +42,6 @@ EMSCRIPTEN_BINDINGS(GenOptions)
         .property("hwMaxActiveLightSources", &mx::GenOptions::hwMaxActiveLightSources)
         .property("hwNormalizeUdimTexCoords", &mx::GenOptions::hwNormalizeUdimTexCoords)
         .property("hwWriteAlbedoTable", &mx::GenOptions::hwWriteAlbedoTable)
+        .property("hwWriteEnvPreConvolution", &mx::GenOptions::hwWriteEnvPreConvolution)
         ;
 }

source/MaterialXGenGlsl/GlslShaderGenerator.cpp

@@ -595,6 +595,13 @@ void GlslShaderGenerator::emitPixelStage(const ShaderGraph& graph, GenContext& c
         emitLineBreak(stage);
     }
 
+    // Emit environment pre-convolution code
+    if (context.getOptions().hwWriteEnvPreConvolution)
+    {
+        // TODO: Implement this
+        assert(false);
+    }
+
     // Set the include file to use for uv transformations,
     // depending on the vertical flip flag.
     if (context.getOptions().fileTextureVerticalFlip)

source/MaterialXGenMsl/MslShaderGenerator.cpp

@@ -1063,6 +1063,13 @@ void MslShaderGenerator::emitPixelStage(const ShaderGraph& graph, GenContext& co
             emitLineBreak(stage);
         }
 
+        // Emit environment pre-convolution code
+        if (context.getOptions().hwWriteEnvPreConvolution)
+        {
+            emitLibraryInclude("pbrlib/genglsl/lib/mx_pre_convolve_environment.glsl", context, stage);
+            emitLineBreak(stage);
+        }
+
         // Set the include file to use for uv transformations,
         // depending on the vertical flip flag.
         if (context.getOptions().fileTextureVerticalFlip)
@@ -1108,6 +1115,10 @@ void MslShaderGenerator::emitPixelStage(const ShaderGraph& graph, GenContext& co
         {
             emitLine(outputSocket->getVariable() + " = float4(mx_generate_dir_albedo_table(), 1.0)", stage);
         }
+        else if (context.getOptions().hwWriteEnvPreConvolution)
+        {
+            emitLine(outputSocket->getVariable() + " = float4(mx_pre_convolve_environment(), 1.0)", stage);
+        }
         else
         {
             // Add all function calls.

source/MaterialXGenShader/GenOptions.h

@@ -90,6 +90,7 @@ class MX_GENSHADER_API GenOptions
         hwMaxActiveLightSources(3),
         hwNormalizeUdimTexCoords(false),
         hwWriteAlbedoTable(false),
+        hwWriteEnvPreConvolution(false),
         hwImplicitBitangents(true),
         emitColorTransforms(true)
     {
@@ -174,6 +175,11 @@ class MX_GENSHADER_API GenOptions
     /// Defaults to false.
     bool hwWriteAlbedoTable;
 
+    /// Enables the generation of the pre-convolved environment map, commonly called the "LD" term
+    /// in the split sum approximation.
+    /// Defaults to false.
+    bool hwWriteEnvPreConvolution;
+
     /// Calculate fallback bitangents from existing normals and tangents
     /// inside the bitangent node.
     bool hwImplicitBitangents;

source/MaterialXGenShader/HwShaderGenerator.cpp

@@ -359,6 +359,12 @@ ShaderPtr HwShaderGenerator::createShader(const string& name, ElementPtr element
         psPrivateUniforms->add(Type::INTEGER, HW::T_ALBEDO_TABLE_SIZE, Value::createValue<int>(64));
     }
 
+    // Add uniforms for the environment pre-convolution
+    if (context.getOptions().hwWriteEnvPreConvolution)
+    {
+        psPrivateUniforms->add(Type::FILENAME, HW::T_ENV_RADIANCE);
+    }
+
     // Create uniforms for the published graph interface
     for (ShaderGraphInputSocket* inputSocket : graph->getInputSockets())
     {

source/MaterialXRender/Util.cpp

@@ -76,6 +76,26 @@ ShaderPtr createAlbedoTableShader(GenContext& context,
     return shader;
 }
 
+ShaderPtr createEnvPreConvolutionShader(GenContext& context,
+                                        DocumentPtr stdLib,
+                                        const string& shaderName)
+{
+    // Construct a dummy nodegraph.
+    DocumentPtr doc = createDocument();
+    doc->importLibrary(stdLib);
+    NodeGraphPtr nodeGraph = doc->addNodeGraph();
+    NodePtr constant = nodeGraph->addNode("constant");
+    OutputPtr output = nodeGraph->addOutput();
+    output->setConnectedNode(constant);
+
+    // Generate the shader
+    GenContext tableContext = context;
+    tableContext.getOptions().hwWriteEnvPreConvolution = true;
+    ShaderPtr shader = createShader(shaderName, tableContext, output);
+
+    return shader;
+}
+
 ShaderPtr createBlurShader(GenContext& context,
                            DocumentPtr stdLib,
                            const string& shaderName,

source/MaterialXRender/Util.h

@@ -47,6 +47,11 @@ MX_RENDER_API ShaderPtr createAlbedoTableShader(GenContext& context,
                                   DocumentPtr stdLib,
                                   const string& shaderName);
 
+/// Create a shader that generates a pre-convolution map for environment lighting.
+MX_RENDER_API ShaderPtr createEnvPreConvolutionShader(GenContext& context,
+                                  DocumentPtr stdLib,
+                                  const string& shaderName);
+
 /// Create a blur shader, using the given standard libraries for code generation.
 MX_RENDER_API ShaderPtr createBlurShader(GenContext& context,
                            DocumentPtr stdLib,

source/MaterialXRenderMsl/MslPipelineStateObject.mm

@@ -608,6 +608,9 @@ int GetStrideOfMetalType(MTLDataType type)
                         samplingProperties.vaddressMode = ImageSamplingProperties::AddressMode::CLAMP;
                         samplingProperties.filterType = ImageSamplingProperties::FilterType::LINEAR;
 
+                        // This `bindImage` call is actually uploading the image on the first call,
+                        // when the image doesn't exist on the GPU. We need to upload custom mips
+                        // during this stage.
                         static_cast<MaterialX::MetalTextureHandler*>
                         (imageHandler.get())->bindImage(env.second, samplingProperties);
                         bindTexture(renderCmdEncoder, (unsigned int)arg.index, env.second, imageHandler);

source/MaterialXView/RenderPipeline.h

@@ -50,7 +50,7 @@ class RenderPipeline
     virtual void bakeTextures() = 0;
 
     virtual void updateAlbedoTable(int tableSize) = 0;
-    virtual mx::ImagePtr convolveEnvironment(mx::ImagePtr envMip0) = 0;
+    virtual mx::ImagePtr convolveEnvironment() = 0;
     virtual std::shared_ptr<void> createTextureBaker(unsigned int width,
                                                      unsigned int height,
                                                      mx::Image::BaseType baseType) = 0;

source/MaterialXView/RenderPipelineGL.cpp

@@ -110,7 +110,7 @@ void GLRenderPipeline::updateAlbedoTable(int tableSize)
     glDrawBuffer(GL_BACK);
 }
 
-mx::ImagePtr GLRenderPipeline::convolveEnvironment(mx::ImagePtr envMip0)
+mx::ImagePtr GLRenderPipeline::convolveEnvironment()
 {
     // TODO: Implement this.
     return nullptr;

source/MaterialXView/RenderPipelineGL.h

@@ -31,7 +31,7 @@ class GLRenderPipeline : public RenderPipeline
     mx::ImageHandlerPtr createImageHandler() override;
     mx::MaterialPtr     createMaterial() override;
     void updateAlbedoTable(int tableSize) override;
-    mx::ImagePtr convolveEnvironment(mx::ImagePtr envMip0) override;
+    mx::ImagePtr convolveEnvironment() override;
     std::shared_ptr<void> createTextureBaker(unsigned int width,
                                              unsigned int height,
                                              mx::Image::BaseType baseType) override;

source/MaterialXView/RenderPipelineMetal.h

@@ -40,7 +40,7 @@ class MetalRenderPipeline : public RenderPipeline
     mx::ImageHandlerPtr createImageHandler() override;
     mx::MaterialPtr     createMaterial() override;
     void updateAlbedoTable(int tableSize) override;
-    mx::ImagePtr convolveEnvironment(mx::ImagePtr envMip0) override;
+    mx::ImagePtr convolveEnvironment() override;
     void renderFrame(void* color_texture, int shadowMapSize, const char* dirLightNodeCat) override;
     void bakeTextures() override;
     mx::ImagePtr getFrameImage() override;

source/MaterialXView/RenderPipelineMetal.mm

@@ -155,15 +155,19 @@
     }
 }
 
-mx::ImagePtr MetalRenderPipeline::convolveEnvironment(mx::ImagePtr envMip0)
+mx::ImagePtr MetalRenderPipeline::convolveEnvironment()
 {
     auto& genContext    = _viewer->_genContext;
     auto& lightHandler  = _viewer->_lightHandler;
     auto& imageHandler  = _viewer->_imageHandler;
     mx::MetalTextureHandlerPtr mtlImageHandler = std::dynamic_pointer_cast<mx::MetalTextureHandler>(imageHandler);
 
+    mx::ImagePtr envMip0 = lightHandler->getEnvRadianceMap();
     int w = envMip0->getWidth();
     int h = envMip0->getHeight();
+
+    mtlImageHandler->createRenderResources(envMip0, true); // Turn mipmaps off
+
     mx::ImagePtr outTex = mx::Image::create(w, h, 4, mx::Image::BaseType::UINT8, true);
     mtlImageHandler->createRenderResources(outTex, true);
     id<MTLTexture> metalTex = mtlImageHandler->getAssociatedMetalTexture(outTex);
@@ -196,35 +200,12 @@
         [desc setStencilAttachment:nil];
 
         MTL(beginEncoder(desc));
+        [MTL(renderCmdEncoder) setDepthStencilState:MTL_DEPTHSTENCIL_STATE(opaque)];
 
         // Create shader.
-        mx::Color3 col;
-        switch (i)
-        {
-            case 0:
-                col = mx::Color3(1.0f, 0.0f, 0.0f);
-                break;
-            case 1:
-                col = mx::Color3(0.0f, 1.0f, 0.0f);
-                break;
-            case 2:
-                col = mx::Color3(0.0f, 0.0f, 1.0f);
-                break;
-            case 3:
-                col = mx::Color3(1.0f, 0.0f, 1.0f);
-                break;
-            case 4:
-                col = mx::Color3(1.0f, 1.0f, 0.0f);
-                break;
-            case 5:
-                col = mx::Color3(0.0f, 1.0f, 1.0f);
-                break;
-            case 6:
-                col = mx::Color3(1.0f, 1.0f, 1.0f);
-                break;
-        }
-        mx::ShaderPtr hwShader = mx::createConstantShader(genContext, _viewer->_stdLib, "__MIP0__" + std::to_string(i), col);
+        mx::ShaderPtr hwShader = mx::createEnvPreConvolutionShader(genContext, _viewer->_stdLib, "__ENV_PRE_CONVOLUTION__");
         mx::MslMaterialPtr material = mx::MslMaterial::create();
+
         try
         {
             material->generateShader(hwShader);
@@ -235,12 +216,9 @@
             return nullptr;
         }
 
-        // Render albedo table.
+        framebuffer->bind(desc);
         material->bindShader();
-        if (material->getProgram()->hasUniform(mx::HW::ALBEDO_TABLE_SIZE))
-        {
-            material->getProgram()->bindUniform(mx::HW::ALBEDO_TABLE_SIZE, mx::Value::createValue(128));
-        }
+
         material->getProgram()->prepareUsedResources(
                         MTL(renderCmdEncoder),
                         _viewer->_identityCamera,

source/MaterialXView/Viewer.cpp

@@ -503,7 +503,8 @@ void Viewer::loadEnvironmentLight()
         }
     }
 
-    mx::ImagePtr mippedEnvRadianceMap = _renderPipeline->convolveEnvironment(envRadianceMap);
+    _lightHandler->setEnvRadianceMap(envRadianceMap);
+    mx::ImagePtr mippedEnvRadianceMap = _renderPipeline->convolveEnvironment();
     _imageHandler->releaseRenderResources(envRadianceMap);
     // Release any existing environment maps and store the new ones.
     _imageHandler->releaseRenderResources(_lightHandler->getEnvRadianceMap());

source/PyMaterialX/PyMaterialXGenShader/PyGenOptions.cpp

@@ -38,6 +38,7 @@ void bindPyGenOptions(py::module& mod)
         .def_readwrite("hwNormalizeUdimTexCoords", &mx::GenOptions::hwNormalizeUdimTexCoords)
         .def_readwrite("hwAmbientOcclusion", &mx::GenOptions::hwAmbientOcclusion)        
         .def_readwrite("hwWriteAlbedoTable", &mx::GenOptions::hwWriteAlbedoTable)
+        .def_readwrite("hwWriteEnvPreConvolution", &mx::GenOptions::hwWriteEnvPreConvolution)
         .def_readwrite("hwImplicitBitangents", &mx::GenOptions::hwImplicitBitangents)
         .def_readwrite("emitColorTransforms", &mx::GenOptions::emitColorTransforms)
         .def(py::init<>());