PR #11306: [XLA:GPU] Fix not unique name issue in sanitize constant pass
PR #11306: [XLA:GPU] Fix not unique name issue in sanitize constant pass

Imported from GitHub PR https://github.com/openxla/xla/pull/11306

Fix `Instruction name is not unique` error reported by JAX UT [`ShardMapTest.test_matmul_reduce_scatter`](https://github.com/google/jax/blob/jaxlib-v0.4.24/tests/shard_map_test.py#L149-L162) in XLA:GPU.

### Background
Error message:
```
!ContainsKey(instruction_names, instruction->name()) Instruction name is not unique: param_1
```
This error is reported after `PrepareHloModuleForIrEmitting()`. The non-unique name `param_1` is generated from 2 different passes `GpuSanitizeConstantNames` and `FusionWrapper`. The related HLO changes are as follows:

1. Original HLO, there's no `param_1` and only got `param_0` in `async_computation`:
```ll
%main.19_spmd (param: s32[4,4], param.1: s32[4,8]) -> s32[2,8] {
  %param = s32[4,4]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %param.1 = s32[4,8]{1,0} parameter(1), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %dot.1 = s32[4,8]{1,0} dot(s32[4,4]{1,0} %param, s32[4,8]{1,0} %param.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=int32]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=158}
  %reduce-scatter-start = ((s32[4,8]{1,0}), s32[2,8]{1,0}) reduce-scatter-start(s32[4,8]{1,0} %dot.1), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"collective_backend_config":{"is_sync":false,"no_parallel_custom_call":false}}
  ROOT %reduce-scatter-done = s32[2,8]{1,0} reduce-scatter-done(((s32[4,8]{1,0}), s32[2,8]{1,0}) %reduce-scatter-start), metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
%async_computation (param_0: s32[4,8]) -> s32[2,8] {
  %param_0 = s32[4,8]{1,0} parameter(0)
  ROOT %reduce-scatter.2 = s32[2,8]{1,0} reduce-scatter(s32[4,8]{1,0} %param_0), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
```

2. `param_0` was changed to `param_1` after `GpuSanitizeConstantNames` pass:
```ll
%main.19_spmd (param: s32[4,4], param.1: s32[4,8]) -> s32[2,8] {
  %param = s32[4,4]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %param.1 = s32[4,8]{1,0} parameter(1), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %dot.1 = s32[4,8]{1,0} dot(s32[4,4]{1,0} %param, s32[4,8]{1,0} %param.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=int32]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=158}
  %reduce-scatter-start = ((s32[4,8]{1,0}), s32[2,8]{1,0}) reduce-scatter-start(s32[4,8]{1,0} %dot.1), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"collective_backend_config":{"is_sync":false,"no_parallel_custom_call":false}}
  ROOT %reduce-scatter-done = s32[2,8]{1,0} reduce-scatter-done(((s32[4,8]{1,0}), s32[2,8]{1,0}) %reduce-scatter-start), metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
%async_computation (param_1: s32[4,8]) -> s32[2,8] {
  %param_1 = s32[4,8]{1,0} parameter(0)
  ROOT %reduce-scatter.2 = s32[2,8]{1,0} reduce-scatter(s32[4,8]{1,0} %param_1), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
```

3. Another `param_1` was generated after `FusionWrapper` pass:
```ll
%async_computation (param_1: s32[4,8]) -> s32[2,8] {
  %param_1 = s32[4,8]{1,0} parameter(0)
  ROOT %reduce-scatter.2 = s32[2,8]{1,0} reduce-scatter(s32[4,8]{1,0} %param_1), channel_id=1, replica_groups={{0,2},{1,3},{4,6},{5,7}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
%wrapped_dot_computation (param_0.1: s32[4,4], param_1: s32[4,8]) -> s32[4,8] {
  %param_0.1 = s32[4,4]{1,0} parameter(0)
  %param_1 = s32[4,8]{1,0} parameter(1)
  ROOT %dot.2 = s32[4,8]{1,0} dot(s32[4,4]{1,0} %param_0.1, s32[4,8]{1,0} %param_1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=int32]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=158}
}

%main.19_spmd (param: s32[4,4], param.1: s32[4,8]) -> s32[2,8] {
  %param = s32[4,4]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %param.1 = s32[4,8]{1,0} parameter(1), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="jit(fwd)/jit(main)/shard_map[mesh=Mesh(\'x\': 2, \'y\': 2) in_names=({0: (\'x\',), 1: (\'y\',)}, {0: (\'x\',)}) out_names=({0: (\'x\', \'y\')},) check_rep=True rewrite=True auto=frozenset()]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=161}
  %dot.1 = s32[4,8]{1,0} dot(s32[4,4]{1,0} %param, s32[4,8]{1,0} %param.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=None preferred_element_type=int32]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=158}
  %reduce-scatter-start = ((s32[4,8]{1,0}), s32[2,8]{1,0}) reduce-scatter-start(s32[4,8]{1,0} %dot.1), channel_id=1, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, dimensions={0}, to_apply=%region_0.7, metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}, backend_config={"operation_queue_id":"0","wait_on_operation_queues":[],"collective_backend_config":{"is_sync":false,"no_parallel_custom_call":false}}
  ROOT %reduce-scatter-done = s32[2,8]{1,0} reduce-scatter-done(((s32[4,8]{1,0}), s32[2,8]{1,0}) %reduce-scatter-start), metadata={op_name="jit(fwd)/jit(main)/jit(shmap_body)/reduce_scatter[axis_name=y scatter_dimension=0 axis_index_groups=None axis_size=2 tiled=True]" source_file="/home/sdp/tenglu/intel-jax/tests/shard_map_test.py" source_line=159}
}
```

### Root cause
`GpuSanitizeConstantNames` runs before `FusionWrapper` and may change HLO instruction name by a [local name uniquer](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_sanitize_constant_names.cc#L37). There're 2 issues here:

* The original HLO instruction name may be changed unexpectedly even though [the pass does not intend to do so](https://github.com/openxla/xla/blob/main/xla/service/gpu/gpu_sanitize_constant_names.cc#L50-L51).
* The global name uniquer of HLO module isn't aware of this change and may wrongly assign the same name to another HLO instruction in other passes, e.g. `FusionWrapper`.

`param_0` is changed to `param_1` by `GpuSanitizeConstantNames` unexpectedly.

### Solution
Only record HLO instruction names in the local name uniquer of `GpuSanitizeConstantNames` and do not change the original names. It exactly follows the pass design.
Copybara import of the project:

--
c39bc256957db681cfe1706fb716986da89edf3a by Lu Teng <teng.lu@intel.com>:

Fix not unique name issue in sanitize constant pass.

Merging this change closes #11306

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11306 from Intel-tensorflow:tenglu/fix_unique_name c39bc256957db681cfe1706fb716986da89edf3a
