EnzymeAD / EnzymeAD/Enzyme-JAX

Last remaining all-gathers

Open
#770 6 comments 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.