Generating MLIR without JAX¯

[arul-lm]

#loc2 = loc(unknown)
module @zml attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  func.func @main(%arg0: tensor<128x128xf16> {mhlo.layout_mode = "default", mhlo.sharding = "{devices=[4,1]<=[4]}"} loc(unknown), %arg1: tensor<128x128xf16> {mhlo.layout_mode = "default", mhlo.sharding = "{devices=[4,1]<=[4]}"} loc(unknown)) -> (tensor<128x128xf16> {mhlo.layout_mode = "default", mhlo.sharding = "{devices=[4,1]<=[4]}"}) {
    %0 = stablehlo.custom_call @Sharding(%arg0) {backend_config = "", mhlo.sharding = "{devices=[1,4]<=[4]}"} : (tensor<128x128xf16>) -> tensor<128x128xf16> loc(#loc3)
    %1 = stablehlo.custom_call @Sharding(%arg1) {backend_config = "", mhlo.sharding = "{devices=[4,1]<=[4]}"} : (tensor<128x128xf16>) -> tensor<128x128xf16> loc(#loc3)
    %2 = stablehlo.dot_general %0, %1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<128x128xf16>, tensor<128x128xf16>) -> tensor<128x128xf16> loc(#loc6)
    %3 = stablehlo.custom_call @Sharding(%2) {backend_config = "", mhlo.sharding = "{devices=[4,1]<=[4]}"} : (tensor<128x128xf16>) -> tensor<128x128xf16> loc(#loc3)
    return %3 : tensor<128x128xf16> loc(#loc1)
  } loc(#loc1)
} loc(#loc5)
#loc = loc("module.zig":84:39)
#loc1 = loc("module.zig":339:39)
#loc3 = loc("tensor.zig":187:44)
#loc4 = loc("tensor.zig":1179:47)
#loc5 = loc("main"(#loc))
#loc6 = loc("dot({m=128,k=128!,f16},{k=128!,n=128,f16},contracting={ { 1, 0 } },batching={  }"(#loc4))

I'm generating MLIR from zml project. How can I take this and pass to sdy_opt? Unfortunately, I won't be able to use JAX to generate sdy.sharding attributes

[tomnatan30]

Hey!

I see you are generating mhlo.sharding and @Sharding custom calls, this works with GSPMD, you need to instead generate sdy.sharding attributes with sdy::TensorShardingAttr and sdy.sharding_constraint instead of the custom calls.

Not sure how ZML is generating the module, but in either python or c++ you are able to insert sdy ops and attrs. Take a look at our dialect documentation - https://github.com/openxla/shardy/blob/main/docs/sdy_dialect.md. Perhaps it would also be useful to see the difference between how jax lowers to gspmd or shardy - https://github.com/openxla/shardy/blob/main/docs/getting_started_jax.ipynb

You can also look at how JAX inserts sdy ops and attrs in python - https://github.com/jax-ml/jax/blob/main/jax/_src/named_sharding.py#L314, https://github.com/jax-ml/jax/blob/main/jax/_src/interpreters/mlir.py#L2616

Finally, there is an alternative. We have an mlir pipeline in XLA to convert gspmd shardings and custom calls to sdy dialect -
https://github.com/openxla/xla/blob/main/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc#L656. But I wouldn't recommend using it as it's limited and not maintained. Instead I would recommend generating sdy IR directly.

Hope this helps!