maleadt
commented
1 month ago Hmm, this is a little disconcerting: Even a very simple Cartesian kernel inhibits a very significant slowdown.
Reduced from our broadcast implementation:
using CUDA, KernelAbstractions, Chairmarks
function main()
a = CuArray{Float32}(undef, 512, 1000)
bc = Broadcast.broadcasted(identity, 0f0)
bc = Broadcast.instantiate(Broadcast.Broadcasted(bc.f, bc.args, axes(a)))
print("Old: ")
display(@b CUDA.@sync copyto_old!(a, bc))
print("New: ")
display(@b CUDA.@sync copyto_new!(a, bc))
end
@inline function copyto_old!(dest::AbstractArray, bc)
axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
isempty(dest) && return dest
bc = Broadcast.preprocess(dest, bc)
function broadcast_kernel(dest, bc)
i = (blockIdx().x-1) * blockDim().x + threadIdx().x
@inbounds if i <= length(dest)
I = CartesianIndices(dest)[i]
dest[I] = bc[I]
end
return
end
kernel = @cuda launch=false broadcast_kernel(dest, bc)
config = launch_configuration(kernel.fun)
threads = min(length(dest), config.threads)
blocks = cld(length(dest), threads)
kernel(dest, bc; threads, blocks)
return dest
end
@inline function copyto_new!(dest::AbstractArray, bc)
axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
isempty(dest) && return dest
bc = Broadcast.preprocess(dest, bc)
@kernel function broadcast_kernel(dest, bc)
I = @index(Global, Cartesian)
@inbounds dest[I] = bc[I]
end
broadcast_kernel(get_backend(dest))(dest, bc; ndrange=size(dest))
return dest
endWithout CUDA.@sync, i.e., measuring launch overhead:
julia> main()
Old: 3.199 μs (12 allocs: 240 bytes)
New: 5.192 μs (58 allocs: 1.859 KiB)With CUDA.@sync, i.e., measuring execution time:
julia> main()
Old: 7.746 μs (12 allocs: 240 bytes)
New: 10.940 μs (58 allocs: 1.859 KiB)The overhead scales, e.g., using 4k x 4k inputs instead:
julia> main()
Old: 30.230 μs (12 allocs: 240 bytes)
New: 61.250 μs (58 allocs: 1.859 KiB)Generated code looks pretty bad, with both extra exceptions, branches, and argument mangling:
define ptx_kernel void @old({ i64, i32 } %state, { i8 addrspace(1)*, i64, [2 x i64], i64 } %0, { [1 x float], [2 x [1 x i64]] } %1) local_unnamed_addr {
conversion:
%.fca.0.extract = extractvalue { i8 addrspace(1)*, i64, [2 x i64], i64 } %0, 0
%.fca.3.extract = extractvalue { i8 addrspace(1)*, i64, [2 x i64], i64 } %0, 3
%.fca.0.0.extract = extractvalue { [1 x float], [2 x [1 x i64]] } %1, 0, 0
%2 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
%3 = zext i32 %2 to i64
%4 = call i32 @llvm.nvvm.read.ptx.sreg.ntid.x()
%5 = zext i32 %4 to i64
%6 = mul nuw nsw i64 %3, %5
%7 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
%8 = add nuw nsw i32 %7, 1
%9 = zext i32 %8 to i64
%10 = add nuw nsw i64 %6, %9
%.not = icmp sgt i64 %10, %.fca.3.extract
br i1 %.not, label %L165, label %L30
L30: ; preds = %conversion
%.fca.2.0.extract = extractvalue { i8 addrspace(1)*, i64, [2 x i64], i64 } %0, 2, 0
%.not6 = icmp eq i64 %.fca.2.0.extract, 0
br i1 %.not6, label %fail, label %pass
L165: ; preds = %pass, %conversion
ret void
fail: ; preds = %L30
call fastcc void @gpu_report_exception({ i64, i32 } %state, i64 ptrtoint ([10 x i8]* @exception16 to i64))
call fastcc void @gpu_signal_exception({ i64, i32 } %state)
call void @llvm.trap()
call void @llvm.trap()
call void asm sideeffect "exit;", ""()
unreachable
pass: ; preds = %L30
%11 = add nsw i64 %10, -1
%12 = bitcast i8 addrspace(1)* %.fca.0.extract to float addrspace(1)*
%13 = getelementptr inbounds float, float addrspace(1)* %12, i64 %11
store float %.fca.0.0.extract, float addrspace(1)* %13, align 4
br label %L165
}define ptx_kernel void @new({ i64, i32 } %state, { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, { i8 addrspace(1)*, i64, [2 x i64], i64 } %1, { [1 x float], [2 x [1 x i64]] } %2) local_unnamed_addr {
conversion:
%.fca.0.0.0.0.extract = extractvalue { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, 0, 0, 0, 0
%.fca.0.0.1.0.extract = extractvalue { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, 0, 0, 1, 0
%.fca.1.0.0.0.0.extract = extractvalue { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, 1, 0, 0, 0, 0
%.fca.1.1.0.0.0.extract = extractvalue { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, 1, 1, 0, 0, 0
%.fca.1.1.0.1.0.extract = extractvalue { [1 x [2 x [1 x i64]]], [2 x [1 x [2 x [1 x i64]]]] } %0, 1, 1, 0, 1, 0
%.fca.0.extract = extractvalue { i8 addrspace(1)*, i64, [2 x i64], i64 } %1, 0
%.fca.2.0.extract = extractvalue { i8 addrspace(1)*, i64, [2 x i64], i64 } %1, 2, 0
%.fca.0.0.extract = extractvalue { [1 x float], [2 x [1 x i64]] } %2, 0, 0
%.not = icmp eq i64 %.fca.1.0.0.0.0.extract, 0
br i1 %.not, label %fail, label %pass
L527: ; preds = %pass6, %pass2
ret void
fail: ; preds = %conversion
call fastcc void @gpu_report_exception({ i64, i32 } %state, i64 ptrtoint ([10 x i8]* @exception19 to i64))
call fastcc void @gpu_signal_exception({ i64, i32 } %state)
call void @llvm.trap()
call void @llvm.trap()
call void asm sideeffect "exit;", ""()
unreachable
pass: ; preds = %conversion
%.not15 = icmp eq i64 %.fca.1.1.0.0.0.extract, 0
br i1 %.not15, label %fail1, label %pass2
fail1: ; preds = %pass
call fastcc void @gpu_report_exception({ i64, i32 } %state, i64 ptrtoint ([10 x i8]* @exception19 to i64))
call fastcc void @gpu_signal_exception({ i64, i32 } %state)
call void @llvm.trap()
call void @llvm.trap()
call void asm sideeffect "exit;", ""()
unreachable
pass2: ; preds = %pass
%3 = call i32 @llvm.nvvm.read.ptx.sreg.tid.x()
%4 = zext i32 %3 to i64
%5 = call i32 @llvm.nvvm.read.ptx.sreg.ctaid.x()
%6 = zext i32 %5 to i64
%7 = sdiv i64 %6, %.fca.1.0.0.0.0.extract
%.neg16 = mul i64 %7, %.fca.1.0.0.0.0.extract
%8 = sdiv i64 %4, %.fca.1.1.0.0.0.extract
%9 = add nsw i64 %8, 1
%reass.add17 = add i64 %.neg16, %8
%reass.add = sub i64 %6, %reass.add17
%reass.mul = mul i64 %reass.add, %.fca.1.1.0.0.0.extract
%10 = add nuw nsw i64 %4, 1
%11 = add i64 %10, %reass.mul
%12 = mul i64 %7, %.fca.1.1.0.1.0.extract
%13 = add i64 %9, %12
%14 = icmp sgt i64 %11, 0
%15 = icmp sle i64 %11, %.fca.0.0.0.0.extract
%16 = and i1 %14, %15
%17 = icmp sgt i64 %13, 0
%18 = icmp sle i64 %13, %.fca.0.0.1.0.extract
%19 = and i1 %17, %18
%20 = and i1 %19, %16
br i1 %20, label %pass6, label %L527
pass6: ; preds = %pass2
%21 = add i64 %12, %8
%22 = mul i64 %21, %.fca.2.0.extract
%23 = add i64 %22, %4
%24 = add i64 %23, %reass.mul
%25 = bitcast i8 addrspace(1)* %.fca.0.extract to float addrspace(1)*
%26 = getelementptr inbounds float, float addrspace(1)* %25, i64 %24
store float %.fca.0.0.extract, float addrspace(1)* %26, align 4
br label %L527
}This results in a much higher register usage, from 28 to 64 (at the PTX level).
cc @vchuravy
The switch to KA.jl significantly slowed down several operations.
CUDA.jl:
permudetims,broadcast, and many othershttps://speed.juliagpu.org/changes/?tre=10&rev=6221589f5befec8f6f157a5a5271667dba09d0b6&exe=11&env=1
Metal.jl:
permudetims