ExecuTorch Compiler 상세

개요

ExecuTorch 컴파일러는 PyTorch 모델을 엣지 디바이스에서 효율적으로 실행할 수 있는 .pte 파일로 변환하는 정적(AOT) 컴파일 시스템이다. 기존 Backend 문서가 백엔드 시스템과 딜리게이트 메커니즘에 초점을 맞춘 반면, 본 문서는 컴파일러 파이프라인의 내부 동작에 집중한다.

ExecuTorch 컴파일러의 핵심은 EXIR(ExecuTorch Intermediate Representation)이라는 중간 표현 체계이다. ATen 다이어렉티브 → Edge 다이어렉티브 → Backend 다이어렉티브로 이어지는 3단계 다이어렉티브 변환을 통해, 사용자 모델을 하드웨어에 최적화된 실행 가능한 형태로 낮추(lowering)는 것이 목표이다. 이 과정에서 커스텀 컴파일러 패스, 연산 융합, 메모리 플래닝, 코드 생성 등 다양한 최적화가 수행된다.

핵심 개념

EXIR 다이어렉티브 체계

ExecuTorch는 모델을 단계적으로 낮추기 위해 세 가지 다이어렉티브를 사용한다:

ATen 다이어렉티브

ATen 다이어렉티브는 torch.export()의 직접적인 출력으로, eager 모드 PyTorch 프로그램이 Exported IR 그래프가 되는 첫 번째 단계이다:

import torch

# 모델 정의
model = MyModel().eval()
example_inputs = (torch.randn(1, 3, 224, 224),)

# Export → ATen 다이어렉티브
aten_program = torch.export.export(model, example_inputs)

ATen 다이어렉티브 속성:
- 기능화(functionalization)가 수행되어 텐서 별칭(alias)과 변이(mutation)가 제거됨
- 모든 연산자는 torch.ops.aten 네임스페이스의 ATen 연산자이거나 등록된 커스텀 연산자
- 모든 연산자에 meta kernel이 있어야 함 (입력 shape으로 출력 shape 추론)
- 모든 텐서는 torch.contiguous_format 메모리 포맷 사용
- 동적 dtype, 암묵적 dtype 프로모션, 암묵적 브로드캐스팅 지원

Edge 다이어렉티브

Edge 다이어렉티브는 엣지 디바이스에 유용하지만 하드웨어에 아직 특화되지 않은 중간 단계이다:

from executorch import exir

# ATen → Edge 변환
edge_program: exir.EdgeProgramManager = exir.to_edge(aten_program)

Edge 다이어렉티브 속성:
- 모든 연산자는 dtype 특화된 Edge 연산자 (하드웨어 독립적)
- 모든 Scalar 타입은 Tensor로 변환
- edge.yaml에 dtype 제약 조건 정의
- 커스텀 그래프 변환(Pass) 적용 가능

dtype 제약 조건 예시:

- func: sigmoid
  namespace: edge
  inherits: aten::sigmoid
  type_alias:
    T0: [Bool, Byte, Char, Int, Long, Short]
    T1: [Double, Float]
    T2: [Float]
  type_constraint:
  - self: T0
    __ret_0: T2
  - self: T1
    __ret_0: T1

Backend 다이어렉티브

Backend 다이어렉티브는 선택적 단계로, 백엔드가 그래프를 하드웨어별 연산자, 메타데이터, 딜리게이트된 로어 모듈로 재작성할 때 사용된다:

# Edge → Backend 변환 (백엔드별 파티셔너 사용)
program = to_edge_transform_and_lower(
    aten_program,
    partitioner=[XnnpackPartitioner()]
).to_executorch()

Backend 다이어렉티브 특징:
- Edge 다이어렉티브 이후에만 실행
- 하드웨어별 융합, 로어링 패턴, 딜리게이트 통합에 사용
- 커스텀 연산자와 달리 백엔드 전용 연산자는 이 단계에서만 도입

커스텀 컴파일러 패스

ExecuTorch는 그래프 최적화를 위한 다양한 커스텀 패스를 제공한다:

패스 분류 체계:

ExportPass (Level 1 패스):

from executorch.exir.pass_base import ExportPass

class ReplaceInPlaceReluWithOutOfPlaceReluPass(ExportPass):
    """in-place relu_를 out-of-place relu로 교체"""

    def call_operator(self, op, args, kwargs, meta):
        if op != torch.ops.aten.relu_.default:
            return super().call_operator(op, args, kwargs, meta)
        return super().call_operator(
            Op(torch.ops.aten.relu.default), args, kwargs, meta
        )

# 패스 실행
replace_pass = ReplaceInPlaceReluWithOutOfPlaceReluPass()
new_graph_module = replace_pass(graph_module).graph_module

Subgraph Rewriter (Level 2 패스):

from torch.fx import subgraph_rewriter

def replace_patterns(graph_module):
    def pattern(x, y):
        x = torch.ops.aten.add.Tensor(x, y)
        x = torch.ops.aten.mul.Tensor(x, y)
        return x

    def replacement(x, y):
        return torch.ops.aten.sub.Tensor(x, y)

    return subgraph_rewriter.replace_pattern_with_filters(
        graph_module, pattern, replacement
    )

파티셔너

파티셔너는 모델 그래프를 하드웨어 지원 서브그래프와 CPU 폴백으로 분할한다:

from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
from torch.fx.passes.operator_support import OperatorSupportBase

class AddMulOperatorSupport(OperatorSupportBase):
    def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
        return node.op == "call_function" and node.target in [
            torch.ops.aten.add.Tensor,
            torch.ops.aten.mul.Tensor,
        ]

partitioner = CapabilityBasedPartitioner(graph_module, AddMulOperatorSupport())
partitions = partitioner.propose_partitions()

동작 원리

전체 컴파일러 파이프라인

ExecuTorch 컴파일러의 전체 파이프라인은 다음 단계로 구성된다:

1단계: 모델 캡처 (torch.export)

import torch
from executorch.exir import to_edge_transform_and_lower

# eager 모드 모델 정의
model = MyModel().eval()
example_inputs = (torch.randn(1, 3, 224, 224),)

# Export → ATen 다이어렉티브
aten_program = torch.export.export(model, example_inputs)

2단계: 엣지 변환 (to_edge)

# ATen → Edge 다이어렉티브
# dtype 특화, Scalar → Tensor 변환 수행
edge_program = exir.to_edge(aten_program)

3단계: 그래프 최적화 (커스텀 패스)

# 커스텀 패스를 통한 그래프 변환
edge_program = edge_program.transform(MyCustomPass())

4단계: 하드웨어 로어링 (to_backend)

# 파티셔너를 통한 그래프 분할 및 백엔드 딜리게이트
program = to_edge_transform_and_lower(
    aten_program,
    partitioner=[XnnpackPartitioner()]
)

5단계: ExecuTorch 프로그램 생성 (to_executorch)

# Backend 다이어렉티브 → .pte 파일 생성
executorch_program = program.to_executorch()

메모리 플래닝

ExecuTorch 컴파일러는 실행 시 메모리 할당을 최적화하기 위해 정적 메모리 플래닝을 수행한다:

코드 생성

Backend 다이어렉티브에서 .pte 파일로 변환될 때, ExecuTorch는 다음 코드를 생성한다:

.pte 파일 구조

.pte 파일 구조:
├── 헤더 (magic number, 버전, 체크섬)
├── 세그먼트 테이블
│   ├── 프로그램 세그먼트 (그래프 구조)
│   ├── 커널 세그먼트 (연산자 구현)
│   └── 데이터 세그먼트 (가중치)
├── 메서드 디스크립터
│   ├── 입력/출력 사양
│   ├── 메모리 플래닝 정보
│   └── 실행 순서
└── 백엔드 딜리게이트 정보

비교/분석

컴파일러 파이프라인 비교

커스텀 패스 프레임워크 비교

코드 생성 전략 비교

장단점

장점

1. 단계적 로어링 (Progressive Lowering): ATen → Edge → Backend 3단계를 통해 점진적으로 하드웨어에 특화되어, 각 단계에서 독립적인 최적화가 가능하다
2. IR Spec 기반 검증: 각 다이어렉티브에서 IR Spec을 준수하는지 검증하여, 컴파일 오류를 조기에 발견할 수 있다
3. 유연한 커스텀 패스: ExportPass, Subgraph Rewriter 등 다양한 수준의 패스를 제공하여 개발자가 세밀한 그래프 변환을 수행할 수 있다
4. 정적 메모리 플래닝: 실행 시 동적 할당 없이 미리 메모리를 계획하여, 엣지 디바이스에서 불필요한 오버헤드를 제거한다
5. 초경량 .pte 파일: FlatBuffer 기반의 플랫 바이너리로, 파싱 오버헤드가 최소화된다

단점

1. 컴파일 시간: AOT 컴파일 방식으로 인해 대규모 모델에서 컴파일 시간이 길어질 수 있다
2. 커스텀 커널 제약: EXIR 연산자 세트 외의 커스텀 커널은 별도 등록이 필요하며, Edge 다이어렉티브에서의 dtype 제약을 준수해야 한다
3. 동적 Shape 제한: 정적 메모리 플래닝 특성상 동적 shape 지원이 제한적이다
4. 학습 미지원: 추론 전용 컴파일러로, 학습 그래프는 지원하지 않는다
5. 디버깅 어려움: AOT 컴파일 과정에서 그래프 변환이 다수 발생하여, 원본 모델과의 대응 관계 추적이 어려울 수 있다

관련 기술

참고 문헌

• ExecuTorch 공식 문서 - Export IR
• ExecuTorch 공식 문서 - Custom Compiler Passes
• ExecuTorch GitHub
• ExecuTorch: A Unified PyTorch Solution to Run AI Models On-Device (arXiv:2605.08195)
• Meta Engineering Blog - ExecuTorch

핵심 정리

1. ExecuTorch 컴파일러는 ATen → Edge → Backend 3단계 다이어렉티브를 통해 모델을 점진적으로 낮추는 AOT 컴파일 시스템이다.
2. EXIR은 각 다이어렉티브에서 IR Spec을 준수하도록 설계되어, 컴파일 과정에서의 오류를 조기에 발견할 수 있다.
3. ExportPass, Subgraph Rewriter 등 다양한 커스텀 패스를 통해 그래프 수준의 세밀한 최적화가 가능하다.
4. 정적 메모리 플래닝과 FlatBuffer 기반 .pte 파일 생성으로 엣지 디바이스에 최적화된 실행 환경을 제공한다.
5. 12개 이상의 하드웨어 백엔드를 지원하며, 딜리게이트 시스템을 통해 하드웨어 전환이 용이하다.