JuliaGPU / JuliaGPU/KernelAbstractions.jl

simple batched dot kernel is ~1.7x slower with Const on Titan RTX

Open
#479 18 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
523
Forks
88
Avg merge
1d 11h
Merged PRs (30d)
25

Description

could someone please help me understand why this should the the case? PTX code is similar and the threads/blocks are identical.

using KernelAbstractions, CUDA, BenchmarkTools

L, N = 256, 32768
x = CuArray(rand(Float32, L,N));
y = CuArray(rand(Float32, L,N));
o = CuArray(rand(Float32, N));
    
@kernel function kernel_ka(o, @Const(x), @Const(y))
    k = @index(Global)

    @inbounds begin
        tmp = 0.0f0
        for i=1:size(x,1)
            tmp += x[i,k] * y[i,k]
        end
        o[k] = tmp
    end
end

function kernel_cuda(o, x, y)
    k = threadIdx().x + (blockIdx().x - 1) * blockDim().x

    @inbounds if k<=size(x,2)
        tmp = 0f0
        for i=1:size(x,1)
            tmp += x[i,k] * y[i,k]
        end
        o[k] = tmp
    end
    return nothing
end                    

batched_dot_ka! = kernel_ka(CUDABackend(), 32)
@benchmark CUDA.@sync batched_dot_ka!(o, x, y; ndrange=length(o))
    
batched_dot_cuda! = @cuda name="batched_dot!" launch=false kernel_cuda(o, x, y)
@benchmark CUDA.@sync batched_dot_cuda!(o, x, y; threads=32, blocks=1024)

the above yields the times below for KA and CUDA, respetively, so KA is ~1.7x slower:

BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  382.110 μs …  3.590 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     491.497 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   474.179 μs ± 62.871 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

                           ▆█                                   
  ▁▃▆▆▅▃▄█▆▄▃▃▂▂▂▂▂▂▂▂▄▅▃▂▃████▅▃▃▃▃▂▂▂▂▂▂▂▂▂▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁ ▂
  382 μs          Histogram: frequency by time          641 μs <

 Memory estimate: 1.72 KiB, allocs estimate: 62.

BenchmarkTools.Trial: 10000 samples with 1 evaluation.
 Range (min … max):  287.471 μs …  3.851 ms  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     294.507 μs              ┊ GC (median):    0.00%
 Time  (mean ± σ):   301.527 μs ± 47.135 μs  ┊ GC (mean ± σ):  0.00% ± 0.00%

  ▄██▆▄▂▂▂                                                  ▂▁ ▂
  ███████████▇▇██▆▆█▆▃▅▁▁▄▃▃▄▃▁▄▁▃▁▁▁▁▁▃▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃▁▁▄███ █
  287 μs        Histogram: log(frequency) by time       468 μs <

 Memory estimate: 336 bytes, allocs estimate: 16.

and here is the PTX code:

KA 
julia> @device_code_ptx batched_dot_ka!(o, x, y; ndrange=length(o))
// PTX CompilerJob of MethodInstance for gpu_kernel_ka(::KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.DynamicCheck, Nothing, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, KernelAbstractions.NDIteration.NDRange{1, KernelAbstractions.NDIteration.DynamicSize, KernelAbstractions.NDIteration.StaticSize{(32,)}, CartesianIndices{1, Tuple{Base.OneTo{Int64}}}, Nothing}}, ::CuDeviceVector{Float32, 1}, ::CuDeviceMatrix{Float32, 1}, ::CuDeviceMatrix{Float32, 1}) for sm_75, maxthreads=32

//
// Generated by LLVM NVPTX Back-End
//

.version 8.3
.target sm_75
.address_size 64

	// .globl	_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE // -- Begin function _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE
                                        // @_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE
.visible .entry _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE(
	.param .align 8 .b8 _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_0[16],
	.param .align 8 .b8 _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_1[16],
	.param .align 8 .b8 _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_2[32],
	.param .align 8 .b8 _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_3[40],
	.param .align 8 .b8 _Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_4[40]
)
.maxntid 32, 1, 1
{
	.reg .pred 	%p<4>;
	.reg .b32 	%r<4>;
	.reg .f32 	%f<10>;
	.reg .b64 	%rd<32>;

// %bb.0:                               // %conversion
	ld.param.u64 	%rd17, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_1];
	mov.u32 	%r2, %ctaid.x;
	mov.u32 	%r1, %tid.x;
	add.s32 	%r3, %r1, 1;
	cvt.u64.u32 	%rd18, %r3;
	mul.wide.u32 	%rd6, %r2, 32;
	add.s64 	%rd7, %rd6, %rd18;
	setp.gt.s64 	%p1, %rd7, %rd17;
	@%p1 bra 	$L__BB0_5;
// %bb.1:                               // %L95
	ld.param.u64 	%rd1, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_2];
	ld.param.u64 	%rd3, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_3+16];
	setp.lt.s64 	%p2, %rd3, 1;
	mov.f32 	%f9, 0f00000000;
	@%p2 bra 	$L__BB0_4;
// %bb.2:                               // %L220.preheader
	ld.param.u64 	%rd2, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_3];
	ld.param.u64 	%rd4, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_4];
	ld.param.u64 	%rd5, [_Z13gpu_kernel_ka16CompilerMetadataI11DynamicSize12DynamicCheckv16CartesianIndicesILi1E5TupleI5OneToI5Int64EEE7NDRangeILi1ES0_10StaticSizeI5_32__ES2_ILi1ES3_IS4_IS5_EEEvEE13CuDeviceArrayI7Float32Li1ELi1EES8_IS9_Li2ELi1EES8_IS9_Li2ELi1EE_param_4+16];
	max.s64 	%rd31, %rd3, 0;
	add.s64 	%rd19, %rd7, 4611686018427387903;
	mul.lo.s64 	%rd20, %rd3, %rd19;
	max.s64 	%rd21, %rd5, 0;
	mul.lo.s64 	%rd22, %rd21, %rd19;
	shl.b64 	%rd23, %rd22, 2;
	add.s64 	%rd30, %rd4, %rd23;
	shl.b64 	%rd24, %rd20, 2;
	add.s64 	%rd29, %rd2, %rd24;
	mov.f32 	%f9, 0f00000000;
$L__BB0_3:                              // %L220
                                        // =>This Inner Loop Header: Depth=1
	ld.global.nc.f32 	%f6, [%rd29];
	ld.global.nc.f32 	%f7, [%rd30];
	fma.rn.f32 	%f9, %f6, %f7, %f9;
	add.s64 	%rd31, %rd31, -1;
	add.s64 	%rd30, %rd30, 4;
	add.s64 	%rd29, %rd29, 4;
	setp.ne.s64 	%p3, %rd31, 0;
	@%p3 bra 	$L__BB0_3;
$L__BB0_4:                              // %L381
	cvt.u64.u32 	%rd25, %r1;
	add.s64 	%rd26, %rd6, %rd25;
	shl.b64 	%rd27, %rd26, 2;
	add.s64 	%rd28, %rd27, %rd1;
	st.global.f32 	[%rd28], %f9;
$L__BB0_5:                              // %L388
	ret;
                                        // -- End function
}
CUDA 
julia> @device_code_ptx @cuda name="batched_dot!" launch=false kernel_cuda(o, x, y)
// PTX CompilerJob of MethodInstance for kernel_cuda(::CuDeviceVector{Float32, 1}, ::CuDeviceMatrix{Float32, 1}, ::CuDeviceMatrix{Float32, 1}) for sm_75

//
// Generated by LLVM NVPTX Back-End
//

.version 8.3
.target sm_75
.address_size 64

	// .globl	batched_dot_            // -- Begin function batched_dot_
                                        // @batched_dot_
.visible .entry batched_dot_(
	.param .align 8 .b8 batched_dot__param_0[16],
	.param .align 8 .b8 batched_dot__param_1[32],
	.param .align 8 .b8 batched_dot__param_2[40],
	.param .align 8 .b8 batched_dot__param_3[40]
)
{
	.reg .pred 	%p<6>;
	.reg .b32 	%r<5>;
	.reg .f32 	%f<22>;
	.reg .b64 	%rd<46>;

// %bb.0:                               // %conversion
	ld.param.u64 	%rd25, [batched_dot__param_2+24];
	mov.u32 	%r1, %tid.x;
	add.s32 	%r2, %r1, 1;
	mov.u32 	%r3, %ctaid.x;
	mov.u32 	%r4, %ntid.x;
	mul.wide.u32 	%rd6, %r3, %r4;
	cvt.u64.u32 	%rd26, %r2;
	add.s64 	%rd7, %rd6, %rd26;
	setp.gt.s64 	%p1, %rd7, %rd25;
	@%p1 bra 	$L__BB0_8;
// %bb.1:                               // %L31
	ld.param.u64 	%rd1, [batched_dot__param_1];
	ld.param.u64 	%rd3, [batched_dot__param_2+16];
	setp.lt.s64 	%p2, %rd3, 1;
	mov.f32 	%f19, 0f00000000;
	@%p2 bra 	$L__BB0_7;
// %bb.2:                               // %L49.preheader
	ld.param.u64 	%rd2, [batched_dot__param_2];
	ld.param.u64 	%rd4, [batched_dot__param_3];
	ld.param.u64 	%rd5, [batched_dot__param_3+16];
	max.s64 	%rd8, %rd3, 0;
	add.s64 	%rd9, %rd7, -1;
	mul.lo.s64 	%rd10, %rd3, %rd9;
	max.s64 	%rd28, %rd5, 0;
	mul.lo.s64 	%rd11, %rd28, %rd9;
	and.b64  	%rd12, %rd8, 1;
	setp.eq.s64 	%p3, %rd8, 1;
	mov.f32 	%f19, 0f00000000;
	mov.u64 	%rd45, 0;
	@%p3 bra 	$L__BB0_5;
// %bb.3:                               // %L49.preheader.new
	and.b64  	%rd13, %rd8, 9223372036854775806;
	shl.b64 	%rd30, %rd11, 2;
	add.s64 	%rd31, %rd30, %rd4;
	add.s64 	%rd43, %rd31, 4;
	shl.b64 	%rd32, %rd10, 2;
	add.s64 	%rd33, %rd32, %rd2;
	add.s64 	%rd42, %rd33, 4;
	mov.u64 	%rd45, 0;
	mov.f32 	%f19, 0f00000000;
$L__BB0_4:                              // %L49
                                        // =>This Inner Loop Header: Depth=1
	ld.global.f32 	%f11, [%rd42+-4];
	ld.global.f32 	%f12, [%rd43+-4];
	fma.rn.f32 	%f13, %f11, %f12, %f19;
	ld.global.f32 	%f14, [%rd42];
	ld.global.f32 	%f15, [%rd43];
	fma.rn.f32 	%f19, %f14, %f15, %f13;
	add.s64 	%rd45, %rd45, 2;
	add.s64 	%rd43, %rd43, 8;
	add.s64 	%rd42, %rd42, 8;
	setp.ne.s64 	%p4, %rd13, %rd45;
	@%p4 bra 	$L__BB0_4;
$L__BB0_5:                              // %L148.loopexit.unr-lcssa
	setp.eq.s64 	%p5, %rd12, 0;
	@%p5 bra 	$L__BB0_7;
// %bb.6:                               // %L49.epil.preheader
	add.s64 	%rd34, %rd45, %rd10;
	shl.b64 	%rd35, %rd34, 2;
	add.s64 	%rd23, %rd2, %rd35;
	add.s64 	%rd36, %rd45, %rd11;
	shl.b64 	%rd37, %rd36, 2;
	add.s64 	%rd24, %rd4, %rd37;
	ld.global.f32 	%f16, [%rd23];
	ld.global.f32 	%f17, [%rd24];
	fma.rn.f32 	%f19, %f16, %f17, %f19;
$L__BB0_7:                              // %L148
	cvt.u64.u32 	%rd38, %r1;
	add.s64 	%rd39, %rd6, %rd38;
	shl.b64 	%rd40, %rd39, 2;
	add.s64 	%rd41, %rd40, %rd1;
	st.global.f32 	[%rd41], %f19;
$L__BB0_8:                              // %L156
	ret;
                                        // -- End function
}

vendor-agnostic code is really appealing but i'm not sure i'm willing to pay this much of a performance penalty for it. thanks!

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the supplied Julia benchmark using kernel_ka and kernel_cuda, then inspect their generated output with @device_code_ptx. Compare the launch configuration and PTX paths to identify what causes the timing difference. Done means the performance discrepancy has a confirmed explanation and a concrete scope for any follow-up fix.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.