From 91af33a48dbbf65fe8bbea21e1b8bc7aa59eedd0 Mon Sep 17 00:00:00 2001 From: Mergen Nachin Date: Tue, 20 Jan 2026 10:08:28 -0800 Subject: [PATCH] Fix XNNWorkspaceManager cleanup tests by using unique IDs instead of raw pointers (#16675) Summary: The PerModelModeCleanup and GlobalModeCleanup tests were failing because they compared raw xnn_workspace_t pointers to verify cleanup behavior. However, memory allocators commonly reuse addresses after deallocation - when xnn_release_workspace frees memory and xnn_create_workspace immediately allocates the same size, the allocator may return the same address. This fix: 1. Adds a unique instance ID (uint64_t) to XNNWorkspace using an atomic counter 2. Adds an id() method to retrieve this unique ID 3. Modifies both cleanup tests to compare IDs instead of raw pointers This correctly verifies that new workspace objects are created after cleanup, regardless of memory address reuse by the allocator. Reviewed By: GregoryComer Differential Revision: D90897333 --- backends/xnnpack/runtime/XNNWorkspace.h | 13 +++++++++++- .../test/runtime/test_workspace_manager.cpp | 20 +++++++++++-------- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/backends/xnnpack/runtime/XNNWorkspace.h b/backends/xnnpack/runtime/XNNWorkspace.h index 36596b05089..507953a10ab 100644 --- a/backends/xnnpack/runtime/XNNWorkspace.h +++ b/backends/xnnpack/runtime/XNNWorkspace.h @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -24,7 +25,8 @@ using WorkspacePtr = /// with appropriate synchronization. class XNNWorkspace { public: - XNNWorkspace(WorkspacePtr workspace) : workspace_(std::move(workspace)){}; + XNNWorkspace(WorkspacePtr workspace) + : id_(next_id_++), workspace_(std::move(workspace)){}; XNNWorkspace(const XNNWorkspace&) = delete; XNNWorkspace& operator=(const XNNWorkspace&) = delete; // Not moveable due to std::mutex. @@ -43,6 +45,13 @@ class XNNWorkspace { return workspace_.get(); } + // Returns a unique ID for this workspace instance. This can be used to + // distinguish between different workspace objects even if they happen to + // have the same raw pointer due to memory reuse. + uint64_t id() const { + return id_; + } + static runtime::Result> create() { // Because this class can't be moved, we need to construct it in-place. xnn_workspace_t workspace = nullptr; @@ -60,7 +69,9 @@ class XNNWorkspace { } private: + static inline std::atomic next_id_{0}; std::mutex mutex_; + uint64_t id_; WorkspacePtr workspace_; }; diff --git a/backends/xnnpack/test/runtime/test_workspace_manager.cpp b/backends/xnnpack/test/runtime/test_workspace_manager.cpp index ddb7074a1ce..8d3203f3f40 100644 --- a/backends/xnnpack/test/runtime/test_workspace_manager.cpp +++ b/backends/xnnpack/test/runtime/test_workspace_manager.cpp @@ -166,7 +166,7 @@ TEST_F(XNNWorkspaceManagerTest, PerModelModeCleanup) { workspace_manager_->set_sharing_mode(WorkspaceSharingMode::PerModel); uintptr_t program_id = 12345; - xnn_workspace_t raw_workspace1 = nullptr; + uint64_t workspace1_id = 0; // Create a scope to control the lifetime of workspace1 { @@ -175,8 +175,8 @@ TEST_F(XNNWorkspaceManagerTest, PerModelModeCleanup) { ASSERT_TRUE(workspace1_result.ok()); auto workspace1 = workspace1_result.get(); - // Store the raw pointer for later comparison - raw_workspace1 = workspace1->unsafe_get_workspace(); + // Store the unique ID for later comparison + workspace1_id = workspace1->id(); // Let workspace1 go out of scope and be destroyed } @@ -188,7 +188,9 @@ TEST_F(XNNWorkspaceManagerTest, PerModelModeCleanup) { auto workspace2 = workspace2_result.get(); // Since the previous workspace was destroyed, we should get a new one. - EXPECT_NE(workspace2->unsafe_get_workspace(), raw_workspace1); + // We compare IDs instead of raw pointers because memory allocators may + // reuse addresses after deallocation. + EXPECT_NE(workspace2->id(), workspace1_id); } TEST_F(XNNWorkspaceManagerTest, GlobalModeCleanup) { @@ -197,7 +199,7 @@ TEST_F(XNNWorkspaceManagerTest, GlobalModeCleanup) { workspace_manager_->set_sharing_mode(WorkspaceSharingMode::Global); uintptr_t program_id = 12345; - xnn_workspace_t raw_workspace1 = nullptr; + uint64_t workspace1_id = 0; // Create a scope to control the lifetime of workspace1 { @@ -206,8 +208,8 @@ TEST_F(XNNWorkspaceManagerTest, GlobalModeCleanup) { ASSERT_TRUE(workspace1_result.ok()); auto workspace1 = workspace1_result.get(); - // Store the raw pointer for later comparison - raw_workspace1 = workspace1->unsafe_get_workspace(); + // Store the unique ID for later comparison + workspace1_id = workspace1->id(); // Let workspace1 go out of scope and be destroyed } @@ -219,7 +221,9 @@ TEST_F(XNNWorkspaceManagerTest, GlobalModeCleanup) { auto workspace2 = workspace2_result.get(); // Since the previous workspace was destroyed, we should get a new one. - EXPECT_NE(workspace2->unsafe_get_workspace(), raw_workspace1); + // We compare IDs instead of raw pointers because memory allocators may + // reuse addresses after deallocation. + EXPECT_NE(workspace2->id(), workspace1_id); } TEST_F(XNNWorkspaceManagerTest, SwitchingModes) {