Skip to content

Don't hold the PatternCache lock while freezing the patterns. - #43

Merged
vswagath1989 merged 1 commit into
torch-spyre:mainfrom
lupalby:fix-patterncache-freeze-deadlock
Sep 21, 2026
Merged

vswagath1989 merged 1 commit into
torch-spyre:mainfrom
lupalby:fix-patterncache-freeze-deadlock

Conversation

@lupalby

@lupalby lupalby commented Sep 21, 2026

Copy link
Copy Markdown
Member

What

PatternCache::get held mutex_ across the FrozenRewritePatternSet
construction. This narrows the lock to the map accesses and builds the frozen
set with it released.

Why

Freezing a PDL module runs a pass manager of its own, which hands its parallel
verifier to the MLIRContext's thread pool and waits for it. And
apply-device-patterns is nested over modules and then functions, so MLIR's
asynchronous pass adaptor runs it on that same pool, many at once, all sharing
the one PatternCache the device holds.

StdThreadPool::processTasks(WaitingForGroup) wakes on !Tasks.empty() and
then pops Tasks.front() regardless of which group that task belongs to. So a
pool thread waiting for its own group picks up a sibling function's run of this
pass and stops on this lock -- and is then no longer available to finish the
group the lock's holder is waiting for. Once every thread that could have
finished the freeze is stopped on the lock, nothing finishes and the compile
hangs.

The two shapes seen together in one hung process:

holder:   PatternCache::get (holds mutex_) -> FrozenRewritePatternSet
                                          -> PassManager::run -> verify
                                          -> processTasks         <- waiting on the pool
blocked:  processTasks -> verifyOpAndDominance -> processTasks     <- waiting on its own group
                      -> stolen adaptor task  -> PatternCache::get <- parked on mutex_

SmartMutex<true> being recursive kept this quiet rather than loud: in one
capture the holder had stolen a sibling task, re-entered its own lock, found the
cache still cold and started rebuilding the whole set on top of itself -- two
nested PatternCache::get frames in one stack.

Cost

Two callers that miss together now build a set each, and the first one back
files it with try_emplace while the others take that, so every caller still
gets one set. Building concurrently is safe: getPatterns only reads the device
and clones into a module of its own.

Testing

Reproduced and measured on a downstream MLIR pipeline that runs this pass over
several modules, each holding one function. The failure is a hang until the CI
timeout, on a different test each time, so it is measured as a hang rate over
repeated runs of one suite rather than as a single pass/fail. Same harness, two
uninstrumented builds:

build suite runs hung
before 15 2
after 30 0

Every test passed on every run of the fixed build.

How likely it is depends on how close the pool is to saturation. Widening the
lock window with a sleep does not reproduce it on a 64-core machine, because
idle workers are always left to finish the group; fewer cores makes it more
likely, which is why CI sees it more often than a workstation does.

No regression test: a reappearance shows up as a CI timeout rather than a test
failure, and the interleaving is not reliably forceable.

Follow-up

mutex_ could become non-recursive now that nothing is held across the freeze.
Nothing needs the recursion, and a non-recursive mutex would have made this an
immediate, obvious failure instead of a hang once in several runs.

`PatternCache::get` held `mutex_` across the `FrozenRewritePatternSet`
construction. Freezing a PDL module runs a pass manager of its own, which
hands its parallel verifier to the MLIRContext's thread pool and waits for
it -- and `apply-device-patterns` is nested over modules and then functions,
so MLIR's asynchronous pass adaptor runs it on that same pool, many at once,
all sharing the one cache the device holds.

A pool thread that is waiting for a task group also runs whatever else is
queued while it waits, whatever group that task belongs to. So a thread
waiting for its own group picks up a sibling function's run of this pass and
stops on this lock, which means it is no longer available to finish the
group the lock's holder is waiting for. Once the threads that would have
finished the freeze are all stopped on the lock, nothing finishes: the
compile hangs. `SmartMutex<true>` being recursive hid it, letting the holder
re-enter and rebuild the set on top of itself rather than fail loudly.

Hold the lock only across the map accesses and build with it released. Two
callers that miss together now build a set each and the first one back files
it, which is the whole cost; cloning the patterns only reads the device.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Alberto Mannari <ert@zurich.ibm.com>

@KFAFSP KFAFSP left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I had removed the lock once before and was still able to trigger a race, maybe I did something else wrong. Still, if this works for you, LGTM.

@vswagath1989
vswagath1989 merged commit 29a121d into torch-spyre:main Sep 21, 2026
3 checks passed
@lupalby
lupalby deleted the fix-patterncache-freeze-deadlock branch September 21, 2026 15:03
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants