MLIRContext usage clarification

I’m playing around a bit with how contexts are managed in MLIRSwift. Currently I have a single, static context but for various reasons I’m considering breaking it into several separate contexts, and I’m wondering about the possible correctness and performance implications of such a change.

One concrete example I’m thinking about is having a separate, immortal context for all of the types, attributes, and identifiers associated with a particular dialect (let’s call it dialectContext). In my architecture, these are likely to be used when constructing a Module initialized with a separate context (let’s call it moduleContext). Things like Locations will be created using ModuleContext, while more constant things like types, certain identifiers and attributes will be borrowed from dialectContext (and possibly from multiple dialectContexts from various dialects).

Glancing at the code, it seems like this will work but I’m wondering if there is something I’m missing that might cause problems for a Module with constituents from different Contexts, or if there are performance implications I should think about.

You can’t mix entities with different context: the underlying principle of the context is that it is used to provide an isolated environment (or “context”, hence the name), in which your entire IR is living. This includes types, attributes, location, dialect, identifiers, and every operation that references these.

You may think that you can get away by keeping types/attributes in a context instance and use another instance for locations, however this won’t be sustainable: all the passes and code manipulating the IR assume a single context and will create new entities (type, attributes, location, …) by calling getContext() on the most convenient object at hand (Pass, Operation, Attribute, Location, …), not respecting the segregation you’re trying to get at.

If you’re trying to optimize storage this way, I think a better direction would be to make it possible for multiple context to share the storage for types, attributes, etc. individually (through injection for example). This may have some “interesting” consequences though as the context is fairly complex for multi-threading management purpose.

(also: an permanent storage for Types (etc.) means an ever growing amount of memory that you can’t reclaim)

1 Like

Thanks, yeah I was wondering about how aggressively folks used getContext on random things… oh well.