Fix gather operations on results of other gather operations

With symbolic gather operations, invalid indices might have been pushed/propagated to previous gather nodes in the dependency tree which could produce out-of-bound memory reads if there was an other gather operation in the tree. This commit also propagates the mask throughout the tree to guarantee that no out-of-bound reads will be made.

Author: njroussel

src/op.cpp

@@ -1415,7 +1415,7 @@ static uint32_t jitc_scatter_gather_index(uint32_t source, uint32_t index) {
 
 /// Change all indices/counters in an expression tree to 'new_index'
 static uint32_t jitc_var_reindex(uint32_t var_index, uint32_t new_index,
-                                 uint32_t size) {
+                                 uint32_t mask, uint32_t size) {
     Variable *v = jitc_var(var_index);
 
     if (v->is_data() || (VarType) v->type == VarType::Void)
@@ -1439,7 +1439,16 @@ static uint32_t jitc_var_reindex(uint32_t var_index, uint32_t new_index,
             uint32_t index_2 = v->dep[i];
             if (!index_2)
                 continue;
-            dep[i] = steal(jitc_var_reindex(index_2, new_index, size));
+
+            if (v->kind == VarKind::Gather && i == 2) {
+                // Gather nodes must have their masks replaced rather than reindexed
+                JitBackend backend = (JitBackend) v->backend;
+                Ref default_mask = steal(jitc_var_mask_default(backend, size));
+                dep[i] = steal(jitc_var_and(mask, default_mask));
+            } else {
+                dep[i] = steal(jitc_var_reindex(index_2, new_index, mask, size));
+            }
+
             v = jitc_var(var_index);
             if (!dep[i])
                 return 0; // recursive call failed, give up
@@ -1527,7 +1536,7 @@ uint32_t jitc_var_gather(uint32_t src, uint32_t index, uint32_t mask) {
     // Don't perform the gather operation if the inputs are trivial / can be re-indexed
     if (!result) {
         Ref index_2 = steal(jitc_var_cast(index, VarType::UInt32, 0));
-        Ref src_reindexed = steal(jitc_var_reindex(src, index_2, var_info.size));
+        Ref src_reindexed = steal(jitc_var_reindex(src, index_2, mask, var_info.size));
         if (src_reindexed) {
             // Temporarily hold an extra reference to prevent 'jitc_var_resize' from changing 'src'
             Ref unused = borrow(src_reindexed);

tests/mem.cpp

@@ -213,3 +213,22 @@ TEST_LLVM(14_gather_symbolic_llvm_mask) {
     Float buf_3 = gather<Float>(buf_2, arange<UInt32>(4, 8, 1));
     jit_assert(strcmp(buf_3.str(), "[5, 6, 7, 8]") == 0);
 }
+
+TEST_BOTH(15_gather_symbolic_multiple_mask) {
+    /* A gather expression that is reindexed/rewritten should properly apply its
+     * mask to any previous gather operations it depends on */
+    Float buf_0 = Float(1, 2, 3, 4, 5, 6, 7, 8);
+
+    // true, true, true, false, true
+    Mask mask_1 = (arange<UInt32>(0, 5, 1) % 4) != 0;
+    UInt32 index_1 = arange<UInt32>(0, 5, 1);
+    Float buf_1 = gather<Float>(buf_0, index_1, mask_1);
+
+    Mask mask_2 = Mask(true, true, false, false);
+    UInt32 index_2 = UInt32(0, 1, -1, -1);
+
+    // This gather will reindex, and should apply `mask_2` to the previous
+    // gather, or else it will lookup invalid memory
+    Float buf_2 = gather<Float>(buf_1, index_2, mask_2);
+    jit_assert(strcmp(buf_2.str(), "[1, 2, 0, 0]") == 0);
+}