Share Shading Context when optimizing/jitting a shader…

src/liboslexec/context.cpp

@@ -68,6 +68,25 @@ ShadingContext::~ShadingContext()
     free_dict_resources();
 }
 
+ShadingContext::RestoreState
+ShadingContext::repurposeForJit()
+{
+    process_errors();
+    // Match previous behavior of nullptr value for texture thread info
+    // as if we created a new ShadingContext
+    RestoreState restore_state { texture_thread_info() };
+    texture_thread_info(nullptr);
+    return restore_state;
+}
+
+// Process any errors from JIT and restore the texture_thread_info
+void
+ShadingContext::restoreFromJit(const RestoreState& restore_state)
+{
+    // Restore "this" ShadingContext for execution
+    process_errors();
+    texture_thread_info(restore_state.m_pre_jit_texture_thread_info);
+}
 
 
 bool
@@ -86,14 +105,14 @@ ShadingContext::execute_init(ShaderGroup& sgroup, int threadindex,
     if (sgroup.nlayers()) {
         sgroup.start_running();
         if (!sgroup.jitted()) {
-            auto ctx = shadingsys().get_context(thread_info());
-            shadingsys().optimize_group(sgroup, ctx, true /*do_jit*/);
+            auto restore_state = repurposeForJit();
+            shadingsys().optimize_group(sgroup, this, true /*do_jit*/);
             if (shadingsys().m_greedyjit
                 && shadingsys().m_groups_to_compile_count) {
                 // If we are greedily JITing, optimize/JIT everything now
                 shadingsys().optimize_all_groups();
             }
-            shadingsys().release_context(ctx);
+            restoreFromJit(restore_state);
         }
         if (sgroup.does_nothing())
             return false;
@@ -260,19 +279,16 @@ ShadingContext::Batched<WidthT>::execute_init(
     if (sgroup.nlayers()) {
         sgroup.start_running();
         if (!sgroup.batch_jitted()) {
-            // Matching ShadingContext::execute_init behavior
-            // of grabbing another context.
-            // TODO:  Is this necessary, why can't we just use the
-            // the existing context()?
-            //shadingsys().template batched<WidthT>().jit_group(sgroup, &context());
-            auto ctx = shadingsys().get_context(context().thread_info());
-            shadingsys().template batched<WidthT>().jit_group(sgroup, ctx);
+            auto restore_state = context().repurposeForJit();
+            shadingsys().template batched<WidthT>().jit_group(sgroup,
+                                                              &context());
+
             if (shadingsys().m_greedyjit
                 && shadingsys().m_groups_to_compile_count) {
                 // If we are greedily JITing, optimize/JIT everything now
                 shadingsys().template batched<WidthT>().jit_all_groups();
             }
-            shadingsys().release_context(ctx);
+            context().restoreFromJit(restore_state);
         }
         // To handle layers that were not used but still possibly had
         // render outputs, we always generate a run function even for

src/liboslexec/oslexec_pvt.h

@@ -2518,6 +2518,18 @@ class OSLEXECPUBLIC ShadingContext {
     // wide data offsets should be used
     int batch_size_executed;
     bool execution_is_batched() const { return batch_size_executed != 0; }
+
+
+    struct RestoreState {
+        TextureSystem::Perthread* m_pre_jit_texture_thread_info;
+    };
+
+    // Rather than allocate an additional ShadingContext for JIT
+    // reuse this one by processing any existing errors
+    // and saving off any necessary state to be restored afterwards
+    RestoreState repurposeForJit();
+    // Process any errors from JIT and restore the state
+    void restoreFromJit(const RestoreState&);
 };