Scatter elision cleanup

Dr.Jit can elide scatter operations when their result can no longer be
referenced by any other operations. The logic to do so, and when
reference count decreases are needed, was dispersed throughout
``eval.cpp``. This commit simplifies the underlying code.

Author: wjakob

src/eval.cpp

@@ -104,19 +104,15 @@ static std::vector<VisitedKey> visit_later;
 // ====================================================================
 
 // Don't perform scatters, whose output buffer is found to be unreferenced
-bool jitc_var_maybe_suppress_scatter(uint32_t index, Variable *v, uint32_t depth) {
+bool jitc_elide_scatter(uint32_t index, const Variable *v) {
+    if ((VarKind) v->kind != VarKind::Scatter)
+        return false;
     Variable *target = jitc_var(v->dep[0]);
     Variable *target_ptr = jitc_var(target->dep[3]);
-    if (target_ptr->ref_count != 0 || depth != 0)
-        return false;
-
     jitc_log(Debug, "jit_eval(): eliding scatter r%u, whose output is unreferenced.", index);
-    if (callable_depth == 0)
-        jitc_var_dec_ref(index, v);
-    return true;
+    return target_ptr->ref_count == 0;
 }
 
-
 /// Recursively traverse the computation graph to find variables needed by a computation
 static void jitc_var_traverse(uint32_t size, uint32_t index, uint32_t depth = 0) {
     if (!visited.emplace(size, index, depth).second)
@@ -125,7 +121,7 @@ static void jitc_var_traverse(uint32_t size, uint32_t index, uint32_t depth = 0)
     Variable *v = jitc_var(index);
     switch ((VarKind) v->kind) {
         case VarKind::Scatter:
-            if (jitc_var_maybe_suppress_scatter(index, v, depth))
+            if (jitc_elide_scatter(index, v))
                 return;
             break;
 
@@ -690,8 +686,14 @@ void jitc_eval_impl(ThreadState *ts) {
 
     ts->scheduled.clear();
 
-    for (uint32_t index: ts->side_effects)
-        jitc_var_traverse(jitc_var(index)->size, index);
+    for (uint32_t index: ts->side_effects) {
+        Variable *v = jitc_var(index);
+
+        if (jitc_elide_scatter(index, v))
+            jitc_var_dec_ref(index);
+        else
+            jitc_var_traverse(v->size, index);
+    }
 
     ts->side_effects.clear();
 
@@ -747,7 +749,6 @@ void jitc_eval_impl(ThreadState *ts) {
 
     for (ScheduledGroup &group : schedule_groups) {
         jitc_assemble(ts, group);
-
         jitc_run(ts, group);
     }
 

src/llvm_eval.cpp

@@ -957,6 +957,7 @@ static void jitc_llvm_render(Variable *v) {
                 }
             }
             break;
+
         case VarKind::Scatter:
             if (v->literal)
                 jitc_llvm_render_scatter_reduce(v, a0, a1, a2, a3);