Autodiff를 위한 TypeTree

TypeTree란 무엇입니까?

Enzyme을 위한 메모리 레이아웃 기술자입니다. 타입이 메모리에서 어떻게 구성되는지 Enzyme에 정확히 알려 주어 도함수를 효율적으로 계산할 수 있게 합니다.

구조

#![allow(unused)]
fn main() {
TypeTree(Vec<Type>)

Type {
    offset: isize,  // byte offset (-1 = everywhere)
    size: usize,    // size in bytes
    kind: Kind,     // Float, Integer, Pointer, etc.
    child: TypeTree // nested structure
}
}

예시: fn compute(x: &f32, data: &[f32]) -> f32

입력 0: x: &f32

#![allow(unused)]
fn main() {
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: 0, size: 4, kind: Float,  // Single value: use offset 0
        child: TypeTree::new()
    }])
}])
}

입력 1: data: &[f32]

#![allow(unused)]
fn main() {
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,  // -1 = all elements
        child: TypeTree::new()
    }])
}])
}

출력: f32

#![allow(unused)]
fn main() {
TypeTree(vec![Type {
    offset: 0, size: 4, kind: Float,  // Single scalar: use offset 0
    child: TypeTree::new()
}])
}

왜 필요합니까?

• Enzyme은 LLVM IR로부터 복잡한 타입 레이아웃을 추론할 수 없습니다
• 느린 메모리 패턴 분석을 피할 수 있습니다
• 중첩된 구조체에 대해 올바른 도함수 계산을 가능하게 합니다
• 어떤 바이트가 미분 가능한 값이고 어떤 바이트가 메타데이터인지 Enzyme에 알려 줍니다

Enzyme이 이 정보로 하는 일:

TypeTree가 없는 경우:

; Enzyme sees generic LLVM IR:
define float @distance(ptr %p1, ptr %p2) {
; Has to guess what these pointers point to
; Slow analysis of all memory operations
; May miss optimization opportunities
}

TypeTree가 있는 경우:

define "enzyme_type"="{[-1]:Float@float}" float @distance(
    ptr "enzyme_type"="{[-1]:Pointer, [-1,0]:Float@float}" %p1, 
    ptr "enzyme_type"="{[-1]:Pointer, [-1,0]:Float@float}" %p2
) {
; Enzyme knows exact type layout
; Can generate efficient derivative code directly
}

TypeTree - 오프셋과 -1에 대한 설명

타입 구조

#![allow(unused)]
fn main() {
Type {
    offset: isize, // WHERE this type starts
    size: usize,   // HOW BIG this type is
    kind: Kind,    // WHAT KIND of data (Float, Int, Pointer)
    child: TypeTree // WHAT'S INSIDE (for pointers/containers)
}
}

오프셋 값

일반 오프셋(0, 4, 8 등)

구조체 내의 특정 바이트 위치

#![allow(unused)]
fn main() {
struct Point {
    x: f32, // offset 0, size 4
    y: f32, // offset 4, size 4
    id: i32, // offset 8, size 4
}
}

&Point에 대한 TypeTree(내부 표현):

#![allow(unused)]
fn main() {
TypeTree(vec![
    Type { offset: 0, size: 4, kind: Float },   // x at byte 0
    Type { offset: 4, size: 4, kind: Float },   // y at byte 4
    Type { offset: 8, size: 4, kind: Integer }  // id at byte 8
])
}

생성되는 LLVM

"enzyme_type"="{[-1]:Pointer, [-1,0]:Float@float, [-1,4]:Float@float, [-1,8]:Integer, [-1,9]:Integer, [-1,10]:Integer, [-1,11]:Integer}"

오프셋 -1(특수 값: “모든 곳”)

“이 패턴이 모든 요소에 대해 반복됨“을 의미합니다

예시 1: 직접 배열 [f32; 100](포인터 간접 참조 없음)

#![allow(unused)]
fn main() {
TypeTree(vec![Type {
    offset: -1, // ALL positions
    size: 4,    // each f32 is 4 bytes
    kind: Float, // every element is float
}])
}

생성되는 LLVM: "enzyme_type"="{[-1]:Float@float}"

예시 1b: 배열 참조 &[f32; 100](포인터 간접 참조 있음)

#![allow(unused)]
fn main() {
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, // ALL array elements
        size: 4,    // each f32 is 4 bytes
        kind: Float, // every element is float
    }])
}])
}

생성되는 LLVM: "enzyme_type"="{[-1]:Pointer, [-1,-1]:Float@float}"

오프셋 0,4,8,12...396을 가진 100개의 개별 Type을 나열하는 대신

예시 2: 슬라이스 &[i32]

#![allow(unused)]
fn main() {
// Pointer to slice data
TypeTree(vec![Type {
    offset: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, // ALL slice elements
        size: 4,    // each i32 is 4 bytes
        kind: Integer
    }])
}])
}

LLVM을 생성합니다: "enzyme_type"="{[-1]:Pointer, [-1,-1]:Integer}"

예제 3: 혼합 구조

#![allow(unused)]
fn main() {
struct Container {
    header: i64,        // offset 0
    data: [f32; 1000],  // offset 8, but elements use -1
}
}

#![allow(unused)]
fn main() {
TypeTree(vec![
    Type { offset: 0, size: 8, kind: Integer }, // header
    Type { offset: 8, size: 4000, kind: Pointer,
        child: TypeTree(vec![Type {
            offset: -1, size: 4, kind: Float // ALL array elements
        }])
    }
])
}

핵심 차이점: 단일 값과 배열

단일 값은 정밀도를 위해 오프셋 0을 사용합니다:

• &f32는 오프셋 0에 정확히 하나의 f32 값을 가집니다
• -1(“모든 위치”)을 사용하는 것보다 더 정밀합니다
• 생성 결과: {[-1]:Pointer, [-1,0]:Float@float}

배열은 효율성을 위해 오프셋 -1을 사용합니다:

• &[f32; 100]에는 같은 패턴이 100번 반복됩니다
• -1을 사용하면 100개의 개별 오프셋을 나열하지 않아도 됩니다
• 생성 결과: {[-1]:Pointer, [-1,-1]:Float@float}