Caching and parallel compilation

[anstow]

We are writing a just in time compiler. As part of this we are attempting to hide the compilation overhead by compiling multiple MLIR graphs in different threads at the same time. As far as I can tell this means we need an MLIRContext per compilation thread (if only to ensure that debug info doesn’t get interleaved). We also have a caching layer - if we receive the same module twice we reuse the executable from the first compilation. The caching layer requires graph hashing and graph comparison. The comparison and hashing from mlir::OperationEquivalence appears to depend on the MLIRContext the module was built in (this is because new copies of the dialects are loaded for each MLIRContext). So we’ve been forced to write our own comparison and hashing layer (obviously this is very similar to mlir::OperationEquivalence). I have a couple of questions:

1. Is there a better way of doing the graph comparison? Can I somehow use mlir::OperationEquivalence with graphs from different contexts?
2. Re-loading all the dialects for each thread seems inefficient - is there a way to re-use this work for multiple MLIRContexts?

[mehdi_amini]

@anstow

As part of this we are attempting to hide the compilation overhead by compiling multiple MLIR graphs in different threads at the same time. As far as I can tell this means we need an MLIRContext per compilation thread

MLIR is thread-safe, but not all of the MLIRContext APIs are. It is very possible to do multi-threading with a single context but it requires care.

@anstow

• Is there a better way of doing the graph comparison? Can I somehow use mlir::OperationEquivalence with graphs from different contexts?

Your analysis is correct: OperationEquivalence is relying on the IR entities being in the same MLIRContext. The reason is that the MLIRContext already hashes the Types and Attributes so OperationEquivalence can avoid doing this and hash the pointer used by the DictionaryAttr for example.

@anstow

Re-loading all the dialects for each thread seems inefficient - is there a way to re-use this work for multiple MLIRContexts?

This is a good question: we should look into the cost of dialect initialization and see if a "MLIRContext::clone()` method could be faster (I’m not sure off-hands).
Overall I would be more interested in tracking down the cost of loading dialects and see if we can lower it!

[jackalcooper]

Glad to see this is discussed. I encountered this problem when running multiple pass managers of the same MLIRContext in parallel.