MaterialXOCIO backported to MaterialX

CMakeLists.txt

@@ -44,6 +44,7 @@ option(MATERIALX_BUILD_GEN_MDL "Build the MDL shader generator back-end." ON)
 option(MATERIALX_BUILD_GEN_MSL "Build the MSL shader generator back-end." ON)
 option(MATERIALX_BUILD_RENDER "Build the MaterialX Render modules." ON)
 option(MATERIALX_BUILD_OIIO "Build OpenImageIO support for MaterialXRender." OFF)
+option(MATERIALX_BUILD_OCIO "Build OpenColorIO support for shader generators." OFF)
 option(MATERIALX_BUILD_TESTS "Build unit tests." ON)
 option(MATERIALX_BUILD_BENCHMARK_TESTS "Build benchmark tests." OFF)
 
@@ -135,6 +136,7 @@ mark_as_advanced(MATERIALX_BUILD_GEN_MDL)
 mark_as_advanced(MATERIALX_BUILD_GEN_MSL)
 mark_as_advanced(MATERIALX_BUILD_RENDER)
 mark_as_advanced(MATERIALX_BUILD_OIIO)
+mark_as_advanced(MATERIALX_BUILD_OCIO)
 mark_as_advanced(MATERIALX_BUILD_TESTS)
 mark_as_advanced(MATERIALX_BUILD_BENCHMARK_TESTS)
 mark_as_advanced(MATERIALX_BUILD_SHARED_LIBS)

source/MaterialXGenGlsl/GlslShaderGenerator.cpp

@@ -885,8 +885,12 @@ ShaderNodeImplPtr GlslShaderGenerator::getImplementation(const NodeDef& nodedef,
     }
     else if (implElement->isA<Implementation>())
     {
-        // Try creating a new in the factory.
-        impl = _implFactory.create(name);
+        if (getColorManagementSystem() && getColorManagementSystem()->hasImplementation(name)) {
+            impl = getColorManagementSystem()->createImplementation(name);
+        } else {
+            // Try creating a new in the factory.
+            impl = _implFactory.create(name);
+        }
         if (!impl)
         {
             // Fall back to source code implementation.

source/MaterialXGenShader/CMakeLists.txt

@@ -1,5 +1,9 @@
 set(MATERIALX_MODULE_NAME MaterialXGenShader)
 
+if(MATERIALX_BUILD_OCIO)
+    find_package(OpenColorIO REQUIRED)
+endif()
+
 file(GLOB_RECURSE materialx_source "${CMAKE_CURRENT_SOURCE_DIR}/*.cpp")
 file(GLOB_RECURSE materialx_headers "${CMAKE_CURRENT_SOURCE_DIR}/*.h*")
 
@@ -31,6 +35,13 @@ target_link_libraries(
     MaterialXFormat
     ${CMAKE_DL_LIBS})
 
+if(MATERIALX_BUILD_OCIO)
+    target_compile_definitions(${MATERIALX_MODULE_NAME} PUBLIC MATERIALX_BUILD_OCIO)
+    target_link_libraries(
+        ${MATERIALX_MODULE_NAME}
+        OpenColorIO::OpenColorIO)
+endif()
+
 target_include_directories(${MATERIALX_MODULE_NAME}
     PUBLIC
         $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../>

source/MaterialXGenShader/ColorManagementSystem.h

@@ -64,6 +64,12 @@ class MX_GENSHADER_API ColorManagementSystem
     ShaderNodePtr createNode(const ShaderGraph* parent, const ColorSpaceTransform& transform, const string& name,
                              GenContext& context) const;
 
+    /// Returns true if the CMS can create a shader node implementation for a locally managed CMS transform
+    virtual bool hasImplementation(const string& /*implName*/) const { return false; }
+
+    /// Create an CMS node implementation for a locally managed transform
+    virtual ShaderNodeImplPtr createImplementation(const string& /*implName*/) const { return {}; }
+
   protected:
     /// Protected constructor
     ColorManagementSystem();

source/MaterialXGenShader/Nodes/OpenColorIONode.cpp

@@ -0,0 +1,142 @@
+//
+// Copyright Contributors to the MaterialXOCIO Project
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#include <MaterialXGenShader/Nodes/OpenColorIONode.h>
+#include <MaterialXGenShader/OpenColorIOManagementSystem.h>
+
+#include <MaterialXCore/Interface.h>
+#include <MaterialXGenGlsl/GlslShaderGenerator.h>
+#include <MaterialXGenShader/GenContext.h>
+#include <MaterialXGenShader/Library.h>
+#include <MaterialXGenShader/ShaderNode.h>
+#include <MaterialXGenShader/Shader.h>
+#include <MaterialXGenShader/ShaderStage.h>
+
+#include <OpenColorIO/OpenColorIO.h>
+#include <OpenColorIO/OpenColorTypes.h>
+
+#include <cstring>
+#include <functional>
+#include <memory>
+#include <string>
+
+MATERIALX_NAMESPACE_BEGIN
+
+namespace
+{
+// Internal OCIO strings:
+constexpr const char OCIO_COLOR3[] = "color3";
+
+// Lengths where needed:
+constexpr auto OCIO_COLOR3_LEN = sizeof(OCIO_COLOR3) / sizeof(OCIO_COLOR3[0]);
+
+} // namespace
+
+ShaderNodeImplPtr OpenColorIONode::create()
+{
+    return std::make_shared<OpenColorIONode>();
+}
+
+void OpenColorIONode::initialize(const InterfaceElement& element, GenContext& context)
+{
+    ShaderNodeImpl::initialize(element, context);
+
+    // Single function shared between color3 and color4 nodes, use a custom hash with only the function name.
+    _hash = std::hash<std::string>{}(getFunctionName());
+}
+
+void OpenColorIONode::emitFunctionDefinition(
+    const ShaderNode& /*node*/,
+    GenContext& context,
+    ShaderStage& stage) const
+{
+    if (stage.getName() == Stage::PIXEL)
+    {
+        auto ocioManager = std::dynamic_pointer_cast<OpenColorIOManagementSystem>(context.getShaderGenerator().getColorManagementSystem());
+
+        auto gpuProcessor = ocioManager->getGpuProcessor(getName());
+        OCIO::GpuShaderDescRcPtr shaderDesc = OCIO::GpuShaderDesc::CreateShaderDesc();
+
+        // TODO: Extend to essl and MDL and possibly SLang.
+        if (context.getShaderGenerator().getTarget() == "genglsl") {
+            shaderDesc->setLanguage(OCIO::GPU_LANGUAGE_GLSL_4_0);
+        } else if (context.getShaderGenerator().getTarget() == "genmsl") {
+            shaderDesc->setLanguage(OCIO::GPU_LANGUAGE_MSL_2_0);
+        } else if (context.getShaderGenerator().getTarget() == "genosl") {
+            shaderDesc->setLanguage(OCIO::LANGUAGE_OSL_1);
+        }
+
+        auto functionName = getFunctionName();
+
+        shaderDesc->setFunctionName(functionName.c_str());
+
+        gpuProcessor->extractGpuShaderInfo(shaderDesc);
+
+        stage.addString(shaderDesc->getShaderText());
+        stage.endLine(false);
+    }
+}
+
+void OpenColorIONode::emitFunctionCall(
+    const ShaderNode& node,
+    GenContext& context,
+    ShaderStage& stage) const
+{
+    if (stage.getName() == Stage::PIXEL)
+    {
+        auto functionName = getFunctionName();
+
+        // TODO: Adjust syntax for other languages. OSL definitely does not have a native color4 type
+        //       and requires an add-on struct to be used.
+
+        // The OCIO function uses a vec4 parameter, so:
+        // Function call for color4: vec4 res = func(in);
+        // Function call for color3: vec3 res = func(vec4(in, 1.0)).rgb;
+        // TODO: Handle LUT samplers.
+        bool isColor3 = getName().back() == '3';
+
+        const auto& shadergen = context.getShaderGenerator();
+        shadergen.emitLineBegin(stage);
+
+        const auto* output = node.getOutput();
+        const auto* colorInput = node.getInput(0);
+
+        shadergen.emitOutput(output, true, false, context, stage);
+        shadergen.emitString(" = ", stage);
+
+        shadergen.emitString(functionName + "(", stage);
+        if (isColor3)
+        {
+            shadergen.emitString("vec4(", stage);
+        }
+        shadergen.emitInput(colorInput, context, stage);
+        if (isColor3)
+        {
+            shadergen.emitString(", 1.0)", stage);
+        }
+
+        shadergen.emitString(")", stage);
+
+        if (isColor3)
+        {
+            shadergen.emitString(".rgb", stage);
+        }
+
+        shadergen.emitLineEnd(stage);
+    }
+}
+
+std::string OpenColorIONode::getFunctionName() const
+{
+    auto name = getName();
+
+    // Strip _color3 and _color4 suffixes and impl prefix:
+    size_t startPos = OpenColorIOManagementSystem::IMPL_PREFIX.size();
+    size_t length = name.size() - OCIO_COLOR3_LEN - 1 - startPos;
+
+    return name.substr(startPos, length);
+}
+
+MATERIALX_NAMESPACE_END

source/MaterialXGenShader/Nodes/OpenColorIONode.h

@@ -0,0 +1,48 @@
+//
+// Copyright Contributors to the MaterialX Project
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#ifndef MATERIALX_OCIO_NODE_H
+#define MATERIALX_OCIO_NODE_H
+
+#ifdef MATERIALX_BUILD_OCIO
+/// @file
+/// OCIO node implementation
+
+#include <MaterialXGenShader/ShaderNodeImpl.h>
+#include <MaterialXCore/Interface.h>
+#include <MaterialXGenShader/GenContext.h>
+#include <MaterialXGenShader/Library.h>
+#include <MaterialXGenShader/ShaderNode.h>
+#include <MaterialXGenShader/ShaderStage.h>
+
+#include <string>
+
+namespace mx = MaterialX;
+
+MATERIALX_NAMESPACE_BEGIN
+
+/// GLSL OCIO node implementation. Takes a Maya OCIO shader fragment and
+/// makes it compatible with the shadergen
+class OpenColorIONode : public ShaderNodeImpl
+{
+  public:
+    static ShaderNodeImplPtr create();
+
+    void initialize(const InterfaceElement& element, GenContext& context) override;
+
+    void emitFunctionDefinition(const ShaderNode& node, GenContext& context, ShaderStage& stage)
+        const override;
+
+    void emitFunctionCall(const ShaderNode& node, GenContext& context, ShaderStage& stage)
+        const override;
+
+  private:
+    std::string getFunctionName() const;
+};
+
+MATERIALX_NAMESPACE_END
+
+#endif
+#endif