ml_program support

[silvasean]

Request description

ml_program is being used for certain JAX -> IREE flows and we are looking at using it in Torch-MLIR. Is there a plan for supporting this in StableHLO?

[joker-eph]

What does it mean to "support it"? What would be the difference with a stablehlo.func / stablehlo.graph operation?

[silvasean]

I mostly would be interested in the globals. In Torch-MLIR we still use func.func

[burmako]

Hi Sean! I'd also like to ask what kind of support you have in mind.

From the looks of it, ml_program.global_load, ml_program.global_store and friends work with the builtin tensor type, which is compatible with stablehlo. Although, I see that the ODS for these ml_program ops says AnyType which makes me worried a bit that I might've missed something.

[silvasean]

Yes, I think they are intended to be more general than just tensors. E.g. for storing a scalar i64 for a counter / rng seed / etc.

[burmako]

Sounds like loading scalars will be compatible with StableHLO as well, given that we can use tensor.from_elements there, and that's also on our list of ops to support in #115.

Overall, given that there are straightforward ways of using MLProgram's globals with StableHLO, and that #115 proposes upholding compatibility guarantees for "All ops, attributes and types" in the ml_program dialect, it looks like we have a solid plan for supporting this in StableHLO, so we can probably consider this ticket resolved?

[silvasean]

I would expect that we would close this PR when appropriate docs and specs have been updated to include ml_program in the set of "what is included in StableHLO"