Generating MLIR code with dynamic dimensions via Reactant?

[oschulz]

Reactant may be an awesome way to export, share and preserve scientific models (not just ML, also simulation code, etc.) written in Julia, so they can be used beyond Julia.

One problem I've run into while exploring this is how to generate StableHLO code that doesn't hardcode batch size. For example, when doing a multi-particle simulation, the exported StableHLO shouldn't hardcode the number of particles - batch size should still be flexible depending on user need and target hardware.

StableHLO itself support things like tensor<?x128xf32> right? And maybe the ? could then even be rewritten to a fixed size at the point/time of use? I've read that JAX export supports emitting StableHLO with ? dimensions via polymorphic_shapes, but I'm not an expert. Can/could Reactant support this as well?

This could be very relevant for a new project that I'm part of.

[mofeing]

Reactant arrays do not yet support dynamic dimensions, but we are discussing how they could get supported. Keep in mind that currently many optimizations require static shapes, so arrays with dynamic dimensions won't be as optimized as the ones with static shapes.

[oschulz]

but we are discussing how they could get supported.

Thanks, that would be great!

Keep in mind that currently many optimizations require static shapes

Reactant or XLA optimizations? For the latter, it could be feasible to replace ? by a batch size before use, at least in theory, right? Not by Reactant, but by whatever custom tooling that will read the stored StableHLO, before handing it to XLA?

[mofeing]

Reactant or XLA optimizations?

Enzyme-JAX; i.e. the optimization engine under Reactant. I'm unaware of the status of XLA optimizations for dynamic dims.

For the latter, it could be feasible to replace ? by a batch size before use, at least in theory, right? Not by Reactant, but by whatever custom tooling that will read the stored StableHLO, before handing it to XLA?

Do you mean to compile it first with dyn dim and later transform the code with a given dim value? You would need to check if the value is propagated everywhere.

What are you trying to do?

[oschulz]

I'm trying to - well, more like exploring the feasibility of - preserving and exchanging physics model code in a host-language independent form. But with a batch size that is not fixed, but can be set at the "receiving end", so it speak. So basically I'm looking for a way to have an "insert batch size here" placeholder in MLIR/StableHLO code. I think typically there will be a fixed batch size at the end, at the point of use, and then fancy optimizations can be run by whatever framework is used there. Since all other dimension sizes will be fixed, I thought the dynamic dimension size of StableHLO could then, by convention, be used as the batch size - as there is a legal way to represent that in StableHLO code via ? (from what I understand).

[avik-pal]

There are shape refinement passes implemented upstream in stablehlo https://openxla.org/stablehlo/dynamism#compiler_passes_for_refining_dynamic_programs