AICF Memory IR Design - MCIR Memory Property System

1. Motivation

대부분의 AI 컴파일러는 메모리 최적화를 다음 방식으로 수행한단.

kernel fusion
layout transform
heuristic scheduling

하지만 이러한 접근은 보통 다음 문제를 가진다.

pattern 이 명시적으로 표현되지 않음
optimization legality 가 불명확
compiler reasoning 이 어려움

AICF 는 메모리 최적화를 구조적 property 기반 IR 로 표현한다.

즉 메모리 최적화는 다음과 같이 정의된다.

memory optimization
=
memory-relevant computation properties
+
graph rewrite rules

 

2. Core Concept

MCIR(Memory-Centric IR) 의 핵심 개념은 다음이다.

Operator semantics
+
Memory behavior properties

즉 연산은 단순히

Y = f(X)

로 표현되는 것이 아니라

Y = f(X)
with memory properties

를 가진다.

예

softmax_weighted_sum
property;
	weighted_streaming_reduction

 

3. Memory Properties

MCIR 에서는 메모리 관련 구조를 property system 으로 표현한다.

현재 AICF 는 다음 네 가지 property 를 정의한다.

online_reducible_norm
weighted_streaming_reduction
rematerializable_intermediate
tile_compatible_compute

이 property 들은 패턴 카탈로그와 직접 대응한다.

 

4. Property 1 online_reducible_norm

의미

통계 계산이 streaming reduction 형태로 변환 가능하다.

즉 다음 구조를 가진다.

state update
+
mergeable partial statistics

IR 표현

reduce_norm(x)

properties:
    online_reducible_norm

Legality

statistics state mergeable
update associative

rewrite

two-pass reduction
→ streaming reduction

lowering

global load
→ local accumulator
→ hierarchical merge

 

5. Property 2 weighted_streaming_reduction

의미

정규화된 weighted reductino 이 streaming 방식으로 계산 가능하다

대표 예

softmax(QK^T)V

IR 표현

weighted_reduce(scores, values)

properties:
    weighted_streaming_reduction

legality

rescaling invariant
normalization mergeable

rewrite

materialized attention
→ streaming attention

lowering

tile load
→ local score compute
→ online normalization
→ weighted accumulation

 

6. Property 3 rematerializable_intermediate

의미

intermediate tensor 가 저장 대신 재계산 가능하다

즉

Y = f(X)

이고

f(X)

를 언제든 다시 계산할 수 있다.

IR 표현

op: elementwise_gelu

properties:
    rematerializable_intermediate

legality

pure function
deterministic
side-effect free

rewrite

store intermediate
→ recompute

lowering

inline recompute

 

7. Property 4 tile_compatible_compute

의미

연산이 tile 내부에서 닫힌 실행 구조를 만들 수 있다.

즉

dependency closure within tile

IR 표현

op: matmul

properties:
    tile_compatible_compute

legality

partial accumulation 가능
working-set fits on-chip

rewrite

naive schedule
→ tiled schedule

lowering

shared memory resident kernel

 

8. Property Interaction

중요한 점은 이 property 들이 독립적인 것이 아니라 결합될 수 있다는 것이다.

예

FlashAttention

tile_compatible_compute
+
weighted_streaming_reduction
+
rematerialization

즉 하나의 kernel 이 여러 memory pattern 을 동시에 사용할 수 있다.

 

9. MCIR Graph Example

예: attention graph

QKᵀ
→ softmax
→ weighted sum

기존 IR

matmul
softmax
matmul

MCIR

attention_op

properties:
    weighted_streaming_reduction
    tile_compatible_compute

rewirte

3 operators
-> fused streaming kernel

 

10. Compiler Pipeline

AICF memory optimization pipeline

graph capture
↓
pattern detection
↓
property annotation
↓
graph rewrite
↓
kernel generation

즉 컴파일러는

pattern → property → rewrite

과정을 수행한다.

 

 

11. Why Property-Based IR Matters

이 접근의 장점은 다음과 같다

• legality reasoning
  • rewirte correctness 명확
• compiler extensibility
  • 새 memory pattern 추가 가능
• hardware abstraction
  • archiecture-independent optimization

 

12. Final Perspective

AICF 는 메모리 최적화를 다음처러 ㅁ본다

memory optimization
=
execution structure transformation

즉 단순한 kernel tuning 이 아니라

computation graph
→ memory-aware execution graph

로 변환하는 과정이다.