fix ImageBlock issues

This commit addresses an ImageBlock issue that occurred when ImageBlock
was used with normalize=True, and when the sample footprint reached
beyond the boundary.

It also updates Dr.Jit-Core to fix a loop recording issue (size of side
effects not correctly propagated) found along the way.

src/render/imageblock.cpp

@@ -242,13 +242,14 @@ MI_VARIANT void ImageBlock<Float, Spectrum>::put(const Point2f &pos,
 
         // Compute the number of filter evaluations needed along each axis
         ScalarVector2u count;
+        uint32_t count_max = dr::ceil2int<uint32_t>(2.f * radius);
         if constexpr (!JIT) {
             if (dr::any(pos_0_u > pos_1_u))
                 return;
             count = count_u;
         } else {
             // Conservative bounds must be used in the vectorized case
-            count = dr::ceil2int<uint32_t>(2.f * radius);
+            count = count_max;
             active &= dr::all(pos_0_u <= pos_1_u);
         }
 
@@ -264,16 +265,15 @@ MI_VARIANT void ImageBlock<Float, Spectrum>::put(const Point2f &pos,
                   *weights_y = (Float *) alloca(sizeof(Float) * count.y());
 
             // Evaluate filters weights along the X and Y axes
-
-            for (uint32_t i = 0; i < count.x(); ++i) {
-                new (weights_x + i)
+            for (uint32_t x = 0; x < count.x(); ++x) {
+                new (weights_x + x)
                     Float(JIT ? m_rfilter->eval(rel_f.x())
                               : m_rfilter->eval_discretized(rel_f.x()));
                 rel_f.x() += 1.f;
             }
 
-            for (uint32_t i = 0; i < count.y(); ++i) {
-                new (weights_y + i)
+            for (uint32_t y = 0; y < count.y(); ++y) {
+                new (weights_y + y)
                     Float(JIT ? m_rfilter->eval(rel_f.y())
                               : m_rfilter->eval_discretized(rel_f.y()));
                 rel_f.y() += 1.f;
@@ -283,11 +283,14 @@ MI_VARIANT void ImageBlock<Float, Spectrum>::put(const Point2f &pos,
             if (unlikely(m_normalize)) {
                 Float wx = 0.f, wy = 0.f;
 
-                for (uint32_t i = 0; i < count.x(); ++i)
-                    wx += weights_x[i];
-
-                for (uint32_t i = 0; i < count.y(); ++i)
-                    wy += weights_y[i];
+                Point2f rel_f2 = dr::ceil(pos_0_f) - pos_f;
+                for (uint32_t i = 0; i < count_max; ++i) {
+                    wx += JIT ? m_rfilter->eval(rel_f2.x())
+                              : m_rfilter->eval_discretized(rel_f2.x());
+                    wy += JIT ? m_rfilter->eval(rel_f2.y())
+                              : m_rfilter->eval_discretized(rel_f2.y());
+                    rel_f2 += 1.f;
+                }
 
                 Float factor = dr::detach(wx * wy);
 
@@ -416,17 +419,11 @@ MI_VARIANT void ImageBlock<Float, Spectrum>::put(const Point2f &pos,
                   *weights_y = (Float *) alloca(sizeof(Float) * count);
 
             for (uint32_t i = 0; i < count; ++i) {
-                Float weight_x = m_rfilter->eval(rel_f.x()),
-                      weight_y = m_rfilter->eval(rel_f.y());
-
-                if (unlikely(m_normalize)) {
-                    dr::masked(weight_x, x + i >= size.x()) = 0.f;
-                    dr::masked(weight_y, y + i >= size.y()) = 0.f;
-                }
+                Float weight_x = m_rfilter->eval(rel_f.x());
+                Float weight_y = m_rfilter->eval(rel_f.y());
 
                 new (weights_x + i) Float(weight_x);
                 new (weights_y + i) Float(weight_y);
-
                 rel_f += 1;
             }
 

src/render/tests/test_imageblock.py

@@ -192,3 +192,28 @@ def test04_read(variants_all, filter_name, border, offset, normalize, enable_ad)
                     ref /= dr.sum(weight)
 
             assert dr.allclose(value, ref, atol=1e-5)
+
+
+@pytest.mark.parametrize("coalesce", [ False, True ])
+@pytest.mark.parametrize("normalize", [ False, True ])
+def test05_boundary_effects(variants_vec_rgb, coalesce, normalize):
+    # Check that everything works correctly even when the image block
+    # is smaller than the filter kernel
+    rfilter = mi.load_dict({'type':'gaussian'})
+    ib = mi.ImageBlock(
+        size=(1, 1),
+        offset=(0, 0),
+        channel_count=1,
+        rfilter=rfilter,
+        normalize=normalize,
+        coalesce=coalesce
+    )
+    ib.put(pos=(0.49, 0.49), values=(dr.ones(mi.Float, 1),))
+
+    v1 = 0.9996645373720975
+    v2 = 0.13499982060871019
+    if normalize:
+        ref = v1 / (v1 + 2*v2)**2
+    else:
+        ref = v1
+    assert dr.allclose(ib.tensor().array, ref, rtol=1e-2)