[Graph API] Is compiled_partition thread-safe for concurrent calls on same instance?

[lxq2t]

We observe crashes (double-free, heap corruption) when multiple threads concurrently call execute() on the same compiled_partition instance. The crash occurs inside matmul_t::execute_impl during scratchpad deallocation.

oneDNN version: 3.6.0
Graph contains: MatMul operations
OS: Linux 6.8.0-060800rc6-generic

Can you clarify:

1. Is compiled_partition::execute() guaranteed thread-safe for concurrent calls on the same instance?
2. If not, should we maintain a pool of compiled_partition instances (one per concurrent execution)?
3. Is there an allocator or scratchpad configuration that enables safe concurrent execution? We already tried dnnl::graph::make_engine_with_allocator(dnnl::engine::kind::cpu, 0, allocator); with functions to system malloc/free but without effect.

// Setup
dnnl::graph::graph g(dnnl::engine::kind::cpu);
// ... add matmul operations ...
g.finalize();

auto partitions = g.get_partitions();
auto cp = partitions[0].compile(inputs, outputs, engine);

// Concurrent execution - CRASHES
std::vector<std::thread> threads;
for (int i = 0; i < 4; ++i) {
    threads.emplace_back([&]() {
        for (int iter = 0; iter < 100; ++iter) {
            // Each thread has its own input/output tensors
            std::vector<dnnl::graph::tensor> my_inputs = /* thread-local */;
            std::vector<dnnl::graph::tensor> my_outputs = /* thread-local */;
            
            cp.execute(stream, my_inputs, my_outputs);  // CRASH here
        }
    });
}
for (auto& t : threads) t.join();

Thank you in advance.

[TaoLv]

Thank you for the questions, @lxq2t . Yes, compiled_partition::execute() should be thread-safe for concurrent calls on the same instance. We have a few internal tests to verify that. For example, this one.
For the problem you are seeing, could you please share a full reproducer for us to investigate?

[lxq2t]

We try to do investigation to isolate issue from our application code and confirm it can be reproducible independently.
In case, when we execute different instances of compiled partitions same way, issue is not reproduced.
May be it's unxpected/unsafe to use of graph in re-entrant way?

Estimated cause

1. Thread starts executing graph partition (e.g., SDPA kernel or more complex)
2. Kernel calls parallel_for() via threadpool interface
3. While waiting, thread steals another graph task of same instance of graph from queue
4. Creates nested/re-entrant kernel execution on same thread but with different input context data
5. Internal kernel state (scratchpads, temporaries, threadlocal) appears to get corrupted (not confirmed, but suspected based on crash patterns)

Minimal Reproduction

Example with simple threadpool that can show issue, with graph of typical ViT block:

https://gist.github.com/lxq2t/46d675bf3368686680f0df079d014d1a

In threadpool, protection by in_parallel is done to avoid nested parallel_for.

// Estimated pattern that triggers the bug
void parallel_for(int n, const std::function<void(int,int)>& fn) override {
    for (int i = 0; i < n; ++i) {
        pushTask([&]{ fn(i, n); });
    }
    while (!all_done) {
        executeOne();  // Steal and execute same graph with different context.
    }

Observed Crashes

Error	Frequency
malloc(): corrupted top size	Common
double free or corruption (out)	Occasional
malloc(): unaligned tcache chunk detected	Occasional
SIGSEGV in kernel	Occasional