[Multi-Core & GPU Programming] HW6 - Triton ResNet

"Triton ResNet"

[Objective]

PyTorch, CUDA Library(cuDNN, cuBLAS 등)를 사용하지 않고 Triton Kernel을 이용하여 ResNet18의 주요 연산(Conv2d, ReLU 등)을 직접 구현하기

[Implementation]

① "make info"를 실행하여 ResNet18 Model의 상세 정보 확인(구조, 입출력, Layer 등)

② "make torch"를 실행하여 PyTorch 기준 결과를 확인 → 정답 기준

③ "make triton"을 실행하여 구현한 전체 Triton Kernel의 결과를 확인

· "TritonReLU_ResNet18.py" : "mgp_nn.TritonReLU"가 어떻게 구현되는지 보여주는 Hint File

· "make test-triton-relu" : 구현한 Triton ReLU Kernel만 테스트

▣ ReLU

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26

import triton import triton.language as tl from mgp import empty     @triton.jit def _relu_forward_kernel(     in_ptr, out_ptr, num_elements,     BLOCK_SIZE: tl.constexpr ):     pid = tl.program_id(0)     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)     a = tl.load(in_ptr + offsets, mask=offsets < num_elements, other=0.)     a = tl.where(a > 0, a, 0)     tl.store(out_ptr + offsets, a, mask=offsets < num_elements)     def triton_relu(x, inplace=False):     if inplace:         out = x     else:         out = empty(x.shape, device=x.device, dtype=x.dtype)     n_elements = x.numel()     def grid(meta): return (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )     _relu_forward_kernel[grid](x, out, n_elements, BLOCK_SIZE=1024)     return out Colored by Color Scripter

cs

· 입력 값이 양수면 그대로 두고, 음수면 0으로 변환

8~9 : ReLU 연산을 수행하는 GPU Kernel 인자

    · in_ptr : 입력 Tensor 주소

    · out_ptr : 출력 Tensor 주소

    · num_elements : 전체 원소 개수

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

11 : 현재 실행 중인 Triton Program의 ID

12 : 이번 Program이 처리할 원소들의 Index 생성

13 : 입력 Tensor에서 값을 Load(mask를 이용해 마지막 Block에서 범위를 넘어가는 Index 무시)

14 : ReLU(a) = max(a, 0) → ReLU 연산

15 : ReLU 결과 값을 Output Tensor에 저장(mask를 이용해 유효한 Block의 범위에 저장)

18 : Python에서 "_relu_forward_kernel" 호출("inplace" Default 값은 False)

19~20 : "inplace == True"인 경우 입력 Tensor 자체를 출력으로 사용(입력 Tensor를 Overwrite)

21~22 : "inplace == False"인 경우 새로운 Output Tensor를 생성(shape, device, dtype은 입력과 동일하게 설정)

23 : 입력 Tensor 전체 원소 개수 계산

24 : 실행할 Triton Program 개수 계산 → grid = ceil(n_elements / BLOCK_SIZE)

25 : Triton GPU Kernel 실행(하나의 Triton Program은 최대 1024개의 원소에 대해 동시에 ReLU 연산 수행)

26 : ReLU 결과 Tensor를 반환

** ReLU : 입력값을 나중에 다시 사용할 필요가 없는 경우↑ → 출력 Tensor를 입력 Tensor에 Overwrite

△ Memory 절약 & Memory 접근 감소

▣ BatchNorm2d

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48

import triton import triton.language as tl from mgp import empty, empty_like     @triton.jit def _bn2d_kernel(     x_ptr, y_ptr, total_elements,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr,     weight_ptr, bias_ptr, mean_ptr, var_ptr,     eps: tl.constexpr,     BLOCK_SIZE: tl.constexpr,):       pid = tl.program_id(0)     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)     mask = offsets < total_elements       c = (offsets // (H * W)) % C     x = tl.load(x_ptr + offsets, mask = mask, other = 0)     w = tl.load(weight_ptr + c, mask = mask, other = 0)     b = tl.load(bias_ptr + c, mask = mask, other = 0)     mean = tl.load(mean_ptr + c, mask = mask, other = 0)     var = tl.load (var_ptr + c, mask = mask, other = 1)       y = (x - mean) * tl.rsqrt(var + eps)     y = y * w + b     tl.store(y_ptr + offsets, y, mask = mask)     def triton_bn2d(     input,     weight, bias, running_mean, running_var,     momentum: float = 0.1,     eps: float = 1e-5):       N, C, H, W = input.shape     output = empty_like(input)     total_elements = N * C * H * W     grid = lambda meta:(triton.cdiv(total_elements, meta["BLOCK_SIZE"]),)       _bn2d_kernel[grid](         input, output, total_elements,         C, H, W,         weight, bias, running_mean, running_var,         eps,         BLOCK_SIZE = 1024,)       return output Colored by Color Scripter

cs

· Channel별로 값을 정규화한 후 Scale과 Bias를 적용하여 출력 분포를 안정화해서 Model이 잘 동작하도록 함

8~12 : BatchNorm 계산을 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · total_elements : 전체 출력 원소 개수

    · C / H / W : 입력 Tensor의 Channel / Height / Width

    · weight_ptr / bias_ptr / mean_ptr / var_ptr : BatchNorm Parameter 주소

    · eps : ε(분모가 0이되는 것을 막기 위한 작은 값)

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

14 : 현재 실행 중인 Triton Program의 ID

15 : 현재 Program이 처리할 원소들의 Index 생성

16 : 마지막 Block에서 전체 원소 범위를 넘어가는 Index를 무시

18 : 각 원소가 몇 번째 Channel에 속하는지 계산(같은 Channel 안에 H x W개의 원소 존재)

19 : 입력 Tensor 값을 Load(mask를 이용해서 유효한 범위만 Load)

20 : 해당 Channel의 Weight를 Load(mask를 이용해서 유효한 범위만 Load)

21 : 해당 Channel의 Bias를 Load(mask를 이용해서 유효한 범위만 Load)

22 : 해당 Channel의 Running Mean을 Load(mask를 이용해서 유효한 범위만 Load)

23 : 해당 Channel의 Running Variance를 Load(mask를 이용해서 유효한 범위만 Load)

25 : BatchNorm 정규화

26 : Scale과 Shift 적용

27 : 계산 결과를 Output Tensor에 저장

30 : Python에서 "_bn2d_kernel" 호출

31~32 : BatchNorm2d에 필요한 입력과 Parameter

36 : 입력 Tensor Shape Load

37 : 입력과 같은 shape, device, dtype을 가진 Output Tensor 생성

38 : 전체 원소 개수 계산

39 : 필요한 Triton Program 개수 계산 → grid = ceil(n_elements / BLOCK_SIZE)

41~46 : Triton GPU Kernel 실행(하나의 Triton Program은 최대 1024개의 원소에 대해 동시에 BatchNorm 연산 수행)

47 : BatchNorm 결과 Tensor를 반환

· x : 입력 / μ : Running Mean / σ² : Running Variance / γ : Weight / β : Bias

▣ MaxPool2d

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63

import triton import triton.language as tl from mgp import empty     @triton.jit def _maxpool2d_kernel(     x_ptr, y_ptr, total_elements,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr,     kH: tl.constexpr, kW: tl.constexpr,     sH: tl.constexpr, sW: tl.constexpr,     pH: tl.constexpr, pW: tl.constexpr,     out_H: tl.constexpr, out_W: tl.constexpr,     BLOCK_SIZE: tl.constexpr,):       pid = tl.program_id(0)     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)     mask = offsets < total_elements       ow = offsets % out_W     oh = (offsets // out_W) % out_H     c = (offsets // (out_H * out_W)) % C     n = offsets // (C * out_H * out_W)     max_val = tl.full((BLOCK_SIZE,), -float("inf"), dtype = tl.float32)       for kh in range(kH):         for kw in range(kW):             ih = oh * sH + kh - pH             iw = ow * sW + kw - pW             valid = mask & (ih >= 0) & (ih < H) & (iw >= 0) & (iw < W)             x_offset = ((n * C + c) * H + ih) * W + iw             val = tl.load(x_ptr + x_offset, mask = valid, other = -float("inf"))             max_val = tl.maximum(max_val, val)     tl.store(y_ptr + offsets, max_val, mask = mask)     def triton_maxpool2d(     x,     kernel_size=(2, 2),     stride=(2, 2),     padding=(0, 0)):       N, C, H, W = x.shape     kH, kW = kernel_size     sH, sW = stride     pH, pW = padding     out_H = (H + 2 * pH - kH) // sH + 1     out_W = (W + 2 * pW - kW) // sW + 1       y = empty((N, C, out_H, out_W), device = x.device, dtype = x.dtype)     total_elements = N * C * out_H * out_W     grid = lambda meta:(triton.cdiv(total_elements, meta["BLOCK_SIZE"]),)       _maxpool2d_kernel[grid](         x, y, total_elements,         C, H, W,         kH, kW,         sH, sW,         pH, pW,         out_H, out_W,         BLOCK_SIZE = 256,)       return y Colored by Color Scripter

cs

· 작은 Window 안에서 가장 큰 값 하나를 뽑아 이미지 Feature Map의 크기를 줄이면서 강한 특징을 남김

8~14 : MaxPool2d 계산을 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · total_elements : 전체 출력 원소 개수

    · C / H / W : 입력 Tensor의 Channel / Height / Width

    · kH, kW : Pooling Window 크기

    · sH, sW : Stride

    · pH, pW : Padding

    · out_H, out_W : 출력 Tensor의 Height, Width

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

16 : 현재 실행 중인 Triton Program의 ID

17 : 현재 Program이 처리할 원소들의 Index 생성

18 : 마지막 Block에서 전체 원소 범위를 넘어가는 Index를 무시

20 : 현재 Output 원소의 Width Index

21 : 현재 Output 원소의 Height Index

22 : 현재 Output 원소의 Channel Index

23 : 현재 Output 원소의 Batch Index

24 : 최댓값을 저장할 변수(가장 큰 값을 찾아야 하므로 - ∞부터 시작)

26~27 : Pooling Window 내부 순회

28~29 : 현재 Output 위치(oh, ow)가 참조해야 하는 Input위치(ih, iw)를 계산

30 : Padding 때문에 Input 범위를 벗어나는 위치 제외

31 : Input Tensor [N, C, H, W]에서 Load할 Memory Offset 계산

32 : Input 값 Load(범위 밖인 경우 - ∞로 처리)

33 : 현재까지의 최댓값과 새로운 값을 비교해 더 큰 값 저장

34 : 최종 MaxPool 결과를 Output Tensor에 저장

37 : Python에서 "_maxpool2d_kernel" 호출

43 : 입력 Tensor Shape Load

44~46 : Pooling 설정값 저장

47~48 : 출력 크기 계산

50 : 출력 Tensor 생성

51 : 출력 Tensor의 전체 원소 개수 계산

52 : 필요한 Triton Program 개수 계산 → grid = ceil(n_elements / BLOCK_SIZE)

54~61 : Triton GPU Kernel 실행(하나의 Triton Program은 최대 256개의 원소에 대해 동시에 MaxPool 연산 수행)

63 : MaxPool 결과 Tensor를 반환

▣ AvgPool2d

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62

import triton import triton.language as tl from mgp import empty     @triton.jit def _avgpool2d_kernel(     x_ptr, y_ptr, total_elements,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr,     kH: tl.constexpr, kW: tl.constexpr,     sH: tl.constexpr, sW: tl.constexpr,     out_H: tl.constexpr, out_W: tl.constexpr,     BLOCK_SIZE: tl.constexpr,):       pid = tl.program_id(0)     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)     mask = offsets < total_elements       ow = offsets % out_W     oh = (offsets // out_W) % out_H     c = (offsets // (out_H * out_W)) % C     n = offsets // (C * out_H * out_W)     acc = tl.zeros((BLOCK_SIZE,), dtype = tl.float32)       for kh in range(kH):         for kw in range(kW):             ih = oh * sH + kh             iw = ow * sW + kw             x_offset = ((n * C + c) * H + ih) * W + iw             val = tl.load(x_ptr + x_offset, mask = mask, other = 0)             acc += val     avg = acc / (kH * kW)       tl.store(y_ptr + offsets, avg, mask = mask)     def triton_avgpool2d(     x,     kernel_size=(7, 7),     stride=(1, 1),     padding=(0, 0),):       N, C, H, W = x.shape     kH, kW = kernel_size     sH, sW = stride     pH, pW = padding     out_H = (H + 2 * pH - kH) // sH + 1     out_W = (W + 2 * pW - kW) // sW + 1       y = empty((N, C, out_H, out_W), device = x.device, dtype = x.dtype)     total_elements = N * C * out_H * out_W     grid = lambda meta:(triton.cdiv(total_elements, meta["BLOCK_SIZE"]),)       _avgpool2d_kernel[grid](         x, y, total_elements,         C, H, W,         kH, kW,         sH, sW,         out_H, out_W,         BLOCK_SIZE=256,)       return y Colored by Color Scripter

cs

· Window 내부 값들의 평균 계산

8~13 : AvgPool2d 계산을 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · total_elements : 전체 출력 원소 개수

    · C / H / W : 입력 Tensor의 Channel / Height / Width

    · kH, kW : Pooling Window 크기

    · sH, sW : Stride

    · out_H, out_W : 출력 Tensor의 Height, Width

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

15 : 현재 실행 중인 Triton Program의 ID

16 : 현재 Program이 처리할 원소들의 Index 생성

17 : 마지막 Block에서 전체 원소 범위를 넘어가는 Index를 무시

19 : 현재 Output 원소의 Width Index

20 : 현재 Output 원소의 Height Index

21 : 현재 Output 원소의 Channel Index

22 : 현재 Output 원소의 Batch Index

23 : Pooling Window 내부 값들의 합을 저장할 변수(평균을 구해야 하므로 0부터 시작) ≠ MaxPool2d의 최댓값을 저장할 변수

25~26 : Pooling Window 내부 순회

27~28 : 현재 Output 위치(oh, ow)가 참조해야 하는 Input위치(ih, iw)를 계산

29 : Input Tensor [N, C, H, W]에서 Load할 Memory Offset 계산

30 : Input 값 Load(마지막 Block에서 범위를 넘어가는 Index 무시)

31 : Pooling Window 안의 값을 누적

32 : 누적합을 Window 크기로 나누어 평균 계산

34 : AvgPool 결과를 Output Tensor에 저장

37 : Python에서 "_avgpool2d_kernel" 호출

43 : 입력 Tensor Shape Load

44~46 : Pooling 설정값 저장

47~48 : 출력 크기 계산

50 : 출력 Tensor 생성

51 : 출력 Tensor의 전체 원소 개수 계산

52 : 필요한 Triton Program 개수 계산 → grid = ceil(n_elements / BLOCK_SIZE)

54~60 : Triton GPU Kernel 실행(하나의 Triton Program은 최대 256개의 원소에 대해 동시에 AvgPool 연산 수행)

62 : AvgPool 결과 Tensor를 반환

** 마지막에 수행되는 AvgPool의 Padding = 0

▣ Linear

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44

import triton import triton.language as tl from mgp import empty     @triton.jit def _linear_kernel(     a_ptr, b_ptr, bias_ptr, c_ptr,     M, N, K,     BLOCK_SIZE: tl.constexpr,):       pid = tl.program_id(0)     row = pid // N     col = pid % N     acc = 0.0       for k in range(K):         a_val = tl.load(a_ptr + row * K + k)         b_val = tl.load(b_ptr + k * N + col)         acc += a_val * b_val     bias_val = tl.load(bias_ptr + col)     tl.store(c_ptr + row * N + col, acc + bias_val)     def triton_linear(a, b, bias):     assert a.is_contiguous(), "Matrix A must be contiguous"     assert b.is_contiguous(), "Matrix B must be contiguous"     assert a.shape[1] == b.shape[0], "Incompatible dimensions"     assert b.shape[1] == bias.shape[0], "Incompatible bias dimensions"       M = a.shape[0]     K = a.shape[1]     N = b.shape[1]       c = empty((M, N), device = a.device, dtype = a.dtype)       grid = (M * N,)       _linear_kernel[grid](             a, b, bias,             c, M, N, K,             BLOCK_SIZE = 1,)       return c Colored by Color Scripter

cs

· 행렬곱에 Bias를 더하는 연산으로, 각 Class에 대한 점수를 계산

8~10 : 행렬곱에 Bias를 더하는 연산을 수행하는 GPU Kernel 인자

    · a_ptr : 입력 행렬 A 주소

    · b_ptr : Weight 행렬 B 주소

    · bias_ptr : Bias 주소

    · c_ptr : 출력 행렬 C 주소

    · M : A의 Row 수

    · N : B의 Column 수

    · K : A의 Column 수 = B의 Row 수

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

12 : 현재 실행 중인 Triton Program의 ID

13 : 출력 행렬 C에서 현재 Program이 계산할 Row Index

14 : 출력 행렬 C에서 현재 Program이 계산할 Column Index

15 : 곱셈 결과를 누적할 변수

17 : 행렬 A, B의 공통 차원 K를 순회

18 : 행렬 A의 값 Load

19 : 행렬 B의 값 Load

20 : 행렬곱 누적

21 : 현재 Output Column에 해당하는 Bias Load

22 : 최종 결과를 Output Tensor에 저장

25 : Python에서 "_linear_kernel" 호출

26~27 : 행렬 A, B가 Memory상 연속적으로 저장되어 있는지 확인

28 : 행렬곱이 가능한지 확인 → A : (M, K) / B : (K, N)

29 : Bias 크기가 Output Column 수 N과 같은지 확인

31~33 : 행렬 크기를 Load → A : (M, K) / B : (K, N) / C : (M, N)

35 : 출력 Tensor C 생성

37 : 출력 원소 개수만큼 Triton Program 실행 → 하나의 Program이 C[row, col] 하나를 계산

39~42 : Triton GPU Kernel 실행(하나의 Triton Program은 1개의 원소에 대해 행렬곱에 Bias를 더하는 연산 수행)

44 : 행렬곱에 Bias를 더한 결과 Tensor를 반환

▣ Conv2d - Naive

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68

import triton import triton.language as tl from mgp import zeros     @triton.jit def _conv2d_kernel(     x_ptr, w_ptr, b_ptr, y_ptr, total_elements,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr, K: tl.constexpr,     R: tl.constexpr, S: tl.constexpr, P: tl.constexpr, Q: tl.constexpr,     sH: tl.constexpr, sW: tl.constexpr,     pH: tl.constexpr, pW: tl.constexpr,     dH: tl.constexpr, dW: tl.constexpr,     bias_exist: tl.constexpr,     BLOCK_SIZE: tl.constexpr,):       pid = tl.program_id(0)     offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)     mask = offsets < total_elements       q = offsets % Q     p = (offsets // Q) % P     k = (offsets // (P * Q)) % K     n = offsets // (K * P * Q)     acc = tl.zeros((BLOCK_SIZE,), dtype = tl.float32)       for c in range(C):         for r in range(R):             for s in range(S):                 ih = p * sH + r * dH - pH                 iw = q * sW + s * dW - pW                 valid = mask & (ih >= 0) & (ih < H) & (iw >= 0) & (iw < W)                 x_offset = ((n * C + c) * H + ih) * W + iw                 w_offset = ((k * C + c) * R + r) * S + s                 x_val = tl.load(x_ptr + x_offset, mask = valid, other = 0)                 w_val = tl.load(w_ptr + w_offset, mask = mask, other = 0)                 acc += x_val * w_val       if bias_exist:         bias_val = tl.load(b_ptr + k, mask = mask, other = 0)         acc += bias_val     tl.store(y_ptr + offsets, acc, mask = mask)     def triton_conv2d(input, weight, bias, stride, padding, dilation):     N, C, H, W = input.shape     K, C, R, S = weight.shape     sH, sW = stride     pH, pW = padding     dH, dW = dilation       P = (H + 2 * pH - dH * (R - 1) - 1) // sH + 1     Q = (W + 2 * pW - dW * (S - 1) - 1) // sW + 1     output = zeros((N, K, P, Q), device = input.device, dtype = input.dtype)     total_elements = N * K * P * Q     grid = lambda meta:(triton.cdiv(total_elements, meta["BLOCK_SIZE"]),)     bias_exist = bias is not None       _conv2d_kernel[grid](         input, weight, bias, output, total_elements,         C, H, W, K, R, S, P, Q,         sH, sW,         pH, pW,         dH, dW,         bias_exist = bias_exist,         BLOCK_SIZE = 128,)       return output Colored by Color Scripter

cs

· Convolution 연산 수행 (Output[n, k, p, q] = ∑Input[n, c, ih, iw] x Weight[k, c, r, s])

8~15 : Conv2d 연산을 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · w_ptr : Weight Tensor 주소

    · b_ptr : Bias Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · total_elements : 출력 Tensor 전체 원소 개수

    · C / H / W : 입력 Tensor의 Input Channel 수 / Height / Width → [N, C, H, W]

    · K / R / S : Weight Tensor의 Output Channel 수 / Kernel Height / Kernel Width → [K, C, R, S]

    · K / P / Q : 출력 Tensor의 Output Channel 수 / Height / Width → [N, K, P, Q]

    · sH, sW : Stride

    · pH, pW : Padding

    · dH, dW : Dilation

    · bias_exist : Bias 존재 여부

    · BLOCK_SIZE : 하나의 Triton Program이 처리할 원소 수

17 : 현재 실행 중인 Triton Program의 ID

18 : 현재 Program이 처리할 Output 원소 Index 생성

19 : 마지막 Block에서 범위를 벗어나는 Output Index 무시

21 : Output Width Index

22 : Output Height Index

23 : Output Channel Index

24 : Batch Index

25 : 각 Output 원소의 Convolution 누적합 저장

27~29 : Conv Window 전체 순회 → 한 Output 값을 계산하려면 C x R x S 개의 곱셈-덧셈 필요

30~31 : 현재 Output 위치(p, q)가 참조해야 하는 Input 위치(ih, iw) 계산

32 : Padding 때문에 Input 범위를 벗어나는 위치 제외

33 : Input Tensor [N, C, H, W]에서 Load할 Memory Offset 계산

34 : Weight Tensor [K, C, R, S]에서 Load할 Memory Offset 계산

35 : Input 값 Load(범위 밖인 경우 0으로 처리)

36 : Weight 값 Load

37 : Convolution 곱셈 결과를 누적

39~41 : Bias가 있는 경우만 Output Channel "k"에 해당하는 Bias를 더함(Bias 확인을 안하면 b_ptr = None인 경우 Error 발생)

42 : 최종 Conv2d 결과를 Output Tensor에 저장

45 : Python에서 "_conv2d_kernel" 호출

46 : 입력 Tensor Shape Load

47 : Weight Tensor Shape Load

48~50 : Conv 설정값 저장

52 : Output Height 계산

53 : Output Width 계산

54 : Output Tensor 생성(0으로 초기화)

55 : Output Tensor 전체 원소 개수 계산

56 : 필요한 Triton Program 개수 계산 → grid = ceil(n_elements / BLOCK_SIZE)

57 : Bias가 있는지 확인

59~66 : Triton GPU Kernel 실행(하나의 Triton Program은 최대 128개의 원소에 대해 동시에 Conv 연산 수행)

68 : Conv 결과 Tensor를 반환

> Result

· 하나의 Output 원소를 독립적으로 계산(Ex. Input : [128, 64, 56, 56] / Weight : [64, 64, 3, 3] / Output : [128, 64, 56, 56])

→ Output 원소 하나 계산 : 64 x 3 x 3 = 576번 연산 / Output 원소 개수 : 128 x 64 x 56 x 56

→ 많은 원소가 각각 576번씩 Loop 진행

· 같은 Weight는 여러 Output 위치에서 반복 사용

→ But, 현재 Code는 Output 원소마다 "w_val"에서 같은 Weight를 다시 Load

· 인접한 Output 위치들은 비슷한 Input Window를 공유

→ But, 현재 Code는 각 Output 원소가 "x_val"에서 자기 Input Window를 따로 Load

▣ Conv2d - 3x3 Tiled Kernel

> make info

· ResNet18에서 3x3 Conv가 가장 많이 존재 → 3x3 Conv만 따로 처리하는 Tiled Kernel을 추가하여 빠르게 계산

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64

@triton.jit def _conv2d_3x3_tiled_kernel(     x_ptr, w_ptr, b_ptr, y_ptr, M_total: tl.constexpr,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr,     K: tl.constexpr, P: tl.constexpr, Q: tl.constexpr,     sH: tl.constexpr, sW: tl.constexpr,     pH: tl.constexpr, pW: tl.constexpr,     bias_exist: tl.constexpr,     BLOCK_M: tl.constexpr,     BLOCK_K: tl.constexpr,):       pid_m = tl.program_id(0)     pid_k = tl.program_id(1)     offsets_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)     offsets_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)       q = offsets_m % Q     p = (offsets_m // Q) % P     n = offsets_m // (P * Q)     acc = tl.zeros((BLOCK_M, BLOCK_K), dtype = tl.float32)       for c in range(C):         for r in range(3):             for s in range(3):                 ih = p[:, None] * sH + r - pH                 iw = q[:, None] * sW + s - pW                 x_offset = ((n[:, None] * C + c) * H + ih) * W + iw                 w_offset = ((offsets_k[None, :] * C + c) * 3 + r) * 3 + s                 valid_x = ((offsets_m[:, None] < M_total)                         & (ih >= 0) & (ih < H)                         & (iw >= 0) & (iw < W))                 valid_w = offsets_k[None, :] < K                 x_val = tl.load(x_ptr + x_offset, mask = valid_x, other = 0)                 w_val = tl.load(w_ptr + w_offset, mask = valid_w, other = 0)                 acc += x_val * w_val       if bias_exist:         bias_val = tl.load(b_ptr + offsets_k, mask = offsets_k < K, other=0.0)         acc += bias_val[None, :]       y_offset = ((n[:, None] * K + offsets_k[None, :]) * P + p[:, None]) * Q + q[:, None]     y_mask = (offsets_m[:, None] < M_total) & (offsets_k[None, :] < K)     tl.store(y_ptr + y_offset, acc, mask = y_mask)     def triton_conv2d(input, weight, bias, stride, padding, dilation):     ...       if R == 3 and S == 3 and dH == 1 and dW == 1:         M_total = N * P * Q         grid_tiled = (triton.cdiv(M_total, 16), triton.cdiv(K, 16),)           _conv2d_3x3_tiled_kernel[grid_tiled](             input, weight, bias, output, M_total,             C, H, W, K, P, Q,             sH, sW,             pH, pW,             bias_exist = bias_exist,             BLOCK_M = 16,             BLOCK_K = 16,)     else:         _conv2d_kernel[grid](...)       return output                       Colored by Color Scripter

cs

· 기존 Naive Kernel에 3x3 Conv인 경우만 별도의 Tiled Kernel로 분기하여 실행

· 항상 Kernel Size가 3x3이라고 가정하여 R, S 인자를 Load 하지 않고 range(3)으로 고정

3~10 : 3x3 Conv2d 연산을 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · w_ptr : Weight Tensor 주소

    · b_ptr : Bias Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · M_total : Batch와 Output 위치를 합친 개수(N x P x Q)

    · C / H / W : 입력 Tensor의 Input Channel 수 / Height / Width → [N, C, H, W]

    · K / P / Q : 출력 Tensor의 Output Channel 수 / Height / Width → [N, K, P, Q]

    · sH, sW : Stride

    · pH, pW : Padding

    · bias_exist : Bias 존재 여부

    · BLOCK_M : 한 Program이 저장할 Output 위치 개수

    · BLOCK_K : 한 Program이 저장할 Output Channel 개수

12 : Output 위치 Tile 번호

13 : Output Channel Tile 번호

14 : 현재 Program이 처리할 Output 위치 Index

15 : 현재 Program이 처리할 Output Channel Index

17~19 : offsets_m을 output 위치 좌표로 변경

20 : 누적합 저장

22 : Input Channel 전체 순회

23~24 : 3x3 Kernel Window 내부 순회

25~26 : 현재 Output 위치들이 참조할 Input 위치 계산

27 : Input Tensor [N, C, H, W]에서 Load할 Memory Offset 계산

28 : Weight Tensor [K, C, 3, 3]에서 Load할 Memory Offset 계산

29~31 : Input 위치가 유효한 범위 안에 있는지 확인(Padding에 의해 범위를 벗어나는 위치 제외)

32 : Output Channel Index가 실제 K보다 작은지 확인

33 : Input 값 Load(범위 밖인 경우 0으로 처리)

34 : Weight 값 Load

35 : Conv 곱셈 결과 누적

37~39 : Bias가 있는 경우 각 Output Channel에 해당하는 Bias를 더함

41 : Output Tensor [N, K, P, Q]에서 저장할 Memory Offset 계산

42 : 유효한 Output 위치와 Output Channel에 대해서만 저장하기 위한 Mask

43 : 최종 Conv 결과를 Output Tensor에 저장

49 : Kernel Size가 3x3이고, Dilation이 1인 경우 Tiled Kernel 사용

50 : Batch와 Output Spatial 위치를 하나로 합친 개수

51 : 2차원 Grid 생성(1번째 축 : Output 위치 Tile 개수 / 2번째 축 : Output Channel Tile 개수)

53~60 : 3x3 Tiled Kernel 실행(하나의 Triton Program은 최대 16x16개의 원소에 대해 동시에 Conv 연산 수행)

61~62 : 3x3 Conv가 아닌 경우 기존 Naive Conv Kernel 사용

63 : Conv 결과 Tensor를 반환

· Naive : 하나의 Program이 Output 원소 128개를 계산하며 각 Output 원소가 독립적으로 C x R x S Loop 수행

→ Output[n, k, p, q] 하나하나를 따로 계산(offset = [n, k, p, q])

· 3x3 Tiled Kernel : 하나의 Program이 (Output 위치 16개) x (Output Channel 16개) 계산

→ 한 번에 256개 Output 원소 계산(offsets_m = [n, p, q] & offsets_k = [k])

> Result

· 현재 Tiled Kernel은 Tiling은 했지만 Scalar 방식으로 Update → Register 사용량↑, Broadcast 연산↑, Tensor Core 사용 X

▣ Conv2d - 3x3 Tiled Dot Kernel

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73

@triton.jit def _conv2d_3x3_dot_kernel(     x_ptr, w_ptr, b_ptr, y_ptr, M_total: tl.constexpr,     C: tl.constexpr, H: tl.constexpr, W: tl.constexpr,     K: tl.constexpr, P: tl.constexpr, Q: tl.constexpr,     sH: tl.constexpr, sW: tl.constexpr,     pH: tl.constexpr, pW: tl.constexpr,     bias_exist: tl.constexpr,     BLOCK_M: tl.constexpr,     BLOCK_K: tl.constexpr,     BLOCK_CRS: tl.constexpr):       pid_m = tl.program_id(0)     pid_k = tl.program_id(1)     offsets_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)     offsets_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)       q = offsets_m % Q     p = (offsets_m // Q) % P     n = offsets_m // (P * Q)     acc = tl.zeros((BLOCK_M, BLOCK_K), dtype = tl.float32)       for base in range(0, C * 9, BLOCK_CRS):         offsets_crs = base + tl.arange(0, BLOCK_CRS)         c = offsets_crs // 9         rs = offsets_crs % 9         r = rs // 3         s = rs % 3         ih = p[:, None] * sH + r[None, :] - pH         iw = q[:, None] * sW + s[None, :] - pW         x_offset = ((n[:, None] * C + c[None, :]) * H + ih) * W + iw         w_offset = ((offsets_k[None, :] * C + c[:, None]) * 3 + r[:, None]) * 3 + s[:, None]         x_mask = (                 (offsets_m[:, None] < M_total)                 & (offsets_crs[None, :] < C * 9)                 & (ih >= 0) & (ih < H)                 & (iw >= 0) & (iw < W))         w_mask = (                 (offsets_k[None, :] < K)                 & (offsets_crs[:, None] < C * 9))         x_tile = tl.load(x_ptr + x_offset, mask = x_mask, other = 0)         w_tile = tl.load(w_ptr + w_offset, mask = w_mask, other = 0)         acc += tl.dot(x_tile, w_tile)       if bias_exist:         bias_val = tl.load(b_ptr + offsets_k, mask = offsets_k < K, other=0.0)         acc += bias_val[None, :]       y_offset = ((n[:, None] * K + offsets_k[None, :]) * P + p[:, None]) * Q + q[:, None]     y_mask = (offsets_m[:, None] < M_total) & (offsets_k[None, :] < K)     tl.store(y_ptr + y_offset, acc, mask = y_mask)     def triton_conv2d(input, weight, bias, stride, padding, dilation):     ...       if R == 3 and S == 3 and dH == 1 and dW == 1:         M_total = N * P * Q         grid_tiled = (triton.cdiv(M_total, 16), triton.cdiv(K, 16),)           _conv2d_3x3_dot_kernel[grid_tiled](             input, weight, bias, output, M_total,             C, H, W, K, P, Q,             sH, sW,             pH, pW,             bias_exist = bias_exist,             BLOCK_M = 16,             BLOCK_K = 16,             BLOCK_CRS = 32,)     else:         _conv2d_kernel[grid](...)       return output Colored by Color Scripter

cs

· 기존 Naive Kernel에 3x3 Conv인 경우만 별도의 Tiled Kernel로 분기하여 실행

· 항상 Kernel Size가 3x3이라고 가정하여 R, S 인자를 Load 하지 않고 range(3)으로 고정

** tl.dot( )을 사용해서 Conv를 행렬곱처럼 계산(Like im2col)

3~11 : 3x3 Conv2d 연산을 행렬곱 연산으로 수행하는 GPU Kernel 인자

    · x_ptr : 입력 Tensor 주소

    · w_ptr : Weight Tensor 주소

    · b_ptr : Bias Tensor 주소

    · y_ptr : 출력 Tensor 주소

    · M_total : Batch와 Output 위치를 합친 개수(N x P x Q)

    · C / H / W : 입력 Tensor의 Input Channel 수 / Height / Width → [N, C, H, W]

    · K / P / Q : 출력 Tensor의 Output Channel 수 / Height / Width → [N, K, P, Q]

    · sH, sW : Stride

    · pH, pW : Padding

    · bias_exist : Bias 존재 여부

    · BLOCK_M : 한 Program이 저장할 Output 위치 개수

    · BLOCK_K : 한 Program이 저장할 Output Channel 개수

    · BLOCK_CRS : C x 3 x 3 방향을 몇 개씩 나눠서 tl.dot에 넣을지 결정

13 : Output 위치 Tile 번호

14 : Output Channel Tile 번호

15 : 현재 Program이 처리할 Output 위치 Index

16 : 현재 Program이 처리할 Output Channel Index

18~20 : offsets_m을 output 위치 좌표로 변경

21 : 출력 Tile의 누적합 저장

23 : C x 3 x 3 방향을 BLOCK_CRS개씩 나누어 순회

24 : 현재 Dot Product에 사용할 C x 3 x 3 Index

25 : Input Channel Index 계산

26~28 : 3x3 Kernel 내부 위치 계산

29~30 : 각 Output 위치와 각 Kernel 위치가 참조할 Input 위치 계산

31 : Input Tensor [N, C, H, W]에서 Load할 Memory Offset 계산

32 : Weight Tensor [K, C, 3, 3]에서 Load할 Memory Offset 계산

33~37 : Input 위치가 유효한 범위 안에 있는지 확인(Padding에 의해 범위를 벗어나는 위치 제외)

38~40 : Weight Index가 유효한지 확인

41 : Input Tile Load

42 : Weight Tile Load

43 : Tile 단위 행렬곱 연산 수행

45~47 : Bias가 있는 경우 각 Output Channel에 해당하는 Bias를 더함

49 : Output Tensor [N, K, P, Q]에서 저장할 Memory Offset 계산

50 : 유효한 Output 위치와 Output Channel에 대해서만 저장하기 위한 Mask

51 : 최종 Conv 결과를 Output Tensor에 저장

57 : Kernel Size가 3x3이고, Dilation이 1인 경우 Tiled Kernel 사용

58 : Batch와 Output Spatial 위치를 하나로 합친 개수

59 : 2차원 Grid 생성(1번째 축 : Output 위치 Tile 개수 / 2번째 축 : Output Channel Tile 개수)

61~69 : 3x3 Dot Kernel 실행(하나의 Triton Program은 최대 16x16개를 계산하고, Cx3x3 방향은 32개씩 나누어 내적 결과 누적)

70~71 : 3x3 Conv가 아닌 경우 기존 Naive Conv Kernel 사용

73 : Conv 결과 Tensor를 반환

> Result

· 추론 시간 ≤ 200ms & 정확도 ≥ 75% 조건 달성

· 3x3 Tiled Kernel : Scalar/Broadcast 방식으로 Cx3x3 방향의 MAC 연산을 반복 수행

· 3x3 Dot Kernel : Cx3x3 방향의 연산을 32개씩 묶어서 행렬곱 형태로 수행

▣ Discussion

· ResNet18을 PyTorch로 수행한 결과와 비교해보면 Triton으로 만든 Conv2d가 PyTorch보다 느림

> PyTorch

· NVIDIA의 cuDNN Library를 사용하여 HW 특화 명령을 직접 사용하는 저수준 언어의 최적화 적용 → Tensor Core 활용도↑

· Tensor의 형태, Batch 크기, Kernel 크기에 따라 최적의 Kernel 조합을 동적으로 결정

· 경우에 따라 Kernel Fusion을 활용하여 연산을 하나의 Kernel로 합쳐서 실행

    → 데이터의 재사용성↑ = 데이터를 중간에 다시 Memory에 Read/Write하는 비용↓

> Triton Kernel

· Triton Kernel : 고수준 언어로서 범용성을 가지는 대신 저수준 언어를 사용하는 Library에 비해 Tensor Core 활용도↓

· 고정된 Tiling을 사용하여 Library에 비해 성능↓

· 다단계 연산에서 결과를 Global Memory에 저장하고, 다시 Load하는 Overhead 발생

∴ Conv2d 연산은 PyTorch를 통한 연산의 성능↑

[Reference]

· HW6 - Triton ResNet (MGP) - Yongjun Park

· 11_Triton Indroduction_mgp_2026 (MGP) - Yongjun Park