From 14401613758c6838a517b9bcc3fc0996cd0818df Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Fri, 3 Jul 2026 22:07:17 +0200 Subject: [PATCH 1/4] [Bug] Fix in-tree field byte offset truncation in DLPack export --- quadrants/python/dlpack_funcs.cpp | 8 +++++--- tests/python/quadrants/lang/test_dlpack.py | 22 ++++++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/quadrants/python/dlpack_funcs.cpp b/quadrants/python/dlpack_funcs.cpp index 4fbc7e95ba..88beab7f89 100644 --- a/quadrants/python/dlpack_funcs.cpp +++ b/quadrants/python/dlpack_funcs.cpp @@ -215,14 +215,16 @@ pybind11::capsule field_to_dlpack(Program *program, SNode *snode, int element_nd "'to_dlpack'.") } - int field_in_tree_offset = program->get_field_in_tree_offset(tree_id, snode); + // A SNode tree packs every field of a program stage, so in-tree byte offsets routinely exceed + // 2^31; the offset must stay 64-bit end to end (DLTensor::byte_offset is uint64_t). + std::size_t field_in_tree_offset = program->get_field_in_tree_offset(tree_id, snode); void *raw_ptr = nullptr; DLDeviceType device_type = DLDeviceType::kDLCPU; std::tie(raw_ptr, device_type) = get_raw_ptr(arch, program, tree_device_ptr); - int byte_offset = 0; - if (field_in_tree_offset >= 0) { + uint64_t byte_offset = 0; + if (field_in_tree_offset > 0) { if (torch_supports_byte_offset()) { byte_offset = field_in_tree_offset; } else { diff --git a/tests/python/quadrants/lang/test_dlpack.py b/tests/python/quadrants/lang/test_dlpack.py index 164712d789..518c0499bc 100644 --- a/tests/python/quadrants/lang/test_dlpack.py +++ b/tests/python/quadrants/lang/test_dlpack.py @@ -334,3 +334,25 @@ def test_dlpack_field_memory_allocation_before_to_dlpack(): assert ( second_time_tc == second_time.to_torch(device=second_time_tc.device) ).all(), f"{second_time_tc} != {second_time.to_torch(device=second_time_tc.device)}" + + +@test_utils.test(arch=[qd.cpu]) +@pytest.mark.run_in_serial +@pytest.mark.slow +def test_dlpack_field_offset_past_2gib(): + # Fields share one SNode tree, so the second field's in-tree byte offset exceeds 2^31; the + # DLPack export must carry it without truncation instead of aliasing the tree base. + pad = qd.field(dtype=qd.f32, shape=(560_000_000,)) + victim = qd.field(dtype=qd.f32, shape=(1000,)) + + @qd.kernel + def fill(): + for i in pad: + pad[i] = 1.0 + for i in victim: + victim[i] = qd.cast(i, qd.f32) + + fill() + qd.sync() + victim_tc = torch.utils.dlpack.from_dlpack(victim.to_dlpack()) + np.testing.assert_array_equal(victim_tc.cpu().numpy(), np.arange(1000, dtype=np.float32)) From 99c7081205ef5ce0be5c67762ec7a73414b865ee Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Fri, 3 Jul 2026 22:07:18 +0200 Subject: [PATCH 2/4] [Bug] Fix UnifiedAllocator chunk tail overstating the usable size --- quadrants/rhi/common/unified_allocator.cpp | 2 +- tests/cpp/rhi/common/host_memory_pool_test.cpp | 16 ++++++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/quadrants/rhi/common/unified_allocator.cpp b/quadrants/rhi/common/unified_allocator.cpp index f2bdce1c8a..268052f557 100644 --- a/quadrants/rhi/common/unified_allocator.cpp +++ b/quadrants/rhi/common/unified_allocator.cpp @@ -69,7 +69,7 @@ void *UnifiedAllocator::allocate(std::size_t size, std::size_t alignment, bool e void *ptr = HostMemoryPool::get_instance().allocate_raw_memory(allocation_size); chunk.data = ptr; chunk.head = (void *)((std::size_t)chunk.data + size); - chunk.tail = (void *)((std::size_t)chunk.head + allocation_size); + chunk.tail = (void *)((std::size_t)chunk.data + allocation_size); chunk.is_exclusive = exclusive; QD_ASSERT(chunk.data != nullptr); diff --git a/tests/cpp/rhi/common/host_memory_pool_test.cpp b/tests/cpp/rhi/common/host_memory_pool_test.cpp index d569123850..ef054acd6a 100644 --- a/tests/cpp/rhi/common/host_memory_pool_test.cpp +++ b/tests/cpp/rhi/common/host_memory_pool_test.cpp @@ -34,4 +34,20 @@ TEST(HostMemoryPool, AllocateMemory) { HostMemoryPoolTestHelper::setDefaultAllocatorSize(oldAllocatorSize); } +TEST(HostMemoryPool, ChunkTailMatchesAllocationSize) { + auto oldAllocatorSize = HostMemoryPoolTestHelper::getDefaultAllocatorSize(); + HostMemoryPoolTestHelper::setDefaultAllocatorSize(102400); // 100KB + + HostMemoryPool pool; + + // The first allocation creates a 100KB chunk. The second does not fit in the remaining 40KB, so + // it must open a new chunk instead of being placed past the real end of the first one. + void *ptr1 = pool.allocate(61440, 16); + void *ptr2 = pool.allocate(81920, 16); + + EXPECT_NE((std::size_t)ptr2, (std::size_t)ptr1 + 61440); + + HostMemoryPoolTestHelper::setDefaultAllocatorSize(oldAllocatorSize); +} + } // namespace quadrants::lang From 06dc81a2bf766af71a730221952e251f924cadd3 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Fri, 3 Jul 2026 22:07:18 +0200 Subject: [PATCH 3/4] [Bug] Fix iterator use-after-erase in CachingAllocator::allocate --- quadrants/rhi/llvm/allocator.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/quadrants/rhi/llvm/allocator.cpp b/quadrants/rhi/llvm/allocator.cpp index 2227e7264c..64a1571559 100644 --- a/quadrants/rhi/llvm/allocator.cpp +++ b/quadrants/rhi/llvm/allocator.cpp @@ -43,8 +43,9 @@ uint64_t *CachingAllocator::allocate(LlvmDevice *device, const LlvmDevice::LlvmR ptr_map_.insert(std::make_pair(remaining_head, remaining_sz)); } ret = reinterpret_cast(it_blk->second); - mem_blocks_.erase(it_blk); + // Erase from ptr_map_ first: mem_blocks_.erase(it_blk) invalidates it_blk. ptr_map_.erase(it_blk->second); + mem_blocks_.erase(it_blk); } else { ret = reinterpret_cast(device->allocate_llvm_runtime_memory_jit(params)); From 1fa430ef618bf91e582b5eec4f32684cc9d7f1c0 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Fri, 3 Jul 2026 22:50:59 +0200 Subject: [PATCH 4/4] [Misc] Reflow comments to the 120-column budget --- quadrants/python/dlpack_funcs.cpp | 4 ++-- tests/cpp/rhi/common/host_memory_pool_test.cpp | 4 ++-- tests/python/quadrants/lang/test_dlpack.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/quadrants/python/dlpack_funcs.cpp b/quadrants/python/dlpack_funcs.cpp index 88beab7f89..55504e9992 100644 --- a/quadrants/python/dlpack_funcs.cpp +++ b/quadrants/python/dlpack_funcs.cpp @@ -215,8 +215,8 @@ pybind11::capsule field_to_dlpack(Program *program, SNode *snode, int element_nd "'to_dlpack'.") } - // A SNode tree packs every field of a program stage, so in-tree byte offsets routinely exceed - // 2^31; the offset must stay 64-bit end to end (DLTensor::byte_offset is uint64_t). + // A SNode tree packs every field of a program stage, so in-tree byte offsets routinely exceed 2^31; the offset must + // stay 64-bit end to end (DLTensor::byte_offset is uint64_t). std::size_t field_in_tree_offset = program->get_field_in_tree_offset(tree_id, snode); void *raw_ptr = nullptr; diff --git a/tests/cpp/rhi/common/host_memory_pool_test.cpp b/tests/cpp/rhi/common/host_memory_pool_test.cpp index ef054acd6a..876f7bce96 100644 --- a/tests/cpp/rhi/common/host_memory_pool_test.cpp +++ b/tests/cpp/rhi/common/host_memory_pool_test.cpp @@ -40,8 +40,8 @@ TEST(HostMemoryPool, ChunkTailMatchesAllocationSize) { HostMemoryPool pool; - // The first allocation creates a 100KB chunk. The second does not fit in the remaining 40KB, so - // it must open a new chunk instead of being placed past the real end of the first one. + // The first allocation creates a 100KB chunk. The second does not fit in the remaining 40KB, so it must open a new + // chunk instead of being placed past the real end of the first one. void *ptr1 = pool.allocate(61440, 16); void *ptr2 = pool.allocate(81920, 16); diff --git a/tests/python/quadrants/lang/test_dlpack.py b/tests/python/quadrants/lang/test_dlpack.py index 518c0499bc..88e7c5188f 100644 --- a/tests/python/quadrants/lang/test_dlpack.py +++ b/tests/python/quadrants/lang/test_dlpack.py @@ -340,8 +340,8 @@ def test_dlpack_field_memory_allocation_before_to_dlpack(): @pytest.mark.run_in_serial @pytest.mark.slow def test_dlpack_field_offset_past_2gib(): - # Fields share one SNode tree, so the second field's in-tree byte offset exceeds 2^31; the - # DLPack export must carry it without truncation instead of aliasing the tree base. + # Fields share one SNode tree, so the second field's in-tree byte offset exceeds 2^31; the DLPack export must carry + # it without truncation instead of aliasing the tree base. pad = qd.field(dtype=qd.f32, shape=(560_000_000,)) victim = qd.field(dtype=qd.f32, shape=(1000,))