[shoumikhin]

What is wrong today

From Python, _unsafe_reset_threadpool(n) is the only way to set the CPU thread count. If it runs after a method that uses XNNPACK is loaded, the next runs can crash or hang:

method = Runtime.get().load_program(pte.buffer).load_method("forward")
_unsafe_reset_threadpool(16)
for _ in range(20):
    method.execute([x])  # segfault, abort, or hang

On Linux x86_64 with an 8-thread default pool, a reset to 16 after loading crashed in 10 of 12 runs, and a reset to 4 also in 10 of 12. The Android bindings reset the pool each time a Module is created, so a second Module can hit it too.

Why it happens

The reset destroys the current pthreadpool and creates a new one. But XNNPACK got that pthreadpool from get_pthreadpool() when the method was loaded, and it keeps using it. So the loaded method runs on a freed pool.

Resizing the pool in place is not safe either. XNNPACK sizes its per-thread buffers by the thread count when it plans a method, and it does not plan again when the count changes.

What this change does

A reset no longer destroys a pthreadpool that get_pthreadpool() has handed out. It stays alive with its size, next to the new one:

• XNNPACK methods loaded before the reset keep running on their old pool.
• Methods loaded after the reset, and ExecuTorch's own kernels, use the new pool.
• A reset back to a kept size reuses that pool, so at most one pool per size is kept. _unsafe_destroy_threadpool() frees kept pools too.
• A pool that was never handed out is destroyed as before, so the usual "reset, then load" order costs nothing extra.

The header and the Python docstring now explain this.

Tests

• New ThreadPoolTest.ResetKeepsHandedOutPthreadpoolUsable: after a reset, the held pool still runs work on 2 threads, and a reset back to its size reuses it. It fails without the fix and passes with it.
• Ran extension_threadpool_test, the native CPU engine tests, the Python tests for the pybindings and XNNPACK, and the LLM custom op tests that call the reset, on Linux x86_64. Tests that failed there fail on main too.
• The script above now passes after resets to 16 and to 4. With 20 resets between 2 thread counts and a new method loaded after each, all 21 methods stay correct and the process gains no threads from round to round.

Fixes #24765

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24770

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 42 Pending

As of commit ce80ef8 with merge base 75af78c ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

PR Review in progress

• Gather context (diff, changed files, checklist)
• Run three independent reviewers
• Consolidate findings
• Fact-check findings
• Post review

View job run