Memory-Centric IR Specification - AICF Intermediate Representation

1. Overview

Memory-Centric IR 은 AICF 컴파일러에서 Semantic IR 과 Kernel IR 사이의 중간 표현이다.

MCIR 의 목적은 다음을 표현하는 것이다.

• data movement
• tile lifecycle
• memory residency
• streaming computation
• materialization boundaries

Semantic IR 이 연산 의미를 표현한다면

what to compute

MCIR 은 실행 구조와 데이터 흐름을 표현한다.

how data moves during execution

 

2. Position in Compilation Pipeline

AICF compilation pipeline

Python Model
     │
     ▼
Semantic IR
     │
     ▼
Memory-Centric IR (MCIR)
     │
     ▼
Kernel IR
     │
     ▼
CUDA Kernel

Semantic IR - MCIR 변환은다음을 수행한다.

• executino region 생성
• streaming structure 구성
• tile structure 정의
• memory reuse 모델링

 

3. Core IR Object Model

MCIR 은 Region-based IR 이다.

기본 구조

Region
 ├── Node
 ├── Region
 └── Boundary

핵심 객체

ExecutionRegion
StreamingRegion
TileRegion
ResidencyRegion
MaterializationBoundary
MCNode
MCValue

 

4. MCValue

MCValue 는 MCIR 에서 데이터 객체를 나타낸다.

MCValue
{
    id: ValueID
    shape: TensorShape
    dtype: DataType

    producer: MCNode
    consumers: list<MCNode>

    residency: MemoryLocation
}

Example

value_17
{
    shape: [128, 64]
    dtype: fp16
    residency: shared
}

 

5. MCNode

MCNode 는 computation 또는 data movement 를 표현한다.

MCNode
{
    id: NodeID
    op: MCOpType

    inputs: list<MCValue>
    outputs: list<MCValue>

    attributes: dict
}

MCNode 는 다음 종류가 있다.

LoadTile
StoreTile
Compute
Reduce
UpdateState
Prefetch
Barrier

 

6. Region System

MCIR 은 region hierarchy 를 사용한다.

ExecutionRegion
   └── StreamingRegion
          └── TileRegion

Region 은 다음 구조를 가진다.

Region
{
    id: RegionID
    type: RegionType

    inputs
    outputs

    nodes: list<MCNode>
    subregions: list<Region>
}

 

7. Execution Region

Execution Region 은 semantic computation unit 이다.

• Attention 
• LayerNorm
• Fused elementwise chain

ExecutionRegion
{
    id: region_0

    inputs:
        Q
        K
        V

    outputs:
        O

    subregions:
        StreamingRegion
}

 

8. Streaming Region

StreamRegion 은 streaming computation pipeline 을 표현한다.

StreamingRegion
{
    stream_axis: sequence

    inputs:
        Q
        K
        V

    state:
        softmax_max
        softmax_sum
        accumulator
}

Streaming Region 은 다음 특징을 가진다.

• intermediate tensor 없음
• partial state 유지
• sequential streaming 

 

9. TileRegion

TileRegion 은 tile-based execution 단위이다.

TileRegion
{
    tile_shape:
        m: 128
        n: 64
        k: 64

    nodes:
        LoadTile
        Compute
        Reduce
        UpdateState
}

TileRegion 은 GPU execution block 과 대응된다.

 

10. Residency Region

Residency Region 은 on-chip memory lifetime 을 표현한다.

ResidencyRegion
{
    memory: shared

    values:
        Q_tile
        K_tile
        V_tile

    lifetime:
        tile_iteration
}

가능한 memory location

register
shared
global

 

11. Materialization Boundary

MaterializationBoundary 는 tensor materialization 지점을 나타낸다.

MaterializationBoundary
{
    value: score_matrix
    location: global_memory
}

컴파일러는 가능한 경우 boundary 를 제거한다.

 

12. Example: Attention MCIR

Semantic IR

Scores = MatMul(Q, Kᵀ)
Scores = Softmax(Scores)
Output = MatMul(Scores, V)

MCIR

ExecutionRegion (Attention)

    StreamingRegion

        for q_tile:

            TileRegion

                LoadTile(Q)

                for kv_tile:

                    LoadTile(K)
                    LoadTile(V)

                    Compute(score)

                    UpdateState(softmax)

                    Compute(accumulate)

 

여기서 중요한 점은

score matrix never materialized

 

13. Lowering Contract

MCIR - Kernel IR 변환 규칙

TileRegion	thread block loop
StreamingRegion	outer loop
ResidencyRegion	shared/register allocation
MCNode	instruction sequence

예

TileRegion

->

for block_tile:
	compute

 

14. Example Kernel Structure

MCIR

Streaming Region

Lowered Kernel

for q_tile:

    load Q_tile

    init softmax

    for kv_tile:

        load K_tile
        load V_tile

        compute score
        update softmax
        accumulate

store output

 

15. Transformation Interface

MCIR transforamtion pass

AttentionPatternPass
StreamingConversionPass
TileSchedulingPass
ResidencyAssignmentPass
MaterializationEliminationPass

각 pass 는 MCIR 을 수정한다.

 

16. Advantages

MCIR 을 사용하면 다음이 가능하다

Explicit memory model

메모리 이동이 IR 에서 명시된다.

Streaming computatoin 표현

FlashAttention 같은 구조를 자연스럽게 표현 가능

Architecture-specific optimization

GPU 구조에 맞는 tile schedule 생성

 

17. Summary

Memory-Centric IR 은 AICF 의 핵심 IR 계층이다.

Semantic IR 이 다음을 표현한다면

mathematical computation

MCIR 은 다음을 표현한다.

execution structure
memory movement
tile lifecycle

이를 통해 AICF 는 semantic graph 로부터 hardward-efficient execution pipeline 을 생성하는 compile 을 목표로 한다.