StableHLO CAPI requires building upstream dialects

[GleasonK]

Request description

Currently CAPI builds linalg and arith, etc -- but most users of CAPI (frameworks) only want StableHLO+CHLO+Serialization. Would be great to offer a stablehlo_dialect_capi which only has core APIs.

One possible separation direction:

• stablehlo_dialect_capi: Just StableHLO+CHLO+Serialization APIs, these are the things useful for PJRT frameworks.
• stablehlo_unified_capi: A separate target which builds the dialect CAPI along with the transformation passes / reference interpreter / other useful APIs, likely will be the target used for StableHLO python bindings

This will probably require splitting the serialization APIs into a separate file StablehloDialectApi.h and StablehloUnifiedApi.h. The following functions should go into the dialect API: https://github.com/openxla/stablehlo/blob/main/stablehlo/integrations/c/StablehloApi.h#L23-L117

The reference API and pass registration should go into the unified API target:
https://source.corp.google.com/piper///depot/google3/third_party/stablehlo/stablehlo/integrations/c/StablehloPasses.h

Looking at the CAPI deps today, I'm surprised that this requires upstream dialects, so definitely something to look into.

(Noted by @steeve at recent OpenXLA community meeting)

[cerisier]

Hey @GleasonK,

I'm having a go at this and am nearly ready to submit a draft PR.
Before submitting anything, I would love your thoughts on a few points:

1. API Export vs Re-export

Basically, either:

1. stablehlo_unified_capi target exports StablehloDialectApi.h as well as StablehloUnifiedApi.h (in the hdrs attribute) (but then downstreams users like python bindings have to import the 2 of them)
2. StablehloUnifiedApi.h re-exports StablehloDialectApi.h which becomes the only header to import.

2. Common deps between new StablehloUnifiedApi and StablehloDialectApi

Currently, stablehlo_capi compiles StablehloAttributes, StablehloDialect and StablehloTypes and exports their respective headers which will likely needed by users of the new stablehlo_dialect_capi.

I have several options:

1. Split those further into multiple targets and depend on them on both stablehlo_unified_capi and stablehlo_dialect_capi targets. (But then both targets need to either re-rexports the headers or exports them for manual inclusion downstream).
2. Have stablehlo_dialect_capi depend on them and stablehlo_unified_capi depend on stablehlo_dialect_capi.
3. Have both stablehlo_unified_capi and stablehlo_dialect_capi compile those files and exports their headers.

3. Backward compat for target names

Should I expose aliases so that stablehlo_capi points to stablehlo_unified_capi to preserve backward compat ?

If confusing, I can go ahead and submit what I have and we start from there ?

Thank you in advance 🙏

[GleasonK]

I don't have any major preferences as long as it works.. These CAPIs get a little strange with linking in MLIR but I think this should all be OK, since that linking is handled at a higher level in other projects.

In general I like the layering approach you've mentioned, is seems like:

• Layer 1: StableHLO Dialect (note we really should include serialization APIs in this core target since frameworks that interact with PJRT all should sent VHLO not StableHLO to plugins).
• Layer 2: StableHLO APIs (dialect + interpreter + other mlir dialects + etc)

(2) In that regard I like 2.2 if that layering works "stablehlo_unified_capi depend on stablehlo_dialect_capi". That said if this proves more difficult then having multiple header deps is probably fine. The way we export for other projects today is this file group, so 1 header vs 2 tends to be opaque to users:

• StableHLO:

  stablehlo/BUILD.bazel

  Line 1078 in b62dc66

  filegroup(

• JAX: https://github.com/jax-ml/jax/blob/9107ee4a2226d9d26269d587b75a4206b3696c03/jaxlib/mlir/_mlir_libs/BUILD.bazel#L351

(3) Backward compatibility will likely be important. I'm not clear what this alias means, if this is just an easy build target, then probably ok to break compat and I can update internally, if this will change what submodules APIs are in and would require modifying lots of Py source code that uses StableHLO bindings..then alias is preferred.

[cerisier]

Thanks for the notes ! I'll make a draft PR with the above guidelines !

[cerisier]

I think this can be closed now 😊