EnzymeAD / EnzymeAD/Enzyme-JAX
Last remaining all-gathers
- Dominant language
- MLIR
- Stars
- 131
- Forks
- 53
- Avg merge
- 1d 10h
- Merged PRs (30d)
- 193
Description
```
%tuple.943 = (u32[], u32[], f64[1,24,48]{2,1,0}, f64[1,24,48]{2,1,0}) tuple(%and.213, %shift-right-logical.2, %broadcast_select_fusion.64, %pad_dynamic-slice_fusion.32)
%conditional.87 = (f64[5,46,94]{2,0,1}) conditional(%compare.756, %tuple.945, %tuple.946), true_computation=%region_176.10734_spmd, false_computation=%region_177.10738_spmd
%tuple.951 = (f64[20,24,48]{2,1,0}) tuple(%dynamic-slice_pad_fusion.16)
%tuple.950 = (u32[], u32[], f64[20,24,48]{2,1,0}, f64[20,24,48]{2,1,0})
%region_176.10734_spmd (param.176: (f64[5,23,47])) -> (f64[5,46,94]) {
%param.176 = (f64[5,23,47]{2,1,0}) parameter(0)
%get-tuple-element.87871 = f64[5,23,47]{2,1,0} get-tuple-element(%param.176), index=0
%copy.742 = f64[5,23,47]{1,0,2} copy(%get-tuple-element.87871)
%all-gather = f64[5,23,94]{1,0,2} all-gather(%copy.742), channel_id=847, replica_groups=[2,2]<=[2,2]T(1,0), dimensions={2}, use_global_device_ids=true
%copy.728 = f64[5,23,94]{2,0,1} copy(%all-gather)
%all-gather.1 = f64[5,46,94]{2,0,1} all-gather(%copy.728), channel_id=848, replica_groups=[2,2]<=[4], dimensions={1}, use_global_device_ids=true
ROOT %tuple.6059 = (f64[5,46,94]{2,0,1}) tuple(%all-gather.1)
}
%region_177.10738_spmd (param.177: (f64[4,23,47])) -> (f64[5,46,94]) {
%param.177 = (f64[4,23,47]{2,1,0}) parameter(0)
%constant.7352 = f64[] constant(0)
%get-tuple-element.87872 = f64[4,23,47]{2,1,0} get-tuple-element(%param.177), index=0
%copy.743 = f64[4,23,47]{1,0,2} copy(%get-tuple-element.87872)
%pad.129 = f64[5,23,47]{1,0,2} pad(%copy.743, %constant.7352), padding=1_0x0_0x0_0
%all-gather.2 = f64[5,23,94]{1,0,2} all-gather(%pad.129), channel_id=849, replica_groups=[2,2]<=[2,2]T(1,0), dimensions={2}, use_global_device_ids=true
%copy.729 = f64[5,23,94]{2,0,1} copy(%all-gather.2)
%all-gather.3 = f64[5,46,94]{2,0,1} all-gather(%copy.729), channel_id=850, replica_groups=[2,2]<=[4], dimensions={1}, use_global_device_ids=true
ROOT %tuple.6060 = (f64[5,46,94]{2,0,1}) tuple(%all-gather.3)
}
```
cc @avik-pal
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.