From 84c1a9f8e5a8c413dea56852894541f0cc28193b Mon Sep 17 00:00:00 2001
From: Brecht Van Lommel <brecht@blender.org>
Date: Tue, 4 Feb 2025 23:33:50 +0100
Subject: [PATCH] Fix Metal shader errors with MaterialX 1.39

* mx_microfacet.glsl uses functions from mx_math.metal, so the
  latter must be included first.
* Defining atan(y, x) for GLSL compatibility conflicts with
  mx_math.metal using ::atan(y_over_x). Work around this by tweaking
  the definition.
* textureQueryLevels(u_envRadiance) does not take into account that
  this texture is wrapped in a MetalTexture class.

Signed-off-by: Brecht Van Lommel <brecht@blender.org>
---
 pxr/imaging/hdSt/materialXShaderGen.cpp | 17 +++++++++++++----
 pxr/imaging/hgiMetal/shaderGenerator.mm |  4 +++-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/pxr/imaging/hdSt/materialXShaderGen.cpp b/pxr/imaging/hdSt/materialXShaderGen.cpp
index 58687a7607..8345b6fb2d 100644
--- a/pxr/imaging/hdSt/materialXShaderGen.cpp
+++ b/pxr/imaging/hdSt/materialXShaderGen.cpp
@@ -447,8 +447,17 @@ HdStMaterialXShaderGen<Base>::_EmitMxInitFunction(
         emitLine("u_envIrradiance = HdGetSampler_domeLightFallback()", mxStage);
         emitLine("u_envRadiance = HdGetSampler_domeLightFallback()", mxStage);
         emitLine("#endif", mxStage, false);
+        emitLine("u_envRadianceMips = textureQueryLevels(u_envRadiance)", mxStage);
+    }
+    else {
+        if (std::is_same_v<Base, MaterialX::MslShaderGenerator>) {
+            // Msl has this wrapped in a MetalTexture class
+            emitLine("u_envRadianceMips = textureQueryLevels(u_envRadiance.tex)", mxStage);
+        }
+        else {
+            emitLine("u_envRadianceMips = textureQueryLevels(u_envRadiance)", mxStage);
+        }
     }
-    emitLine("u_envRadianceMips = textureQueryLevels(u_envRadiance)", mxStage);
     Base::emitLineBreak(mxStage);
 
     // Initialize MaterialX Texture samplers with HdGetSampler equivalents
@@ -1309,12 +1318,12 @@ HdStMaterialXShaderGenMsl::_EmitMxFunctions(
     mx::GenContext& mxContext,
     mx::ShaderStage& mxStage) const
 {
-    mx::ShaderGenerator::emitLibraryInclude(
-        "pbrlib/" + mx::GlslShaderGenerator::TARGET
-        + "/lib/mx_microfacet.glsl", mxContext, mxStage);
     mx::ShaderGenerator::emitLibraryInclude(
         "stdlib/" + mx::MslShaderGenerator::TARGET
         + "/lib/mx_math.metal", mxContext, mxStage);
+    mx::ShaderGenerator::emitLibraryInclude(
+        "pbrlib/" + mx::GlslShaderGenerator::TARGET
+        + "/lib/mx_microfacet.glsl", mxContext, mxStage);
     _EmitConstantsUniformsAndTypeDefs(
         mxContext, mxStage,_syntax->getConstantQualifier());
 
diff --git a/pxr/imaging/hgiMetal/shaderGenerator.mm b/pxr/imaging/hgiMetal/shaderGenerator.mm
index f1db1beb1d..6ace769997 100644
--- a/pxr/imaging/hgiMetal/shaderGenerator.mm
+++ b/pxr/imaging/hgiMetal/shaderGenerator.mm
@@ -515,7 +515,9 @@ void _Init(
               "template <typename T>\n"
               "T mod(T y, T x) { return fmod(y, x); }\n\n"
               "template <typename T>\n"
-              "T atan(T y, T x) { return atan2(y, x); }\n\n"
+              "T atan(T y_over_x) { return ::atan(y_over_x); }\n\n"
+              "template <typename T>\n"
+              "T atan(T y, T x) { return ::atan2(y, x); }\n\n"
               "template <typename T>\n"
               "T bitfieldReverse(T x) { return reverse_bits(x); }\n\n"
               "template <typename T>\n"