fix: multi-device execution and sharding [take III]

[avik-pal]

• fixes for MultiGPU runs
• Tests
• API to load output shardings from XLA
• JLL Changes
• Adds @code_mhlo that prints mlir with mhlo dialect
• Remove forced replication of outputs
  • remove linear results sharding info
• use OpSharding to regenerate the outputs
• ConcreteRNumber can now be replicated
• JLL bump

[avik-pal]

dont merge some more fixes needed

[avik-pal]

beautiful error on TPUs

ERROR: INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/service/spmd/spmd_partitioner.cc:5245) !HasReplicatedSharding(hlo->sharding()) || CanSideEffectingHaveReplicatedSharding(hlo) side-effect HLO cannot have a replicated sharding: %cust
om-call.8 = f32[12,2]{1,0} custom-call(f32[12,2]{1,0} %sine.1), custom_call_target="xla.sdy.FuncResultSharding", custom_call_has_side_effect=true, sharding={replicated}, frontend_attributes={xla.sdy.sharding="#sdy.sharding_per_value<[<@mesh, [{}, {}], re
plicated={\"data\", \"model\"}>]>"}, metadata={op_name="custom-call.9"}

[avik-pal]

beautiful error on TPUs

ERROR: INTERNAL: RET_CHECK failure (third_party/tensorflow/compiler/xla/service/spmd/spmd_partitioner.cc:5245) !HasReplicatedSharding(hlo->sharding()) || CanSideEffectingHaveReplicatedSharding(hlo) side-effect HLO cannot have a replicated sharding: %cust
om-call.8 = f32[12,2]{1,0} custom-call(f32[12,2]{1,0} %sine.1), custom_call_target="xla.sdy.FuncResultSharding", custom_call_has_side_effect=true, sharding={replicated}, frontend_attributes={xla.sdy.sharding="#sdy.sharding_per_value<[<@mesh, [{}, {}], re
plicated={\"data\", \"model\"}>]>"}, metadata={op_name="custom-call.9"}

Just to note. This is only on TPUs, the same code works fine on GPUs

[avik-pal]

Remaining changes are only julia side so hopefully no more JLL building needed

[avik-pal]

This is now ready, we just need the JLL build to go through. Once that is done I will test with TPUs

[wsmoses]

very minor but technically if we want to be correct, hlo is the low level post xla IR, and mhlo (aka mlir hlo) and now stablehlo is what we have before

re naming of code_hlo and code_mhlo lol

we don't need to fix here though

[avik-pal]

very minor but technically if we want to be correct, hlo is the low level post xla IR, and mhlo (aka mlir hlo) and now stablehlo is what we have before

@code_mhlo actually still prints the mhlo not the xla IR. But it is useful since all the shardy ops are expanded into custom calls.

module @reactant_fn_test1 attributes {mhlo.frontend_attributes = {xla.sdy.meshes = "{mesh = #sdy.mesh<[\\\22x\\\22=1, \\\22y\\\22=2]>}"}, mhlo.num_partitions = 2 : i64, mhlo.num_replicas = 1 : i64} {
  func.func @main(%arg0: tensor<12x4x16xf32> {mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding<@mesh, [{\\\22x\\\22}, {}, {\\\22y\\\22}]>"}, mhlo.sharding = "{devices=[1,1,2]<=[2]}"}) -> (tensor<f32>, tensor<12x4x1xf32>, tensor<12x4x16xf32> {mhlo.sharding = "{devices=[1,1,2]<=[2]}"}, tensor<12x4x16xf32>) {
    %0 = mhlo.constant dense<1.000000e+00> : tensor<f32>
    %1 = "mhlo.broadcast_in_dim"(%0) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>) -> tensor<16x4x12xf32>
    %2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %3 = "mhlo.transpose"(%arg0) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<12x4x16xf32>) -> tensor<16x4x12xf32>
    %4 = mhlo.add %3, %3 : tensor<16x4x12xf32>
    %5 = mhlo.add %3, %1 : tensor<16x4x12xf32>
    %6 = mhlo.multiply %5, %4 : tensor<16x4x12xf32>
    %7 = mhlo.reduce(%4 init: %2) applies mhlo.add across dimensions = [0, 1, 2] : (tensor<16x4x12xf32>, tensor<f32>) -> tensor<f32>
    %8 = mhlo.reduce(%4 init: %2) applies mhlo.add across dimensions = [0] : (tensor<16x4x12xf32>, tensor<f32>) -> tensor<4x12xf32>
    %9 = "mhlo.transpose"(%8) <{permutation = dense<[1, 0]> : tensor<2xi64>}> : (tensor<4x12xf32>) -> tensor<12x4xf32>
    %10 = mhlo.reshape %9 : (tensor<12x4xf32>) -> tensor<12x4x1xf32>
    %11 = "mhlo.transpose"(%5) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<16x4x12xf32>) -> tensor<12x4x16xf32>
    %12 = "mhlo.transpose"(%6) <{permutation = dense<[2, 1, 0]> : tensor<3xi64>}> : (tensor<16x4x12xf32>) -> tensor<12x4x16xf32>
    %13 = mhlo.custom_call @xla.sdy.FuncResultSharding(%11) {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\\\22x\\\22}, {}, {\\\22y\\\22}]>]>"}} : (tensor<12x4x16xf32>) -> tensor<12x4x16xf32>
    return %7, %10, %13, %12 : tensor<f32>, tensor<12x4x1xf32>, tensor<12x4x16xf32>, tensor<12x4x16xf32>
  }
}

We should definitely rename code_hlo to return the actual HLO module at some point. Though I need to check how to get that without dumping it to a file.

[avik-pal]

This is now ready to go!