Skip to content

Xnnpack disable workspace nonlock #17780

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Froskekongen wants to merge 7 commits into
base: main
Choose a base branch
from
+34 −1
Open

Xnnpack disable workspace nonlock #17780

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
21d9bdd
Remove lock when Disabled Workspace mode is set.
Froskekongen Mar 2, 2026
7f4bb96
Run xnnpack workspace disable/enable lock test
Froskekongen Mar 2, 2026
ccdbca7
Merge branch 'main' into xnnpack-disable-workspace-nonlock
Froskekongen Mar 2, 2026
019327b
Fix lint error
Froskekongen Mar 4, 2026
3871367
Merge branch 'main' into xnnpack-disable-workspace-nonlock
Froskekongen Mar 4, 2026
93e9e76
Merge branch 'pytorch:main' into xnnpack-disable-workspace-nonlock
Froskekongen Mar 5, 2026
1a97263
Merge branch 'main' into xnnpack-disable-workspace-nonlock
Froskekongen Mar 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions backends/xnnpack/runtime/XNNWorkspace.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@ class XNNWorkspace {
XNNWorkspace& operator=(XNNWorkspace&&) = delete;

std::pair<std::unique_lock<std::mutex>, xnn_workspace_t> acquire() {
if (!lock_required_) {
return {std::unique_lock<std::mutex>{}, workspace_.get()};
}
auto lock = std::unique_lock<std::mutex>(mutex_);
return {std::move(lock), workspace_.get()};
}
Expand All @@ -52,6 +55,10 @@ class XNNWorkspace {
return id_;
}

void disable_locking() {
lock_required_ = false;
}
Comment on lines +58 to +60
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The lock_required_ field is a plain bool that is written by disable_locking() and read by acquire(). Although the current single call site in XNNWorkspaceManager.cpp calls disable_locking() before the workspace is shared with any other thread (making it safe in practice), the public disable_locking() method has no documented threading constraints, and calling it on an already-shared workspace while acquire() is being called concurrently from another thread would be an undefined behavior data race. Consider making lock_required_ a std::atomic<bool> to ensure correctness regardless of call order, or at minimum add a comment documenting the thread-safety requirement (that it must only be called before the workspace is shared with other threads).

Copilot uses AI. Check for mistakes.

static runtime::Result<std::shared_ptr<XNNWorkspace>> create() {
// Because this class can't be moved, we need to construct it in-place.
xnn_workspace_t workspace = nullptr;
Expand All @@ -72,6 +79,7 @@ class XNNWorkspace {
static inline std::atomic<uint64_t> next_id_{0};
std::mutex mutex_;
uint64_t id_;
bool lock_required_ = true;
WorkspacePtr workspace_;
};

Expand Down
1 change: 1 addition & 0 deletions backends/xnnpack/runtime/XNNWorkspaceManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ XNNWorkspaceManager::get_or_create_workspace(uintptr_t program_id) const {
return create_result.error();
}

create_result.get()->disable_locking();
Copy link
Member

@GregoryComer GregoryComer Mar 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Can we pass this in as a arg to XNNWorkspace::create?

return create_result.get();
} else if (mode == WorkspaceSharingMode::PerModel) {
return get_or_create_model_workspace(program_id);
Expand Down
2 changes: 1 addition & 1 deletion backends/xnnpack/test/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..)

include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake)

set(_test_srcs runtime/test_xnnexecutor.cpp
set(_test_srcs runtime/test_xnnexecutor.cpp runtime/test_workspace_manager.cpp
${EXECUTORCH_ROOT}/extension/threadpool/test/threadpool_test.cpp
)

Expand Down
24 changes: 24 additions & 0 deletions backends/xnnpack/test/runtime/test_workspace_manager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,18 @@ TEST_F(XNNWorkspaceManagerTest, DisabledMode) {
workspace2->unsafe_get_workspace(), workspace3->unsafe_get_workspace());
}

TEST_F(XNNWorkspaceManagerTest, DisabledModeAcquireDoesNotLock) {
workspace_manager_->set_sharing_mode(WorkspaceSharingMode::Disabled);

auto workspace_result = workspace_manager_->get_or_create_workspace(12345);
ASSERT_TRUE(workspace_result.ok());
auto workspace = workspace_result.get();

auto [lock, ptr] = workspace->acquire();
ASSERT_NE(ptr, nullptr);
EXPECT_FALSE(lock.owns_lock());
}

TEST_F(XNNWorkspaceManagerTest, PerModelMode) {
// In PerModel mode, calls with the same program_id should return the same
// workspace.
Expand Down Expand Up @@ -139,6 +151,18 @@ TEST_F(XNNWorkspaceManagerTest, PerModelMode) {
workspace1->unsafe_get_workspace(), workspace3->unsafe_get_workspace());
}

TEST_F(XNNWorkspaceManagerTest, PerModelAcquireStillLocks) {
workspace_manager_->set_sharing_mode(WorkspaceSharingMode::PerModel);

auto workspace_result = workspace_manager_->get_or_create_workspace(12345);
ASSERT_TRUE(workspace_result.ok());
auto workspace = workspace_result.get();

auto [lock, ptr] = workspace->acquire();
ASSERT_NE(ptr, nullptr);
EXPECT_TRUE(lock.owns_lock());
}
Comment on lines +154 to +164
Copy link

Copilot AI Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The PR adds DisabledModeAcquireDoesNotLock and PerModelAcquireStillLocks to verify locking behavior per mode, but there is no corresponding GlobalAcquireStillLocks test for the Global sharing mode. Since Global mode shares a workspace across all callers and also relies on the mutex to serialize access, adding a test verifying that acquire() in Global mode returns a locked unique_lock would make the test suite complete and symmetric.

Copilot uses AI. Check for mistakes.

TEST_F(XNNWorkspaceManagerTest, GlobalMode) {
// In Global mode, all calls should return the same workspace.
workspace_manager_->set_sharing_mode(WorkspaceSharingMode::Global);
Expand Down
Loading