Primitives is an "open set", MLIR with allow-unregistered-dialects is an open set, but StableHLO isn't

[erick-xanadu]

Hello,

I have defined custom primitives in JAXPR. I also built the pjrt_c_api_cpu_plugin.so because I wanted to experiment with passing custom dialects to a custom compiler.

I've been able to:

• define my own custom primitives
• use the generated python classes from a custom MLIR dialect to define MLIR lowerings from these primitives to the custom primitives. (It is important to note that these MLIR classes do not include the MLIR bindings and therefore are printed with only the generic syntax when looking at the generated MLIR module, type inference also doesn't work, etc...).
• I am also able to set ctx.allow_unregistered_dialects from JAX's MLIR context to allow custom dialects.

However, I believe somewhere during the PJRT and before I am able to compile the MLIR module that I receive, the PJRT rejects the MLIR module because it does not allow unregistered dialects. Specifically I thin kthis happens in the ParseMlirModuleString method in mlir_to_hlo.cc.

I was wondering if it would be possible to either allow unregistered dialects (and then my own compiler can lower these to StableHLO) or allow for MLIR dialect plugins mechanism to allow me to register a custom dialect with the PJRT.

Or maybe this could be an issue on my side? Is there another interface I could use to make sure unregistered dialects are allowed in my MLIR Module? Maybe something in PjRt?

Thanks!