tl;dr We are doing a separate kernel launch for every element of the fusion tree, which is SUPER wasteful especially for otherwise small tensors.
Here's a sample script:
using TensorKit, AMDGPU, Adapt, LinearAlgebra, Printf
function warmup(seconds = 3.0)
X = AMDGPU.rand(Float32, 2048, 2048)
t0 = time()
while time() - t0 < seconds
X * X
end
AMDGPU.synchronize()
return nothing
end
function bench(f, n; reps = 7)
f(); AMDGPU.synchronize()
best = Inf
for _ in 1:reps
AMDGPU.synchronize()
e = @elapsed begin
for _ in 1:n; f(); end
AMDGPU.synchronize()
end
best = min(best, e / n)
end
return best * 1e3
end
nfusiontrees(t) = length(collect(TensorKit.fusiontrees(t)))
nstored(t) = sum(prod(size(b)) for (_, b) in blocks(t))
function report(name, t_cpu, perm)
t_gpu = adapt(ROCArray, t_cpu)
nb, nf, ne = length(collect(blocks(t_cpu))), nfusiontrees(t_cpu), nstored(t_cpu)
c = bench(() -> permute(t_cpu, perm; copy = true), 200)
g = bench(() -> permute(t_gpu, perm; copy = true), 200)
# device allocations issued by ONE permute
AMDGPU.synchronize()
s0 = copy(AMDGPU.alloc_stats)
permute(t_gpu, perm; copy = true)
AMDGPU.synchronize()
d = AMDGPU.alloc_stats - s0
@printf("%-18s trees=%-5d elems=%-6d %-10s CPU %7.3f ms GPU %8.3f ms GPU/CPU %6.2fx allocs/permute=%d\n",
name, nf, ne, Base.format_bytes(ne * 8), c, g, g / c, d.alloc_count)
return nothing
end
println("=== permute cost vs. fusion-tree count (AMDGPU / MI210) ===")
println("(best-of-7 batches of 200; GPU warmed up first)\n")
warmup()
Vp = Vect[SU2Irrep](1//2 => 1)
Vv = Vect[SU2Irrep](0 => 1, 1 => 2)
A = randn(Float64, Vp ⊗ Vp', Vv ⊗ Vv ⊗ Vv' ⊗ Vv')
perm6 = ((2, 3, 5, 6), (1, 4))
report("SU2 iPEPO tensor", A, perm6)
Wp, Wv = ℂ^2, ℂ^7
B = randn(ComplexF64, Wp ⊗ Wp', Wv ⊗ Wv ⊗ Wv' ⊗ Wv')
report("trivial symmetry", B, perm6)
println()
for nmax in 1:3
V = Vect[SU2Irrep](j => 2 for j in 0:nmax)
T = randn(Float64, Vp ⊗ Vp', V ⊗ V ⊗ V' ⊗ V')
report("SU2 jmax=$nmax", T, perm6)
end
println()
Up = Vect[U1Irrep](1//2 => 1, -1//2 => 1)
Uv = Vect[U1Irrep](-1 => 2, 0 => 3, 1 => 2)
C = randn(Float64, Up ⊗ Up', Uv ⊗ Uv ⊗ Uv' ⊗ Uv')
report("U1", C, perm6)Here are the results I see on an MI210:
SU2 iPEPO tensor trees=42 elems=329 2.570 KiB CPU 0.123 ms GPU 2.144 ms GPU/CPU 17.43x allocs/permute=23
trivial symmetry trees=1 elems=9604 75.031 KiB CPU 0.045 ms GPU 0.022 ms GPU/CPU 0.50x allocs/permute=1
SU2 jmax=1 trees=42 elems=672 5.250 KiB CPU 0.095 ms GPU 2.107 ms GPU/CPU 22.11x allocs/permute=23
SU2 jmax=2 trees=323 elems=5168 40.375 KiB CPU 0.726 ms GPU 15.136 ms GPU/CPU 20.84x allocs/permute=121
SU2 jmax=3 trees=1364 elems=21824 170.500 KiB CPU 2.910 ms GPU 57.840 ms GPU/CPU 19.88x allocs/permute=391
U1 trees=70 elems=2226 17.391 KiB CPU 0.053 ms GPU 1.168 ms GPU/CPU 21.95x allocs/permute=1
The GPU is way way slower due to all the kernel launches. The allocs/permute number tells you how many allocations occurred per permute.
I think a good first step would be to try to fuse these into one launch if we can,
tl;dr We are doing a separate kernel launch for every element of the fusion tree, which is SUPER wasteful especially for otherwise small tensors.
Here's a sample script:
Here are the results I see on an MI210:
The GPU is way way slower due to all the kernel launches. The
allocs/permutenumber tells you how many allocations occurred per permute.I think a good first step would be to try to fuse these into one launch if we can,