search

openxla / stablehlo

Backward compatible ML compute opset inspired by HLO/MHLO
Apache License 2.0
357 stars 100 forks source link

ml_program support #182

Open silvasean opened 1 year ago

silvasean commented 1 year ago

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 commented 1 year ago

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

silvasean commented 1 year ago

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

burmako commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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 commented 1 year ago

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"