Expose a StableHLO Dialect CAPI that doesn't requires building upstream dialects

[cerisier]

This is an attempt PR to address #2593

This PR introduces 2 seperate targets:

1. stablehlo_dialect_capi which embeds StableHLO + Serialization.
2. stablehlo_unified_capi which embeds stablehlo_dialect_capi as well as transformation passes (incl linalg) and reference interpreter.

Note:

• I kept CHLO C API as a separate target as it was. I could include it in stablehlo_dialect_capi but since it was already exposed as a separate target and not included in the former stablehlo_capi I didn't do it just yet.

• stablehlo_passes is still included in dialect_capi because it's needed by serialisation.

• I kept an alias stablehlo_capi on stablehlo_unified_capi just so that downstreams don't have to rename the target they depend on. Let me know if this is un-needed.

• I had to extract a small type conversion delcared in linagl_passes but used in stablehlo_passes which forced stablehlo_passes to depend on linagl_passes.
  I named this extracted target stablehlo_type_conversions and put it under stablehlo/transforms/conversions. Let me know if this is not the right location.

Dependency Graph

stablehlo_dialect_capi

stablehlo_unified_capi

Finally I tested integration in XLA and JAX as well as python bindings to validate that it links correctly.

Let me know what you think and if you see further changes to add to this PR :)

[GleasonK]

This looks mostly good! And definitely a good first step. I think the part 2 of this (can be separate and I can try my hand at it) would be splitting the dialect API passes from stablehlo_passes into something else. I.e. I think we only need serialization passes, so no need to build the other passes.

Or maybe the overhead of this dep isn't too bad in which case we don't need to do this... Have you done any analysis on the library size to compare before/after?

Finally I tested integration in XLA and JAX as well as python bindings to validate that it links correctly.

Thank you so much for this. This tends to be the trickier part of these CAPI / Python binding changes :)

Also RE lint failure: ./build_tools/github_actions/lint_clang_format.sh -f should fix it up

[steeve]

@GleasonK tbh I'm pretty OK with keeping the other passes, as long as they don't pull other dialects (or perhaps, LinAlg only). For instance, CHLO -> StableHLO. Or perhaps you count them as serialization passes too ?

[GleasonK]

Let's keep the passes. I was thinking if we should split into a hypothetical "framework_passes" and "plugin_passes" looking at the things in passes this may not be a big deal:

Useful for getting things to StableHLO (pre-serialization):
- ChloLegalizeToStablehloPass 
- ShapeLegalizeToStablehloPass
- StablehloLegalizeDeprecatedOpsPass
- StablehloCompatibilityExpanderPass

Serialization passes:
- VhloLegalizeToStablehloPass
- VhloToVersionPass
- StablehloLegalizeToVhloPass 

====== Could draw the line there for framework passes ======

Graph simplifications (could go on either side, FW / plugin):
- StablehloAggressiveFolderPass
- StablehloAggressiveSimplificationPass

Dynamic shape handling (plugins):
- StablehloCanonicalizeDynamismPass  
- StablehloRefineArgumentsPass 
- StablehloRefineShapesPass 

Quantization handling (for plugins):
- StablehloLegalizeQDQToQuantizedOpPass 
- StablehloLegalizeQuantizedOpToQDQPass
- StablehloLegalizeQuantToMathPass 

Other (numeric precision / HW specific passes):
- StablehloComplexMathExpanderPass 
- StablehloConvertToSignlessPass 
- StablehloLegalizeCompositeToCallPass 

This shouldn't pull in many new MLIR dialects, maybe linalg by the sound of it but maybe not even(?). Happy to defer that decision, can prototype and see if the binary size changes much.

[GleasonK]

LGTM, will need a rebase, and there's a buildifier lint error

[cerisier]

@GleasonK Should be good ! Many thanks for the review.

[cerisier]

Not sure what happened... Buildified again.

[GleasonK]

Thanks! Also I've sent you an OpenXLA org invite so your future PRs won't need approval to run CI