From 1d4a616214c4ea21c681308d51735784fd5412b5 Mon Sep 17 00:00:00 2001 From: Duncan Betts Date: Sun, 20 Sep 2026 13:48:14 +0100 Subject: [PATCH] fix(core): preserve layout state after overflow --- src/core/layout.cpp | 24 ++-------- src/core/layout.h | 20 ++------ tests/cmake/CoreTests.cmake | 3 ++ tests/test_layout.cpp | 96 +++++++++++++++++++++++++++++++++++++ 4 files changed, 106 insertions(+), 37 deletions(-) create mode 100644 tests/test_layout.cpp diff --git a/src/core/layout.cpp b/src/core/layout.cpp index 3047748202..7c42b51c89 100644 --- a/src/core/layout.cpp +++ b/src/core/layout.cpp @@ -86,39 +86,23 @@ std::size_t LayoutBuilder::finish(std::size_t alignment, std::string_view label) return align_up(peak_, alignment, label); } -WorkspaceLayoutBuilder::Scope::Scope(WorkspaceLayoutBuilder& builder) noexcept - : builder_(&builder), saved_cursor_(builder.cursor_) {} - -WorkspaceLayoutBuilder::Scope::~Scope() noexcept { - if (builder_ != nullptr) { builder_->cursor_ = saved_cursor_; } -} - -WorkspaceLayoutBuilder::Scope::Scope(Scope&& other) noexcept - : builder_(other.builder_), saved_cursor_(other.saved_cursor_) { - other.builder_ = nullptr; -} - Tensor WorkspaceLayoutBuilder::alloc(DType dtype, std::initializer_list shape, std::size_t alignment) { Tensor tensor(nullptr, dtype, shape); - cursor_ = align_up(cursor_, alignment, "workspace layout"); - cursor_ = checked_add(cursor_, tensor.bytes(), "workspace layout"); - if (cursor_ > peak_) { peak_ = cursor_; } + (void)layout_.add(tensor.bytes(), alignment, "workspace layout"); return tensor; } DeviceSpan WorkspaceLayoutBuilder::alloc_bytes(std::size_t bytes, std::size_t alignment) { if (bytes == 0) { return {}; } - cursor_ = align_up(cursor_, alignment, "workspace layout"); - cursor_ = checked_add(cursor_, bytes, "workspace layout"); - if (cursor_ > peak_) { peak_ = cursor_; } + (void)layout_.add(bytes, alignment, "workspace layout"); return DeviceSpan{nullptr, bytes}; } -WorkspaceLayoutBuilder::Scope WorkspaceLayoutBuilder::scope() noexcept { return Scope(*this); } +WorkspaceLayoutBuilder::Scope WorkspaceLayoutBuilder::scope() noexcept { return layout_.scope(); } std::size_t WorkspaceLayoutBuilder::peak_bytes(std::size_t alignment) const { - return align_up(peak_, alignment, "workspace layout"); + return layout_.finish(alignment, "workspace layout"); } } // namespace ninfer diff --git a/src/core/layout.h b/src/core/layout.h index 20e0a4f5d9..d673e617fd 100644 --- a/src/core/layout.h +++ b/src/core/layout.h @@ -62,23 +62,10 @@ class LayoutBuilder { // Dry-run counterpart of WorkspaceArena. Target allocation helpers can run against this builder // and the real arena, including identical nested scope lifetimes, without maintaining byte -// formulas. +// formulas. LayoutBuilder owns its cursor, alignment checks, and peak accounting. class WorkspaceLayoutBuilder { public: - class Scope { - public: - ~Scope() noexcept; - Scope(const Scope&) = delete; - Scope& operator=(const Scope&) = delete; - Scope(Scope&& other) noexcept; - Scope& operator=(Scope&&) = delete; - - private: - friend class WorkspaceLayoutBuilder; - explicit Scope(WorkspaceLayoutBuilder& builder) noexcept; - WorkspaceLayoutBuilder* builder_ = nullptr; - std::size_t saved_cursor_ = 0; - }; + using Scope = LayoutBuilder::Scope; [[nodiscard]] Tensor alloc(DType dtype, std::initializer_list shape, std::size_t alignment = 256); @@ -87,8 +74,7 @@ class WorkspaceLayoutBuilder { [[nodiscard]] std::size_t peak_bytes(std::size_t alignment = 256) const; private: - std::size_t cursor_ = 0; - std::size_t peak_ = 0; + LayoutBuilder layout_; }; } // namespace ninfer diff --git a/tests/cmake/CoreTests.cmake b/tests/cmake/CoreTests.cmake index e0e4a652c7..a9d0d4c22b 100644 --- a/tests/cmake/CoreTests.cmake +++ b/tests/cmake/CoreTests.cmake @@ -16,6 +16,9 @@ ninfer_add_test(ninfer_tensor_test SOURCES "${CMAKE_CURRENT_LIST_DIR}/../t ninfer_add_test(ninfer_arena_test SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_arena.cpp" LIBRARIES ninfer_core) +ninfer_add_test(ninfer_layout_test SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_layout.cpp" + LIBRARIES ninfer_core) + ninfer_add_test(ninfer_materialization_budget_test SOURCES "${CMAKE_CURRENT_LIST_DIR}/../test_materialization_budget.cpp" LIBRARIES ninfer_core) diff --git a/tests/test_layout.cpp b/tests/test_layout.cpp new file mode 100644 index 0000000000..90a29377d2 --- /dev/null +++ b/tests/test_layout.cpp @@ -0,0 +1,96 @@ +#include "core/layout.h" + +#include +#include +#include +#include + +namespace { + +int failures = 0; + +void expect(bool condition, const char* message) { + if (condition) { return; } + ++failures; + std::cerr << message << '\n'; +} + +template +void expect_throws(Operation operation, const char* message) { + try { + operation(); + } catch (const Exception&) { return; } + expect(false, message); +} + +void test_workspace_scopes() { + ninfer::WorkspaceLayoutBuilder layout; + expect(layout.peak_bytes() == 0, "empty workspace must need no storage"); + const auto tensor = layout.alloc(ninfer::DType::BF16, {3, 5}); + expect(tensor.data == nullptr && tensor.bytes() == 30 && tensor.ne[0] == 3 && tensor.ne[1] == 5, + "dry-run tensor must retain shape and dtype without storage"); + expect(layout.peak_bytes(1) == 30, "first tensor needs 30 bytes"); + { + auto outer = layout.scope(); + const auto bytes = layout.alloc_bytes(17, 64); + expect(bytes.data == nullptr && bytes.bytes == 17, + "dry-run byte span must not own storage"); + expect(layout.peak_bytes(1) == 81, "aligned allocation ends at byte 81"); + { + auto inner = layout.scope(); + auto moved = std::move(inner); + (void)layout.alloc_bytes(11, 128); + expect(layout.peak_bytes(1) == 139, "nested allocation ends at byte 139"); + } + (void)layout.alloc_bytes(100, 1); + expect(layout.peak_bytes(1) == 181, "moved inner scope must restore byte 81"); + } + (void)layout.alloc_bytes(200, 1); + expect(layout.peak_bytes(1) == 230, "outer scope must restore byte 30"); + expect(layout.peak_bytes() == 256, "final estimate rounds the peak to its requested alignment"); + + expect_throws( + [&] { + auto scope = layout.scope(); + (void)layout.alloc_bytes(100, 1); + throw std::runtime_error("unwind workspace scope"); + }, + "test exception must propagate through scope cleanup"); + expect(layout.peak_bytes(1) == 330, "scope unwinding must preserve the peak"); + (void)layout.alloc_bytes(101, 1); + expect(layout.peak_bytes(1) == 331, "scope unwinding must restore the cursor"); +} + +void test_workspace_boundaries() { + ninfer::WorkspaceLayoutBuilder layout; + (void)layout.alloc_bytes(3, 1); + const auto empty = layout.alloc_bytes(0, 3); + expect(empty.data == nullptr && empty.bytes == 0 && layout.peak_bytes(1) == 3, + "zero-byte scratch must remain a no-op, including its alignment"); + expect_throws([&] { (void)layout.alloc_bytes(1, 3); }, + "non-power-of-two alignment must fail"); + expect_throws([&] { (void)layout.peak_bytes(0); }, + "zero final alignment must fail"); + expect_throws( + [&] { (void)layout.alloc_bytes(std::numeric_limits::max(), 8); }, + "allocation end overflow must fail"); + (void)layout.alloc_bytes(1, 1); + expect(layout.peak_bytes(1) == 4, "failed allocation must not consume alignment padding"); + + ninfer::WorkspaceLayoutBuilder maximum; + (void)maximum.alloc_bytes(std::numeric_limits::max(), 1); + expect(maximum.peak_bytes(1) == std::numeric_limits::max(), + "maximum representable unaligned size must remain representable"); + expect_throws([&] { (void)maximum.peak_bytes(2); }, + "final alignment overflow must fail"); + expect_throws([&] { (void)maximum.alloc_bytes(1, 2); }, + "cursor alignment overflow must fail"); +} + +} // namespace + +int main() { + test_workspace_scopes(); + test_workspace_boundaries(); + return failures == 0 ? 0 : 1; +}