본문 바로가기

AI Compiler framework

AI Compiler 프로젝트 개발 재시작 방향

1. 개발 목표 재정의

이 프로젝트의 핵심 목표는 개별 CUDA 커널을 빠르게 구현하는 것이 아니다.

목표는 모델을 구성하는 연산의 의미와 특성을 분석하고, 이를 기반으로 그래프를 변환·축소·최적화하는 AI Compiler의 최소 구조를 만드는 것이다.

따라서 개발의 출발점은 커널 최적화가 아니라 다음 질문이어야 한다.

이 연산은 입력 정보를 어떻게 변환하고, 어떤 정보를 보존하거나 제거하며, 다른 연산과 어떻게 결합될 수 있는가?


2. 연산자 정의부터 시작

먼저 소수의 핵심 연산자를 선정한다.

  • Linear / MatMul
  • Add / Bias
  • ReLU
  • Dropout
  • Reshape / Transpose
  • Softmax

각 연산자는 단순 실행 코드가 아니라 다음 정보를 가져야 한다.

  • 입력과 출력의 shape
  • dtype
  • 수학적 의미
  • 선형성 여부
  • 가역성 여부
  • 정보 붕괴 여부
  • 파라미터 수
  • 연산량과 메모리 접근량
  • 결합 가능한 연산
  • 제거하거나 단순화할 수 있는 조건
  • 수치적으로 안전한 변환 조건

이 메타데이터가 이후 그래프 최적화의 근거가 된다.


3. 최소 프레임워크 구현

초기 프레임워크는 다음 기능만 제공한다.

Tensor

  • shape
  • dtype
  • device
  • data 또는 buffer 참조

Operator

  • 입력과 출력 정의
  • shape inference
  • 실행 함수
  • 연산 특성 메타데이터

Graph

  • 연산 노드 등록
  • Tensor 간 연결
  • 위상 정렬
  • 그래프 출력 및 검증

Executor

  • 그래프 순차 실행
  • CPU 기준 구현
  • 중간 Tensor 값 확인

처음부터 CUDA Graph, Autograd, 동적 분기, 메모리 풀까지 포함하지 않는다.

우선 작은 그래프를 정확하게 표현하고 실행할 수 있어야 한다.


4. 첫 번째 최적화 패스 구현

최소 그래프가 동작하면 간단한 rewrite부터 구현한다.

예시:

  • 연속된 Linear 연산 결합
  • MatMul + Bias → Linear
  • Linear + ReLU → FusedLinearReLU
  • 불필요한 Reshape 제거
  • 연속 Transpose 제거
  • 상수 연산 미리 계산
  • 사용되지 않는 노드 제거
  • 동일한 연산 결과 재사용

초기에는 성능 향상보다 그래프를 분석하고 변환하는 구조 자체를 만드는 것이 중요하다.


5. 표현력과 붕괴 분석 연결

각 연산에 다음과 같은 분석 값을 부여한다.

  • 표현 공간을 확장하는가
  • 표현 공간을 유지하는가
  • 표현 공간을 제한하는가
  • 입력 정보 일부를 비가역적으로 제거하는가
  • 파라미터 수 대비 실제 활성화되는 경로는 어느 정도인가

예를 들어 ReLU는 음수 영역을 제거하므로 입력 상태 일부를 동일한 0으로 붕괴시킨다.

Dropout은 학습 중 임의의 표현 경로를 제거해 과도한 표현 공간 탐색을 제한한다.

Linear는 출력 차원과 rank에 따라 표현 가능한 부분공간을 확장하거나 축소한다.

이러한 정보는 이후 다음 최적화에 활용할 수 있다.

  • 과도하게 큰 Linear 축소
  • 저랭크 근사
  • 중복 표현 경로 제거
  • 불필요한 확장 후 붕괴 패턴 탐지
  • 모델 구조 단순화 후보 제안

6. 개발 순서

1단계

Tensor, Operator, Graph의 최소 인터페이스를 정의한다.

2단계

Linear, Add, ReLU만으로 작은 MLP 그래프를 실행한다.

3단계

각 연산의 메타데이터와 의미 분석 구조를 추가한다.

4단계

그래프 rewrite와 pattern matching을 구현한다.

5단계

연산량, 메모리량, 파라미터 수, 붕괴 특성을 기반으로 최적화 후보를 출력한다.

6단계

CPU 구현이 안정된 뒤 CUDA backend와 기존 커널을 연결한다.


7. 당장 구현할 최소 결과물

첫 번째 목표는 다음 코드가 동작하는 상태다.

Input
  ↓
Linear
  ↓
ReLU
  ↓
Linear
  ↓
Output

이 그래프에 대해 프레임워크가 다음 정보를 출력하도록 한다.

Node 0: Linear
Input shape: [1, 10]
Output shape: [1, 20]
Parameters: 220
Properties: linear, trainable

Node 1: ReLU
Input shape: [1, 20]
Output shape: [1, 20]
Properties: nonlinear, irreversible, collapse-producing

Node 2: Linear
Input shape: [1, 20]
Output shape: [1, 5]
Parameters: 105
Properties: linear, trainable

그다음 컴파일러가 다음과 같은 분석을 제공하도록 한다.

Total parameters: 325
Potential collapse point: ReLU
Linear expansion: 10 → 20
Linear contraction: 20 → 5
Optimization candidates:
- inspect effective rank of first Linear
- measure ReLU activation ratio
- consider hidden dimension reduction

결론

개발을 다시 시작할 때는 CUDA 커널이나 개별 연산 성능부터 들어가지 않는다.

먼저 연산을 표현하고, 연산의 의미를 기록하고, 그래프를 분석하고 변환할 수 있는 최소 컴파일러 구조를 만든다.

커널은 그 구조가 안정된 이후에 연결되는 실행 수단이다.

현재 가장 먼저 구현해야 할 것은 다음 세 가지다.

Operator 정의
Graph 표현
최소 실행기

이 세 가지가 만들어지면 기존에 구현한 CUDA 커널, 메모리 플래너, Graph Capture 기능도 다시 명확한 위치를 가지게 된다.