TorchInductor / torch.compile 단권화

1 추상의 정체 왜 존재

• eager overhead (Python dispatch, allocator) 제거 → kernel fusion
• graph 추출은 partial — graph break 허용이 TF graph와 차별점이다
• AOT Autograd로 backward도 같은 IR에서 컴파일된다

2 5단계 Pipeline ★ Dy·AO·FX·IN·CG

Python fn (eager)
      │  ① TorchDynamo
      │      · CPython frame intercept
      │      · guard 생성
      ▼
  FX Graph  (torch operator level)
      │  ② AOT Autograd
      │      · fwd + bwd 분리
      │      · joint graph → partitioner
      ▼
  FX Graph  (aten / prims, 2개: fwd, bwd)
      │  ③ Decomposition
      │      · _refs / _decomp 적용
      ▼
  Inductor Input IR (lowered aten)
      │  ④ Inductor
      │      · Loops/Pointwise/Reduction
      │      · Scheduler fusion
      ▼
  Triton Python  /  C++/OpenMP
      │  ⑤ Codegen → ptxas / gcc
      ▼
  PTX + SASS  /  .so

source: torch/_dynamo/convert_frame.py · torch/_functorch/aot_autograd.py · torch/_inductor/compile_fx.py

3 단계별 책임 매트릭스

4 Entry API

import torch
def f(x, y):
    return (x + y).relu() * 2

g = torch.compile(f,
       backend="inductor",
       mode="reduce-overhead",
       fullgraph=False,
       dynamic=None)
g(x, y)   # compile on first call
g(x, y)   # guard hit → cached kernel

5 mode preset

6 Cache 위치 JIT 산출물

7 Inductor ↔ 다른 경로

• Triton 컴파일러 내부 ↗ V11 — Inductor는 Triton의 소비자다
• MLIR lowering 일반 ↗ V12
• XLA / TVM 비교 ↗ V14
• distributed training path ↗ V17

1 Frame Evaluation 가로채기 ★

• target: CPython 3.8+ · PEP 523 _PyInterpreterState_SetEvalFrameFunc
• 대상: tensor-producing op만 trace, Python 제어흐름은 유지
• 결과: FX Graph + guard set + residual bytecode

source: torch/_dynamo/convert_frame.py (v2.x)

2 Symbolic Evaluator

Python frame
  bytecode = co_code
  for each op:
    if op produces Tensor:
      → trace into FX
      → record Guard(expr)
    elif op needs concrete value:
      → graph break
      → fallback to eager
    else:
      → emit residual bytecode

3 FakeTensor & shape prop

• trace 시 실제 tensor 대신 FakeTensor (shape+dtype+device만)
• 실제 compute 0회 · metadata만 흘림
• dynamic shape은 SymInt로 상징화 ↗ §15

4 Guard 시스템 기본 ★

각 guard는 C++로 lowering되어 call 당 ~ns overhead (v2.3+)

5 Graph break 조건

• Python data-dependent control (if x.item() > 0)
• 지원 안 된 builtin (print, input 등)
• .numpy() / CPU scalar 추출
• unsupported torch.* op (목록 있음)
• 사용자 @torch._dynamo.graph_break()

6 explain API 진단

from torch._dynamo import explain
exp = explain(f)(x, y)
print(exp.graph_count,
      exp.graph_break_count,
      exp.break_reasons)

반환: graph 수, break 수, 각 break의 (op, reason, stack)

7 Recompile 트리거

8 cache_size_limit

1 존재 이유 ★

• eager autograd: backward는 runtime에 autograd engine이 구성 → fusion 대상 불가
• AOT: backward를 FX Graph로 미리 구성 → fusion · decomp 적용 가능

source: torch/_functorch/aot_autograd.py

2 Joint Graph

inputs:   (x1, x2, ..., xk)          primals
labels:   (l1, ..., lm)               (loss inputs)
outputs:  (y1, ..., yn)               forward outputs
            + (g1, ..., gk)           backward grads

joint_fn(primals, tangents) -> (outs, grad_ins)
  = torch.autograd.Function.apply
     but traced into ONE FX graph

3 aot_function API

from torch._functorch.aot_autograd import aot_function
compiled = aot_function(
   fn,
   fw_compiler=my_fw,   # FX -> callable
   bw_compiler=my_bw,
   partition_fn=default_partition,
   decompositions=core_aten)

4 Partitioner ★

• default_partition: autograd와 동일한 activation 저장
• min_cut_rematerialization_partition: recompute로 saved tensor 감소 (checkpointing) ↗ V17 §7

5 Saved tensor 결정

6 functionalization

• in-place op (add_, relu_) → out-of-place로 재작성
• view는 copy로 구체화되지 않고 alias info로 전달
• 이유: fusion/decomp가 pure functional IR을 가정

7 aot_module nn.Module

compiled = aot_module(
   mod,
   fw_compiler=inductor_fw,
   bw_compiler=inductor_bw)

• module param/buffer는 graph input으로 승격
• state-free functional form으로 compile

8 Inference-only 경로

• torch.no_grad() 또는 inference_mode이면 AOT가 forward graph만 생성
• partitioner 호출 없음, bwd compiler 미사용

9 세 경로 대조

1 torch.fx.Graph ★

• Python 객체 — C++ IR 아님
• pretty-printable, print(gm.graph)로 즉시 관측
• transform은 Python rewriter로 작성

2 Node 종류 전체표

3 Node 필드

• op, target, args, kwargs, name
• meta['val'] — FakeTensor (shape/dtype/device)
• meta['stack_trace'] — 원본 코드 위치
• users, all_input_nodes — DAG edge

4 예시: f(x) = relu(x·W + b)

def f(x, W, b):
    y = x @ W
    z = y + b
    return z.relu()

gm = torch.fx.symbolic_trace(f)
print(gm.graph)

graph():
  %x : [num_users=1] = placeholder[target=x]
  %W : [num_users=1] = placeholder[target=W]
  %b : [num_users=1] = placeholder[target=b]
  %matmul : call_function[target=torch.matmul](args=(%x, %W))
  %add    : call_function[target=operator.add](args=(%matmul, %b))
  %relu   : call_method[target=relu](args=(%add,))
  return (%relu,)

5 변환 API

# 1) node-by-node walk
for n in gm.graph.nodes:
    if n.op == "call_function" and \
       n.target is operator.add:
        with gm.graph.inserting_after(n):
            new = gm.graph.call_function(
                torch.add, n.args)
            n.replace_all_uses_with(new)
            gm.graph.erase_node(n)
gm.recompile()

• replace_pattern — subgraph rewrite
• Interpreter — 재실행으로 meta 재계산
• Transformer — node 변환 framework

6 GraphModule

생성된 source는 gm.code로 확인 — eager fallback 및 디버깅의 핵심.

7 Symbolic Trace vs Dynamo

8 Pass 배치 순서

1. Dynamo → FX (torch op level)
2. AOT Autograd → fwd+bwd FX
3. Decomposition → prims/aten core
4. Inductor lowering → Inductor IR ↗ §7

1 Op의 두 층위 ★

source: torch/_refs · torch/_prims · torch/_decomp

2 canonical form

• Inductor가 케이스 폭발을 피하도록 → 줄어든 op 집합만 처리
• 수학적 동치이되 수치 비트단위 동일은 보장하지 않음

3 _refs · _decomp 역할

• torch._refs.* — op의 reference Python 구현 (prims 기반)
• torch._decomp.* — aten→aten 분해 (derivative-safe)
• core_aten_decompositions() — Inductor가 기본으로 받는 dict

4 예시 before/after ★ softmax

BEFORE (torch op level, Dynamo 직후):
  %a : placeholder[target=x]
  %sm = call_function[target=aten.softmax](
          args=(%a, -1))
  return (%sm,)

AFTER decomp (core/prims level):
  %a  = placeholder[target=x]
  %mx = call_function[target=aten.amax](
          args=(%a, [-1], True))
  %sub = call_function[target=aten.sub](
          args=(%a, %mx))
  %ex  = call_function[target=aten.exp](
          args=(%sub,))
  %sm  = call_function[target=aten.sum](
          args=(%ex, [-1], True))
  %div = call_function[target=aten.div](
          args=(%ex, %sm))
  return (%div,)

stable softmax: subtract amax for overflow safety ↗ V09 §9

5 Decomp 전후 장단점

6 Decomp 등록

from torch._decomp import register_decomposition
from torch import Tensor

@register_decomposition(aten.mse_loss)
def mse_loss(x: Tensor, t: Tensor,
             reduction: int = 1):
    d = (x - t)
    loss = d * d
    if reduction == 1:
        return loss.mean()
    return loss

7 Decomp 우선순위

1. user custom decomp
2. _refs (prims 기반)
3. _decomp (aten→aten)
4. decomp 없음 → Inductor fallback (extern kernel)

8 Opaque op 유지

• aten::mm, aten::convolution — decomp 대신 extern kernel (cuBLAS/cuDNN) 호출로 남김
• Inductor max-autotune 모드는 Triton matmul 후보도 생성

1 Entry point ★

source: torch/_inductor/compile_fx.py :: compile_fx_inner

2 입력 불변조건

• graph는 functional (in-place 없음)
• 모든 node에 meta['val'] (FakeTensor) 존재
• op은 core aten + prims + 약간의 허용 set
• dynamic shape은 SymInt로 표현

3 FakeTensor

4 Shape Propagation

for node in graph:
  fake_args = [a.meta['val'] for a in node.inputs]
  out_fake  = node.target(*fake_args)     # FakeTensor op
  node.meta['val'] = out_fake             # shape/dtype 기록

• FakeTensor 연산은 meta kernel 호출 — 실제 memory 0
• dispatch는 Meta key로 라우팅

5 Symbolic Shape dynamic

6 Layout Inference

• input stride 관측 → contiguous/channels_last/custom 분기
• 출력 layout은 make_contiguous_strides_for(shape) 기본
• layout 결정은 fusion 가능성에 영향

7 Lowering 개요

for node in gm.graph.nodes:
  lower = LOWERINGS[node.target]
  ir_out = lower(*input_ir_tensors,
                 **kwargs)
  env[node] = ir_out                # TensorBox

• op별 lowering 함수는 dict에 등록
• 결과는 TensorBox (wrap Loops IR) ↗ §7

8 Fallback 정책

9 backend 라우팅 분기

• FakeTensor.device == "cuda" → Triton codegen
• == "cpu" → C++/OpenMP codegen ↗ §10
• mixed device → split subgraph

1 IR 계층 ★

source: torch/_inductor/ir.py

2 핵심 클래스 전체표

3 Pointwise 예시

def relu_add(a, b):
    def inner(idx):
        x = a.load(idx)
        y = b.load(idx)
        return ops.maximum(x + y, 0)
    return Pointwise(
       device=a.device,
       dtype=a.dtype,
       inner_fn=inner,
       ranges=a.size)

• inner_fn(idx)는 symbolic index → value
• compose는 함수 합성으로 이루어짐 (== vertical fusion 기반) ↗ §8

4 Reduction 예시

Reduction(
   device=x.device,
   dst_dtype=torch.float32,
   src_dtype=x.dtype,
   inner_fn=lambda idx, ridx:
       x.load(idx + ridx),
   ranges=[M],
   reduction_ranges=[N],
   reduction_type="sum")

• non-reduced axis = ranges
• reduced axis = reduction_ranges
• type: sum · prod · max · min · argmax · argmin · welford

5 ops namespace

ops는 backend에 의해 재해석 (Triton ops vs C++ ops)

6 Buffer 두 종류

• InputBuffer — graph input / placeholder
• ComputedBuffer — Inductor가 생성할 tensor
• ExternKernel — cuBLAS/cuDNN 호출 node

7 Realization 시점

1 Scheduler 역할 ★

source: torch/_inductor/scheduler.py :: Scheduler

2 Fusion 종류

3 Fusion의 이득

• kernel launch 1회 → overhead ↓
• intermediate buffer 제거 → memory BW ↓
• register에서 값 재사용 → L2/HBM traffic ↓

4 가능/불가능 조건표 ★

5 Fusion 제한 요인

1. shared memory 용량 (persistent reduction)
2. register 압박 — 너무 많이 합치면 spill
3. iteration space 불일치
4. dtype·device 경계
5. 명시적 realize() hint

6 알고리즘 개요

nodes = topo_sort(IR graph)
for n in nodes:
  for pred in n.deps:
    if can_fuse(pred, n):
      fuse(pred, n)            # vertical
for n in nodes:
  for sibling in same_range(n):
    if can_fuse_h(n, sibling):
      fuse_h(n, sibling)       # horizontal

• cost model은 memory bytes 기반 heuristic
• loop-order-aware (loop_ordering_after_fusion)

7 반복 합침

1 Triton Backend 진입 ★

source: torch/_inductor/codegen/triton.py

2 생성 코드 skeleton

@triton.jit
def triton_poi_fused_add_relu_0(
    in_ptr0, in_ptr1, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
  xoffset = tl.program_id(0) * XBLOCK
  xindex = xoffset + tl.arange(0, XBLOCK)
  xmask = xindex < xnumel
  x0 = tl.load(in_ptr0 + xindex, xmask)
  x1 = tl.load(in_ptr1 + xindex, xmask)
  r  = tl.maximum(x0 + x1, 0.0)
  tl.store(out_ptr0 + xindex, r, xmask)

• kernel name: triton_{kind}_fused_{ops}_{hash}
• kind: poi pointwise · red reduction · per persistent

3 Block Size 결정

autotune config는 triton_heuristics.py에서 후보 생성. size_hint에 따라 분기.

4 Reduction 형태

• persistent: reduction 차원을 한 번에 register에 → no-loop · R ≤ ~2048
• looped: reduction loop + accumulator · R 제한 없음
• split: 두 kernel (partial + final)
• 선택 기준: reduction numel, shared memory 용량

5 출력 코드 구성

call_fn:
  args: input ptr, output ptr, shape
  grid: launch config
  stream: current CUDA stream
  triton_kernel[grid](... , XBLOCK=256)

6 Horizontal fusion in codegen

• 동일 grid로 여러 output을 한 kernel에 작성
• 서로 다른 output buffer에 tl.store 반복
• launch overhead 상각, I/O 재사용

7 Matmul Template

• prologue: load_scale · cast 등 fuse
• epilogue: bias · gelu · layernorm 등 fuse
• 후보는 MI300/A100/H100별로 다름

8 Kernel 이름 규칙

이름으로 Inductor가 만든 어떤 kind인지 바로 읽힘.

1 존재 이유

source: torch/_inductor/codegen/cpp.py · cpp_wrapper.py

2 출력 코드 skeleton

extern "C" void kernel(
    const float* in0,
    const float* in1,
    float* out,
    long N) {
  #pragma omp parallel for
  for(long i=0; i<N; i+=16){
    auto a = at::vec::Vectorized<float>
                ::loadu(in0+i);
    auto b = at::vec::Vectorized<float>
                ::loadu(in1+i);
    (a + b).relu().store(out+i);
  }
}

3 Vectorization 경로

at::vec::Vectorized<T>가 ISA 추상화.

4 OpenMP 병렬화

• outer loop에 #pragma omp parallel for
• thread 수 = torch.get_num_threads()
• reduction은 omp parallel reduction(+:acc)

5 Reduction scheme

6 Build path

.py source
  → write /tmp/torchinductor_*/xx.cpp
  → compile (gcc/clang) -> .so
  → dlopen -> ctypes fn ptr
  → call via cpp_wrapper

7 cpp_wrapper 모드

• config.cpp_wrapper = True → Python wrapper까지 C++로
• dispatch overhead 제거 → deployment에 유리
• AOT export 경로 (torch.export + AOTInductor)의 기반

8 Inductor CPU 한계

1 목적 ★

source: torch/_inductor/cudagraph_trees.py

2 Inductor 활성화 조건

• mode="reduce-overhead" 또는 mode="max-autotune"
• 또는 config.triton.cudagraphs = True
• CUDA device에 한함

3 Graph capture 조건

4 Static Input Copy

• Inductor는 graph 실행 시 static input buffer에 입력을 복사한 뒤 replay
• copy 비용 vs launch 절약 trade-off
• config.triton.cudagraph_skip_dynamic_graphs로 우회

5 CUDA Graph Trees

default ON in v2.2+ when mode uses cudagraph

6 Re-capture 트리거

1. input shape 변화 → 새 graph
2. new tensor object → pointer mismatch → re-capture
3. warmup 단계 (첫 2~3회는 plain run)

7 적용 유/불리 상황

8 비활성 패턴

• NCCL collective (동기화 포함) — stream-aware collective 필요 ↗ V15 §15
• .cpu() / .item() 호출
• dynamic allocator가 새 block을 요청

9 진단 env

1 설정 경로

• torch._inductor.config.* — Python attr
• TORCHINDUCTOR_* env — env 변수 매핑
• 우선순위: env > Python attr > default

source: torch/_inductor/config.py

2 주요 bool flag

3 Size-related

4 Debug flag ★

5 Cache dir 구조

/tmp/torchinductor_$USER/
  └ <fxgraph_hash>/
      ├ output_code.py        # Python wrapper
      ├ fx_graph_readable.py  # pre-lowering
      ├ fx_graph_runnable.py  # re-executable
      ├ fx_graph_transformed.py
      ├ triton_<hash>.py      # kernel source
      └ fx_graph.aot_inductor.so (cpp_wrapper)

6 Shape policy

자세한 동작은 ↗ §15 Guard.

7 Trace dump

• TORCHINDUCTOR_TRACE=1 → 각 kernel call에 NVTX range
• Nsight Systems로 load · kernel별 분리 관측 가능 ↗ V18 §12

8 권장 조합

1 backend 인자

2 기본 목록

3 언제 eager?

• graph break 원인만 확인 → eager
• bwd 오류 원인만 확인 → aot_eager
• Inductor 버그 격리 → aot_eager로 bisect

4 Custom backend signature

from torch._dynamo import register_backend

@register_backend
def my_be(gm, example_inputs):
    # gm: torch.fx.GraphModule
    # example_inputs: list[FakeTensor]
    return lambda *args: gm(*args)

• signature: (gm, example_inputs) → callable
• callable은 eager-equivalent이어야 함
• 내부에서 AOT Autograd를 호출할지 선택

5 AOT Autograd 조합

from functorch.compile import \
    aot_module_simplified

def my_be(gm, inputs):
    return aot_module_simplified(
        gm, inputs,
        fw_compiler=my_compile,
        bw_compiler=my_compile)

backend이 fwd/bwd 모두 컴파일하려면 AOT를 명시 호출.

6 backend 비교 축

상세 비교는 ↗ V14 §15.

7 list_backends()

import torch
print(torch._dynamo.list_backends())
# 현재 설치된 backend 이름 목록

8 minifier 경로

• TORCHDYNAMO_REPRO_AFTER=dynamo
• TORCHDYNAMO_REPRO_AFTER=aot
• 버그를 minimal FX graph로 재현 → fx_graph_runnable.py로 산출

1 왜 필요한가

2 3 종류 impl

3 custom_op API

@torch.library.custom_op(
  "my::scaled_add",
  mutates_args=())
def scaled_add(
  x: torch.Tensor,
  y: torch.Tensor,
  alpha: float) -> torch.Tensor:
    return x + alpha * y

4 Fake / Meta impl ★

@scaled_add.register_fake
def _(x, y, alpha):
    torch._check(x.shape == y.shape)
    return torch.empty_like(x)

• fake impl은 shape/dtype만 돌려줌 — 실제 compute 없음
• FakeTensor 경로에서 호출 (Dynamo·Inductor·AOT)
• 없으면 trace 실패 → eager fallback

5 Backward impl

def setup_ctx(ctx, inputs, output):
    x, y, alpha = inputs
    ctx.alpha = alpha

def bwd(ctx, grad):
    return (grad,
            ctx.alpha * grad,
            None)

scaled_add.register_autograd(
    bwd, setup_context=setup_ctx)

6 Inductor lowering 등록

from torch._inductor.lowering import register_lowering

@register_lowering(
    torch.ops.my.scaled_add)
def _(x, y, alpha):
    return x + alpha * y    # IR ops

• 등록 시 Inductor가 op을 inline해 fusion 가능
• 미등록이면 extern kernel로 유지 (fusion 불가)

7 등록 계층 요약

8 torch.library 네임스페이스

• Library("my", "DEF") — 새 op 선언
• Library("aten", "IMPL") — 기존 aten override
• torch.library.opcheck(fn, args) — 등록 정합성 검사

1 Guard의 책임

2 Static vs Symbolic

3 SymInt / SymFloat

• SymInt = symbolic integer, 대수 연산 지원
• backing state: ShapeEnv가 constraint 보관
• shape 관련 비교는 guard 생성의 주요 원천

4 Guard 생성 시점

trace:
  if x.shape[0] == 1024:     # guard: s0 == 1024  (static)
  if x.shape[0] > 0:         # guard: s0 > 0        (constraint)
  y = torch.empty(x.shape)   # no guard (shape-passthrough)
  z = x.view(-1, 8)          # guard: s0 mod 8 == 0

source: torch/fx/experimental/symbolic_shapes.py

5 dynamic=True 효과

• 모든 shape을 SymInt로 시작
• constraint는 필요 시만 추가
• 재컴파일 훨씬 감소
• 단, fusion 기회 약간 감소

6 mark_dynamic / mark_static

torch._dynamo.mark_dynamic(x, 0)
# axis 0만 symbolic, 나머지 static
torch._dynamo.mark_static(y, 1)
# axis 1은 반드시 특정 값으로 컴파일

• mixed 정책: 일부 axis만 dynamic
• 초기 호출에서 한 번 호출 → Dynamo hint

7 재컴파일 감소 전략

1. padding/bucketing으로 shape을 고정된 bucket 집합에 정렬
2. dynamic=True 또는 mark_dynamic으로 symbolic화
3. cache_size_limit 조정
4. graph break 제거 (재컴파일이 graph 단위임)

8 관측 API

9 ShapeEnv

1 TORCH_LOGS 카테고리

여러 카테고리 동시: TORCH_LOGS="+inductor,output_code"

2 TORCH_COMPILE_DEBUG=1

• 모든 중간 산출물을 cache dir에 덤프
• fx_graph · output_code · fusion decisions 포함
• debug_trace/ 하위 디렉토리에 step-by-step

3 Inductor dump 경로 ★

$TORCHINDUCTOR_CACHE_DIR or /tmp/torchinductor_$USER/

  <fxhash>/
    ├ fx_graph_readable.py  ← Dynamo 직후
    ├ fx_graph_transformed.py ← decomp 후
    ├ fx_graph_runnable.py  ← 단독 실행 가능
    ├ output_code.py        ← Python wrapper + triton src
    ├ triton_poi_*.py       ← per-kernel source
    └ debug_trace/
        ├ 0_before_pre_grad_graph.py
        ├ 1_after_decomp_graph.py
        ├ 2_after_post_grad_graph.py
        └ *_ir_pre/post_fusion.txt

4 fx_graph_runnable 활용

# 단독 재현
import torch
exec(open("fx_graph_runnable.py").read())
# args0, mod 변수 사용 가능
mod(*args0)

버그를 최소 FX로 축소할 때 유용. minifier 출력과 동일 형식.

5 소스 디렉토리 지도

6 Trace 연계

• NVTX marker로 각 kernel launch를 표시
• Nsight Systems timeline에서 triton_poi_* 이름으로 식별
• roofline·stall 분석은 ↗ V18

7 흔한 증상 → 힌트

1 Pipeline 단계 × debug 카드

2 단계 산출 IR 요약

3 Fusion 의사결정 1-liner

• pointwise+pointwise → vertical (기본)
• pointwise+reduction → producer 흡수
• reduction+reduction → persistent에서만
• matmul 전후 → template prologue/epilogue
• scatter/mutation → fuse 금지

4 재컴파일 최소화 5원칙

1. shape bucketing
2. dynamic=True 또는 mark_dynamic
3. Python 상수 arg 최소
4. graph break 제거
5. cache_size_limit 확대

5 mode 결정 트리

inference · decode loop  → reduce-overhead
training · 큰 matmul       → max-autotune
디버깅·정확도 bisect       → aot_eager
graph break 원인 찾기      → eager

6 Custom op 등록 체크리스트

• □ kernel impl (CUDA/CPU)
• □ register_fake
• □ register_autograd (학습 시)
• □ register_lowering (fusion 원할 때)
• □ opcheck로 정합성 확인

7 핵심 env 6종

8 타 권 연결

• Triton 내부 pipeline·layout ↗ V11
• MLIR 하위 lowering 일반 ↗ V12
• XLA/TVM 비교 매트릭스 ↗ V14 §15
• distributed training 경로 ↗ V17
• roofline·stall 분석 ↗ V18

9 한 줄 요약