From 0e19c2f16a4d4d383fb597b19d927d8be14b9dec Mon Sep 17 00:00:00 2001 From: Tobias Rittig Date: Mon, 8 Sep 2025 15:41:51 +0200 Subject: [PATCH 1/5] Support dynamic local arrays and structs in dr::Local (C++) --- include/drjit/local.h | 172 +++++++++++++++++++++++++++++++++------- tests/local_ext.cpp | 63 +++++++++++++-- tests/test_local_ext.py | 88 +++++++++++++++++++- 3 files changed, 287 insertions(+), 36 deletions(-) diff --git a/include/drjit/local.h b/include/drjit/local.h index 7e2fe176a..3ac3c9a5e 100644 --- a/include/drjit/local.h +++ b/include/drjit/local.h @@ -16,8 +16,18 @@ NAMESPACE_BEGIN(drjit) -template struct Local { - using Index = uint32_t; +template , + typename SFINAE = int> +struct Local { + static constexpr size_t Size = Size_; + static_assert(Size != Dynamic, "Scalar local arrays are only fixed size. " + "If you meant to instantiate a JIT variant " + "or a DRJIT_STRUCT you may have forgotten to " + "add the Index template parameter."); + using Value = Value_; + using Index = Index_; using Mask = bool; Local() = default; @@ -57,50 +67,156 @@ template struct Local { Value m_data[Size]; }; -template struct Local> { - static constexpr JitBackend Backend = backend_v; - using Index = uint32_array_t; - using Mask = mask_t; - - Local() { - m_index = jit_array_create(Backend, Value::Type, 1, Size); +/** + * \brief Local memory implemented on top of drjit-core jit_array_* + * \details The array `value` of static or dynamic width will be used + * to initialize the entries of local memory with length `Size`. + * `Size` can be drjit::Dynamic, in which case a call to resize will + * be required before usage. + */ +template +struct Local || (is_array_v && is_drjit_struct_v)>> +{ + static constexpr JitBackend Backend = backend_v; + static constexpr size_t Size = Size_; + using Value = Value_; + using Index = Index_; + using Mask = mask_t; + + /** + * \brief Allocate local memory + * \param value optional inital value (also used when resizing dynamic memory) + */ + Local(Value value = empty()) + : m_size(Size == Dynamic ? 1 : Size), m_value(value) { + initialize(); } - Local(const Value &value) { - uint32_t tmp = jit_array_create(Backend, Value::Type, 1, Size); - m_index = jit_array_init(tmp, value.index()); - jit_var_dec_ref(tmp); + ~Local() { + for (uint32_t index : m_arrays) + jit_var_dec_ref(index); } - - ~Local() { jit_var_dec_ref(m_index); } Local(const Local &) = delete; Local(Local &&l) { - m_index = l.m_index; - l.m_index = 0; + m_arrays.swap(l.m_arrays); + l.m_arrays.clear(); } Local &operator=(const Local &) = delete; Local &operator=(Local &&l) { - jit_var_dec_ref(m_index); - m_index = l.m_index; - l.m_index = 0; + for (uint32_t index : m_arrays) + jit_var_dec_ref(index); + m_arrays.swap(l.m_arrays); + l.m_arrays.clear(); } Value read(const Index &offset, const Mask &active = true) const { - return Value::steal(jit_array_read(m_index, offset.index(), active.index())); + Value result; + size_t counter = 0; + auto callback = [&](auto &result, auto &&callback) -> void { + using T = std::decay_t; + if constexpr (is_jit_v && depth_v == 1) { + if (counter >= m_arrays.size()) + jit_raise("Local::read(): internal error, ran out of " + "variable arrays!"); + result = T::steal(jit_array_read(m_arrays[counter++], + offset.index(), active.index())); + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(result), [&](auto &result) { + callback(result, callback); + }); + } + }; + callback(result, callback); + + if (counter != m_arrays.size()) + jit_raise( + "Local::read(): internal error, did not access all variable " + "arrays!"); + + return result; } void write(const Index &offset, const Value &value, const Mask &active = true) { - uint32_t new_index = jit_array_write(m_index, offset.index(), - value.index(), active.index()); - jit_var_dec_ref(m_index); - m_index = new_index; + size_t counter = 0; + auto callback = [&](auto &value, auto &&callback) -> void { + using T = std::decay_t; + if constexpr (is_jit_v && depth_v == 1) { + if (counter >= m_arrays.size()) + jit_raise("Local::write(): internal error, ran out of " + "variable arrays!"); + + if (value.index_ad()) + jit_raise("Local memory writes are not differentiable. You " + "must use 'drjit.detach()' to disable gradient " + "tracking of the written value."); + + uint32_t result = + jit_array_write(m_arrays[counter], offset.index(), + value.index(), active.index()); + jit_var_dec_ref(m_arrays[counter]); + m_arrays[counter++] = result; + + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(value), + [&](auto &value) { callback(value, callback); }); + } + }; + callback(value, callback); + + if (counter != m_arrays.size()) + jit_raise( + "Local.write(): internal error, did not access all variable " + "arrays!"); } - size_t size() const { return Size; } + size_t size() { return m_size; } + + /** + * Reserve a new array of `length` and discard any current contents + */ + void resize(size_t size) { + for (uint32_t index : m_arrays) + jit_var_dec_ref(index); + m_arrays.clear(); + m_size = size; + initialize(); + } private: - uint32_t m_index; -}; + void initialize() { + auto callback = [&](auto &value, auto &&callback) -> void { + using T = std::decay_t; + if constexpr (is_jit_v && depth_v == 1) { + uint32_t result; + if (!value.empty()) { + uint32_t i1 = value.index(); + size_t width = jit_var_size(i1); + uint32_t i2 = jit_array_create( + backend_v, var_type>::value, + width, m_size); + result = jit_array_init(i2, i1); + jit_var_dec_ref(i2); + } else { + result = jit_array_create( + backend_v, var_type>::value, + 1, m_size); + } + m_arrays.push_back(result); + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(value), + [&](auto &value) { callback(value, callback); }); + } + }; + callback(m_value, callback); + } + size_t m_size; + Value m_value; + vector m_arrays; +}; NAMESPACE_END(drjit) diff --git a/tests/local_ext.cpp b/tests/local_ext.cpp index 02eade5df..45d68615f 100644 --- a/tests/local_ext.cpp +++ b/tests/local_ext.cpp @@ -6,15 +6,68 @@ namespace nb = nanobind; namespace dr = drjit; +using nb::literals::operator""_a; + +template +auto bind_local(nb::module_ &m, const dr::string& name) { + auto c = nb::class_(m, name.c_str()) + .def(nb::init<>()) + .def(nb::init()) + .def("__len__", &Local::size) + .def("read", &Local::read, "index"_a, "active"_a = true) + .def("write", &Local::write, "offset"_a, "value"_a, "active"_a = true); + + if constexpr (dr::is_jit_v) + c = c.def("resize", &Local::resize); + + return c; +} + template void bind(nb::module_ &m) { using UInt32 = dr::uint32_array_t; + using Local10 = dr::Local; + using LocalDyn = dr::Local; + + bind_local(m, "Local10"); + + if constexpr (dr::is_jit_v) + bind_local(m, "LocalDyn"); + + struct MyStruct + { + Float value; + UInt32 priority; + DRJIT_STRUCT(MyStruct, value, priority) + + MyStruct(int i) : value(i), priority(i) {} + }; + + auto mystruct = nb::class_(m, "MyStruct") + .def(nb::init<>()) + .def(nb::init()) + .def_rw("value", &MyStruct::value) + .def_rw("priority", &MyStruct::priority); + + nb::handle u32; + if constexpr (dr::is_array_v) + u32 = nb::type(); + else + u32 = nb::handle((PyObject *) &PyLong_Type); + + nb::handle f32; + if constexpr (dr::is_array_v) + f32 = nb::type(); + else + f32 = nb::handle((PyObject *) &PyFloat_Type); + + nb::dict fields; + fields["value"] = f32; + fields["priority"] = u32; + mystruct.attr("DRJIT_STRUCT") = fields; - m.def("lookup", [](UInt32 offset, Float value, UInt32 offset2) { - dr::Local local(3.f); - local.write(offset, value); - return local.read(offset2); - }); + using LocalStruct10 = dr::Local; + bind_local(m, "LocalStruct10"); } NB_MODULE(local_ext, m) { diff --git a/tests/test_local_ext.py b/tests/test_local_ext.py index ba458f728..c738f6e07 100644 --- a/tests/test_local_ext.py +++ b/tests/test_local_ext.py @@ -11,11 +11,93 @@ def get_pkg(t): return m.cuda elif backend == dr.JitBackend.Invalid: return m.scalar + +def is_constant_valued(local, value): + for i in range(len(local)): + assert dr.allclose(local.read(i), value) @pytest.test_arrays('float32,shape=(*)') -def test01_lookup(t): +def test01_initialization(t): pkg = get_pkg(t) + initial = 25.4 - assert dr.all(pkg.lookup(4, 5.0, 2) == 3.0) - assert dr.all(pkg.lookup(4, 5.0, 4) == 5.0) + l10 = pkg.Local10() + assert len(l10) == 10 + + with pytest.raises(AssertionError): + is_constant_valued(l10, initial) + + l10 = pkg.Local10(initial) + assert len(l10) == 10 + + is_constant_valued(l10, initial) + + +@pytest.test_arrays('float32,is_jit,shape=(*)') +def test02_dynamic_initialization(t): + pkg = get_pkg(t) + initial = dr.full(t, 25.4, 15) + + ldyn = pkg.LocalDyn(initial) + assert len(ldyn) == 1 + is_constant_valued(ldyn, initial) + + ldyn.resize(20) + assert len(ldyn) == 20 + is_constant_valued(ldyn, initial) + + +@pytest.test_arrays('float32,shape=(*)') +def test03_write_read(t): + pkg = get_pkg(t) + width = 20 + + local = pkg.Local10() + + for i in range(len(local)): + value = dr.arange(t, i, i+width) + if dr.backend_v(t) == dr.JitBackend.Invalid: + value = value[0] + local.write(i, value) + + sum = dr.zeros(t, width) + for i in range(len(local)): + sum += local.read(i) + + expected = dr.sum(dr.arange(t, len(local))) + (dr.arange(t, width) * len(local)) + + if dr.backend_v(t) == dr.JitBackend.Invalid: + sum = sum[0] + expected = expected[0] + + assert dr.allclose(sum, expected) + + +@pytest.test_arrays('float32,shape=(*)') +def test04_struct(t): + pkg = get_pkg(t) + width = 20 if dr.backend_v(t) == dr.JitBackend else 1 + + values = dr.ones(pkg.MyStruct, width) + local = pkg.LocalStruct10(values) + + def validate_index(idx, value): + struct = local.read(idx) + assert dr.width(struct) == width + assert dr.allclose(struct.value, value) + assert dr.allclose(struct.priority, value) + + for i in range(10): + validate_index(i, 1) + + values = dr.zeros(pkg.MyStruct, width) + local.write(0, values) + + validate_index(0, 0) + for i in range(1,10): + validate_index(i, 1) + + if dr.backend_v(t) == dr.JitBackend: + with pytest.raises(RuntimeError, match="out of bounds"): + validate_index(10, 1) From e2089cd5e95ae839580af17790148b45aa980555 Mon Sep 17 00:00:00 2001 From: Tobias Rittig Date: Mon, 15 Sep 2025 18:07:50 +0200 Subject: [PATCH 2/5] #Fix nested lambda implementations in `Local` --- include/drjit/local.h | 153 ++++++++++++++++++++++-------------------- 1 file changed, 82 insertions(+), 71 deletions(-) diff --git a/include/drjit/local.h b/include/drjit/local.h index 3ac3c9a5e..eb206934b 100644 --- a/include/drjit/local.h +++ b/include/drjit/local.h @@ -67,6 +67,84 @@ struct Local { Value m_data[Size]; }; + +NAMESPACE_BEGIN(detail) + +template +void init_impl(const T &value, const size_t size, vector& arrays) { + if constexpr (is_jit_v && depth_v == 1) { + uint32_t result; + if (!value.empty()) { + uint32_t i1 = value.index(); + size_t width = jit_var_size(i1); + uint32_t i2 = jit_array_create( + backend_v, var_type>::value, + width, size); + result = jit_array_init(i2, i1); + jit_var_dec_ref(i2); + } else { + result = jit_array_create( + backend_v, var_type>::value, + 1, size); + } + arrays.push_back(result); + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(value), + [&](auto &v) { init_impl(v, size, arrays); }); + } +} + +template +void read_impl(T &result, + const uint32_t &offset, + const uint32_t &active, + const vector &arrays, + size_t &counter) { + if constexpr (is_jit_v && depth_v == 1) { + if (counter >= arrays.size()) + jit_raise("Local::read(): internal error, ran out of " + "variable arrays!"); + result = T::steal(jit_array_read(arrays[counter++], offset, active)); + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(result), [&](auto &r) { + read_impl(r, offset, active, arrays, counter); + }); + } +} + +template +void write_impl(const uint32_t &offset, + const T &value, + const uint32_t &active, + vector &arrays, + size_t &counter) { + if constexpr (is_jit_v && depth_v == 1) { + if (counter >= arrays.size()) + jit_raise("Local::write(): internal error, ran out of " + "variable arrays!"); + + if (value.index_ad()) + jit_raise("Local memory writes are not differentiable. You " + "must use 'drjit.detach()' to disable gradient " + "tracking of the written value."); + + uint32_t result = + jit_array_write(arrays[counter], offset, value.index(), active); + jit_var_dec_ref(arrays[counter]); + arrays[counter++] = result; + + } else if constexpr (is_traversable_v) { + /// Recurse and try again if the object is traversable + traverse_1(fields(value), + [&](auto &v) { write_impl(offset, v, active, arrays, counter); }); + } +} + +NAMESPACE_END(detail) + + /** * \brief Local memory implemented on top of drjit-core jit_array_* * \details The array `value` of static or dynamic width will be used @@ -90,7 +168,7 @@ struct Local()) : m_size(Size == Dynamic ? 1 : Size), m_value(value) { - initialize(); + detail::init_impl(m_value, m_size, m_arrays); } ~Local() { @@ -113,22 +191,7 @@ struct Local void { - using T = std::decay_t; - if constexpr (is_jit_v && depth_v == 1) { - if (counter >= m_arrays.size()) - jit_raise("Local::read(): internal error, ran out of " - "variable arrays!"); - result = T::steal(jit_array_read(m_arrays[counter++], - offset.index(), active.index())); - } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable - traverse_1(fields(result), [&](auto &result) { - callback(result, callback); - }); - } - }; - callback(result, callback); + detail::read_impl(result, offset.index(), active.index(), m_arrays, counter); if (counter != m_arrays.size()) jit_raise( @@ -140,31 +203,7 @@ struct Local void { - using T = std::decay_t; - if constexpr (is_jit_v && depth_v == 1) { - if (counter >= m_arrays.size()) - jit_raise("Local::write(): internal error, ran out of " - "variable arrays!"); - - if (value.index_ad()) - jit_raise("Local memory writes are not differentiable. You " - "must use 'drjit.detach()' to disable gradient " - "tracking of the written value."); - - uint32_t result = - jit_array_write(m_arrays[counter], offset.index(), - value.index(), active.index()); - jit_var_dec_ref(m_arrays[counter]); - m_arrays[counter++] = result; - - } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable - traverse_1(fields(value), - [&](auto &value) { callback(value, callback); }); - } - }; - callback(value, callback); + detail::write_impl(offset.index(), value, active.index(), m_arrays, counter); if (counter != m_arrays.size()) jit_raise( @@ -182,38 +221,10 @@ struct Local void { - using T = std::decay_t; - if constexpr (is_jit_v && depth_v == 1) { - uint32_t result; - if (!value.empty()) { - uint32_t i1 = value.index(); - size_t width = jit_var_size(i1); - uint32_t i2 = jit_array_create( - backend_v, var_type>::value, - width, m_size); - result = jit_array_init(i2, i1); - jit_var_dec_ref(i2); - } else { - result = jit_array_create( - backend_v, var_type>::value, - 1, m_size); - } - m_arrays.push_back(result); - } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable - traverse_1(fields(value), - [&](auto &value) { callback(value, callback); }); - } - }; - callback(m_value, callback); - } - size_t m_size; Value m_value; vector m_arrays; From 8354a1bd0bcd2f681b416a541867665c3319a1d5 Mon Sep 17 00:00:00 2001 From: Tobias Rittig Date: Mon, 22 Sep 2025 17:44:04 +0200 Subject: [PATCH 3/5] -Remove usage of /// in local.h Fix the docs for `traverse` accordingly --- docs/cpp.rst | 6 +++--- include/drjit/local.h | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/cpp.rst b/docs/cpp.rst index 9079640d2..2a7e879b0 100644 --- a/docs/cpp.rst +++ b/docs/cpp.rst @@ -494,11 +494,11 @@ trees: template void visit_jit_pairs(T &v0, T &v1) { if constexpr (dr::is_jit_v && dr::depth_v == 1) { - /// Do something with 'v0' and 'v1' + // Do something with 'v0' and 'v1' } else if constexpr (dr::is_traversable_v) { - /// Recurse and try again if the object is traversable + // Recurse and try again if the object is traversable dr::traverse_2( - /// Extract the fields of 'v0' and 'v1' + // Extract the fields of 'v0' and 'v1' dr::fields(v0), dr::fields(v1), // .. and call the following lambda function on them [&](auto &x, auto &y) { visit_jit_pairs(x, y); } diff --git a/include/drjit/local.h b/include/drjit/local.h index eb206934b..ec9c9f034 100644 --- a/include/drjit/local.h +++ b/include/drjit/local.h @@ -89,7 +89,7 @@ void init_impl(const T &value, const size_t size, vector& arrays) { } arrays.push_back(result); } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable + // Recurse and try again if the object is traversable traverse_1(fields(value), [&](auto &v) { init_impl(v, size, arrays); }); } @@ -107,7 +107,7 @@ void read_impl(T &result, "variable arrays!"); result = T::steal(jit_array_read(arrays[counter++], offset, active)); } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable + // Recurse and try again if the object is traversable traverse_1(fields(result), [&](auto &r) { read_impl(r, offset, active, arrays, counter); }); @@ -136,7 +136,7 @@ void write_impl(const uint32_t &offset, arrays[counter++] = result; } else if constexpr (is_traversable_v) { - /// Recurse and try again if the object is traversable + // Recurse and try again if the object is traversable traverse_1(fields(value), [&](auto &v) { write_impl(offset, v, active, arrays, counter); }); } From 9c87a29672cfb67271207d35a78254b158d66080 Mon Sep 17 00:00:00 2001 From: Tobias Rittig Date: Wed, 1 Oct 2025 14:27:46 +0200 Subject: [PATCH 4/5] Cleanup move constructors in dr::Local (C++) --- include/drjit/local.h | 22 ++++++++-------------- 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/include/drjit/local.h b/include/drjit/local.h index ec9c9f034..6b73633eb 100644 --- a/include/drjit/local.h +++ b/include/drjit/local.h @@ -39,15 +39,9 @@ struct Local { ~Local() = default; Local(const Local &) = delete; - Local(Local &&l) { - for (size_t i = 0; i < Size; ++i) - m_data[i] = l.m_data[i]; - } + Local(Local &&l) = default; Local &operator=(const Local &) = delete; - Local &operator=(Local &&l) { - for (size_t i = 0; i < Size; ++i) - m_data[i] = l.m_data[i]; - } + Local &operator=(Local &&l) = default; Value read(const Index &offset, const Mask &active = true) const { if (active) @@ -176,16 +170,16 @@ struct Local Date: Mon, 6 Oct 2025 12:13:10 +0200 Subject: [PATCH 5/5] Make dr::Local (C++) loop compatible Meaning it can be wrapped in a DRJIT_STRUCT itself --- include/drjit/local.h | 19 ++++++++--- tests/local_ext.cpp | 73 ++++++++++++++++++++++++++++++++++++++++- tests/test_local_ext.py | 19 +++++++++-- 3 files changed, 103 insertions(+), 8 deletions(-) diff --git a/include/drjit/local.h b/include/drjit/local.h index 6b73633eb..76d73e468 100644 --- a/include/drjit/local.h +++ b/include/drjit/local.h @@ -38,9 +38,9 @@ struct Local { } ~Local() = default; - Local(const Local &) = delete; + Local(const Local &) = default; Local(Local &&l) = default; - Local &operator=(const Local &) = delete; + Local &operator=(const Local &) = default; Local &operator=(Local &&l) = default; Value read(const Index &offset, const Mask &active = true) const { @@ -169,10 +169,21 @@ struct Local +#include #include +#include +#include namespace nb = nanobind; namespace dr = drjit; @@ -20,12 +23,75 @@ auto bind_local(nb::module_ &m, const dr::string& name) { if constexpr (dr::is_jit_v) c = c.def("resize", &Local::resize); + m.def(("test_" + name + "_loop").c_str(), []() { + auto initial = Local(); + auto counter = int32_t(0); + + if constexpr (dr::is_jit_v) + initial.resize(10); + + dr::tie(initial, counter) = dr::while_loop( + dr::make_tuple(initial, counter), + [](const Local& l, const int32_t& i) { + DRJIT_MARK_USED(l); + return i < 5; + }, + [](Local& l, int32_t& i) { + l.write(i, dr::full(i)); + auto written = l.read(i); + DRJIT_MARK_USED(written); + i += 1; + } + ); + for(unsigned int i = 0; i < 5; ++i) { + auto value = initial.read(i); + if(dr::any(value != dr::full(i))) { + jit_raise("Index %d doesn't match %s", i, dr::string(value).c_str()); + } + } + }); + + m.def(("test_" + name + "_loop_struct").c_str(), []() { + auto initial = Local(); + auto counter = int32_t(0); + + if constexpr (dr::is_jit_v) + initial.resize(10); + + + struct LoopState { + Local l; + int32_t counter; + + DRJIT_STRUCT(LoopState, l, counter) + } ls = { initial, counter }; + + dr::tie(ls) = dr::while_loop( + dr::make_tuple(ls), + [](const LoopState &ls) { return ls.counter < 5; }, + [](LoopState &ls) { + ls.l.write(ls.counter, dr::full(ls.counter)); + auto written = ls.l.read(ls.counter); + DRJIT_MARK_USED(written); + ls.counter += 1; + } + ); + for(unsigned int i = 0; i < 5; ++i) { + auto value = ls.l.read(i); + if(dr::any(value != dr::full(i))) { + jit_raise("Index %d doesn't match %s", i, dr::string(value).c_str()); + } + } + }); + return c; } template void bind(nb::module_ &m) { using UInt32 = dr::uint32_array_t; + using Bool = dr::bool_array_t; + using Local10 = dr::Local; using LocalDyn = dr::Local; @@ -40,6 +106,10 @@ void bind(nb::module_ &m) { UInt32 priority; DRJIT_STRUCT(MyStruct, value, priority) + Bool operator!=(const MyStruct& other) const { + return priority != other.priority; + } + MyStruct(int i) : value(i), priority(i) {} }; @@ -47,7 +117,8 @@ void bind(nb::module_ &m) { .def(nb::init<>()) .def(nb::init()) .def_rw("value", &MyStruct::value) - .def_rw("priority", &MyStruct::priority); + .def_rw("priority", &MyStruct::priority) + .def(nb::self != nb::self); nb::handle u32; if constexpr (dr::is_array_v) diff --git a/tests/test_local_ext.py b/tests/test_local_ext.py index c738f6e07..8d31fb2a0 100644 --- a/tests/test_local_ext.py +++ b/tests/test_local_ext.py @@ -11,7 +11,7 @@ def get_pkg(t): return m.cuda elif backend == dr.JitBackend.Invalid: return m.scalar - + def is_constant_valued(local, value): for i in range(len(local)): assert dr.allclose(local.read(i), value) @@ -66,11 +66,11 @@ def test03_write_read(t): sum += local.read(i) expected = dr.sum(dr.arange(t, len(local))) + (dr.arange(t, width) * len(local)) - + if dr.backend_v(t) == dr.JitBackend.Invalid: sum = sum[0] expected = expected[0] - + assert dr.allclose(sum, expected) @@ -101,3 +101,16 @@ def validate_index(idx, value): if dr.backend_v(t) == dr.JitBackend: with pytest.raises(RuntimeError, match="out of bounds"): validate_index(10, 1) + +@pytest.test_arrays('float32,shape=(*)') +def test05_loop(t): + pkg = get_pkg(t) + pkg.test_Local10_loop() + pkg.test_Local10_loop_struct() + + if dr.backend_v(t) == dr.JitBackend: + pkg.test_LocalDyn_loop() + pkg.test_LocalDyn_loop_struct() + + pkg.test_LocalStruct10_loop() + pkg.test_LocalStruct10_loop_struct()