Add partial and optional support for OCIO to MaterialX

CMakeLists.txt

@@ -45,6 +45,7 @@ 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_TESTS "Build unit tests." OFF)
+option(MATERIALX_BUILD_OCIO "Build OpenColorIO support for shader generators." OFF)
 option(MATERIALX_BUILD_BENCHMARK_TESTS "Build benchmark tests." OFF)
 
 option(MATERIALX_BUILD_SHARED_LIBS "Build MaterialX libraries as shared rather than static." 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_BENCHMARK_TESTS)
 mark_as_advanced(MATERIALX_BUILD_SHARED_LIBS)
 mark_as_advanced(MATERIALX_BUILD_MONOLITHIC)

libraries/stdlib/genosl/lib/vector4_extra_ops.osl

@@ -0,0 +1,45 @@
+// Adds some syntactic sugar allowing mixing vector4 and color4 as
+// arguments of some binary operators used by OCIO transform code.
+
+vector4 __operator__mul__(matrix m, vector4 v)
+{
+  return vector4(v.x * m[0][0] + v.y * m[0][1] + v.z * m[0][2] + v.w * m[0][3], 
+                 v.x * m[1][0] + v.y * m[1][1] + v.z * m[1][2] + v.w * m[1][3], 
+                 v.x * m[2][0] + v.y * m[2][1] + v.z * m[2][2] + v.w * m[2][3], 
+                 v.x * m[3][0] + v.y * m[3][1] + v.z * m[3][2] + v.w * m[3][3]);
+}
+
+vector4 __operator__mul__(color4 c, vector4 v)
+{
+  return vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a) * v;
+}
+
+vector4 __operator__mul__(vector4 v, color4 c)
+{
+  return v * vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a);
+}
+
+vector4 __operator__sub__(color4 c, vector4 v)
+{
+  return vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a) - v;
+}
+
+vector4 __operator__add__(vector4 v, color4 c)
+{
+  return v + vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a);
+}
+
+vector4 __operator__add__(color4 c, vector4 v)
+{
+  return vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a) + v;
+}
+
+vector4 pow(color4 c, vector4 v)
+{
+  return pow(vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a), v);
+}
+
+vector4 max(vector4 v, color4 c)
+{
+  return max(v, vector4(c.rgb.r, c.rgb.g, c.rgb.b, c.a));
+}

source/MaterialXGenGlsl/GlslShaderGenerator.cpp

@@ -776,8 +776,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/MaterialXGenMdl/MdlShaderGenerator.cpp

@@ -307,8 +307,12 @@ ShaderNodeImplPtr MdlShaderGenerator::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/MaterialXGenMsl/MslShaderGenerator.cpp

@@ -1269,8 +1269,12 @@ ShaderNodeImplPtr MslShaderGenerator::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

@@ -11,3 +11,9 @@ mx_add_library(MaterialXGenShader
         MaterialXCore
     EXPORT_DEFINE
         MATERIALX_GENSHADER_EXPORTS)
+
+if(MATERIALX_BUILD_OCIO)
+    find_package(OpenColorIO REQUIRED)
+    target_link_libraries(${TARGET_NAME} PUBLIC OpenColorIO::OpenColorIO)
+    target_compile_definitions(${TARGET_NAME} PUBLIC MATERIALX_BUILD_OCIO)
+endif()

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,209 @@
+//
+// Copyright Contributors to the MaterialX Project
+// SPDX-License-Identifier: Apache-2.0
+//
+
+#ifdef MATERIALX_BUILD_OCIO
+
+#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";
+constexpr const char COLOR4_SUFFIX[] = "_color4_temp";
+
+// 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<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.
+        bool isOSL = false;
+        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);
+            isOSL = true;
+        }
+
+        auto functionName = getFunctionName();
+
+        shaderDesc->setFunctionName(functionName.c_str());
+
+        gpuProcessor->extractGpuShaderInfo(shaderDesc);
+
+        string shaderText = shaderDesc->getShaderText();
+
+        // For OSL, we need to extract the function from the shader OCIO creates.
+        if (isOSL)
+        {
+            const ShaderGenerator& shadergen = context.getShaderGenerator();
+            shadergen.emitLibraryInclude("stdlib/genosl/lib/vector4_extra_ops.osl", context, stage);
+            shadergen.emitLineBreak(stage);
+            auto startpos = shaderText.find(string{"color4 "} + shaderDesc->getFunctionName());
+            if (startpos != string::npos)
+            {
+                auto endpos = shaderText.find(string{"outColor = "} + shaderDesc->getFunctionName(), startpos);
+                if (endpos != string::npos)
+                {
+                    shaderText = shaderText.substr(startpos, endpos - startpos);
+                }
+            }
+        }
+
+        stage.addString(shaderText);
+        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.
+        // TODO: Handle LUT samplers.
+        const bool isColor3 = getName().back() == '3';
+
+        const auto& shadergen = context.getShaderGenerator();
+        shadergen.emitLineBegin(stage);
+
+        const auto* output = node.getOutput();
+        const auto* colorInput = node.getInput(0);
+
+        if (context.getShaderGenerator().getTarget() == "genosl")
+        {
+            // For OSL, since swizzling the output of a function is not allowed, we need:
+            // Function call for color4: color4 res = func(in);
+            // Function call for color3:
+            //    color4 res_color4 = func(color4(in, 1.0));
+            //    color res = res_color4.rgb;
+            if (isColor3)
+            {
+                shadergen.emitString("color4 " + output->getVariable() + COLOR4_SUFFIX + " = ", stage);
+                shadergen.emitString(functionName + "(color4(", stage);
+                shadergen.emitInput(colorInput, context, stage);
+                shadergen.emitString(", 1.0))", stage);
+                shadergen.emitLineEnd(stage);
+                shadergen.emitLineBegin(stage);
+                shadergen.emitOutput(output, true, false, context, stage);
+                shadergen.emitString(" = " + output->getVariable() + COLOR4_SUFFIX + ".rgb", stage);
+                shadergen.emitLineEnd(stage);
+            }
+            else
+            {
+                shadergen.emitOutput(output, true, false, context, stage);
+                shadergen.emitString(" = ", stage);
+                shadergen.emitString(functionName + "(", stage);
+                shadergen.emitInput(colorInput, context, stage);
+                shadergen.emitString(")", stage);
+                shadergen.emitLineEnd(stage);
+            }
+        }
+        else
+        {
+            // 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;
+            shadergen.emitOutput(output, true, false, context, stage);
+            shadergen.emitString(" = ", stage);
+
+            shadergen.emitString(functionName + "(", stage);
+            if (isColor3)
+            {
+                if (context.getShaderGenerator().getTarget() == "genglsl")
+                {
+                    shadergen.emitString("vec4(", stage);
+                }
+                else if (context.getShaderGenerator().getTarget() == "genmsl")
+                {
+                    shadergen.emitString("float4(", 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);
+        }
+    }
+}
+
+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
+
+#endif

source/MaterialXGenShader/Nodes/OpenColorIONode.h

@@ -0,0 +1,44 @@
+//
+// 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>
+
+MATERIALX_NAMESPACE_BEGIN
+
+/// OCIO node implementation. Takes an OCIO GpuProcessor and
+/// uses it to inject shadergen code.
+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:
+    string getFunctionName() const;
+};
+
+MATERIALX_NAMESPACE_END
+
+#endif
+#endif