JuliaGPU / JuliaGPU/CUDA.jl

Improve vectorization of (shared) memory accesses

Open
#68 12 comments 0 reactions 0 assignees View on GitHub
cuda kernels performance
Dominant language
Julia
Stars
1.4k
Forks
281
Avg merge
1d 7h
Merged PRs (30d)
30

Description

Given the following MWE:

```c
#define BLOCK_SIZE 5

__global__ void kernel() {
__shared__ float dia[BLOCK_SIZE][BLOCK_SIZE];
__shared__ float peri_col[BLOCK_SIZE][BLOCK_SIZE];

int idx = threadIdx.x;
for (int i = 0; i < BLOCK_SIZE; i++) {
for (int j = 0; j < i; j++)
peri_col[idx][i] -= dia[j][i];
peri_col[idx][i] /= dia[i][i];
}
}
```

`nvcc` generates PTX with loads clustered together:

```ptx
.version 6.2
.target sm_35
.address_size 64

// .globl kernel
// kernel_dia has been demoted
// kernel_peri_col has been demoted

.visible .entry kernel(

)
{
.reg .f32 %f<36>;
.reg .b32 %r<4>;
// demoted variable
.shared .align 4 .b8 kernel_dia[100];
// demoted variable
.shared .align 4 .b8 kernel_peri_col[100];

mov.u32 %r1, %tid.x;
mov.u32 %r2, kernel_peri_col;
mad.lo.s32 %r3, %r1, 20, %r2;
ld.shared.f32 %f1, [%r3];
ld.shared.f32 %f2, [kernel_dia];
div.rn.f32 %f3, %f1, %f2;
ld.shared.f32 %f4, [%r3+4];
ld.shared.f32 %f5, [kernel_dia+4];
ld.shared.f32 %f6, [kernel_dia+24];
ld.shared.f32 %f7, [%r3+8];
ld.shared.f32 %f8, [kernel_dia+8];
ld.shared.f32 %f9, [kernel_dia+28];
ld.shared.f32 %f10, [kernel_dia+48];
ld.shared.f32 %f11, [%r3+12];
ld.shared.f32 %f12, [kernel_dia+12];
ld.shared.f32 %f13, [kernel_dia+32];
ld.shared.f32 %f14, [kernel_dia+52];
ld.shared.f32 %f15, [kernel_dia+72];
ld.shared.f32 %f16, [%r3+16];
ld.shared.f32 %f17, [kernel_dia+16];
ld.shared.f32 %f18, [kernel_dia+36];
ld.shared.f32 %f19, [kernel_dia+56];
ld.shared.f32 %f20, [kernel_dia+76];
ld.shared.f32 %f21, [kernel_dia+96];
st.shared.f32 [%r3], %f3;
sub.f32 %f22, %f4, %f5;
div.rn.f32 %f23, %f22, %f6;
st.shared.f32 [%r3+4], %f23;
sub.f32 %f24, %f7, %f8;
sub.f32 %f25, %f24, %f9;
div.rn.f32 %f26, %f25, %f10;
st.shared.f32 [%r3+8], %f26;
sub.f32 %f27, %f11, %f12;
sub.f32 %f28, %f27, %f13;
sub.f32 %f29, %f28, %f14;
div.rn.f32 %f30, %f29, %f15;
st.shared.f32 [%r3+12], %f30;
sub.f32 %f31, %f16, %f17;
sub.f32 %f32, %f31, %f18;
sub.f32 %f33, %f32, %f19;
sub.f32 %f34, %f33, %f20;
div.rn.f32 %f35, %f34, %f21;
st.shared.f32 [%r3+16], %f35;
ret;
}
```

Note the pessimistic alignment of 4 for the shared memory arrays. I guess this follows the C spec, but it rules out vectorization. However, when ptxas assembles this code, it knows the memory layout and "physical" alignment of the shared memory arrays, and as a result nicely vectorizes this code:

```sass
MOV R1, c[0x0][0x44];
S2R R0, SR_TID.X;
MOV32I R3, 0x64;
IMAD R0, R0, 0x14, R3;
LDS R10, [RZ];
LDS R5, [R0];
CAL 0x170;

LDS.128 R4, [RZ];
LDS.128 R8, [0x10];
LDS R20, [R0+0x4];
LDS.128 R12, [0x30];
LDS R18, [0x60];
LDS.64 R16, [0x20];
LDS.64 R2, [0x48];

LDS R19, [R0+0x8];
FADD R5, R20, -R5;
LDS R9, [R0+0xc];
LDS R4, [R0+0x10];
STS [R0], R25;
CAL 0x170;
STS [R0+0x4], R25;
```

We tackle this differently, using LLVM's Load Store Vectorizer by specifying a much more optimistic alignment that enables vectorization (https://github.com/JuliaGPU/CUDAnative.jl/pull/204):

```julia
const BLOCK_SIZE = 5

function kernel()
dia = @cuStaticSharedMem(Float32, (BLOCK_SIZE,BLOCK_SIZE))
peri_col = @cuStaticSharedMem(Float32, (BLOCK_SIZE,BLOCK_SIZE))

index = threadIdx().x
for i = 1:BLOCK_SIZE
for j = 1:i-1
peri_col[i, index] -= dia[i, j]
end
peri_col[i, index] /= dia[i, i]
end
end
```

```ptx
.version 6.0
.target sm_35
.address_size 64

// .globl ptxcall_kernel_2
// shmem1 has been demoted
// shmem2 has been demoted

.visible .entry ptxcall_kernel_2(
.param .align 8 .b8 ptxcall_kernel_2_param_0[24],
.param .u64 ptxcall_kernel_2_param_1
)
{
.reg .f32 %f<36>;
.reg .b32 %r<3>;
.reg .b64 %rd<4>;
// demoted variable
.shared .align 16 .b8 shmem1[100];
// demoted variable
.shared .align 16 .b8 shmem2[100];
mov.u32 %r1, %tid.x;
mul.lo.s32 %r2, %r1, 5;
mul.wide.u32 %rd1, %r2, 4;
mov.u64 %rd2, shmem2;
add.s64 %rd3, %rd2, %rd1;
ld.shared.f32 %f1, [%rd3];
ld.shared.v4.f32 {%f2, %f3, %f4, %f5}, [shmem1];
div.rn.f32 %f6, %f1, %f2;
st.shared.f32 [%rd3], %f6;
ld.shared.f32 %f7, [%rd3+4];
sub.f32 %f8, %f7, %f3;
ld.shared.f32 %f9, [shmem1+24];
div.rn.f32 %f10, %f8, %f9;
st.shared.f32 [%rd3+4], %f10;
ld.shared.f32 %f11, [%rd3+8];
sub.f32 %f12, %f11, %f4;
ld.shared.f32 %f13, [shmem1+28];
sub.f32 %f14, %f12, %f13;
ld.shared.f32 %f15, [shmem1+48];
div.rn.f32 %f16, %f14, %f15;
st.shared.f32 [%rd3+8], %f16;
ld.shared.f32 %f17, [%rd3+12];
sub.f32 %f18, %f17, %f5;
ld.shared.f32 %f19, [shmem1+32];
sub.f32 %f20, %f18, %f19;
ld.shared.f32 %f21, [shmem1+52];
sub.f32 %f22, %f20, %f21;
ld.shared.v2.f32 {%f23, %f24}, [shmem1+72];
div.rn.f32 %f25, %f22, %f23;
st.shared.f32 [%rd3+12], %f25;
ld.shared.f32 %f26, [%rd3+16];
ld.shared.f32 %f27, [shmem1+16];
sub.f32 %f28, %f26, %f27;
ld.shared.f32 %f29, [shmem1+36];
sub.f32 %f30, %f28, %f29;
ld.shared.f32 %f31, [shmem1+56];
sub.f32 %f32, %f30, %f31;
sub.f32 %f33, %f32, %f24;
ld.shared.f32 %f34, [shmem1+96];
div.rn.f32 %f35, %f33, %f34;
st.shared.f32 [%rd3+16], %f35;
ret;
}
```

```sass

MOV R1, c[0x0][0x44];
S2R R0, SR_TID.X;
LDS.128 R4, [RZ];
IMUL32I R0, R0, 0x5;
ISCADD R0, R0, 0x70, 0x2;
LDS R2, [R0];
CAL 0x178;

LDS R2, [R0+0x4];
STS [R0], R13;
LDS R4, [0x18];
FADD R2, -R5, R2;
CAL 0x178;
LDS R2, [R0+0x8];
STS [R0+0x4], R13;

LDS R3, [0x1c];
FADD R2, -R6, R2;
LDS R4, [0x30];
FADD R2, R2, -R3;
CAL 0x178;
LDS R2, [R0+0xc];
STS [R0+0x8], R13;

LDS R3, [0x20];
FADD R2, -R7, R2;
LDS R6, [0x34];
LDS.64 R4, [0x48];
FADD R2, R2, -R3;
FADD R2, R2, -R6;
CAL 0x178;

STS [R0+0xc], R13;
LDS R2, [R0+0x10];
LDS R3, [0x10];
LDS R4, [0x24];
LDS R6, [0x38];
FADD R2, R2, -R3;
FADD R2, R2, -R4;

LDS R4, [0x60];
FADD R2, R2, -R6;
FADD R2, -R5, R2;
CAL 0x178;
STS [R0+0x10], R13;
EXIT;
FCHK.DIVIDE P0, R2, R4;
```

While our PTX contains vectorized loads now (16, vs 20 with nvcc), the resulting SASS still performs more loads (remaining at 16, as we don't cluster our loads for ptxas to vectorize, vs only 12). This is due to ptxas vectorizing load chains with gaps, eg. `LDS.128 R8, [0x10]` corresponding with loading elements 4 to 7 from `kernel_dia`, while `kernel_dia+20` isn't actually used by the code (and `R9` isn't either in the SASS code). LLVM refuses to vectorize these accesses, emitting 3 loads instead (it should have still been able to vectorize accesses to elements 6 and 7, looking into that right now).

So this probably needs a fix to the Load Store Vectorizer to vectorize chains with gaps.
Alternatively, we could cluster loads like nvcc, but the existing load/store clustering DAG mutation isn't quite as aggressive and as such only enables `ptxas` to merge a couple of loads (maybe there's some missing bits since NVPTX hasn't been using the machine scheduler before):

```patch
@@ NVPTXTargetMachine.cpp

class NVPTXPassConfig : public TargetPassConfig {
+ ScheduleDAGInstrs *
+ createMachineScheduler(MachineSchedContext *C) const override {
+ ScheduleDAGMILive *DAG = createGenericSchedLive(C);
+ DAG->addMutation(createLoadClusterDAGMutation(DAG->TII, DAG->TRI));
+ DAG->addMutation(createStoreClusterDAGMutation(DAG->TII, DAG->TRI));
+ return DAG;
+ }
};

@@ NVPTXSubtarget.h

class NVPTXSubtarget : public NVPTXGenSubtargetInfo {
+ bool enableMachineScheduler() const override { return true; }
};
```

cc @jlebar
reduced from rodinia/lud
ref https://lists.llvm.org/pipermail/llvm-dev/2018-June/124209.htm

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.