Week 2: PyTorch Eager Mode

PyTorch + NPU 온라인 모임 #2 | 2024-12-11

소개

이번 강의는 PyTorch 기초의 첫 번째 시간으로, Eager mode1에 대해 다룹니다.

오늘 다룰 주제들

x = torch.matmul(y, z) 한 줄을 실행하면 내부에서 정확히 무슨 일이 일어날까요? 이 질문에서 출발해, 55단계에 달하는 call stack을 끝까지 따라가고 그 과정에서 등장한 PyTorch internal 개념을 정리합니다.

  • torch.matmul의 call stack 추적 - CUDA GEMM 진입점에 breakpoint를 걸어 얻은 55단계 call stack 분석
    • Python → C++ 바인딩 → matmul dispatch → mm의 autograd kernel → CUDA kernel redispatch → cuBLAS 진입 순으로 전체 경로를 따라가기
    • dispatcher가 dispatch key set을 계산해 kernel을 고르고 redispatch를 반복하는 구조
    • call stack의 상당 부분은 반복되는 dispatch logic이거나 native_functions.yaml, derivatives.yaml에서 자동 생성된 코드 → 55단계 깊이만큼 사람이 관리하는 코드가 많은 것은 아님
  • Eager mode의 구성 요소 - 예제에서 등장한 개념들을 세 축으로 나눠 정리
    • Tensor: TensorImpl·Storage로 이어지는 내부 구조, view·clone 같은 lifecycle
    • Operator: dispatcher와 dispatch key(set)가 동작하는 메커니즘
    • Runtime: Device, Stream, Event, DeviceGuard, Allocator 등 저수준 인프라
  • PyTorch build - 지금까지 본 코드가 실제 라이브러리가 되는 과정
    • 소스 빌드 과정과 주요 환경변수, 빌드 산출물(shared object, Python module)
    • ATen·C10·자동 생성 코드로 나뉘는 소스 디렉토리 구조

torch.matmul을 수행하면 어떤 일이 일어나는가?

PyTorch에서 x = torch.matmul(y, z)를 실행하면 내부적으로 어떤 일이 벌어질까요? 파이썬에서 import torch를 한 뒤 텐서를 만들어 연산을 수행하면 결과가 나오고, 이를 프린트해 볼 수 있습니다. 그런데 그 과정에서 실제로 어떤 일이 일어나는지를 파악하는 가장 효과적인 방법 중 하나는 call stack을 확인해 보는 것입니다.

Call Stack 분석

아래는 GPU에서 matmul이 수행되도록 설정하고 CUDA GEMM2 진입점에 breakpoint를 건 뒤 torch.matmul을 수행했을 때의 call stack snapshot입니다. 실제 수행 지점까지 많은 중간 단계를 거칩니다.

이 snapshot은 Python 3.13.0(conda) 환경의 debug 빌드에서 캡처한 것입니다(commit과 빌드 옵션은 기록되지 않았습니다). 이 문서가 기준으로 삼는 v2.14.0 소스와 대조해 보면 표에 등장하는 심볼 이름·파일 경로·전체 dispatch 흐름은 그대로 유효하지만, 파일 내 라인 번호는 버전에 따라 이동합니다(예: Dispatcher.h의 redispatch는 v2.14.0 기준 843행). 또한 최근 버전에서는 cuBLAS 호출 직전에 TunableOp 체크와 backend 선택(cuBLAS/cuBLASLt) 레이어가 있어, gemm<float> → gemm_internal<float> → gemm_internal_cublas<float>(v2.14.0 기준 CUDABlas.cpp 1421 → 1226 → 969행) 순으로 스택이 1~2단 더 깊어집니다.

#0 at::cuda::blas::gemm<float> at aten/src/ATen/cuda/CUDABlas.cpp:1127
#1 operator at aten/src/ATen/native/cuda/Blas.cpp:457
#2 operator at aten/src/ATen/native/cuda/Blas.cpp:457
#3 at::native:: at aten/src/ATen/native/cuda/Blas.cpp:457
#4 at::native::structured_mm_out_cuda::impl at aten/src/ATen/native/cuda/Blas.cpp:604
#5 at:: at build/aten/src/ATen/RegisterCUDA.cpp:12173
#6 c10::impl::detail::WrapFunctionIntoFunctor_<... at aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#7 c10::impl::wrap_kernel_functor_unboxed_<... at aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:468
#8 c10::callUnboxedKernelFunction<... at aten/src/ATen/core/boxing/KernelFunction_impl.h:53
#9 c10::KernelFunction::call<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/boxing/KernelFunction_impl.h:105
#10 c10::Dispatcher::redispatch<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/dispatch/Dispatcher.h:714
#11 c10::TypedOperatorHandle<at::Tensor at aten/src/ATen/core/dispatch/Dispatcher.h:536
#12 at::_ops::mm::redispatch at build/aten/src/ATen/Operators_3.cpp:4010
#13 at::redispatch::mm at build/aten/src/ATen/RedispatchFunctions.h:5217
#14 operator at torch/csrc/autograd/generated/VariableType_3.cpp:13455
#15 torch::autograd::VariableType:: at torch/csrc/autograd/generated/VariableType_3.cpp:13456
#16 c10::impl::detail::WrapFunctionIntoFunctor_<... at aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#17 c10::impl::wrap_kernel_functor_unboxed_<... at aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:485
#18 c10::callUnboxedKernelFunction<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/boxing/KernelFunction_impl.h:53
#19 c10::KernelFunction::call<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/boxing/KernelFunction_impl.h:105
#20 c10::Dispatcher::callWithDispatchKeySlowPath<... at aten/src/ATen/core/dispatch/Dispatcher.h:661
#21 c10::Dispatcher::call<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/dispatch/Dispatcher.h:680
#22 c10::TypedOperatorHandle<at::Tensor at aten/src/ATen/core/dispatch/Dispatcher.h:531
#23 at::_ops::mm::call at build/aten/src/ATen/Operators_3.cpp:4003
#24 at::Tensor::mm at build/aten/src/ATen/core/TensorBody.h:2999
#25 at::native::_matmul_impl at aten/src/ATen/native/LinearAlgebra.cpp:2031
#26 at::native::matmul at aten/src/ATen/native/LinearAlgebra.cpp:2181
#27 at:: at build/aten/src/ATen/RegisterCompositeImplicitAutograd.cpp:2774
#28 c10::impl::detail::WrapFunctionIntoFunctor_<... at aten/src/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h:13
#29 c10::impl::wrap_kernel_functor_unboxed_<... at aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h:468
#30 c10::callUnboxedKernelFunction<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/boxing/KernelFunction_impl.h:53
#31 c10::KernelFunction::call<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/boxing/KernelFunction_impl.h:105
#32 c10::Dispatcher::callWithDispatchKeySlowPath<... at aten/src/ATen/core/dispatch/Dispatcher.h:661
#33 c10::Dispatcher::call<at::Tensor, at::Tensor const&, at::Tensor const&> at aten/src/ATen/core/dispatch/Dispatcher.h:680
#34 c10::TypedOperatorHandle<at::Tensor at aten/src/ATen/core/dispatch/Dispatcher.h:531
#35 at::_ops::matmul::call at build/aten/src/ATen/Operators_4.cpp:3192
#36 at::Tensor::matmul at build/aten/src/ATen/core/TensorBody.h:2899
#37 operator at torch/csrc/autograd/generated/python_torch_functions_0.cpp:4909
#38 torch::autograd::THPVariable_matmul at torch/csrc/autograd/generated/python_torch_functions_0.cpp:4911
#39 cfunction_call at /usr/local/src/conda/python-3.13.0/Objects/methodobject.c:540
#40 _PyObject_MakeTpCall at /usr/local/src/conda/python-3.13.0/Objects/call.c:242
#41 _PyEval_EvalFrameDefault at /usr/local/src/conda/python-3.13.0/Python/generated_cases.c.h:813
#42 PyEval_EvalCode at /usr/local/src/conda/python-3.13.0/Python/ceval.c:596
#43 run_eval_code_obj at /usr/local/src/conda/python-3.13.0/Python/pythonrun.c:1323
#44 run_mod at /usr/local/src/conda/python-3.13.0/Python/pythonrun.c:1408
#45 pyrun_file at /usr/local/src/conda/python-3.13.0/Python/pythonrun.c:1241
#46 _PyRun_SimpleFileObject at /usr/local/src/conda/python-3.13.0/Python/pythonrun.c:490
#47 _PyRun_AnyFileObject at /usr/local/src/conda/python-3.13.0/Python/pythonrun.c:77
#48 pymain_run_file_obj at /usr/local/src/conda/python-3.13.0/Modules/main.c:409
#49 pymain_run_file at /usr/local/src/conda/python-3.13.0/Modules/main.c:428
#50 pymain_run_python at /usr/local/src/conda/python-3.13.0/Modules/main.c:696
#51 Py_RunMain at /usr/local/src/conda/python-3.13.0/Modules/main.c:775
#52 Py_BytesMain at /usr/local/src/conda/python-3.13.0/Modules/main.c:829
#53 __libc_start_call_main at ../sysdeps/nptl/libc_start_call_main.h:58
#54 __libc_start_main_impl at ../csu/libc-start.c:360

Call depth가 무려 55단계나 됨 뭔가 엄청 복잡한 일이 일어나고 있는 것처럼 보임

Eager mode에서는 breakpoint를 잡은 시점에서 이전 상태를 파악하기 어려운 경우도 많지만, torch.matmul 같은 경우에는 call stack을 보면 어떤 과정을 거쳐 어떤 일이 일어났는지 거의 다 파악할 수 있습니다.

왜 이렇게 복잡한가?

Call stack이 55단계나 되는 이유는 크게 세 가지로 나눌 수 있습니다.

1. torch.matmul 자체의 복잡성

torch.matmul은 input tensor의 shape에 따라 서로 다른 연산으로 분기합니다:

Input 조합호출되는 연산
1D × 1Dtorch.dot
2D × 1Dtorch.mv
1D × 2Dunsqueeze 후 torch.mm
2D × 2Dtorch.mm
N-D × 2-D 또는 N-D × 1-D보통 batch 차원을 fold3하여 torch.mm / torch.mv
양쪽 다 batch 차원batch 차원 broadcast 후 torch.bmm

2. Language boundary 전환

PyTorch는 Python → C++ → CUDA(또는 다른 kernel language)로 이어지는 다층 구조입니다. 언어 경계를 넘을 때마다 binding과 wrapper 코드가 필요하므로 call stack이 길어집니다.

3. 다양한 실행 시나리오의 동적 결정

같은 matmul 호출이라도 런타임 상황에 따라 실행 경로가 달라집니다:

결정 축가능한 경우
DeviceCPU, CUDA, XPU, MPS, …
Autogradbackward graph 생성 필요 vs 불필요
Tracingtorch.jit.trace / torch.compile 중 vs 아님

이처럼 “지금 이 호출에서 어떤 구현(kernel)을 실행해야 하는가”를 런타임에 결정하여 해당 kernel을 호출하는 과정을 PyTorch에서는 dispatch라고 부릅니다. 여기서 kernel은 dispatch key 하나에 등록된 함수입니다. 실제 연산을 수행하는 CUDA 코드도, autograd 기록만 하고 연산은 넘기는 wrapper도 dispatcher 입장에서는 모두 kernel입니다. 그리고 위 표의 각 경우 하나하나를 나타내는 태그가 dispatch key입니다. 예를 들어 CPU, CUDA 같은 device도, Autograd 같은 기능도 각각 하나의 dispatch key로 표현됩니다.

한 번의 op 호출에는 여러 관심사가 동시에 걸려 있을 수 있습니다. “CUDA tensor이면서 gradient 추적이 필요한” 호출이라면 CUDA key와 Autograd key가 둘 다 유효합니다. 이때 dispatcher라는 컴포넌트는 다음 순서로 움직입니다.

  1. dispatch key set 계산: 입력 tensor들이 가진 key에 tracing mode 같은 thread-local 상태의 key를 합쳐 dispatch key set을 만듭니다.
  2. 최고 우선순위 key 선택: set에서 우선순위가 가장 높은 key(예: Autograd)의 kernel을 호출합니다.
  3. 준비 작업 후 redispatch: 그 kernel은 자기 관심사의 준비 작업(예: backward graph 노드 생성)을 하고, 실행 도중에 dispatcher를 다시 호출합니다(redispatch). 이미 처리한 key는 제외하도록 표시됩니다.
  4. 다음 key의 kernel 실행: 남은 key 중 우선순위가 가장 높은 key(예: CUDA)의 kernel이 실제 연산을 수행하고 반환합니다.

redispatch는 kernel 실행이 끝난 뒤가 아니라 실행 도중에 일어나는 중첩 호출입니다. 안쪽 kernel이 결과를 반환하면 바깥 kernel이 마무리 작업(예: 결과 tensor에 backward history 기록)을 이어갑니다. 이 중첩만큼 call stack이 깊어집니다.

위 표는 대표적인 예시이며, 실제 dispatch key는 이보다 많습니다. 각 key는 device, autograd, tracing 같은 관심사를 하나씩 독립적으로 맡습니다. 한 호출에 걸린 key마다 dispatch → kernel 실행 → redispatch가 반복되므로 call stack이 깊어집니다. dispatcher의 내부 동작은 아래 “PyTorch Dispatcher” 절에서 다룹니다.

Week 1에 소개했던 PyTorch의 주요 특성과의 연관

지난주 “PyTorch 2.0” 섹션의 주요 특성 중 오늘 다루는 Eager mode call stack과 연관된 항목은 다음과 같습니다. 회색 항목은 이번 주에는 해당하지 않습니다.

특성이번 주 연관
NumPy-like experienceO - interactive mode로 텐서 연산을 곧바로 실행
Heterogeneous computingO - Device별 dispatch의 기반
MPI-like distributed programming modelX - 이번 주 범위 아님
Integration with compute libraries / ML compilerO - cuBLAS 등 외부 라이브러리 호출
Three language layers (Python → C++ → kernel)O - call stack에서 language boundary 전환이 직접 드러남
Define-by-run with TorchDynamoX - graph mode는 다음 주
Codegen을 적극적으로 활용O - call stack의 상당 부분이 자동 생성 코드
Various backend integration pointsO - dispatch target으로서의 backend 등록

Codegen 활용은 Week 1에서 빠져 있던 항목으로, 이번 주에 새로 추가됩니다. PyTorch 코드베이스에는 도구가 자동 생성하는 코드가 많습니다. operator마다 같은 기능의 코드가 반복되므로 사람이 직접 쓰지 않고 생성 도구가 만들며, automated differentiation도 이 코드 생성으로 구현됩니다.

Python vs. PyTorch Internal (C++)

Call stack은 밑에서 위로 진입합니다. 크게 두 영역으로 나뉘며, 오늘은 위쪽의 PyTorch internal (C++) 영역을 중심으로 살펴봅니다.

PyTorch Internal (C++): #0 ~ #38
#Call Stack역할
#0at::cuda::blas::gemm<float> at CUDABlas.cpp:1127cuBLAS GEMM 호출 - 실제 GPU 행렬곱 수행 진입점
#1-3operator at native/cuda/Blas.cpp:457CUDA mm 내부 lambda - cuBLAS 호출을 감싸는 wrapper
#4at::native::structured_mm_out_cuda::impl at native/cuda/Blas.cpp:604CUDA mm 구현체 - structured kernel의 device별 impl
#5at:: at build/.../RegisterCUDA.cpp:12173CUDA backend 등록 - native_functions.yaml에서 생성된 CUDA dispatch 등록 코드 (자동생성)
#6-9WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hdispatch logic - 등록된 kernel functor를 unboxed 호출로 변환하여 실행
#10c10::Dispatcher::redispatch at Dispatcher.h:714redispatch - Autograd kernel이 실행 도중 다음 dispatch key(CUDA)로 재전달
#11c10::TypedOperatorHandle at Dispatcher.h:536redispatch의 typed handle wrapper
#12at::_ops::mm::redispatch at build/.../Operators_3.cpp:4010mm redispatch 진입점 - op별 생성된 C++ redispatch 진입점 (자동생성)
#13at::redispatch::mm at build/.../RedispatchFunctions.h:5217mm redispatch convenience 함수 (자동생성)
#14operator at generated/VariableType_3.cpp:13455Autograd kernel 내부 lambda (자동생성)
#15torch::autograd::VariableType::mm at generated/VariableType_3.cpp:13456Autograd kernel - backward graph 세팅 후 forward 실행, gradient 필요 시 기록 (자동생성; 실제 심볼은 VariableType 아래 익명 namespace 안에 생성됨)
#16-19WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hdispatch logic - mm의 Autograd key에 대한 kernel 실행
#20c10::Dispatcher::callWithDispatchKeySlowPath at Dispatcher.h:661profiler 경로 - RecordFunction callback이 등록되어 있고 op이 observed일 때만 타며, kernel 호출을 RecordFunction guard로 감쌈. kernel 선택은 이미 Dispatcher::call에서 끝남
#21c10::Dispatcher::call at Dispatcher.h:680dispatch 진입 - mm op에 대한 첫 번째 dispatch (Autograd key 선택)
#22c10::TypedOperatorHandle at Dispatcher.h:531dispatch의 typed handle wrapper
#23at::_ops::mm::call at build/.../Operators_3.cpp:4003mm C++ 진입점 - op별 생성된 dispatch 진입점, Dispatcher 호출 (자동생성)
#24at::Tensor::mm at build/.../TensorBody.h:2999at::Tensor의 mm 메서드 - mm::call로 위임
#25at::native::_matmul_impl at LinearAlgebra.cpp:2031matmul 구현체 - input shape 분석 후 적절한 연산(mm, bmm 등) 선택
#26at::native::matmul at LinearAlgebra.cpp:2181matmul 진입 - _matmul_impl로 위임
#27at:: at build/.../RegisterCompositeImplicitAutograd.cpp:2774matmul kernel 등록 - device/autograd 무관한 composite kernel로 등록 (자동생성)
#28-31WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hdispatch logic - matmul의 CompositeImplicitAutograd key에 대한 kernel 실행
#32c10::Dispatcher::callWithDispatchKeySlowPath at Dispatcher.h:661profiler 경로 - RecordFunction callback이 등록되어 있고 op이 observed일 때만 타며, kernel 호출을 RecordFunction guard로 감쌈. kernel 선택은 이미 Dispatcher::call에서 끝남
#33c10::Dispatcher::call at Dispatcher.h:680dispatch 진입 - matmul op에 대한 dispatch
#34c10::TypedOperatorHandle at Dispatcher.h:531dispatch의 typed handle wrapper
#35at::_ops::matmul::call at build/.../Operators_4.cpp:3192matmul C++ 진입점 - op별 생성된 dispatch 진입점 (자동생성)
#36at::Tensor::matmul at build/.../TensorBody.h:2899at::Tensor의 matmul 메서드 - matmul::call로 위임
#37operator at generated/python_torch_functions_0.cpp:4909Python binding 내부 lambda - Python 인자 파싱 후 C++ 호출 (자동생성)
#38torch::autograd::THPVariable_matmul at generated/python_torch_functions_0.cpp:4911Python → C++ 진입점 - torch.matmul() 호출 시 CPython이 최초로 진입하는 C++ 함수 (자동생성)
CPython / libc: #39 ~ #54 (오늘은 스킵)
#Call Stack역할
#39cfunction_call at Objects/methodobject.c:540CPython이 C 확장 함수(PyCFunction)를 호출하는 진입점
#40_PyObject_MakeTpCall at Objects/call.c:242CPython 호출 프로토콜 - tp_call 슬롯을 통한 callable 객체 호출
#41_PyEval_EvalFrameDefault at Python/generated_cases.c.h:813CPython 바이트코드 인터프리터 - CALL 명령어 처리 중
#42PyEval_EvalCode at Python/ceval.c:596컴파일된 code object를 프레임에서 평가
#43-51run_eval_code_obj → Py_RunMain at Python/pythonrun.c ~ Modules/main.cPython 런타임 초기화 및 스크립트 파일 실행 체인
#52Py_BytesMain at Modules/main.c:829Python 프로세스 시작점 - python3 바이너리의 main 함수
#53-54__libc_start_main at libc_start_call_main.h:58libc 진입 - OS가 프로세스를 시작하는 최하위 레벨

PyTorch Eager Mode: High-Level Architecture

op 호출 하나가 거치는 Eager mode 구성 요소는 다음과 같습니다.

Front-End가 Dispatcher를 호출하고, 중간 kernel(예: Autograd)이 Dispatcher로 redispatch한 뒤, Device backend kernel이 Device runtime으로 연산을 수행한다 Front-End Dispatcher 중간 kernel (예: Autograd) Device backend kernel Device runtime ① dispatch ② redispatch ③ dispatch
  1. ① dispatch: Dispatcher가 key set에서 우선순위가 가장 높은 key의 kernel을 호출
  2. ② redispatch: 중간 kernel이 실행 도중, 처리한 key를 제외한 set으로 Dispatcher를 다시 호출
  3. ③ dispatch: 남은 key 중 backend key의 kernel을 호출. backend kernel은 끝단이라 되돌아가지 않음

다이어그램의 Front-End는 사용자가 호출하는 API 표면에서 Dispatcher 직전까지의 구간을 가리킵니다. torch.matmul(y, z)를 호출하면 Python API에서 시작해, 자동 생성된 Python binding(THPVariable_matmul, call stack #38)을 거쳐 op별 C++ 진입점(at::_ops::matmul::call, #35 — call stack 표에서 “C++ 진입점”으로 표시된 부분)에 도달합니다. PyTorch 문서의 “C++ Frontend”는 libtorch의 torch::nn API를 가리키는 별개 용어이므로 여기서는 쓰지 않습니다. Front-End는 Python 인자를 C++ 타입으로 변환해 op과 함께 Dispatcher에게 전달하는 데서 끝납니다. 어떤 kernel이 실행될지는 Front-End가 아니라 Dispatcher가 결정하며, Dispatcher는 인자에서 dispatch key set을 계산하고 그중 우선순위가 가장 높은 key를 선택하여 해당 kernel을 호출합니다.

앞 절에서 본 대로 kernel은 실행 도중에 다시 Dispatcher를 호출하여 redispatch합니다. 예를 들어 Autograd kernel이 backward graph 노드를 준비한 뒤 자신의 key를 제외하고 redispatch하면, Dispatcher가 남은 key 중 우선순위가 가장 높은 key(예: CUDA)의 device backend kernel을 선택합니다. 이렇게 Dispatcher와 kernel을 오가며 autograd, backend 등 여러 kernel이 한 op 호출에 차례로 개입합니다. device backend kernel은 Device Runtime을 사용해 실제 연산을 수행합니다. 중간 kernel은 Autograd 하나가 아니라 tracing, matmul → mm 분해(CompositeImplicitAutograd) 등 여러 층일 수 있으며, 각 층이 같은 방식으로 redispatch합니다. 중간 층이 하나도 없으면 Dispatcher가 바로 backend kernel을 호출합니다.

Call Stack의 주요 특성

자동 생성된 코드

Call stack의 상당 부분이 자동 생성된 코드입니다. 코드 생성은 빌드 타임에 일어나고, 생성된 코드는 사람이 작성한 코드와 함께 빌드되어 바이너리와 Python 패키지가 됩니다.

자동 생성된 코드 목록 (10개 항목)
#Call Stack역할
#5at:: at build/.../RegisterCUDA.cpp:12173CUDA backend dispatch 등록 - gen.py가 native_functions.yaml의 dispatch 항목으로부터 생성
#12-13at::_ops::mm::redispatch at build/.../Operators_3.cpp:4010mm op의 redispatch 진입점 - Autograd → CUDA 전환 시 사용
#14-15torch::autograd::VariableType::mm at generated/VariableType_3.cpp:13456mm의 Autograd kernel - gen_variable_type.py가 derivatives.yaml로부터 생성
#23at::_ops::mm::call at build/.../Operators_3.cpp:4003mm op의 C++ dispatch 진입점 - Dispatcher 호출
#27at:: at build/.../RegisterCompositeImplicitAutograd.cpp:2774matmul kernel 등록 - device 무관 composite kernel
#35at::_ops::matmul::call at build/.../Operators_4.cpp:3192matmul op의 C++ dispatch 진입점
#37-38torch::autograd::THPVariable_matmul at generated/python_torch_functions_0.cpp:4911Python → C++ binding - gen_python_functions.py가 native_functions.yaml로부터 생성

자동 생성 파일의 구조

코드 생성 도구(gen.py, gen_variable_type.py, gen_python_functions.py 등)는 native_functions.yaml과 derivatives.yaml을 입력으로 받아 다음 파일을 생성합니다.

build/aten/src/ATen ← gen.py (input: native_functions.yaml)

  • ops/{operator}.h - op별 public C++ API 함수(at::matmul 등) 선언. dispatcher 진입점 구조체(at::_ops::*)는 ops/{operator}_ops.h에 선언
  • core/TensorBody.h - at::Tensor 클래스의 메서드 선언
  • Operators.cpp - op별 call/redispatch 함수 구현
  • Register{backend}.cpp - backend별 kernel 등록 코드. PyTorch가 초기화될 때 이 파일들이 각 backend의 kernel을 Dispatcher에 등록합니다. {backend}에는 CUDA, CPU, XPU 등 다양한 종류가 들어갈 수 있습니다.

torch/include/torch/csrc/autograd/generated ← tools/autograd의 생성 스크립트들 (gen_autograd.py가 orchestration; input: derivatives.yaml, native_functions.yaml)

  • Functions.h, python_functions.h, python_return_types.h
  • variable_factories.h, VariableType.h, ViewFuncs.h

torch/csrc/autograd/generated

  • VariableType.cpp ← gen_variable_type.py (input: derivatives.yaml, native_functions.yaml)
  • Functions.cpp ← gen_autograd_functions.py (input: derivatives.yaml)
  • python_functions.cpp ← gen_autograd_functions.py (input: derivatives.yaml)
  • python_torch_functions.cpp ← gen_python_functions.py (input: native_functions.yaml)
핵심 입력 파일
  • native_functions.yaml - PyTorch가 기본으로 제공하는 op들의 스펙을 기술한 파일입니다. 각 op의 이름, 인자, 반환 타입, 지원하는 backend별 구현 함수 등이 정의되어 있고, ATen op의 정의와 등록이 이 파일에서 출발합니다. ATen을 위한 기술(description)입니다.
  • derivatives.yaml - 각 op의 autograd 미분 규칙을 정의한 파일입니다.

생성된 코드는 빌드 타임에 사람이 직접 작성한 코드와 함께 컴파일됩니다. PyTorch 프로세스가 시작되면 초기화 과정에서 생성된 등록 코드가 각 backend의 kernel을 Dispatcher에 등록(registration)합니다. 이 등록이 끝나야 Dispatcher가 런타임에 올바른 kernel을 찾아 호출합니다.

Dispatch Logic

Dispatch logic은 call stack에서 반복적으로 등장하는 template 기반 패턴입니다. Dispatcher는 operator의 tensor-like argument와 thread-local state에서 dispatch key set을 구성하고, 우선순위가 가장 높은 key에 등록된 kernel을 조회해 호출합니다.

Dispatch logic 반복 패턴 (6회)
#Call Stack역할
#6-9WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hCUDA kernel dispatch - CUDA backend로 등록된 mm kernel을 functor로 감싸서 실행
#10-13Dispatcher::redispatch → mm::redispatch at Dispatcher.h → Operators_3.cppRedispatch - Autograd kernel이 실행 도중 다음 dispatch key(CUDA)로 mm을 재전달
#16-19WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hAutograd kernel dispatch - Autograd key로 등록된 mm의 VariableType kernel 실행
#20-23Dispatcher::call → mm::call → Tensor::mm at Dispatcher.h → Operators_3.cppmm dispatch 진입 - dispatch key set에서 최우선 key(Autograd) 선택 후 kernel 호출
#28-31WrapFunctionIntoFunctor → KernelFunction::call at boxing/*.hmatmul kernel dispatch - CompositeImplicitAutograd key로 등록된 matmul kernel 실행
#32-35Dispatcher::call → matmul::call → Tensor::matmul at Dispatcher.h → Operators_4.cppmatmul dispatch 진입 - runtime key(예: AutogradCUDA) slot에 등록된 kernel 호출. matmul은 CompositeImplicitAutograd alias key로 등록되어 이 kernel이 여러 slot을 채움

matmul 호출 한 번에 dispatch가 세 번 일어납니다.

PyTorch Dispatcher

Dispatcher는 입력 tensor 등을 보고 dispatch table(Autograd, Tracing, XLA, CUDA, CPU)에서 function pointer 하나를 고른다. vtable과 달리 runtime 확장, multiple dispatch, TLS, boxing/unboxing을 지원한다

Dispatcher는 C++의 virtual function table(V-table)과 비슷한 개념이지만(아래 설명은 Edward Z. Yang의 2020년 블로그 글을 따릅니다), 훨씬 더 복잡한 요구사항을 다룹니다. 같은 matmul이라도 CPU에서 실행될 때, CUDA에서 실행될 때, autograd가 켜져 있을 때 등 서로 다른 컨텍스트에서 다르게 동작해야 하는 모든 경우를 처리하기 위해, 별도의 데이터 구조와 알고리즘으로 설계되었습니다.

Dispatch table에는 각 key에 해당하는 function pointer가 들어 있습니다. 한 호출에서 여러 key가 동시에 유효할 수 있고, 이 key들의 집합이 dispatch key set입니다. Dispatcher는 key set에서 우선순위가 가장 높은 key의 kernel을 호출하고, 호출된 kernel은 실행 도중 자기 key를 제외한 set으로 redispatch합니다(앞 절 참조).

상세 설명: http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/

Call Stack을 단계별로 구분

개발자가 직접 작성한 부분

55단계 중 상당수는 자동 생성된 wrapper와 반복되는 dispatcher logic입니다. 이를 빼면 개발자가 직접 작성한 코드는 두 곳입니다. matmul이 input shape에 따라 연산을 선택하는 코드(#25-26)와, CUDA에서 mm operator를 실행하는 device-specific 구현(#0-4)입니다.

개발자가 직접 작성한 코드 vs 자동 생성/dispatch 코드
#Call Stack구분
#0at::cuda::blas::gemm<float> at CUDABlas.cpp:1127CUDA로 구현된 mm - cuBLAS GEMM 호출
#1-3operator at native/cuda/Blas.cpp:457CUDA mm 내부 구현
#4at::native::structured_mm_out_cuda::impl at native/cuda/Blas.cpp:604CUDA mm 구현체
#5at:: at build/.../RegisterCUDA.cpp:12173(자동생성)
#6-9dispatch logic
#10-13redispatch logic
#14-15VariableType (autograd)(자동생성)
#16-19dispatch logic
#20-24Dispatcher::call → Tensor::mm
#25at::native::_matmul_impl at LinearAlgebra.cpp:2031matmul 구현체 - input shape에 따라 mm, bmm 등으로 분기
#26at::native::matmul at LinearAlgebra.cpp:2181matmul 진입점
#27at:: at RegisterCompositeImplicitAutograd.cpp:2774(자동생성)
#28-35dispatch logic
#36at::Tensor::matmul
#37-38THPVariable_matmul(자동생성, Python binding)
#39-54CPython / libc

전체 과정 요약

Call Stack 전체 과정: 8단계
단계Call Stack역할
1#54-52: __libc_start_main → Py_BytesMainlibc가 Python 프로세스를 실행
2#51-39: Py_RunMain → cfunction_callPython이 PyTorch를 Python 수준에서 처리
3#38-36: THPVariable_matmul → Tensor::matmulPython에서 C++ code로 넘어가는 과정 - 자동 생성된 Python binding을 통해 진입
4#35-27: matmul::call → RegisterCompositeImplicitAutogradmatmul kernel을 dispatch하는 과정 - Dispatcher가 runtime key slot에 채워진 CompositeImplicitAutograd kernel 선택
5#26-24: matmul → _matmul_impl → Tensor::mmmatmul(mm)을 처리하는 과정 - input shape 분석 후 2D×2D이므로 mm을 호출
6#23-15: mm::call → VariableType::mmmm의 autograd kernel을 dispatch하여 수행하는 과정 - backward graph 세팅 후 redispatch
7#14-5: mm::redispatch → RegisterCUDAmm의 CUDA kernel을 dispatch하여 수행하는 과정 - Autograd kernel이 실행 도중 CUDA key로 redispatch
8#4-0: structured_mm_out_cuda::impl → gemm<float>cuBLAS로 진입하는 과정 - 최종적으로 GPU에서 행렬곱 수행

Python에서 C++ code로 넘어가는 과정

#36 at::Tensor::matmul
    at build/aten/src/ATen/core/TensorBody.h:2899
#37 operator
    at torch/csrc/autograd/generated/python_torch_functions_0.cpp:4909
#38 torch::autograd::THPVariable_matmul
    at torch/csrc/autograd/generated/python_torch_functions_0.cpp:4911
#39 cfunction_call
    at .../Objects/methodobject.c:540
#40 _PyObject_MakeTpCall
    at .../Objects/call.c:242

PyTorch는 성능과 확장성을 위해 Python binding을 자체 코드 생성으로 만들고, dispatcher로 CPU, CUDA, XPU 같은 여러 backend를 지원합니다.

generate_code.py가 native_functions.yaml을 읽고 Python/C++ 간 인터페이스 코드를 생성합니다:

  • python_torch_functions.cpp: torch 함수의 정의 (예: torch.matmul은 THPVariable_matmul로 연결)
  • python_variable_methods.cpp: torch.Tensor에 대한 method 정의

참고: https://discuss.pytorch.org/t/how-are-python-bindings-created/46453

native_functions.yaml

native_functions.yaml은 사람이 직접 작성하는 파일로, ATen library의 operator를 정의하고 등록하는 전체 딕셔너리 역할을 합니다. GitHub 저장소나 PyTorch 소스를 빌드하면 확인할 수 있습니다. 이 파일은 dispatcher와 밀접하게 연관되어 backend별 kernel과 autograd 관련 기능들을 설정할 수 있으며, Python 인터페이스와 C++ 구현도 이 파일을 기준으로 연결됩니다. generate_code.py 같은 코드 생성 도구들이 이 파일의 정보를 바탕으로 실제 코드를 생성합니다.

- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor
  device_check: NoCheck         # device 검사를 비활성화
  structured_delegate: add.out  # add.out의 정의를 상속
  variants: function, method    # namespace 함수 + Tensor 메서드
  dispatch:
    SparseCPU, SparseCUDA, SparseMPS, SparseMeta, SparseXPU: add_sparse
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta, SparseCsrXPU: add_sparse_csr
    MkldnnCPU: mkldnn_add
    ZeroTensor: add_zerotensor
    NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA, NestedTensorXPU: NestedTensor_add_Tensor
  tags: [core, pointwise]
  • func: add 함수의 signature 정의 - 이름, 인자, 반환 타입을 지정
  • device_check: NoCheck: device 검사를 비활성화 - 전달받은 모든 tensor가 동일한 device에 있는지 검사하지 않도록 함
  • structured_delegate: add.out: add.out의 정의를 상속받음
  • variants: function (namespace 안에서 호출, 예: torch.add()) / method (Tensor의 메서드로 호출, 예: a.add())
  • dispatch: dispatch key 또는 backend별 impl 함수 지정 - Sparse, MKL-DNN, NestedTensor 등 각 backend에 맞는 구현 함수를 매핑

matmul kernel을 dispatch하는 과정

#35 at::_ops::matmul::call           ← Op의 C++ 진입점
#34 c10::TypedOperatorHandle<...>
#33 c10::Dispatcher::call<...>       ← Dispatch logic
#32 c10::Dispatcher::callWithDispatchKeySlowPath<...>
#31 c10::KernelFunction::call<...>
#30 c10::callUnboxedKernelFunction<...>
#29 c10::impl::wrap_kernel_functor_unboxed_<...>
#28 c10::impl::detail::WrapFunctionIntoFunctor_<...>
#27 at:: at RegisterCompositeImplicitAutograd.cpp  ← matmul kernel
#26 at::native::matmul

먼저 Op의 C++ 진입점(at::_ops::matmul::call, #35)이 Dispatcher를 호출합니다. Dispatcher가 인자에서 dispatch key set을 계산해 dispatch table에서 해당 kernel을 찾아 호출합니다. 이 과정을 거쳐 실제 matmul kernel이 실행됩니다.

#27에서 호출되는 matmul kernel은 RegisterCompositeImplicitAutograd.cpp에 등록된 kernel입니다.

CompositeImplicitAutograd란?

matmul은 input tensor의 shape에 따라 dot, mv, mm, bmm 등 미분 가능한 함수들의 시퀀스로 표현될 수 있습니다. 이 경우 matmul 자체에 대한 미분 규칙을 별도로 정의할 필요가 없습니다. 내부에서 호출되는 각 함수들이 이미 autograd kernel을 갖고 있으므로, 전체 미분은 이들의 합성으로 자동 처리됩니다.

이러한 방식을 CompositeImplicitAutograd라고 합니다. CompositeImplicitAutograd는 alias key라서 runtime에 key set에서 직접 선택되지 않습니다. 등록 시 dispatcher가 이 kernel을 CPU, CUDA, AutogradCPU, AutogradCUDA 같은 여러 runtime key slot에 채워 넣습니다. 그래서 실제 dispatch는 예컨대 AutogradCUDA key로 일어나고, 실행되는 함수는 같은 composite kernel입니다. 이 kernel은 forward 계산을 그대로 수행하며, autograd 전용 kernel은 따로 없습니다. Forward 실행 중 내부에서 호출되는 개별 op들이 각자의 autograd kernel로 backward graph를 구성하므로, 별도의 미분 정의 없이 backward가 계산됩니다. 등록 코드는 코드 생성 도구가 만든 RegisterCompositeImplicitAutograd.cpp에 있습니다.

matmul을 처리하는 과정

matmul operator는 aten/src/ATen/native/LinearAlgebra.cpp에 구현되어 있습니다. matmul은 input tensor의 차원에 따라 수행할 연산을 고릅니다. 이번 예제는 2D × 2D이므로 mm을 호출해 2D 행렬 곱셈을 dispatch합니다.

#24 at::Tensor::mm at build/aten/src/ATen/core/TensorBody.h:2999
#25 at::native::_matmul_impl at aten/src/ATen/native/LinearAlgebra.cpp:2031  ← mm 호출
#26 at::native::matmul at aten/src/ATen/native/LinearAlgebra.cpp:2181

torch.matmul은 input tensor의 차원에 따라 다른 연산을 수행합니다:

  • If both tensors are 1-dimensional, the dot product (scalar) is returned.
  • If both arguments are 2-dimensional, the matrix-matrix product is returned. (이번 예제에 해당)
  • If the first argument is 1-dimensional and the second argument is 2-dimensional, a 1 is prepended to its dimension for the purpose of the matrix multiply. After the matrix multiply, the prepended dimension is removed.
  • If the first argument is 2-dimensional and the second argument is 1-dimensional, the matrix-vector product is returned.
  • If both arguments are at least 1-dimensional and at least one argument is N-dimensional (where N > 2), then a batched matrix multiply is returned. If the first argument is 1-dimensional, a 1 is prepended to its dimension for the purpose of the batched matrix multiply and removed after. If the second argument is 1-dimensional, a 1 is appended to its dimension for the purpose of the batched matrix multiply and removed after. The first N-2 dimensions of each argument, the batch dimensions, are broadcast (and thus must be broadcastable).

출처: torch.matmul - PyTorch Documentation

mm의 autograd kernel을 dispatch

Dispatch logic이 한 번 더 수행되어 mm을 처리할 kernel을 고르고, autograd kernel(#15)이 호출됩니다. 이 kernel은 자동 생성된 코드로, 빌드하면 torch/csrc/autograd/generated/VariableType_3.cpp에 들어 있습니다. 숫자 3은 생성 코드가 길어 여러 파일로 나눈 결과입니다. 이 dispatch는 앞의 matmul dispatch와 같은 방식으로 일어납니다.

#15 torch::autograd::VariableType:: at .../VariableType_3.cpp:13456  ← Autograd kernel
#16-19 (dispatch logic: WrapFunctionIntoFunctor → KernelFunction::call)
#20    c10::Dispatcher::callWithDispatchKeySlowPath<...>
#21    c10::Dispatcher::call<...>
#22    c10::TypedOperatorHandle<...>
#23    at::_ops::mm::call at build/.../Operators_3.cpp:4003
#24    at::Tensor::mm at build/.../TensorBody.h:2999

Autograd kernel이란?

GPU에서 연산을 실행하면 CUDA kernel이 사용되고, 학습을 위해 함수가 호출되면 동적 계산 그래프(dynamic computation graph)가 생성되어야 합니다. 이 두 기능(CUDA 실행과 autograd)은 dispatch table에서 적절히 연결됩니다. Dispatch key set이 생성되면 우선순위가 가장 높은 key가 선택되어 실행되는데, autograd key가 CUDA key보다 우선순위가 높으므로 autograd kernel이 먼저 실행됩니다.

Autograd kernel은 실제 backward 계산을 수행하지 않습니다. 대신 forward 계산을 하면서 동적 계산 그래프를 업데이트하는 역할을 합니다. 구체적으로는:

  • compute_requires_grad() 같은 조건을 판별하여, 필요한 경우 backward 계산을 위한 그래프 노드를 추가
  • Forward 함수 실행 전후에 sanity check 수행
  • 실제 연산은 dispatch table에서 다음 key(예: CUDA)로 redispatch하여 수행

이는 PyTorch의 define-by-run 방식과 관련이 있습니다. 그래프를 미리 정의하고 실행하는 define-and-run 방식과 달리, PyTorch는 eager mode에서 연산을 하나하나 즉시 실행합니다. 이 과정에서 각 연산이 수행될 때마다 autograd kernel이 동적 계산 그래프를 자동으로 생성하고, 후속 backward 계산을 위한 히스토리를 기록합니다.

Autograd kernel의 실제 생성 코드는 아래 “Operator > Kernel” 절에서 v2.14.0 codegen 산출물로 살펴봅니다. gradient가 필요한 경우 backward graph 노드를 준비한 뒤 forward 연산을 redispatch하고, 반환된 결과에 backward history를 기록하는 구조입니다.

mm의 CUDA kernel을 dispatch

#4  at::native::structured_mm_out_cuda::impl
    at aten/src/ATen/native/cuda/Blas.cpp:604
#5  at:: at build/aten/src/ATen/RegisterCUDA.cpp:12173
#6-9  (dispatch logic)
#10 c10::Dispatcher::redispatch<...>   ← Redispatch
#11 c10::TypedOperatorHandle<...>
#12 at::_ops::mm::redispatch
#13 at::redispatch::mm
#14-15 (from autograd VariableType)

이 구간의 용어

  • Redispatch: Autograd kernel이 실행 도중 자기 key를 제외하고 다음 dispatch key로 재전달
  • RegisterCUDA.cpp: native_functions.yaml에서 생성된 CUDA backend 등록 코드

RegisterCUDA.cpp와 관련된 native_functions.yaml

RegisterCUDA.cpp에 등록되는 CUDA kernel이 native_functions.yaml에서 어떻게 정의되어 있는지를 보면, mm op의 구조를 이해할 수 있습니다. mm은 mm.out으로 delegate하고, mm.out은 structured: True로 정의되어 있어 structured kernel 방식으로 구현됩니다. dispatch 항목에서 CUDA: mm_out_cuda로 지정된 부분이 코드 생성 도구에 의해 RegisterCUDA.cpp에 등록되는 CUDA backend kernel입니다.

- func: mm(Tensor self, Tensor mat2) -> Tensor
  structured_delegate: mm.out     # mm.out의 정의를 상속
  variants: function, method
  dispatch:
    SparseCPU, SparseCUDA, SparseMPS, SparseXPU: _sparse_mm
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta, SparseCsrXPU: _sparse_csr_mm
  tags: core

- func: mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!)
  structured: True                # ← structured kernel 방식으로 구현
  dispatch:
    CPU: mm_out_cpu
    CUDA: mm_out_cuda             # ← 이것이 RegisterCUDA.cpp에 등록됨
    MTIA: mm_out_mtia
    MPS: mm_out_mps
    XPU: mm_out_xpu
    SparseCPU, SparseCUDA, SparseMPS, SparseXPU: _sparse_mm_out
    SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta, SparseCsrXPU: _sparse_csr_mm_out

cuBLAS로 진입하는 과정

Call stack의 최하위에서 cuBLAS 호출이 발생합니다. 여기서 structured는 native_functions.yaml의 structured: True 태그를 의미합니다. Structured kernel은 입력 검증, output tensor 준비, 실제 kernel 구현의 역할을 일정한 형태로 분리합니다. 입력 검증과 output shape 결정은 사람이 TORCH_META_FUNC(mm)에 작성하고(LinearAlgebra.cpp), 이를 output 할당·kernel 호출과 엮는 wrapper class(structured_mm_out_cuda)는 code generator가 생성합니다. 개발자는 meta 함수와 정규화된 tensor를 받는 TORCH_IMPL_FUNC 구현만 작성합니다.

#0 at::cuda::blas::gemm<float>
   at aten/src/ATen/cuda/CUDABlas.cpp:1127
...
#4 at::native::structured_mm_out_cuda::impl
   at aten/src/ATen/native/cuda/Blas.cpp:604 # ← structured kernel

#4의 structured_mm_out_cuda::impl이 CUDA mm의 structured kernel 구현체입니다. 이 함수에서 cuBLAS 호출까지 내려가는 과정을 코드로 따라가면 다음과 같습니다:

1단계: TORCH_IMPL_FUNC 매크로. structured kernel의 진입점. mm은 bias가 없는 addmm의 특수한 경우이므로, beta=0으로 addmm_out_cuda_impl에 위임합니다.

// aten/src/ATen/native/cuda/Blas.cpp
TORCH_IMPL_FUNC(mm_out_cuda)(const Tensor& self, const Tensor& mat2,
    const Tensor& result) {
    // NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
    addmm_out_cuda_impl(const_cast<Tensor&>(result), result, self, mat2, 0, 1);
}

2단계: addmm_out_cuda_impl. 입력 검증 후 cuBLAS 호출 준비

Tensor& addmm_out_cuda_impl(Tensor& result, const Tensor& self,
    const Tensor& mat1, const Tensor& mat2, const Scalar& beta,
    const Scalar& alpha, Activation activation=Activation::None,
    bool disable_addmm_cuda_lt_override=false) {
  // Shape checks {
  // Make sure to keep addmm_cuda below in sync with this code; it
  // preflights a check to try to avoid actually needing to call
  // expand().
  TORCH_CHECK(mat1.dim() == 2 && mat2.dim() == 2, "tensors must be 2-D");
  TORCH_CHECK(
    mat1.dtype() == mat2.dtype(),
    "expected mat1 and mat2 to have the same dtype, but got: ",
    mat1.dtype(), " != ", mat2.dtype()
  )
  ...
}

3단계: cuBLAS GEMM 호출. addmm_out_cuda_impl은 조건에 따라 일반 cuBLAS 경로인 at::cuda::blas::gemm(위 call stack의 #0이 이 경로) 또는 cublasLt 기반의 gemm_and_bias를 호출합니다. 후자는 launchGemmAndBiasCublasLt가 감싸며, TunableOp이 켜져 있으면 그쪽으로 먼저 분기합니다.

// aten/src/ATen/native/cuda/Blas.cpp (v2.14.0, 341-380행)
template <typename scalar_t, typename res_scalar_t = scalar_t>
bool launchGemmAndBiasCublasLt(
    // args contains result which is modified
    cublasCommonArgs& args,
    const std::optional<Tensor>& self,
    const Scalar& alpha,
    Activation activation = Activation::None
) {
  // self_ptr == nullptr implies ignore bias epilogue
  // and use standard gemm-like API.
  const auto* self_ptr = self.has_value() ? self.value().const_data_ptr<scalar_t>() : static_cast<const scalar_t*>(nullptr);

  const auto tuning_ctx = at::cuda::tunable::getTuningContext();
  if (tuning_ctx->IsTunableOpEnabled()) {
    launchTunableGemmAndBias<scalar_t>(
      args, alpha, self_ptr, activation_to_gemm_and_blas_arg(activation)
    );
    return true;
  }

  return at::cuda::blas::gemm_and_bias<scalar_t, res_scalar_t>(
    args.transa == 't',
    args.transb == 't',
    args.m,
    args.n,
    args.k,
    alpha.to<at::opmath_type<scalar_t>>(),
    args.mata->const_data_ptr<scalar_t>(),
    args.lda,
    args.matb->const_data_ptr<scalar_t>(),
    args.ldb,
    self_ptr,
    args.result->data_ptr<res_scalar_t>(),
    args.result_ld,
    activation_to_gemm_and_blas_arg(activation)
  );
}

cuBLAS 경로는 cublasSgemm/cublasGemmEx를, cuBLASLt 경로는 cublasLtMatmul을 호출합니다.

참고: https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cuda/Blas.cpp

Eager Mode에 대한 Detail

Eager mode는 Tensor, Operator, Runtime 세 요소로 나눌 수 있습니다. 앞의 torch.matmul 예제에서 등장한 개념을 이 순서로 정리합니다.

Tensor

Tensor는 Python과 C++ 양쪽에서 같은 객체로 다뤄져야 하고, GPU 등 여러 device의 memory 할당도 고려해야 합니다. 이 절은 operator의 입출력인 tensor의 내부 구성을 다룹니다.

Tensor의 구조

data

grad

intrusive_ptr

at::Tensor

c10::TensorImpl

c10::Storage

metadata

AutogradMeta

at::Tensor (gradient)

c10::StorageImpl

참고: PyTorch internals (Edward Z. Yang)

PyTorch 내부에서 대부분의 구현은 C++로 되어 있습니다. Python의 torch.Tensor는 내부적으로 at::Tensor라는 C++ 객체에 대응되며, 이 객체는 c10::TensorImpl이라는 구현체(impl)를 통해 대부분의 정보를 관리합니다.

  • at::Tensor - c10::TensorImpl을 intrusive_ptr4로 가리키는 얇은 wrapper
  • c10::TensorImpl - 실제 구현체로, 핵심 정보를 가짐:
    • c10::Storage: 실제 data가 저장되는 공간.
      • 예를 들어 GPU에 메모리가 할당되면, 그 메모리 덩어리가 이 StorageImpl과 1:1로 매핑됨.
      • intrusive_ptr로 관리되므로 여러 tensor가 같은 storage를 공유할 수 있음. view는 이 공유를 이용함
      • view를 생성하면 새 storage가 아닌 기존 storage를 공유하면서 metadata만 다른 tensor가 만들어짐
    • metadata - sizes, strides, storage offset, dtype, device. Stride를 이용해 tensor[i][j]를 data_ptr[i * stride[0] + j * stride[1]]로 접근
    • AutogradMeta - gradient, grad_fn 등 autograd 관련 bookkeeping 정보를 tensor 단위로 저장. grad가 필요 없는 tensor는 null로 두는 최적화가 있으나, requires_grad가 True였다가 False로 돌아온 경우처럼 default 상태의 AutogradMeta가 남아 있을 수도 있어 “항상 null”은 아님(TensorImpl.h 주석)

Tensor의 Lifecycle

Tensor의 lifecycle은 생성, 소멸, 복사의 세 가지 과정으로 나눌 수 있습니다.

생성 - scratch로 새로 만드는 경우:

  • aten/src/ATen/native/TensorFactories.cpp에 구현된 함수들을 통해 생성
    • empty(), zeros(), ones() 등

생성 - 기존 tensor로부터 파생되는 경우:

기존 tensor로부터 파생될 때는 별도의 storage를 갖는 경우와 기존 tensor의 storage를 공유하는 경우로 나뉩니다. Storage를 공유하는 쪽이 view입니다. 새 storage를 만들지 않고 intrusive_ptr로 reference만 공유하므로 reference count가 증가합니다.

Storage를 새로 만드는 경우

  • at::native::clone()(TensorFactories.cpp): empty_strided(또는 empty_like)로 새 storage를 잡은 뒤 copy_()로 data 복제
  • 별도의 c10::Storage를 가진 at::Tensor
  • 산술 연산(add, matmul 등)의 결과 tensor도 같은 방식으로 새 storage를 가짐

Storage를 기존 tensor와 공유하는 경우 (view)

  • aten/src/ATen/native/TensorShape.cpp에 구현된 함수로 view 생성
  • at::native::alias_with_sizes_and_strides() 등
  • view(), transpose() 같은 shape 조작 op이 이 경우에 해당
  • at::Tensor는 새로 만들지만 c10::Storage는 intrusive_ptr로 공유

소멸:

Tensor가 소멸된다고 해서 storage까지 바로 소멸되지는 않습니다.

at::Tensor의 소멸:

  • Python에서 torch.Tensor 객체가 제거되는 시점에 소멸
    • 명시적으로 객체를 제거 (del keyword 사용)
    • 또는 garbage collection에 의해 자연스럽게 소멸
  • Python 측 객체가 소멸되면 C++ at::Tensor와 내부 포인터들도 함께 해제
  • THPVariable_dealloc이 entry point (구버전에서는 THPVariable_subclass_dealloc)

c10::Storage의 소멸:

  • intrusive_ptr의 reference count가 감소하지만, 다른 tensor(view 등)가 공유 중이면 바로 소멸되지 않음
  • Reference count가 0이 되는 시점에 StorageImpl이 소멸하고 memory는 allocator에 반환됨. CUDA의 caching allocator는 이때 cudaFree를 부르지 않고 block을 pool에 보관해 재사용하므로, driver로의 실제 반환은 torch.cuda.empty_cache() 시점

Tensor를 가지고 할 수 있는 일들

Tensor로 할 수 있는 일들은 모두 operator(op)에 해당하며, op의 동작 방식에 따라 다음과 같이 분류할 수 있습니다.

분류예시
operator 수행 결과를 tensor에 직접 반영 → in-place operatoradd_(), transpose_()
operator 수행 결과로 새로운 tensor가 생성실제 연산이 일어나는 경우 → 새로운 storage가 생김matmul(), add(), 산술 연산
기존 data의 재배치만 일어나는 경우 - view만 생기는 경우view(), transpose()
기존 data의 재배치만 일어나는 경우 - 새로운 storage가 생기는 경우clone(), non-contiguous tensor의 contiguous()

단, 경계에 있는 op들도 있습니다. reshape()은 가능하면 view를 반환하지만 view로 표현할 수 없는 경우(예: non-contiguous tensor) copy를 만들고, contiguous()는 이미 contiguous한 tensor에 대해서는 복사 없이 자기 자신을 그대로 반환합니다.

In-place operator

>>> a = torch.tensor([[1., 2.], [3., 4.]])
>>> b = torch.tensor([[5., 6.], [7., 8.]])
>>> a.add_(b)
tensor([[ 6.,  8.], [10., 12.]])
>>> a.transpose_(0, 1)
tensor([[ 6., 10.], [ 8., 12.]])

새로운 storage 생성 (연산)

>>> c = torch.matmul(a, b)
>>> print(c)
tensor([[100., 116.], [124., 144.]])
>>> d = torch.add(c, 10)
>>> print(d)
tensor([[110., 126.], [134., 154.]])

View만 생성 (storage 공유)

>>> e = torch.reshape(d, (4, 1))
>>> print(e)
tensor([[110.], [126.], [134.], [154.]])
>>> f = torch.transpose(e, 0, 1)
>>> print(f)
tensor([[110., 126., 134., 154.]])

새로운 storage 생성 (재배치)

>>> g = d.clone()      # 새 storage에 복제
>>> t = d.t()          # non-contiguous view
>>> h = t.contiguous() # 새 storage에 재배치
>>> t.is_contiguous(), h.is_contiguous()
(False, True)

View와 Storage 분리

View는 같은 data를 해석만 달리해 변환 시간과 memory를 아낍니다:

  • 같은 data를 다른 shape으로 해석 (reshape)
  • 특정 차원을 반복해 접근 (broadcast)

TensorIterator는 elementwise op용 helper로, broadcast와 stride 처리를 대신해 주므로 kernel은 원소 단위 lambda만 쓰면 됩니다:

at::Tensor my_add_cuda(const at::Tensor& ta, const at::Tensor& tb) {
    auto result = at::empty_like(ta);
    auto iter = at::TensorIteratorConfig()
        .add_output(result)
        .add_const_input(ta)
        .add_const_input(tb)
        .build();
    at::native::gpu_kernel(iter, []GPU_LAMBDA(float a, float b) {
        return a + b;
    });
    return result;
}

view 계열 op은 metadata(size, stride, offset)만 바꾸고 storage를 공유합니다. 단 reshape와 flatten은 view로 표현할 수 없으면(예: transposed tensor) 새 storage에 복사합니다.

Operator무엇을 바꾸는가Storage예시 (x3 = torch.arange(24).reshape(2, 3, 4))
view()shape항상 공유 (불가능하면 에러)x3.view(6, 4)
reshape()shape가능하면 공유, 아니면 복사x3.reshape(2, 12)
flatten()차원 병합가능하면 공유, 아니면 복사x3.flatten(1, 2)
transpose() / permute()차원 순서항상 공유x3.permute(2, 0, 1)
squeeze() / unsqueeze()크기 1 차원 제거·추가항상 공유x3.unsqueeze(0)
expand() / broadcast_to()크기 1 차원을 stride 0으로 확장항상 공유x3[:, :1, :].expand(2, 3, 4)
as_strided()size와 stride 직접 지정항상 공유torch.as_strided(x3, (2, 2), (1, 2))

Operator

Operator 절에서 다루는 개념은 다음과 같습니다:

  • Kernel - 실제 연산을 수행하는 구현체
  • Dispatcher - op과 컨텍스트에 맞는 kernel을 찾아 호출하는 메커니즘
  • Dispatch key - 어떤 kernel을 선택할지 결정하는 키
  • Dispatch key set - 여러 dispatch key의 집합
  • Registration - kernel을 dispatcher에 등록하는 과정

참고: http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/

Kernel

Generated Autograd Kernel for add (PyTorch v2.14.0 codegen 실행 결과)

아래 코드는 v2.14.0 소스에서 autograd codegen(python -m tools.autograd.gen_autograd)을 직접 실행해 얻은 실제 생성 코드입니다. 생성 파일은 VariableTypeEverything.cpp가 아니라 10개로 샤딩된 VariableType_0.cpp ~ VariableType_9.cpp 중 VariableType_2.cpp이며, 같은 파일의 등록부에서 m.impl("add.Tensor", TORCH_FN(VariableType::add_Tensor));로 dispatcher에 연결됩니다.

// In VariableType_2.cpp (v2.14.0 codegen 산출물, 일부 블록 생략)

at::Tensor add_Tensor(c10::DispatchKeySet ks, const at::Tensor & self, const at::Tensor & other, const at::Scalar & alpha) {
  auto& self_ = unpack(self, "self", 0);
  auto& other_ = unpack(other, "other", 1);
  [[maybe_unused]] auto _any_requires_grad = compute_requires_grad( self, other );

  [[maybe_unused]] auto _any_has_forward_grad_result = (isFwGradDefined(self) || isFwGradDefined(other));
  c10::intrusive_ptr<AddBackward0> grad_fn;
  if (_any_requires_grad) {
    grad_fn = c10::make_intrusive<AddBackward0>();
    grad_fn->set_next_edges(collect_next_edges( self, other ));
    grad_fn->alpha = alpha;
    grad_fn->other_scalar_type = other.scalar_type();
    grad_fn->self_scalar_type = self.scalar_type();
  }
  #ifndef NDEBUG
  auto self__storage_saved =
    self_.has_storage() ? ::std::optional<Storage>(self_.storage()) : ::std::nullopt;
  c10::intrusive_ptr<TensorImpl> self__impl_saved;
  if (self_.defined()) self__impl_saved = self_.getIntrusivePtr();
  auto other__storage_saved =
    other_.has_storage() ? ::std::optional<Storage>(other_.storage()) : ::std::nullopt;
  c10::intrusive_ptr<TensorImpl> other__impl_saved;
  if (other_.defined()) other__impl_saved = other_.getIntrusivePtr();
  #endif
  auto _tmp = ([&]() {
    at::AutoDispatchBelowADInplaceOrView guard;
    return at::redispatch::add(ks & c10::after_autograd_keyset, self_, other_, alpha);
  })();
  auto result = std::move(_tmp);
  #ifndef NDEBUG
  // storage/impl sanity check — TORCH_INTERNAL_ASSERT 기반 검증 (생략)
  #endif
  if (grad_fn) {
      set_history(flatten_tensor_args( result ), grad_fn);
  }
  // forward-mode AD (fw_grad) 처리 — _any_has_forward_grad_result인 경우
  // tangent를 계산해 result._set_fw_grad(...)로 기록 (생략)
  if (grad_fn) {
      fire_node_creation_hooks(grad_fn);   // 2.14에서 추가
  }
  return result;
}
참고: old style 생성 코드 (PyTorch 1.x, VariableTypeEverything.cpp 시절)
// In VariableTypeEverything.cpp

Tensor add_Tensor(const Tensor & self, const Tensor & other, Scalar alpha) {
    auto& self_ = unpack(self, "self", 0);
    auto& other_ = unpack(other, "other", 1);
    std::shared_ptr<AddBackward0> grad_fn;
    if (compute_requires_grad( self, other )) {
        grad_fn = std::shared_ptr<AddBackward0>(new AddBackward0(), deleteNode);
        grad_fn->set_next_edges(collect_next_edges( self, other ));
        grad_fn->alpha = alpha;
    }
    #ifndef NDEBUG
    c10::optional<Storage> self__storage_saved =
        self_.has_storage() ? c10::optional<Storage>(self_.storage()) : c10::nullopt;
    c10::intrusive_ptr<TensorImpl> self__impl_saved;
    if (self_.defined()) self__impl_saved = self_.getIntrusivePtr();
    c10::optional<Storage> other__storage_saved =
        other_.has_storage() ? c10::optional<Storage>(other_.storage()) : c10::nullopt;
    c10::intrusive_ptr<TensorImpl> other__impl_saved;
    if (other_.defined()) other__impl_saved = other_.getIntrusivePtr();
    #endif
    auto tmp = ([&]() {
        at::AutoNonVariableTypeMode non_var_type_mode(true);
        return at::add(self_, other_, alpha);
    })();
    auto result = std::move(tmp);
    #ifndef NDEBUG
    if (self__storage_saved.has_value())
        AT_ASSERT(self__storage_saved.value().is_alias_of(self_.storage()));
    if (self__impl_saved) AT_ASSERT(self__impl_saved == self_.getIntrusivePtr());
    if (other__storage_saved.has_value())
        AT_ASSERT(other__storage_saved.value().is_alias_of(other_.storage()));
    if (other__impl_saved) AT_ASSERT(other__impl_saved == other_.getIntrusivePtr());
    #endif
    if (grad_fn) {
        set_history(flatten_tensor_args( result ), grad_fn);
    }
    return result;
}

위 참고의 old style 코드와 비교하면 구조(그래프 노드 생성 → guard 아래에서 forward 실행 → history 기록)는 동일하지만, 세부가 달라졌습니다:

  • Guard: at::AutoNonVariableTypeMode → at::AutoDispatchBelowADInplaceOrView (1.10에서 deprecated 후 교체)
  • Forward 호출: at::add(...) → at::redispatch::add(ks & c10::after_autograd_keyset, ...) — dispatch key set을 명시적으로 전달하며 자신(autograd)의 key를 제외하고 redispatch
  • grad_fn이 std::shared_ptr → c10::intrusive_ptr, 검증이 AT_ASSERT → TORCH_INTERNAL_ASSERT, c10::optional → ::std::optional
  • Forward-mode AD 지원 코드가 추가됨
  • 2.14: set_history 뒤에 fire_node_creation_hooks(grad_fn) 호출이 추가됨. 새 autograd node가 만들어질 때 등록된 hook을 부르는 경로로, 2.13 codegen에는 없음

Registration API

TORCH_LIBRARY에서 m.def로 neg schema를 정의하고, TORCH_LIBRARY_IMPL(aten, CPU)에서 m.impl로 CPU kernel을, TORCH_LIBRARY_IMPL(_, Tracer)에서 m.fallback으로 tracing handler를 등록하는 코드

Dispatch table의 함수 포인터가 어떻게 등록되는지를 담당하는 것이 operator registration API입니다. 이 API와 상호작용하는 세 가지 주요 방법이 있습니다:

  1. def - operator의 schema를 정의
  2. impl - 특정 dispatch key에 대한 구현(kernel)을 등록
  3. fallback - 특정 dispatch key의 모든 operator에 대한 기본 handler를 등록
  • Kernel을 Dispatcher에 등록하는 API
  • Register{Backend}.cpp는 gen.py가 native_functions.yaml로부터 생성

참고: http://blog.ezyang.com/2020/09/lets-talk-about-the-pytorch-dispatcher/

Dispatcher

dispatch key(Autograd, Tracing, XLA, CUDA, CPU)별 function pointer 칸을 가진 dispatch table

Dispatcher는 모든 operator마다 함수 포인터 테이블(dispatch table)을 유지합니다. 각 dispatch key는 PyTorch의 cross-cutting concern(백엔드, autograd, tracing 등)에 대응하며, 위 다이어그램에서 CPU, CUDA, XLA 같은 백엔드뿐 아니라 autograd, tracing 같은 상위 개념에 대한 dispatch entry도 볼 수 있습니다. Dispatcher는 입력 tensor와 기타 정보를 기반으로 dispatch key를 계산한 뒤, 해당 테이블의 함수 포인터로 indirect jump합니다.

C++ vtable과의 비교

Dispatch table은 C++의 virtual table과 비슷하지만 다음 세 가지가 다릅니다:

C++ vtablePyTorch Dispatch Table
할당 단위클래스(class)별operator별 - 새 operator 추가 시 새 dispatch table만 할당하면 됨
dispatch 기준첫 번째 인자(this)만 고려모든 Tensor 인자(multiple dispatch) + thread-local state(TLS) 고려
boxing/unboxing없음calling convention의 일부로 boxing/unboxing 지원

PyTorch는 새 subclass보다 새 operator를 추가하는 방식으로 확장되므로, class 단위 vtable보다 operator 단위 dispatch table이 맞습니다. 대신 dispatch key는 자유롭게 늘릴 수 없고, 새 dispatch key를 추가하려면 PyTorch core에 패치를 제출해야 합니다. out-of-tree backend용으로는 PrivateUse1~PrivateUse3 key가 예약되어 있어, NPU 같은 새 device는 core 패치 없이 이 key에 kernel을 등록해 붙일 수 있습니다(torch.utils.rename_privateuse1_backend).

역사적으로 PyTorch는 원래 virtual method를 사용하여 dynamic dispatch를 구현했으나, vtable이 제공하는 것 이상의 유연성이 필요해지면서 현재의 dispatcher로 재구현했습니다.

Dispatch Key와 Dispatch Key Set

Tensor 입력, Local Include, Global의 key set을 합친 뒤 Local Exclude를 빼고, 남은 key 중 첫 번째 key를 고른다 (Autograd 위치는 2020년 기준)

위 그림은 원 슬라이드(2020년) 기준이라 Autograd가 Global set 쪽에 그려져 있습니다. 현재 PyTorch에서는 Autograd key가 tensor의 dispatch key set에 포함됩니다.

Dispatch table에서 어떤 dispatch key로 index할지를 결정하는 핵심 추상화가 dispatch key set입니다. Dispatch key set은 dispatch key들의 bitset으로, 여러 소스에서 가져온 dispatch key set을 union하고(일부는 mask out), 최종 dispatch key set에서 우선순위가 가장 높은 key를 선택하여 dispatch합니다.

Dispatch key set의 소스
소스설명
Tensor input각 입력 tensor가 자신의 dispatch key set을 기여 (예: CPU tensor → CPU key)
Local include settensor와 무관한 modal 기능 (예: tracing) - thread-local로 특정 scope 내에서 on/off
Global set항상 고려되는 dispatch key (참고: Autograd key는 현재 tensor의 dispatch key set에 포함되어 있음 - 과거에는 global set에 있었음)
Local exclude setdispatch에서 제외할 key를 지정. handler가 자신의 key를 처리한 후 mask off하여 재처리를 방지하는 패턴에 사용

Dispatch Key Set이 정해지는 과정

1단계: 입력 tensor들로부터 dispatch key set 수집

template <typename... Args>
DispatchKeySet multi_dispatch_key_set(const Args&... args) {
  return MultiDispatchKeySet().apply(args...).ts;
}

2단계: TLS의 include/exclude set을 합치고 fallthrough key를 제거해 최종 dispatch key set 계산

template<class... Args>
DispatchKeySet getDispatchKeySetUnboxed(const Args&... args) const {
  auto ks = detail::multi_dispatch_key_set(args...);
  return getDispatchKeySetFromRawDispatchKeySet(ks);
}

C10_ALWAYS_INLINE DispatchKeySet getDispatchKeySetFromRawDispatchKeySet(
    DispatchKeySet ks,
    DispatchKeySet key_mask = DispatchKeySet(DispatchKeySet::FULL)) const {
  c10::impl::LocalDispatchKeySet tls =
      c10::impl::tls_local_dispatch_key_set();
  auto dispatch_keys = (ks | tls.included_) - tls.excluded_;   // TLS 합산
  if (requiresBitsetPerBackend_) {
    auto backend_idx = dispatch_keys.getBackendIndex();
    return dispatch_keys & nonFallthroughKeysPerBackend_[backend_idx] & key_mask;
  } else {
    return dispatch_keys & nonFallthroughKeys_ & key_mask;
  }
}

Runtime

Runtime은 실제 device에서 연산을 수행하기 위한 저수준 인프라를 제공합니다. 주요 구성 요소는 Device, Stream, Event, DeviceGuard, Allocator, Generator입니다.

Device

Tensor가 저장되어 있는 device를 표현합니다.

  • Device는 device type과 index로 unique하게 구분
  • Python 수준에서는 string으로 표현 (예: cuda:0)

c10::Device는 value type으로 해석해야 합니다:

  • 두 c10::Device object가 같은 device type과 index를 가지면 → 같은 concrete device

Device index 해석:

  • Negative: current device를 의미
  • Non-negative: 특정한 concrete한 device
  • Device type이 CPU면 device index는 -1 또는 0만 허용 (CPU는 하나뿐이므로, torch.device('cpu')는 -1)

(c10/core/Device.h)

sycl::device& get_raw_device(DeviceIndex device) {
  initDevicePoolCallOnce();
  check_device_index(device);
  return *gDevicePool.devices[device];
}

Unique한 c10::Device는 gDevicePool에 들어있는 sycl::device와 1:1 관계 (XPU 예시)

Stream

Heterogeneous computing에서의 실행 모델에서, host가 device로 command/data를 전달하면 device가 이를 처리합니다. Host는 device 완료를 기다리지 않고 asynchronous하게 동작합니다.

Stream은 host가 device로 전달한 command가 담긴 queue입니다:

  • Host는 stream에 command를 삽입
  • Device는 stream에서 FIFO 순서로 command를 처리
  • 각 device마다 stream pool이 존재

device 2 · stream pool

device 1 · stream pool

command 삽입

host

stream (FIFO)

stream (FIFO)

stream (FIFO)

stream (FIFO)

(c10/core/Stream.h | XPU의 sycl::queue에 대응)

// Return whether all asynchronous work previously enqueued on this stream
// has completed running on the device.
bool Stream::query() const {
  impl::VirtualGuardImpl impl{device_.type()};  // ← DeviceGuard를 통해 device specific 구현체와 연결
  return impl.queryStream(*this);
}
  • c10/core의 구현체는 device agnostic
  • Device specific 구현체와의 연결은 DeviceGuard를 통해 이루어짐
Stream API
  • query(): stream에 삽입된 모든 command가 처리되었는지 여부를 반환
    • c10::Stream::query() → c10::impl::VirtualGuardImpl::queryStream → c10::impl::DeviceGuardImplInterface::queryStream (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::queryStream → c10::xpu::XPUStream::query → sycl::queue에 ext_oneapi_empty 호출해 queue가 비어있는지 확인
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::queryStream → c10::cuda::CUDAStream::query → DeviceGuard를 선언하고 cudaStreamQuery를 호출해 stream의 command가 모두 처리되었는지 확인
  • synchronize(): host가 stream 안의 모든 command가 처리될 때까지 대기
    • c10::Stream::synchronize → c10::impl::VirtualGuardImpl::synchronizeStream → c10::impl::DeviceGuardImplInterface::synchronizeStream (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::synchronizeStream → c10::xpu::XPUStream::synchronize → sycl::queue에 wait_and_throw를 호출해 queue에 삽입된 모든 command가 처리될 때까지 대기
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::synchronizeStream → c10::cuda::CUDAStream::synchronize → DeviceGuard를 선언하고 c10::cuda::stream_synchronize를 호출 → cudaStreamSynchronize를 호출해 stream의 모든 command가 처리될 때까지 대기
  • wait(event):
    • c10::Stream::wait → event.block(stream)

Event

Event는 device progress를 확인하거나 stream 간의 dependence를 제어하기 위해 사용합니다.

Stream 2EventStream 1HostStream 2EventStream 1Hostrecord command 삽입was_marked_for_recording = truewait 명령 삽입(record 완료까지 대기)완료: query() = true생성event.record(stream1)event.block(stream2)앞선 command 처리record command 처리wait 해제, 이후 command 처리
event.record(stream)
  • Stream에 record command를 삽입
  • Record command가 처리되어야 이 event를 기다리는 stream이 재시작할 수 있음
  • was_marked_for_recording이 false → true로 변경
event.block(stream)
  • Record command가 처리될 때까지 대기하라는 wait 명령(CUDA의 cudaStreamWaitEvent)을 stream에 삽입
  • Stream은 wait 명령에 도달하면 record command의 처리 여부를 확인
  • Stream과 event의 device type은 같아야 하지만 device index는 다를 수 있음
  • 같은 device index여도 deadlock이 발생하지 않음 (wait 명령이 처리되기 전까지 계속 command를 처리해 record command를 처리할 수 있으므로)

(c10/core/Event.h | XPU: sycl::event 대응)

Event 구현 방식

Event는 InlineEvent<VirtualGuardImpl>을 멤버로 소유하며, 실제 동작은 이 멤버를 통해 위임됩니다:

// c10/core/Event.h
struct Event final {
    impl::InlineEvent<impl::VirtualGuardImpl> impl_;  // ← 구현을 멤버로 소유
};
// c10/core/impl/InlineEvent.h
template <typename T>
struct InlineEvent final {
    void recordOnce(const Stream& stream);
    void record(const Stream& stream);
    void block(const Stream& stream);
    bool query();
    void synchronize();

    void* event_ = nullptr;
    T backend_;                          // VirtualGuardImpl - DeviceGuard에 의존
    DeviceType device_type_;
    DeviceIndex device_index_ = -1;
    EventFlag flag_ = EventFlag::PYTORCH_DEFAULT;
    bool was_marked_for_recording_ = false;
};

Device specific한 부분은 T backend_(VirtualGuardImpl)를 거쳐 DeviceGuard에 의존합니다:

Event API 상세 (호출 체인과 XPU/CUDA 구현)
  • 다양한 getter들: device(), device_type(), device_index(), flag(), was_marked_for_recording(), eventId()
  • record(stream): 주어진 stream에 record command를 삽입. record command가 처리되어야 이 event를 기다리는 stream이 재시작할 수 있음
    • stream과 event의 device type은 같아야 함 (event의 device index는 최초 record 시점에 stream의 index를 따라 설정됨)
    • c10::Event::record → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::record → c10::impl::VirtualGuardImpl::record → c10::impl::DeviceGuardImplInterface::record (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::record → stream의 sycl::queue에 ext_oneapi_submit_barrier로 barrier를 제출하고 그 barrier가 돌려주는 sycl::event를 보관함. host는 기다리지 않음. enable_timing인 event는 submit_profiling_tag를 대신 사용
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::record → 주어진 stream의 device로 setDevice하고 cudaEventRecord를 호출. 원래 device 복원은 2.14부터 c10::make_scope_exit로 처리되어 예외가 나도 되돌아감
    • c10::impl::DeviceGuardImplInterface::record를 수행한 후 was_marked_for_recording_을 true로 변경
  • recordOnce(stream): 아직 record된 적이 없을 때만 record. Event 자체가 thread-safe하지 않으므로 여러 thread에서 동시에 부르면 record가 여러 번 실행될 수 있음
    • c10::Event::recordOnce → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::recordOnce → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::record
    • was_marked_for_recording_이 false일 때만 record 실행
  • block(stream): event에 record를 호출했을 때 삽입된 record command가 처리될 때까지 대기하라는 wait 명령을 stream에 삽입. stream은 wait 명령을 처리할 때 record command의 처리 여부를 확인함
    • stream과 event의 device type은 같아야 하지만 device index는 다를 수 있음
    • 같은 device index여도 wait 명령이 stream에서 처리되기 전까지 계속 command를 처리해 record command를 처리할 수 있으므로 deadlock이 발생하지 않음
    • c10::Event::block → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::block → c10::impl::VirtualGuardImpl::block → c10::impl::DeviceGuardImplInterface::block (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::block → stream의 sycl::queue에 sycl::event list를 argument로 ext_oneapi_submit_barrier를 호출. ext_oneapi_submit_barrier는 cudaStreamWaitEvent과 같은 역할
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::block → cudaStreamWaitEvent를 호출. 주어진 stream은 event의 record command가 처리될 때까지 대기함
  • query(): event에 record를 호출했을 때 삽입된 record command가 처리되었는지 여부를 반환. record를 호출한 적이 없으면 항상 true를 반환
    • c10::Event::query → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::query → c10::impl::VirtualGuardImpl::queryEvent → c10::impl::DeviceGuardImplInterface::queryEvent (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::queryEvent → sycl::event의 command_execution_status가 event_command_status::complete인지 여부를 반환
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::queryEvent → event에 대응되는 CUDA event에 cudaEventQuery를 호출해 event의 record command가 처리되었는지 여부를 반환
  • synchronize(): host가 event에 record를 호출했을 때 삽입된 record command가 처리될 때까지 대기
    • c10::Event::synchronize → c10::impl::InlineEvent<c10::impl::VirtualGuardImpl>::synchronize → c10::impl::VirtualGuardImpl::synchronizeEvent → c10::impl::DeviceGuardImplInterface::synchronizeEvent (virtual)
      • XPU 구현: c10::xpu::impl::XPUGuardImpl::synchronizeEvent → sycl::event에 wait_and_throw를 호출
      • CUDA 구현: c10::cuda::impl::CUDAGuardImpl::synchronizeEvent → CUDA event에 cudaEventSynchronize를 호출해 event의 record command가 처리될 때까지 대기

DeviceGuard

DeviceGuard는 특정한 device, stream, event의 안전한 사용을 보장하는 장치입니다:

  • 정해진 scope에서 특정 device를 set하고, scope을 벗어나면 원래 device로 reset
  • RAII (Resource Acquisition Is Initialization) 디자인 패턴
  • 중첩 사용 가능. 안쪽 guard가 끝나면 바깥 guard의 device로, 바깥 guard가 끝나면 원래 device로 복원
g1이 current device를 1에서 2로, 안쪽 g2가 3으로 바꾸고, scope를 벗어날 때마다 2, 1 순으로 복원된다 g1이 current device를 1에서 2로, 안쪽 g2가 3으로 바꾸고, scope를 벗어날 때마다 2, 1 순으로 복원된다
DeviceGuard 사용 예제: CUDA stream의 query() 구현
bool query() const {
  DeviceGuard guard{stream_.device()};
  cudaError_t err = C10_CUDA_ERROR_HANDLED(cudaStreamQuery(stream()));

  if (err == cudaSuccess) {
    return true;
  } else if (err != cudaErrorNotReady) {
    C10_CUDA_CHECK(err);
  } else {
    // ignore and clear the error if not ready
    (void)cudaGetLastError();
  }

  return false;
}

사용 예제: current device가 stream_.device()로 세팅되고, guard가 소멸되며 원복

구현 방식: DeviceGuard → InlineDeviceGuard → XPUGuardImpl (device-specific)

1단계: DeviceGuard. RAII wrapper로 guard_.reset_device()를 호출:

void reset_device(at::Device device) {
  guard_.reset_device(device);
}

2단계: InlineDeviceGuard. VirtualGuardImpl을 통해 device-specific 구현체로 위임:

template <typename U = T>
typename std::enable_if_t<std::is_same_v<U, VirtualGuardImpl>> reset_device(
    at::Device device,
    const impl::DeviceGuardImplInterface* impl = nullptr) {
  auto index = device.index();
  if (index == -1)
    return;
  if (device.type() == original_device_.type()) {
    AT_ASSERT(impl == nullptr || impl->type() == device.type());
    impl_.setDevice(device);

3단계: XPUGuardImpl (device-specific). 실제 device 전환 수행:

void setDevice(Device d) const override {
  TORCH_CHECK(d.is_xpu(), "Expected a XPU device, but got ", d);
  c10::xpu::set_device(d.index());
}

VirtualGuardImpl API

분류API
Device 관련exchangeDevice, getDevice, setDevice, uncheckedSetDevice, deviceCount, getDeviceCapability, synchronizeDevice
Stream 관련getStream, getNewStream, getDefaultStream, getStreamFromGlobalPool, exchangeStream, getStreamNativeHandle
Event 관련record, block, queryEvent, synchronizeEvent, elapsedTime, destroyEvent
Stream 동기화queryStream, synchronizeStream, isStreamCapturing
기타recordDataPtrOnStream

Allocator

Device memory를 관리하는 structure입니다. 대부분의 중요한 함수들이 virtual로 정의되어 있습니다:

  • virtual DataPtr allocate(size_t n)

  • DataPtr clone(const void* data, std::size_t n) - non-virtual, allocate와 copy_data를 조합해 구현

  • virtual bool is_simple_data_ptr(const DataPtr& data_ptr)

  • virtual DeleterFnPtr raw_deleter()

  • virtual void copy_data(void* dest, const void* src, std::size_t count)

  • void raw_allocate(size_t n) / void raw_deallocate(void* ptr) - non-virtual. allocate와 raw_deleter를 조합한 편의 함수

  • Device specific allocator는 이 structure를 상속받아 구현

  • DataPtr: device에 잡힌 메모리를 표현하기 위한 class (device와 그 device 내 주소에 해당하는 pointer를 보유)

(c10/core/Allocator.h)

Random Number Generator

  • Pseudo Random Number Generator (PRNG) engine을 사용하는 interface
  • 사용자는 seed를 제공하거나 상태를 확인할 수 있음

두 층으로 나뉩니다. 사용자가 다루는 at::Generator는 device agnostic한 wrapper이고, device별 구현은 c10::GeneratorImpl을 상속해 작성합니다.

at::Generator (aten/src/ATen/core/Generator.h) - impl_에 위임하는 wrapper

  • void set_current_seed(uint64_t seed) / uint64_t current_seed() const / uint64_t seed()
  • void set_offset(uint64_t offset) / uint64_t get_offset() const
  • void set_state(const at::Tensor& new_state) / at::Tensor get_state() const
  • const c10::intrusive_ptr<c10::GeneratorImpl>& getIntrusivePtr() const
  • Generator clone() const - impl_->clone()을 감싸 반환
  • philox_state(uint64_t increment) - 2.14 신규. CUDA Graph capture 중에도 Philox seed/offset을 안전하게 예약·조회

c10::GeneratorImpl (c10/core/GeneratorImpl.h) - device별로 구현하는 pure virtual

  • virtual void set_current_seed(uint64_t seed) = 0 / virtual uint64_t current_seed() const = 0 / virtual uint64_t seed() = 0
  • virtual void set_offset(uint64_t offset) = 0 / virtual uint64_t get_offset() const = 0
  • virtual void set_state(const c10::TensorImpl& new_state) = 0 / virtual c10::intrusive_ptr<c10::TensorImpl> get_state() const = 0
  • c10::intrusive_ptr<GeneratorImpl> clone() const - non-virtual. 내부에서 virtual GeneratorImpl* clone_impl() const = 0을 호출하므로 device별 구현은 clone_impl만 제공

PyTorch Build

Build System: PEP 517 + scikit-build-core

PyTorch 2.14부터 빌드는 pyproject.toml에 선언된 PEP 517 backend scikit-build-core(build-backend = "scikit_build_core.build")가 수행합니다. pip과 python -m build는 setup.py를 실행하지 않습니다. 오랫동안 쓰이던 python setup.py develop / install은 2.14에서 deprecation shim으로만 남았고, 단계적으로 제거됩니다.

setup.py 제거 일정 (v2.14.0 setup.py 머리 주석 기준):

PyTorch 버전python setup.py <command> 동작
2.14–2.15install/develop은 pip install --no-build-isolation -v [-e] .로 forwarding, 나머지는 안내 메시지와 함께 실패
2.16–2.17모든 명령이 안내 메시지와 함께 실패
2.18setup.py 파일 자체가 제거됨

대체 명령어 (shim이 출력하는 안내문과 동일):

기존 (deprecated)대체 명령
python setup.py developspin develop 또는 python -m pip install --no-build-isolation -v -e .
python setup.py installspin install 또는 python -m pip install --no-build-isolation -v .
python setup.py bdist_wheelpython -m build --wheel --no-isolation
python setup.py sdistpython -m build --sdist
  • spin은 scientific-python 프로젝트들이 쓰는 developer CLI로, pip install --group dev에 포함됩니다. 2.14의 [tool.spin.commands]에는 develop, editable, install, clean, lint, fixlint, quicklint, quickfix와 regenerate_* 명령이 있습니다.
  • 환경변수를 통한 빌드 커스터마이즈(DEBUG=1, USE_CUDA=0, MAX_JOBS=... 등)는 대체 명령에서도 그대로 동작합니다.
  • 배경: pytorch/pytorch#152276, 전환 PR #180247·#180248

How to Build from Source

v2.14.0 README “From Source” 절의 Linux 절차입니다.

# Get the PyTorch Source
git clone https://github.com/pytorch/pytorch
cd pytorch
# 기존 checkout을 업데이트하는 경우에도 필요
git submodule sync
git submodule update --init --recursive

# Install Dependencies (공통)
pip install --group dev
# Linux
pip install mkl-static mkl-include
# (선택) CUDA + LAPACK 지원이 필요한 경우 magma 설치 (conda 환경에서 실행, CUDA 버전 지정)
.ci/docker/common/install_magma_conda.sh 12.4
# (선택) torch.compile (inductor/triton)을 쓰는 경우. Intel GPU면 먼저 export USE_XPU=1
make triton

# (선택) 빌드 옵션은 환경변수로 지정
export MAX_JOBS=16
export USE_CUDA=1   # CUDA 비활성화: USE_CUDA=0, ROCm: USE_ROCM=0, Intel GPU: USE_XPU=0

# CMake가 현재 환경의 라이브러리를 찾을 수 있도록 prefix 지정
export CMAKE_PREFIX_PATH="${CONDA_PREFIX:-'$(dirname $(which conda))/../'}:${CMAKE_PREFIX_PATH}"
# venv 사용 시: export CMAKE_PREFIX_PATH="${VIRTUAL_ENV}:${CMAKE_PREFIX_PATH}"

# Install PyTorch (editable). 아래 두 명령은 같은 일을 함
python -m pip install --no-build-isolation -v -e .
spin develop
  • Python 3.10 이상, C++20 지원 컴파일러 필요 (Linux: gcc 11.3.0+)
  • 디스크 공간 10GB 이상, 초기 빌드 시 30~60분 소요 (이후 재빌드는 훨씬 빠름)
  • conda가 필수는 아니며 uv 등으로 만든 venv에서도 빌드 가능 (CUDA, MKL처럼 pip으로 받을 수 없는 의존성만 시스템에 있으면 됨)
  • editable 설치라 Python code는 source 형태로 남아있음
  • https://github.com/pytorch/pytorch?tab=readme-ov-file#from-source

Build Configuration

빌드 옵션은 환경변수로 지정하고, cmake/EnvVarForwarding.cmake가 이를 CMake cache 변수로 넘깁니다. 예전에 setup.py가 -D로 넘기던 로직이 이 파일로 옮겨 왔고, 환경변수 목록 문서도 이 파일 머리 주석에 있습니다.

Forwarding 규칙 (파일 머리 주석 기준):

  • BUILD_, USE_, CMAKE_로 시작하거나 EXITCODE로 끝나는 환경변수는 같은 이름의 CMake cache 변수로 전달됨
  • 그 외에는 _ENV_ALIASES(예: CUDNN_LIB_DIR → CUDNN_LIBRARY)와 _ENV_PASSTHROUGH(예: BLAS, CUDNN_LIBRARY, TORCH_CUDA_ARCH_LIST) 목록에 있는 것만 전달됨
  • 목록에 없는 값은 CMake 옵션(-D / cmake.define)으로 직접 지정
  • 일부는 이 파일이 아니라 다른 곳에서 읽음. DEBUG/REL_WITH_DEB_INFO/MAX_JOBS는 pyproject.toml, NCCL_*는 cmake/Modules/FindNCCL.cmake, CUDA_HOME/CUDA_PATH는 CMake의 CUDAToolkit 탐지

cache 지속성은 변수 종류에 따라 다릅니다.

  • 접두 변수(BUILD_/USE_/CMAKE_)는 configure할 때마다 현재 환경변수 값으로 cache를 FORCE 덮어씁니다(_envfwd_apply). 환경변수를 바꾸고 다시 빌드하면 반영됩니다.
  • alias/passthrough 목록 변수와 -D로 직접 넣은 값은 cache에 없을 때만 채워집니다. 한 번 들어가면 환경변수를 바꿔도 반영되지 않으므로 ccmake build로 cache를 직접 수정하거나 build/CMakeCache.txt(또는 build/ 전체)를 삭제해야 합니다.
  • build type과 compiler는 항상 다시 적용됩니다.

일반 설정 환경변수

변수설명처리 위치
DEBUG디버그 심볼과 함께 최적화 없이 빌드 (-O0 -g)pyproject.toml의 [[tool.scikit-build.overrides]]가 cmake.build-type = "Debug"로 매핑
REL_WITH_DEB_INFO최적화와 디버그 심볼과 함께 빌드위와 같은 방식으로 RelWithDebInfo
USE_CUSTOM_DEBINFO지정된 파일들에 대해서만 디버그 정보 포함 ("file1.cpp;file2.cpp")접두 forwarding
BUILD_TEST테스트 빌드 활성화접두 forwarding
MAX_JOBS컴파일 병렬 수[tool.scikit-build.env]가 CMAKE_BUILD_PARALLEL_LEVEL로 alias. 직접 지정한 CMAKE_BUILD_PARALLEL_LEVEL이 우선
CFLAGSC와 C++ 파일 컴파일 플래그 (CXXFLAGS 설정 시 C++는 그쪽을 따름)CMake / scikit-build-core가 직접 읽음
CC / CXX사용할 C/C++ 컴파일러[tool.scikit-build.env]에 등재되어 scikit-build-core의 sysconfig 컴파일러 주입을 끄고 PATH에서 탐지

setup.py와 함께 제거된 변수 (2.14부터 동작하지 않음):

변수대체 방법
CMAKE_FRESHbuild/ 디렉토리를 삭제하고 다시 빌드
CMAKE_ONLY대체 없음 (setup.py 전용 디버깅 기능이었음)
USE_NINJACMAKE_GENERATOR=Ninja로 generator 지정 (ninja가 있으면 기본값)

DEBUG_CUDA는 이 표에 속하지 않습니다. 원래부터 환경변수가 아니라 CMake 옵션(-DDEBUG_CUDA=1)이었고, 예전 setup.py 주석이 환경변수처럼 잘못 적어 두었던 것을 2.14 문서가 바로잡았습니다.

기능 토글 환경변수

USE_*/BUILD_*는 접두 규칙으로 모두 forwarding됩니다. 자주 쓰는 것만 추립니다.

변수설명
USE_CUDACUDA 빌드 활성화/비활성화
USE_CUDNNcuDNN 빌드 활성화/비활성화
USE_CUSPARSELTcuSPARSELt 빌드 활성화/비활성화
USE_CUDSScuDSS 빌드 활성화/비활성화
USE_CUFILEcuFile 빌드 활성화/비활성화
USE_FBGEMMFBGEMM 빌드 활성화/비활성화
USE_KINETOlibkineto (프로파일링) 활성화/비활성화
USE_NUMPYNumPy 빌드 활성화/비활성화
USE_MKLDNNMKLDNN 사용 활성화/비활성화
USE_NNPACKNNPACK 빌드 활성화/비활성화
USE_DISTRIBUTED분산 학습 (c10d, gloo, mpi) 빌드 활성화/비활성화
USE_GLOO / USE_MPI / USE_TENSORPIPE개별 분산 backend 활성화/비활성화
USE_OPENMPOpenMP 병렬화 활성화/비활성화
USE_FLASH_ATTENTIONFlash Attention 빌드 활성화/비활성화
USE_MEM_EFF_ATTENTIONMemory Efficient Attention 빌드 활성화/비활성화
USE_SYSTEM_LIBS시스템 제공 third-party 라이브러리 사용. 개별 USE_SYSTEM_* 토글로 확장됨
BUILD_LIBTORCH_WHL / BUILD_PYTHON_ONLYlibtorch만 wheel로 빌드 / 별도 libtorch에 대해 Python wheel만 빌드

버전 및 아키텍처 관련

변수설명처리 위치
PYTORCH_BUILD_VERSION / PYTORCH_BUILD_NUMBERwheel 버전 지정tools/metadata의 version metadata provider (forwarding 아님)
TORCH_CUDA_ARCH_LIST빌드할 CUDA 아키텍처 지정 (예: "8.0;9.0")passthrough
TORCH_XPU_ARCH_LIST빌드할 XPU 아키텍처 지정 (예: "ats-m150,lnl-m")passthrough
PYTORCH_ROCM_ARCH빌드할 AMD GPU 타겟 지정 (예: "gfx900;gfx906")cmake/public/utils.cmake가 환경에서 읽음

외부 라이브러리 관련

변수설명처리 위치
CUDA_HOME (Linux/macOS) / CUDA_PATH (Windows)CUDA 설치 위치CMake의 CUDAToolkit 탐지
CUDNN_LIBRARY, CUDNN_INCLUDE_DIR, CUDNN_ROOTcuDNN 설치 위치passthrough (CUDNN_LIB_DIR은 CUDNN_LIBRARY의 alias)
NCCL_ROOT, NCCL_LIB_DIR, NCCL_INCLUDE_DIRNCCL 설치 위치cmake/Modules/FindNCCL.cmake가 환경에서 읽음
BLAS사용할 BLAS 라이브러리 (MKL, Eigen, ATLAS, FlexiBLAS, OpenBLAS). 없으면 빌드 실패passthrough
MIOPEN_PATHMIOpen 설치 루트 (예전 MIOPEN_LIB_DIR 등은 더 이상 사용 안 함)LoadHIP.cmake가 환경에서 읽음

Build의 최종결과물

Python package를 설치하면 shared object와 Python module들이 함께 설치됩니다.

Shared objects

libtorch_python.so  libtorch.so  libtorch_xxx.so  libc10_xxx.so  libc10.so

Python modules

_awaits _custom_op _decomp _dispatch _dynamo _export _functorch _higher_order_ops _inductor _lazy _library _logging _native _numpy _prims _prims_common _refs _strobelight _subclasses _vendor accelerator amp ao autograd backends compiler contrib cpu cuda distributed distributions export fft func futures fx jit linalg masked monitor mps mtia multiprocessing nativert nested nn numa onnx optim package profiler quantization signal sparse special testing utils xpu

Eager Mode의 Shared Object들

libtorch_python.so

libtorch.so

libtorch_cuda.so

libtorch_cpu.so

libtorch_xpu.so

libc10_cuda.so

libc10_xpu.so

libc10.so

libcudart.so

libsycl.so

  • libtorch.so는 여러 so를 한꺼번에 묶는 wrapper 역할. 실구현체는 들어있지 않음
  • libtorch_cuda.so/libc10_cuda.so는 USE_CUDA=1, libtorch_xpu.so/libc10_xpu.so는 USE_XPU=1로 빌드했을 때만 생김
  • libtorch_cuda.so, libtorch_xpu.so는 CPU fallback 등을 위해 libtorch_cpu.so를 직접 링크함 (caffe2/CMakeLists.txt의 target_link_libraries(torch_cuda PUBLIC torch_cpu_library ...))
  • 2.14 CUDA wheel의 torch/lib에는 이 밖에 libtorch_cuda_linalg.so, libtorch_nvshmem.so, libcaffe2_nvrtc.so, libshm.so, libtorch_global_deps.so도 들어 있음

Eager Mode의 소스 코드 구조

주요 라이브러리

  • ATen (A Tensor C++ Library)
    • Tensor 연산 및 수학적 작업을 정의한 library
    • 사용자에게 노출되는 고수준 API
    • Python integration point
    • namespace: at
  • C10 (Caffe2 and ATen)
    • metadata/memory 관리 및 op dispatch 담당
    • 내부적으로 사용되는 저수준 API
    • Backend integration point
    • namespace: c10
  • Generated files
    • 주로 op별 public C++ API(at::*, at::Tensor::*)와 dispatcher 진입점(at::_ops::*), autograd kernels, dispatch registration 코드
    • native_functions.yaml와 derivatives.yaml로부터 생성

참고: https://discuss.pytorch.org/t/whats-the-difference-between-aten-and-c10/114034/3

직접 작성된 코드 vs. 생성된 코드

Source Code

직접 작성된 코드

생성된 코드

aten/src/ATen

c10

core

native

cuda, xpu, ...

core

cuda, xpu, ...

build/aten/src/ATen

torch

include/torch/csrc/
autograd/generated

csrc/autograd/generated

Directory Structure: aten (중요한 것들 위주)

aten/src/ATen/

  • benchmarks/: ATen benchmark code
  • core/: Tensor 관련 핵심 선언과 구현
    • boxing/, dispatch/, op_registration/: operator dispatch 관련 구현
  • cuda/, hip/, metal/, mps/, vulkan/, xpu/: device dependent 구현 (runtime 등)
    • CUDAEvent, XPUDevice 등
  • cpu/, cudnn/, miopen/, mkl/: device 디렉터리와 나란히 있는 계산 library 전용 디렉터리 (vectorized CPU kernel, cuDNN·MIOpen·MKL wrapper)
  • ops/: per-operator header(ops/{op}.h, ops/{op}_ops.h)가 설치되는 곳. 아래 Generated Files 절과 연결됨
  • native/: operator 구현체들, native_functions.yaml이 여기 포함됨
    • Device independent한 operator 구현체는 바로 밑에 포함
    • Device dependent한 operator 구현체는 sub directory 안에 포함
      • cpu, cuda, cudnn, hip
  • templates/: native_functions.yaml을 통해 gen source code들의 뼈대가 포함
  • test/: ATen test code

Directory Structure: c10

c10/

  • benchmark/: C10 benchmark code
  • core/: PyTorch 핵심 구현과 선언, aten이 상속받아 구현
    • Device 관련: Allocator, Device, DeviceGuard, Event, Stream 등
    • Dispatch 관련: DispatchKey, DispatchKeyset 등
    • Tensor 관련: TensorImpl, StorageImpl 등
  • cuda/, xpu/, hip/, metal/, mobile/: 특정 device의 전용 구현
    • CUDACachingAllocator, XPUStream 등
  • macros/: compiler에서 해석할 inline attributes 관련 macro
    • Shared library visibility 조절, warning 강제 무시 등
  • test/: C10 test code
  • util/: helper class/method들
    • Environments, exception, backtrace 등
    • Array, BFloat16, Complex 등

Directory Structure: Generated Files

build/aten/src/ATen/ : gen.py (input: native_functions.yaml)

  • ops/{operator}.h
  • core/TensorBody.h
  • Operators.cpp
  • Register{backend}.cpp

torch/include/torch/csrc/autograd/generated/ : tools/autograd의 생성 스크립트들 (gen_autograd.py가 orchestration; input: derivatives.yaml, native_functions.yaml)

  • Functions.h
  • python_functions.h
  • python_return_types.h
  • variable_factories.h
  • VariableType.h
  • ViewFuncs.h

torch/csrc/autograd/generated/

  • VariableType.cpp (gen_variable_type.py)
  • Functions.cpp (gen_autograd_functions.py)
  • python_functions.cpp (gen_autograd_functions.py)
  • python_torch_functions.cpp - gen_python_functions.py (input: native_functions.yaml)

Footnotes

  1. 연산을 코드에 작성된 순서대로 즉시 실행하는 define-by-run 방식의 실행 모델. 그래프를 먼저 만들어 두고 나중에 실행하는 graph mode와 대비되는 개념으로, Week 1의 “Eager vs. Graph Mode”에서 소개했습니다. ↩

  2. General Matrix Multiply. BLAS(Basic Linear Algebra Subprograms)에서 정의하는 행렬 곱셈 연산으로, NVIDIA GPU용 구현이 call stack 끝에서 만나게 될 cuBLAS입니다. Week 1의 supercomputing 배경(BLAS/cuBLAS)에서 소개했습니다. ↩

  3. 큰 tensor의 batch 차원을 앞쪽 행 차원에 접어 GEMM 한 번으로 처리하는 것. copy 없이 view로 접을 수 없으면 should_fold가 bmm 경로를 택합니다. 양쪽이 3D이고 한쪽 batch 크기가 1인 경우, 그 tensor가 gradient를 요구하면 batch 차원을 squeeze해 fold 경로로 보냅니다(backward 메모리 절약). ↩

  4. c10::intrusive_ptr은 reference count를 별도 control block이 아니라 가리키는 객체 안에 두는(intrusive) smart pointer입니다. std::shared_ptr보다 포인터 하나 크기로 가볍고, raw pointer에서 다시 감쌀 수 있어 Python 바인딩과 C++ 사이를 오가기 쉽습니다. TensorImpl, StorageImpl이 이 방식으로 관리되며, view가 storage를 공유한다는 것은 같은 StorageImpl의 count를 하나 올린다는 뜻입니다. ↩