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 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:
- Original HLO, there's no
param_1 and only got param_0 in async_computation:
%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}
}
param_0 was changed to param_1 after GpuSanitizeConstantNames pass:
%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}
}
- Another
param_1 was generated after FusionWrapper pass:
%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. There're 2 issues here:
- The original HLO instruction name may be changed unexpectedly even though the pass does not intend to do so.
- 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