Skip to content

Commit 6960460

Browse files
committed
docs: typetree in autodiff
1 parent 18121a9 commit 6960460

File tree

2 files changed

+164
-0
lines changed

2 files changed

+164
-0
lines changed

src/SUMMARY.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,7 @@
108108
- [Installation](./autodiff/installation.md)
109109
- [How to debug](./autodiff/debugging.md)
110110
- [Autodiff flags](./autodiff/flags.md)
111+
- [Type Trees](./autodiff/type-trees.md)
111112

112113
# Source Code Representation
113114

src/autodiff/type-trees.md

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
# TypeTrees for Autodiff
2+
3+
## What are TypeTrees?
4+
Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.
5+
6+
## Structure
7+
```rust
8+
TypeTree(Vec<Type>)
9+
10+
Type {
11+
offset: isize, // byte offset (-1 = everywhere)
12+
size: usize, // size in bytes
13+
kind: Kind, // Float, Integer, Pointer, etc.
14+
child: TypeTree // nested structure
15+
}
16+
```
17+
18+
## Example: `fn compute(x: &f32, data: &[f32]) -> f32`
19+
20+
**Input 0: `x: &f32`**
21+
```rust
22+
TypeTree(vec![Type {
23+
offset: -1, size: 8, kind: Pointer,
24+
child: TypeTree(vec![Type {
25+
offset: -1, size: 4, kind: Float,
26+
child: TypeTree::new()
27+
}])
28+
}])
29+
```
30+
31+
**Input 1: `data: &[f32]`**
32+
```rust
33+
TypeTree(vec![Type {
34+
offset: -1, size: 8, kind: Pointer,
35+
child: TypeTree(vec![Type {
36+
offset: -1, size: 4, kind: Float, // -1 = all elements
37+
child: TypeTree::new()
38+
}])
39+
}])
40+
```
41+
42+
**Output: `f32`**
43+
```rust
44+
TypeTree(vec![Type {
45+
offset: -1, size: 4, kind: Float,
46+
child: TypeTree::new()
47+
}])
48+
```
49+
50+
## Why Needed?
51+
- Enzyme can't deduce complex type layouts from LLVM IR
52+
- Prevents slow memory pattern analysis
53+
- Enables correct derivative computation for nested structures
54+
- Tells Enzyme which bytes are differentiable vs metadata
55+
56+
## What Enzyme Does With This Information:
57+
58+
Without TypeTrees:
59+
```llvm
60+
; Enzyme sees generic LLVM IR:
61+
define float @distance(i8* %p1, i8* %p2) {
62+
; Has to guess what these pointers point to
63+
; Slow analysis of all memory operations
64+
; May miss optimization opportunities
65+
}
66+
```
67+
68+
With TypeTrees:
69+
```llvm
70+
define "enzyme_type"="{[]:Float@float}" float @distance(
71+
ptr "enzyme_type"="{[]:Pointer}" %p1,
72+
ptr "enzyme_type"="{[]:Pointer}" %p2
73+
) {
74+
; Enzyme knows exact type layout
75+
; Can generate efficient derivative code directly
76+
}
77+
```
78+
79+
# TypeTrees - Offset and -1 Explained
80+
81+
## Type Structure
82+
83+
```rust
84+
Type {
85+
offset: isize, // WHERE this type starts
86+
size: usize, // HOW BIG this type is
87+
kind: Kind, // WHAT KIND of data (Float, Int, Pointer)
88+
child: TypeTree // WHAT'S INSIDE (for pointers/containers)
89+
}
90+
```
91+
92+
## Offset Values
93+
94+
### Regular Offset (0, 4, 8, etc.)
95+
**Specific byte position within a structure**
96+
97+
```rust
98+
struct Point {
99+
x: f32, // offset 0, size 4
100+
y: f32, // offset 4, size 4
101+
id: i32, // offset 8, size 4
102+
}
103+
```
104+
105+
TypeTree for `&Point` (internal representation):
106+
```rust
107+
TypeTree(vec![
108+
Type { offset: 0, size: 4, kind: Float }, // x at byte 0
109+
Type { offset: 4, size: 4, kind: Float }, // y at byte 4
110+
Type { offset: 8, size: 4, kind: Integer } // id at byte 8
111+
])
112+
```
113+
114+
Generates LLVM:
115+
```llvm
116+
"enzyme_type"="{[]:Float@float}"
117+
```
118+
119+
### Offset -1 (Special: "Everywhere")
120+
**Means "this pattern repeats for ALL elements"**
121+
122+
#### Example 1: Array `[f32; 100]`
123+
```rust
124+
TypeTree(vec![Type {
125+
offset: -1, // ALL positions
126+
size: 4, // each f32 is 4 bytes
127+
kind: Float, // every element is float
128+
}])
129+
```
130+
131+
Instead of listing 100 separate Types with offsets `0,4,8,12...396`
132+
133+
#### Example 2: Slice `&[i32]`
134+
```rust
135+
// Pointer to slice data
136+
TypeTree(vec![Type {
137+
offset: -1, size: 8, kind: Pointer,
138+
child: TypeTree(vec![Type {
139+
offset: -1, // ALL slice elements
140+
size: 4, // each i32 is 4 bytes
141+
kind: Integer
142+
}])
143+
}])
144+
```
145+
146+
#### Example 3: Mixed Structure
147+
```rust
148+
struct Container {
149+
header: i64, // offset 0
150+
data: [f32; 1000], // offset 8, but elements use -1
151+
}
152+
```
153+
154+
```rust
155+
TypeTree(vec![
156+
Type { offset: 0, size: 8, kind: Integer }, // header
157+
Type { offset: 8, size: 4000, kind: Pointer,
158+
child: TypeTree(vec![Type {
159+
offset: -1, size: 4, kind: Float // ALL array elements
160+
}])
161+
}
162+
])
163+
```

0 commit comments

Comments
 (0)