Add lazytrace option to avoid combinatoric optixTrace inlining (#1815)

OptiX inlines all calls to optixTrace at module compilation time. Due to OSL's
lazy layer evaluation, if layer A calls layer B M times, and layer B calls
layer C N times, that can lead to M*N inlines of optixTrace if layer C contains
a trace operation.

In practice we've observed single trace ops being inlined hundreds of times,
leading to minutes-long shader compilations.

This patch adds a new option, lazytrace, to run layers with trace ops
unconditionally at the start of shader evaluation. This costs some potential
performance in cases where the trace layer would never be evaluated, but
removes all the compilation penalties that the inlining was presenting.

---------

Signed-off-by: Chris Hellmuth <chellmuth@gmail.com>

Author: chellmuth

src/cmake/testing.cmake

@@ -299,6 +299,7 @@ macro (osl_add_all_tests)
                 isconstant
                 layers layers-Ciassign layers-entry layers-lazy layers-lazyerror
                 layers-nonlazycopy layers-repeatedoutputs
+                lazytrace
                 length-reg linearstep
                 logic loop luminance-reg
                 matrix matrix-reg matrix-arithmetic-reg

src/include/OSL/oslexec.h

@@ -228,6 +228,8 @@ class OSLEXECPUBLIC ShadingSystem {
     ///                              output connections (1). For debugging.
     ///    int lazyerror          Run layers lazily even if they have error
     ///                              ops after optimization (1).
+    ///    int lazytrace          Run layers lazily even if they have trace
+    ///                              operations (1). Recommend 0 for OptiX.
     ///    int lazy_userdata      Retrieve userdata lazily (0).
     ///    int userdata_isconnected  Should interpolated=1 params (that may
     ///                              receive userdata) return true from

src/liboslexec/instance.cpp

@@ -35,6 +35,7 @@ ShaderInstance::ShaderInstance(ShaderMaster::ref master, string_view layername)
     , m_outgoing_connections(false)
     , m_renderer_outputs(false)
     , m_has_error_op(false)
+    , m_has_trace_op(false)
     , m_merged_unused(false)
     , m_last_layer(false)
     , m_entry_layer(false)

src/liboslexec/llvm_gen.cpp

@@ -187,6 +187,13 @@ BackendLLVM::llvm_run_connected_layers(Symbol& sym, int symindex, int opnum,
         const Connection& con(inst()->connection(c));
         // If the connection gives a value to this param
         if (con.dst.param == symindex) {
+            // Non-lazy layers are run upfront directly via llvm_call_layer.
+            // Eliding them here doesn't change the semantics of execution,
+            // but it will prevent optixTrace calls from being repeatedly
+            // inlined when lazytrace=0.
+            if (!group()[con.srclayer]->run_lazily())
+                continue;
+
             // already_run is a set of layers run for this particular op.
             // Just so we don't stupidly do several consecutive checks on
             // whether we ran this same layer. It's JUST for this op.
@@ -3022,6 +3029,12 @@ LLVMGEN(llvm_gen_trace)
     };
     llvm::Value* r = rop.ll.call_function("osl_trace", args);
     rop.llvm_store_value(r, Result);
+
+    // Mark the instance as containing a trace call.
+    // With lazytrace=0, we will want to flag the instance
+    // for eager execution.
+    rop.inst()->has_trace_op(true);
+
     return true;
 }
 

src/liboslexec/oslexec_pvt.h

@@ -663,6 +663,7 @@ class ShadingSystemImpl {
     }
     bool countlayerexecs() const { return m_countlayerexecs; }
     bool lazy_userdata() const { return m_lazy_userdata; }
+    bool lazy_trace() const { return m_lazy_trace; }
     bool userdata_isconnected() const { return m_userdata_isconnected; }
     int profile() const { return m_profile; }
     bool no_noise() const { return m_no_noise; }
@@ -871,6 +872,7 @@ class ShadingSystemImpl {
     bool m_lazyunconnected;       ///< Run lazily even if not connected?
     bool m_lazyerror;             ///< Run lazily even if it has error op
     bool m_lazy_userdata;         ///< Retrieve userdata lazily?
+    bool m_lazy_trace;            ///< Run lazily even if it has trace call
     bool m_userdata_isconnected;  ///< Userdata params isconnected()?
     bool m_clearmemory;           ///< Zero mem before running shader?
     bool m_debugnan;              ///< Root out NaN's?
@@ -1261,6 +1263,9 @@ class ShaderInstance {
     bool has_error_op() const { return m_has_error_op; }
     void has_error_op(bool val) { m_has_error_op = val; }
 
+    bool has_trace_op() const { return m_has_trace_op; }
+    void has_trace_op(bool val) { m_has_trace_op = val; }
+
     /// Should this instance only be run lazily (i.e., not
     /// unconditionally)?
     bool run_lazily() const
@@ -1284,6 +1289,10 @@ class ShaderInstance {
         // Shaders with error ops are not lazy unless lazyerror is on.
         if (!shadingsys().m_lazyerror && has_error_op())
             return false;
+
+        if (!shadingsys().lazy_trace() && has_trace_op())
+            return false;
+
         return true;
     }
 
@@ -1512,6 +1521,7 @@ class ShaderInstance {
     bool m_outgoing_connections;         ///< Any outgoing connections?
     bool m_renderer_outputs;             ///< Any outputs params render outputs?
     bool m_has_error_op;                 ///< Any error ops in the code?
+    bool m_has_trace_op;                 ///< Any trace ops in the code?
     bool m_merged_unused;                ///< Unused because of a merge
     bool m_last_layer;                   ///< Is it the group's last layer?
     bool m_entry_layer;                  ///< Is it an entry layer?

src/liboslexec/shadingsys.cpp

@@ -1045,6 +1045,7 @@ ShadingSystemImpl::ShadingSystemImpl(RendererServices* renderer,
     , m_lazyunconnected(true)
     , m_lazyerror(true)
     , m_lazy_userdata(false)
+    , m_lazy_trace(true)
     , m_userdata_isconnected(false)
     , m_clearmemory(false)
     , m_debugnan(false)
@@ -1580,6 +1581,7 @@ ShadingSystemImpl::attribute(string_view name, TypeDesc type, const void* val)
     ATTR_SET("lazyglobals", int, m_lazyglobals);
     ATTR_SET("lazyunconnected", int, m_lazyunconnected);
     ATTR_SET("lazyerror", int, m_lazyerror);
+    ATTR_SET("lazytrace", int, m_lazy_trace);
     ATTR_SET("lazy_userdata", int, m_lazy_userdata);
     ATTR_SET("userdata_isconnected", int, m_userdata_isconnected);
     ATTR_SET("clearmemory", int, m_clearmemory);
@@ -1769,6 +1771,7 @@ ShadingSystemImpl::getattribute(string_view name, TypeDesc type, void* val)
     ATTR_DECODE("lazylayers", int, m_lazylayers);
     ATTR_DECODE("lazyglobals", int, m_lazyglobals);
     ATTR_DECODE("lazyunconnected", int, m_lazyunconnected);
+    ATTR_DECODE("lazytrace", int, m_lazy_trace);
     ATTR_DECODE("lazy_userdata", int, m_lazy_userdata);
     ATTR_DECODE("userdata_isconnected", int, m_userdata_isconnected);
     ATTR_DECODE("clearmemory", int, m_clearmemory);
@@ -2404,6 +2407,7 @@ ShadingSystemImpl::getstats(int level) const
     BOOLOPT(lazyunconnected);
     BOOLOPT(lazyerror);
     BOOLOPT(lazy_userdata);
+    BOOLOPT(lazy_trace);
     BOOLOPT(userdata_isconnected);
     BOOLOPT(clearmemory);
     BOOLOPT(debugnan);

testsuite/lazytrace/main_shader.osl

@@ -0,0 +1,11 @@
+// Copyright Contributors to the Open Shading Language project.
+// SPDX-License-Identifier: BSD-3-Clause
+// https://github.com/AcademySoftwareFoundation/OpenShadingLanguage
+
+shader main_shader(int x = 0)
+{
+    printf("executing main shader\n");
+    if (u == 99.5230942) {
+        printf("  x = %i\n", x);
+    }
+}

testsuite/lazytrace/ref/out.txt

@@ -0,0 +1,10 @@
+Compiled main_shader.osl -> main_shader.oso
+Compiled trace_shader.osl -> trace_shader.oso
+Connect trace_layer.x to main_layer.x
+executing main shader
+
+Connect trace_layer.x to main_layer.x
+executing trace shader
+  Hit!
+executing main shader
+

testsuite/lazytrace/run.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+
+# Copyright Contributors to the Open Shading Language project.
+# SPDX-License-Identifier: BSD-3-Clause
+# https://github.com/AcademySoftwareFoundation/OpenShadingLanguage
+
+shader_commands = " ".join([
+    "-shader trace_shader trace_layer",
+    "-shader main_shader main_layer",
+    "-connect trace_layer x main_layer x",
+])
+
+# Run once with default (lazytrace=1), and once explicitly disabled
+command += testshade(shader_commands)
+command += testshade("{} --options 'lazytrace=0'".format(shader_commands))

testsuite/lazytrace/trace_shader.osl

@@ -0,0 +1,15 @@
+// Copyright Contributors to the Open Shading Language project.
+// SPDX-License-Identifier: BSD-3-Clause
+// https://github.com/AcademySoftwareFoundation/OpenShadingLanguage
+
+shader trace_shader(output int x = 0)
+{
+    printf("executing trace shader\n");
+    vector dir = vector(1.0, 1.0, 1.0) - P;
+
+    int result = trace(P, dir);
+    if (result) {
+        printf("  Hit!\n");
+        x = 1;
+    }
+}