PyTorch 내부 구조 10 - CUDA Extension으로 Custom Kernel 연결하기
CUDA kernel을 PyTorch operator로 만들려면 kernel 코드뿐 아니라 tensor contract와 runtime semantics를 함께 맞춰야 한다
kernel만 빠르면 끝나는 것이 아니다
CUDA extension을 붙일 때 많은 사람이 kernel 코드 자체에만 집중한다. 하지만 실제 operator로 쓰려면 다음이 모두 맞아야 한다.
- input/output shape contract
- dtype / device validation
- stream usage
- error handling
- autograd 연결 가능성
즉 extension 작업은 kernel engineering과 framework integration을 동시에 요구한다.
좋은 습관
- 먼저 Python reference 구현으로 semantics를 고정한다
- 그다음 C++ / CUDA로 같은 계약을 재현한다
- 마지막에 backward와 edge case를 붙인다
다음 글에서는 operator schema와 meta 함수처럼, custom op를 "PyTorch다운 방식"으로 등록하는 데 필요한 요소를 본다.