[AD] Fix missing masks on symbolic calls inside loops

src/extra/call.cpp

@@ -581,12 +581,29 @@ struct CallOp : public dr::detail::CustomOpBase {
         for (size_t i = 0; i < args.size(); ++i)
             m_args.push_back_borrow((uint32_t) args[i]);
 
+        m_mask_stack = jit_var_mask_peek(backend);
+        if (!m_mask_stack) {
+            // Find wavefront size, sizes are guaranteed to be compatible here
+            size_t size = jit_var_size(index);
+
+            size_t mask_size = jit_var_size(index);
+            size = mask_size > size ? mask_size : size;
+
+            for (uint64_t arg_i : args) {
+                size_t arg_size = jit_var_size((uint32_t)arg_i);
+                size = arg_size > size ? arg_size : size;
+            }
+
+            m_mask_stack = jit_var_mask_default(backend, size);
+        }
+
         m_name_op = "Call: " + m_name;
     }
 
     ~CallOp() {
         jit_var_dec_ref(m_index);
         jit_var_dec_ref(m_mask);
+        jit_var_dec_ref(m_mask_stack);
         if (m_cleanup)
             m_cleanup(m_payload);
     }
@@ -597,6 +614,7 @@ struct CallOp : public dr::detail::CustomOpBase {
 
     /// Implements f(arg..., grad(arg)...) -> grad(rv)...
     void forward() override {
+        scoped_push_mask mask_guard(m_backend, m_mask_stack);
         std::string name = m_name + " [ad, fwd]";
 
         index64_vector args, rv;
@@ -630,6 +648,8 @@ struct CallOp : public dr::detail::CustomOpBase {
     /// Implements f(arg..., grad(rv)...) -> grad(arg) ...
     void backward() override {
         scoped_isolation_guard isolation_guard;
+        scoped_push_mask mask_guard(m_backend, m_mask_stack);
+
         std::string name = m_name + " [ad, bwd]";
 
         index64_vector args, rv;
@@ -766,7 +786,7 @@ struct CallOp : public dr::detail::CustomOpBase {
     std::string m_name, m_name_op;
     const char *m_variant;
     const char *m_domain;
-    uint32_t m_index, m_mask;
+    uint32_t m_index, m_mask, m_mask_stack;
     size_t m_callable_count;
     index32_vector m_args;
     index64_vector m_args2;

tests/test_switch.py

@@ -679,3 +679,42 @@ def f2(a, _):
         result = dr.switch(index, funcs, a, b)
         assert dr.allclose(result, [3, 6, 3, 4])
         assert dr.allclose(buf1.grad, [2, 2, 0, 0])
+
+
+# Loop mask should implicitly be applied to calls (with AD too)
+@pytest.mark.parametrize("symbolic_config", [(True, True), (False, True), (False, False)])
+@pytest.test_arrays('float32,diff,shape=(*)')
+@dr.syntax()
+def test18_apply_loop_mask(t, capsys, symbolic_config):
+    dr.set_flag(dr.JitFlag.Debug, True)
+    Float = t
+    UInt32 = dr.uint32_array_t(t)
+    with dr.scoped_set_flag(dr.JitFlag.SymbolicLoops, symbolic_config[0]):
+        with dr.scoped_set_flag(dr.JitFlag.SymbolicCalls, symbolic_config[1]):
+            def f(value):
+                return value * 2
+            def g(value):
+                return Float(0)
+
+            switch_idx = UInt32([0, 1, 123]) # idx=2 is invalid
+
+            val = Float(2)
+            dr.enable_grad(val)
+            dr.set_grad(val, 1)
+
+            # Every idx does one iteration of the loop, except idx=2
+            count = UInt32(0, 0, 1)
+            out = Float(1)
+            while count < 1:
+                out = dr.switch(switch_idx, [f, g], val)
+                count = count + 1
+
+            dr.forward_to(out)
+
+            # idx=2 output values should be ignored
+            assert dr.allclose(out[0], 4)
+            assert dr.allclose(out[1], 0)
+            assert dr.allclose(out.grad[0], 2)
+            assert dr.allclose(out.grad[1], 0)
+            transcript = capsys.readouterr().err
+            assert "Attempted to invoke callable with index 123, but this value must be strictly smaller than 2" not in transcript