Repository: tensorflow
Issue: MHLO -> HLO does not respect sharding, inserting a tuple without sharding
Labels: stat:awaiting tensorflower,type:bug,comp:runtime,comp:xla

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

source

### Tensorflow Version

2e896fbe1e0ea4df33fbcfe780a1036f431b4e89

### Custom Code

Yes

### OS Platform and Distribution

Linux Ubunto 20.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/Compiler version

10.3

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

Converting an MHLO program to HLO that is fully annotated with shardings, results in HLO that has a tuple instruction that is without sharding.

Input MLIR `sharding-not-respected-mhlo-to-hlo.mlir`:
```mlir
func.func @main(%arg0: tensor<2x2xi32> {mhlo.sharding = "{devices=[2,1]0,1}"}) -> (tensor<2x2xi32> {mhlo.sharding = "{devices=[2,1]0,1}"}) {
  %0 = mhlo.add %arg0, %arg0 {mhlo.sharding = "{devices=[2,1]0,1}"} : tensor<2x2xi32>
  return %0 : tensor<2x2xi32>
}
```

Command:
```
xla-translate -mlir-hlo-to-hlo-text sharding-not-respected-mhlo-to-hlo.mlir
```

Result:
```hlo
HloModule main, entry_computation_layout={(s32[2,2]{1,0})->s32[2,2]{1,0}}

ENTRY %main.5 (Arg_0.1: s32[2,2]) -> s32[2,2] {
  %Arg_0.1 = s32[2,2] parameter(0), sharding={devices=[2,1]0,1}
  %add.2 = s32[2,2] add(s32[2,2] %Arg_0.1, s32[2,2] %Arg_0.1), sharding={devices=[2,1]0,1}, metadata={source_file="sharding-not-respected-mhlo-to-hlo.mlir" source_line=2}
  %tuple.3 = (s32[2,2]) tuple(s32[2,2] %add.2)
  ROOT %get-tuple-element.4 = s32[2,2] get-tuple-element((s32[2,2]) %tuple.3), index=0, sharding={devices=[2,1]0,1}
}
```

You can see that a tuple instruction has been inserted that has no sharding annotation.
```hlo
  %tuple.3 = (s32[2,2]) tuple(s32[2,2] %add.2)
```

This [test](https://github.com/tensorflow/tensorflow/blob/2e896fbe1e0ea4df33fbcfe780a1036f431b4e89/tensorflow/compiler/xla/translate/mhlo_to_hlo/tests/sharding.mlir#L21) expects a tuple without sharding. Is this really the expected behavior?
For example this tuple instruction would cause a problem if you want to do SPMD partitioning. Then the partitioner would insert an unwanted all-gather instruction.

If you do a conversion without sharding of the same MLIR
```mlir
func.func @main(%arg0: tensor<2x2xi32>) -> tensor<2x2xi32> {
  %0 = mhlo.add %arg0, %arg0 : tensor<2x2xi32>
  return %0 : tensor<2x2xi32>
}
```
You get a nicer result without the redundant `tuple` and `get-tuple-element` instructions
```hlo
HloModule main, entry_computation_layout={(s32[2,2]{1,0})->s32[2,2]{1,0}}

ENTRY %main.3 (Arg_0.1: s32[2,2]) -> s32[2,2] {
  %Arg_0.1 = s32[2,2] parameter(0)
  ROOT %add.2 = s32[2,2] add(s32[2,2] %Arg_0.1, s32[2,2] %Arg_0.1), metadata={source_file="sharding-not-respected-mhlo-to-hlo.mlir" source_line=2}
}
```

I can see two solutions here. 
1. Make the tuple instruction inherit the correct sharding.
2. Return directly the result of the `add` operation.

### Standalone code to reproduce the issue

```shell
See the description.
```


### Relevant log output

_No response_</details>