Conv2D (NCHW, group=1) 구현/연산

핵심 아이디어

• 입력 X[N,Cin,H,W]를 im2col로 A=[H_out*W_out, K](K = Cin*Kh*Kw)로 펴고,
• 가중치 W[Cout,Cin,Kh,Kw]를 pack해 B=[K, Cout]로 만들고,
• Y_tmp = A @ B = [HWo, Cout]를 계산 후
• 최종 출력 버퍼 y_n의 연속 메모리(NCHW)에 맞도록 전치해 [Cout, HWo]로 기록,
• bias는 채널(row) 기준으로 브로드캐스트 더하기.

이 흐름이 forward/backward에서 하나의 축 약속으로 일관되게 유지된다.

---

텐서 레이아웃 & 기호

• 입력/출력: NCHW 연속 (row-major)
  • 한 배치 내 출력 y_n 메모리는 [Cout, HWo] row-major (채널별로 Ho*Wo가 연속)
• 보조 차원:
  • HWo = H_out * W_out
  • K = Cin * Kh * Kw
• im2col K축 전개 순서(매우 중요): (ci, kh, kw), kw가 최속(fastest)
  → k = ((ci * Kh) + kh) * Kw + kw

---

Forward 경로

1) 출력 공간 크기

H_out = (H + 2*pad_h - dil_h*(Kh-1) - 1) / stride_h + 1
W_out = (W + 2*pad_w - dil_w*(Kw-1) - 1) / stride_w + 1
HWo   = H_out * W_out
K     = Cin * Kh * Kw

2) im2col

• 입력 X[n] → A = im2col(X[n])
  A shape: [HWo, K] (row-major)
  열 인덱스 k는 (ci, kh, kw)로 전개(kw fastest).

3) W pack

• 원형: W[co, ci, kh, kw]
• pack: W_KC[K, Cout]로 변환 (kw fastest, (ci,kh,kw) 순)
  • 커널: pack_w_oihw_to_KC
  • 인덱싱: idx_in = ((co*Cin+ci)*Kh+kh)*Kw + kw
    idx_out = k*Cout + co

4) GEMM

Y_tmp = A [HWo, K] @ W_KC [K, Cout]  -> [HWo, Cout]

5) 출력 버퍼로 전치

• 최종 출력 버퍼 y_n은 NCHW에서 한 배치의 메모리가 [Cout, HWo]로 쓰이도록 잡혀 있음.
• 따라서 transpose(Y_tmp [HWo,Cout]) -> y_n [Cout, HWo].

6) Bias

• B[Cout]를 채널(row) 축으로 브로드캐스트:
• 커널: add_bias_rows(y_n, B, Cout, HWo)

• 주의: gemm_run의 bias 옵션은 끄고(혼동 방지), 항상 별도 커널로 더한다.

for co in [0..Cout):
  for hw in [0..HWo):
    y_n[co, hw] += B[co]

Backward 경로

목표: (dW, dB, dX) = ∂L/∂(W,B,X) for L = sum(Y * dY).

레이아웃 불변식
Forward에서 최종 Y를 [Cout,HWo]로 썼으므로, backward로 들어오는 dY도 [Cout, HWo] row-major다.

0) 공통 준비

• im2col(X[n]) -> A = [HWo, K]
• W_CK = [Cout, K] (forward와 동일 전개 (ci,kh,kw); 커널 pack_w_oihw_to_CK)
• dWpack 누적 버퍼: [Cout, K] (제로 초기화)

1) dB

• 축: 채널(co)별로 공간(HWo) 전체 합
• 커널: reduce_db_rows_kernel(gy = [Cout, HWo])

dB[co] = Σ_{hw} dY[co, hw]

2) dW

• 공식: dWpack = dY^T @ A
  • 여기서 dY는 [Cout, HWo] → 전치 필요 없음
    tA = dY[Cout, HWo], tB = A[HWo, K]

dWpack_batch = dY[Cout, HWo] @ A[HWo, K] -> [Cout, K]
dWpack += dWpack_batch      (배치 누적)

• 배치 루프 끝나면 dWpack[Cout,K] → dW[O,I,H,W]로 언팩 누적
  커널: unpack_ck_to_oihw_add

3) dX

• 공식: dX_col = dY_hw @ W_CK
  • 여기서 dY는 [Cout, HWo]이고, A는 [HWo, K]였으니
    dY_hw = transpose(dY) = [HWo, Cout]가 필요

dY_HT = transpose(dY [Cout, HWo]) = [HWo, Cout]
dX_col = dY_HT[HWo, Cout] @ W_CK[Cout, K] -> [HWo, K]
dX[n] = col2im(dX_col) -> [Cin, H, W]

수치미분 체크(테스트 요약)

• Forward: NumPy conv2d_ref와 np.allclose.
• Backward: 스칼라 손실 L = sum(Y * dY)에 대해,
  • dW_num = ∂L/∂W (수치미분) vs. 구현 dW
  • dB_num = ∂L/∂B (수치미분) vs. 구현 dB
• 우리가 맞춘 후:
  • forward: 랜덤 W/B, bias-only, one-hot 모두 True
  • backward: dW, dB numeric True

---

흔한 함정 & 디버깅 체크리스트

1. K 전개 순서 불일치
   • im2col과 W pack의 K축 전개 순서가 다르면 랜덤 W에서만 오차 폭발.
   • 이번 구현의 약속: (ci, kh, kw), kw가 최속.
2. 출력 메모리 해석 오류
   • NCHW 연속에서 한 배치 y_n은 [Cout, HWo] row-major처럼 써야 한다.
   • GEMM 결과가 [HWo, Cout]이면 전치해서 쓴다.
3. dY 레이아웃 착각
   • Forward 저장 레이아웃을 기준으로 backward의 dY도 해석해야 한다.
     이번 케이스: [Cout, HWo].
   • dB, dW, dX 각 경로에서 맞는 축으로 합/전치 여부 확인.
4. bias 브로드캐스트 축
   • bias는 채널 축(co) 기준.
   • [Cout, HWo]에 row-wise add.
5. transpose 런처의 (M,N) 정의
   • 우리의 런처는 in[M,N] → out[N,M].
   • 어디서나 M/N 순서 맞는지 확인(특히 dX 경로).
6. atomic add / 누적 초기화
   • dB/dW 누적 전에 zero-init 되어야 한다.
   • reduce에서 atomicAdd 사용.

---

경계값/옵션 처리

• stride/pad/dilation:
  im2col에서 좌표 h_in = ho*sH - pH + kh*dH, w_in = wo*sW - pW + kw*dW;
  경계 밖은 0.
• groups: 현재 groups=1만.
  groups 확장은 채널/가중치/행렬 분할(블록 GEMM)로 처리 가능.
• Bias=None: bias 경로 스킵.

---

성능 팁(선택)

• 전치 제거: gemm_run이 임의 stride를 허용해 tY를 [Cout, HWo] stride {HWo,1}로 직접 쓸 수 있으면 forward에서 Y_tmp/전치를 없앨 수 있음.
• W pack 캐싱: 훈련/추론에서 W가 고정이면 W_KC/W_CK 캐시.
• 블로킹/공유메모리: im2col/col2im/pack에서 블록 크기와 coalescing 최적화.

 

수학적 요약

• Forward:
   

A = im2col(X[n])          # [HWo, K]
B = pack(W)               # [K, Cout]
Y_tmp = A @ B             # [HWo, Cout]
y_n = transpose(Y_tmp)    # [Cout, HWo]
y_n += B_bias (row-wise)

• Backward:

dB = reduce_rows(dY, axis=HWo)
dWpack += dY @ A          # [Cout, K]
dW = unpack(dWpack)       # [Cout, Cin, Kh, Kw]

dY_HT = transpose(dY)     # [HWo, Cout]
dX_col = dY_HT @ W_CK     # [HWo, K]
dX = col2im(dX_col)

---

마무리

• 축 약속을 하나만 정하고(여기선 (ci,kh,kw) with kw fastest)
  forward 저장 레이아웃([Cout,HWo])을 기준으로 backward의 모든 경로를 정렬.
• 이 두 가지 원칙만 지키면 conv2d는 일관되게 맞는다.
• 지금 코드는 그 원칙대로 정리되어 있고, forward/ backward numeric 검증까지 통과.