EnzymeAD / EnzymeAD/Reactant.jl
Bad sharding comm
- Dominant language
- Julia
- Stars
- 370
- Forks
- 74
- Avg merge
- 18h 47m
- Merged PRs (30d)
- 30
Description
before optimizations:
```julia
using Reactant
@assert length(Reactant.devices()) >= 2
N = min((length(Reactant.devices()) ÷ 2) * 2, 8)
@show N
mesh = Sharding.Mesh(reshape(Reactant.devices()[1:N], 2, :), (:x, :y))
sharding = Sharding.NamedSharding(mesh, (:x, :y))
@show mesh
x = Reactant.to_rarray(reshape(collect(Int, 1:(1024*12)), 1024, 12); sharding)
function rotate(x)
y = x[1:100, :]
x[1:924, :] = x[101:1024, :]
x[925:1024, :] = y
nothing
end
@show @code_hlo shardy_passes=:to_mhlo_shardings optimize=false rotate(x)
@show @code_hlo shardy_passes=:to_mhlo_shardings rotate(x)
@show @code_xla shardy_passes=:to_mhlo_shardings rotate(x)
function pad(x)
return Reactant.Ops.pad(x, eltype(x)(0); low=[5, 0], high=[15, 0], interior=[0, 0])
end
@show @code_hlo shardy_passes=:to_mhlo_shardings optimize=false pad(x)
@show @code_hlo shardy_passes=:to_mhlo_shardings pad(x)
@show @code_xla shardy_passes=:to_mhlo_shardings pad(x)
function dus(x, y)
x[6:size(x, 1)-15, :] = y
nothing
end
y = Reactant.to_rarray(ones(Int, 1024-20, 12); sharding)
@show @code_hlo shardy_passes=:to_mhlo_shardings optimize=false dus(x, y)
@show @code_hlo shardy_passes=:to_mhlo_shardings dus(x, y)
@show @code_xla shardy_passes=:to_mhlo_shardings dus(x, y)
function dus2(x, y)
x[6:size(x, 1)-15, 2:11] = y
nothing
end
y = Reactant.to_rarray(ones(Int, 1024-20, 10); sharding)
@show @code_hlo shardy_passes=:to_mhlo_shardings optimize=false dus2(x, y)
@show @code_hlo shardy_passes=:to_mhlo_shardings dus2(x, y)
@show @code_xla shardy_passes=:to_mhlo_shardings dus2(x, y)
```
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.