- PyTorch는 텐서 표현, 자동 미분, 디스패치, 코드 생성이 맞물린 큰 C++ 코드베이스라서 내부 구조를 알아야 커널과 연산자를 안전하게 확장할 수 있음
- 텐서는 실제 데이터와 별도로 size, stride, offset, dtype, device를 들고 있으며, 이 분리 덕분에 복사 없이 다양한 view를 만들 수 있음
- 연산 호출은 Python 바인딩에서 C++로 넘어간 뒤 autograd, device·layout, dtype 디스패치를 차례로 거쳐 CPU·CUDA 커널에 도달함
- 코드베이스에서 핵심적으로 볼 곳은
torch/,torch/csrc/,aten/,c10/이며, 실제 연산자와 커널 작업은 대부분 ATen을 중심으로 이뤄짐 - 커널을 추가하려면 스키마, 에러 검사, output/inplace/functional 변형, dtype 디스패치, 병렬화, 데이터 접근 도구를 함께 이해해야 하며 새 TH 코드는 피하는 편이 좋음
텐서의 물리 표현과 stride
- PyTorch의 중심 데이터 구조는 Tensor이며, n차원 데이터와 함께 크기, 원소 타입, 장치 같은 메타데이터를 가짐
- 일반적인 텐서 표현은 원소를 메모리에 연속 배치하고, 논리적 차원 정보는 별도의 size 메타데이터로 기록하는 방식임
- stride는 논리적 인덱스를 물리 메모리 위치로 바꾸는 규칙임
- 각 차원의 인덱스에 해당 차원의 stride를 곱하고 모두 더해 실제 위치를 계산함
tensor[1, 0]같은 접근은 이 계산을 거쳐 연속 배열 안의 위치로 변환됨
TensorAccessor는 포인터 연산을 직접 다루지 않고도 이런 인덱싱 계산을 처리하게 해주는 편의 클래스임
view, storage, Tensor의 내부 분리
- PyTorch의 view는 stride를 활용해 복사 없이 다른 논리적 텐서를 표현함
tensor[1, :]는 두 번째 행을 나타내는 새 view를 반환하지만 실제 데이터는 복사하지 않음- view에서 데이터를 수정하면 같은 물리 메모리를 공유하는 원본 텐서에도 반영됨
- 열 view처럼 물리 메모리에서 원소가 연속되지 않은 경우에도 stride 조정으로 표현할 수 있음
- 첫 번째 열은 원소 사이에 간격이 있으므로 stride를 1이 아니라 2로 둘 수 있음
- Stride Visualizer로 stride 표현이 만들 수 있는 view를 실험해볼 수 있음
- view를 지원하려면 사용자에게 보이는 Tensor와 실제 데이터를 담는 Storage를 분리해야 함
- 여러 Tensor가 같은 Storage를 공유할 수 있음
- Storage는 dtype과 물리 크기를 정의함
- 각 Tensor는 size, stride, offset으로 물리 메모리를 어떻게 논리적으로 해석할지 정함
- view가 원래 물리 메모리를 붙잡고 있으면, 원본 메모리를 해제하기 위해 view를 복사해 기존 메모리와 분리해야 함
- 당시 PyTorch는 Storage 개념을 별도로 두지만, view를 base tensor가 뒷받침하는 tensor로 정의해 Storage 간접 계층을 없애는 방향도 관심사였음
디스패치와 확장 지점
torch.mm같은 연산 호출에는 추상적으로 두 단계의 디스패치가 있음- 첫 번째는 device type과 layout 기반 디스패치임
- CPU Tensor인지 CUDA Tensor인지, strided Tensor인지 sparse Tensor인지에 따라 구현이 달라짐
- 이 단계는 동적 디스패치이며,
libcaffe2.so와libcaffe2_gpu.so처럼 별도 라이브러리에 있는 커널로 넘어가기 위해 필요함
- 두 번째는 dtype 기반 디스패치임
float곱셈과int곱셈은 다른 코드가 필요하므로, 커널이 지원하는 dtype에 대해 switch 형태로 분기함
- 첫 번째는 device type과 layout 기반 디스패치임
- PyTorch Tensor 확장의 핵심 축은 device, layout, dtype임
- device는 물리 메모리가 있는 위치를 나타내며 CPU, NVIDIA GPU의
cuda, AMD GPU의hip, TPU의xla등이 포함됨 - layout은 물리 메모리를 논리적으로 해석하는 방식이며, strided Tensor가 일반적이고 sparse Tensor는 indices와 data 텐서 쌍을 사용함
- dtype은 각 원소에 저장되는 값을 나타내며 float, integer, quantized integer 등이 가능함
- device는 물리 메모리가 있는 위치를 나타내며 CPU, NVIDIA GPU의
- Tensor 확장은 이 세 축의 조합으로 표현되지만, 모든 조합에 실제 커널이 존재하는 것은 아님
- Tensor 기능을 확장하는 다른 방법은 PyTorch Tensor를 감싸는 래퍼 클래스를 만드는 것임
- 래퍼 클래스는 PyTorch 본체 밖에서 개발할 수 있음
- autograd backward pass 중 해당 tensor를 전달해야 한다면 래퍼보다 실제 Tensor 확장이 필요함
- sparse tensor는 embedding이 만드는 gradient를 sparse로 유지해야 하므로 Python 객체 래퍼가 아니라 진짜 Tensor 확장이어야 함
- Tensor 구조체는 자주 호출되는 기본 연산에서 가상 디스패치를 피하기 위해 고정 layout을 선호함
- TensorImpl struct는 모든 tensor-like 객체가 공통으로 갖는 필드 prefix를 포함함
- strided Tensor에 특히 중요한 필드는 main struct에 남겨둠
- sparse Tensor의 indices와 values 같은 커스텀 필드는 per-Tensor suffix에 저장됨
autograd와 코드 경로
- PyTorch가 처음 출시될 때의 구별점은 Tensor 위의 자동 미분이었음
- PyTorch는 역방향 자동 미분(reverse-mode automatic differentiation)을 구현함
- forward 계산을 거꾸로 따라가며 gradient를 계산함
loss에서 시작해grad_loss,grad_next_h2처럼 역방향 값을 계산함- PyTorch는 이런 값을 보통
grad라고 부르지만, 엄밀히는 vector가 왼쪽에서 곱해진 Jacobian임
- autograd는 소스 코드를 생성하는 소스 간 변환을 하지 않음
- forward 연산을 수행할 때 backward에 필요한 메타데이터를 저장함
loss.backward()가 호출되면 이 메타데이터를 사용함
- 당시 구조에서는 Tensor를 감싸는 Variable이 있고, Variable이 AutogradMeta를 저장함
- C++에서는 Variable-Tensor merge 작업이 진행 중이었음
- autograd가 들어가면 디스패치 경로도 달라짐
- CPU나 CUDA 구현으로 가기 전에 Variable 디스패치가 먼저 발생함
- Variable 디스패치는 Variable을 풀고 실제 구현을 호출한 뒤, 결과를 다시 Variable로 감싸고 backward용 메타데이터를 기록함
- 한번 non-Variable Tensor 영역으로 들어가면 함수 반환 전까지 Variable 영역으로 다시 들어가지 않음
PyTorch 코드베이스와 커널 작성
- PyTorch 코드베이스에서 실무적으로 중요한 디렉터리는 네 곳임
torch/: 사용자가 import하는 Python 모듈이 들어 있음torch/csrc/: Python과 C++ 사이의 바인딩, autograd engine, JIT compiler, C++ frontend 등을 포함함aten/: “A Tensor Library”의 약자로 Tensor 연산을 구현하는 C++ 라이브러리이며, 커널 코드를 찾을 가능성이 높은 곳임c10/: Tensor와 Storage 데이터 구조의 실제 구현을 포함한 핵심 추상화가 들어 있음
torch.add호출은 대략 Python 인자 파싱, Variable 디스패치, device type/layout 디스패치, 실제 커널 순서로 진행됨- Python 쪽에는
torch._C.VariableFunctions.add같은 함수가 노출됨 - C++ 구현은 자동 생성되며,
PythonArgParser로args와kwargs에서 C++ 객체를 꺼냄 - 이후
dispatch_add가 호출되고, GIL을 해제한 뒤 C++ Tensor 메서드가 실행됨 - 반환된 Tensor는 다시
PyObject로 감싸짐
- Python 쪽에는
- Tensor의
add메서드 자체에서는 아직 가상 디스패치가 일어나지 않으며, inline 메서드가Type객체의 가상 메서드를 호출함Type은 동적 디스패치를 위한 장치로 쓰임- 예시에서는 모든 device type에 같은
add구현을 쓰기 때문에TypeDefault로 dispatch됨 - 당시
Typedispatch를 다른 메커니즘으로 대체하는 작업도 진행 중이었음
- PyTorch 커널은 보통 다음 요소로 구성됨
- 코드 생성과 Python 바인딩 생성을 위한 schema
- 입력 Tensor의 차원과 속성을 확인하는 에러 검사
- 결과 Tensor 할당
- dtype별 커널로 넘어가기 위한 dtype 디스패치
- CPU 멀티코어 활용을 위한 병렬화
- 실제 데이터 접근과 계산
- 연산자 schema는 native README에 형식이 정리되어 있으며, 필요하면 derivatives.yaml에 derivative를 정의해야 함
- 에러 검사는 저수준
TORCH_CHECK또는 고수준TensorArg기반 API로 할 수 있음TORCH_CHECK는 boolean 조건과 에러 메시지 구성 요소를 받아 조건이 거짓이면 오류를 냄TensorArg는 argument 이름 같은 메타데이터를 포함하고,checkDim()같은 함수로 사용자 친화적인 오류 메시지를 제공함
- PyTorch 연산자는 종종 세 가지 변형을 함께 작성해야 함
abs_out: 사전 할당된 출력에 쓰며out=키워드 인자를 구현함abs_: inplace 연산임abs: 일반 functional 버전임- 보통
abs_out이 핵심 구현이고 나머지는 얇은 래퍼지만, 경우에 따라 각각 특화 구현이 필요함
- dtype 디스패치에는
AT_DISPATCH_ALL_TYPES매크로를 사용함- 대상 Tensor의 dtype과 lambda를 받아 dtype별 특수화를 생성함
- 어떤 dtype을 지원할지도 이 매크로 계열이 결정함
- Dispatch.h에 사용 사례별 선택 지침이 있음
- 데이터 접근에는 상황별 도구를 사용함
TensorAccessor: 차원 수와 dtype을 template parameter로 고정하고, runtime 검사 뒤 unchecked access를 제공하며 stride를 올바르게 처리함PackedTensorAccessor: CUDA launch로 accessor를 넘길 때 유용함TensorIterator: pointwise 연산처럼 규칙적인 접근에서 broadcasting과 type promotion을 자동 처리함Vec256: CPU vector instruction을 사용하는 커널 작성을 돕고,binary_kernel_vec같은 helper가 vectorized path와 나머지 일반 instruction 처리를 도와줌
- 레거시 TH 스타일 커널은 C 스타일 코드, 수동 참조 카운팅,
generic/디렉터리의#define scalar_t기반 다중 컴파일을 사용함- 새 TH 코드를 추가하기보다 기존 TH 함수를 ATen으로 port하는 작업이 유용함
- PyTorch C++ 개발에서는 workflow 효율이 중요함
- 많이 포함되는 header, 특히 CUDA 파일에 포함되는 header를 수정하면 rebuild가 매우 길어질 수 있음
- CI는 별도 설정 없이 테스트할 수 있지만 신호를 받기까지 1~2시간이 걸릴 수 있음
- 특정 CI 설정 문제는 Docker image를 로컬에서 실행해 재현할 수 있음
- CONTRIBUTING guide의 ccache 설정이 권장됨
- CUDA build는 느리기 때문에 노트북보다 CPU와 RAM이 충분한 서버에서 빌드하는 편이 좋음
- 기여를 시작할 곳으로는 issue tracker가 적합함
triaged라벨은 적어도 한 명의 PyTorch 개발자가 초기 평가를 했다는 뜻임- high priority, autograd, small 같은 라벨로 이슈를 찾을 수 있음
- 코드 작성 외에도 문서 개선, 버그 재현, RFC 논의 참여가 기여 활동에 포함됨