Matthias Krüger
c29fb2e57e
Rollup merge of #144197 - KMJ-007:type-tree, r=ZuseZ4
...
TypeTree support in autodiff
# TypeTrees for Autodiff
## What are TypeTrees?
Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.
## Structure
```rust
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
}
```
## Example: `fn compute(x: &f32, data: &[f32]) -> f32`
**Input 0: `x: &f32`**
```rust
TypeTree(vec![Type {
offset: -1, size: 8, kind: Pointer,
child: TypeTree(vec![Type {
offset: -1, size: 4, kind: Float,
child: TypeTree::new()
}])
}])
```
**Input 1: `data: &[f32]`**
```rust
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()
}])
}])
```
**Output: `f32`**
```rust
TypeTree(vec![Type {
offset: -1, size: 4, kind: Float,
child: TypeTree::new()
}])
```
## Why Needed?
- Enzyme can't deduce complex type layouts from LLVM IR
- Prevents slow memory pattern analysis
- Enables correct derivative computation for nested structures
- Tells Enzyme which bytes are differentiable vs metadata
## What Enzyme Does With This Information:
Without TypeTrees (current state):
```llvm
; 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
}
```
With TypeTrees (our implementation):
```llvm
define "enzyme_type"="{[]:Float@float}" float ``@distance(``
ptr "enzyme_type"="{[]:Pointer}" %p1,
ptr "enzyme_type"="{[]:Pointer}" %p2
) {
; Enzyme knows exact type layout
; Can generate efficient derivative code directly
}
```
# TypeTrees - Offset and -1 Explained
## Type Structure
```rust
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)
}
```
## Offset Values
### Regular Offset (0, 4, 8, etc.)
**Specific byte position within a structure**
```rust
struct Point {
x: f32, // offset 0, size 4
y: f32, // offset 4, size 4
id: i32, // offset 8, size 4
}
```
TypeTree for `&Point` (internal representation):
```rust
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
])
```
Generates LLVM:
```llvm
"enzyme_type"="{[]:Float@float}"
```
### Offset -1 (Special: "Everywhere")
**Means "this pattern repeats for ALL elements"**
#### Example 1: Array `[f32; 100]`
```rust
TypeTree(vec![Type {
offset: -1, // ALL positions
size: 4, // each f32 is 4 bytes
kind: Float, // every element is float
}])
```
Instead of listing 100 separate Types with offsets `0,4,8,12...396`
#### Example 2: Slice `&[i32]`
```rust
// 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
}])
}])
```
#### Example 3: Mixed Structure
```rust
struct Container {
header: i64, // offset 0
data: [f32; 1000], // offset 8, but elements use -1
}
```
```rust
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
}])
}
])
```
2025-09-28 18:13:11 +02:00
..
2025-03-01 22:11:14 +01:00
2024-06-06 19:16:53 +02:00
2024-06-06 19:16:53 +02:00
2024-06-06 19:16:53 +02:00
2024-06-18 14:25:59 -04:00
2025-06-19 07:38:01 +08:00
2025-06-16 23:04:33 +00:00
2025-09-05 16:18:49 +02:00
2025-06-18 18:39:25 +08:00
2025-06-18 18:39:25 +08:00
2025-07-10 13:55:04 +08:00
2024-08-10 18:06:10 +00:00
2025-06-23 12:38:35 -07:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-07-23 08:32:55 +02:00
2025-09-19 04:11:35 +00:00
2025-07-22 14:27:48 +02:00
2024-06-07 11:12:04 +02:00
2025-06-24 09:06:49 +00:00
2025-06-12 15:10:12 +02:00
2024-12-28 11:58:32 +08:00
2025-05-09 19:54:43 +08:00
2024-07-23 10:31:21 -04:00
2024-07-23 10:31:21 -04:00
2025-07-10 13:55:04 +08:00
2024-08-10 18:06:10 +00:00
2025-07-04 17:27:24 +02:00
2024-07-19 12:27:47 -04:00
2024-07-19 12:27:47 -04:00
2025-07-10 13:55:04 +08:00
2024-09-22 19:11:29 -04:00
2025-07-10 13:55:04 +08:00
2025-07-10 13:55:04 +08:00
2024-07-09 14:30:57 -04:00
2025-06-12 15:10:12 +02:00
2024-09-05 08:43:38 +00:00
2024-06-14 13:59:44 +02:00
2024-09-05 08:43:38 +00:00
2024-09-05 08:43:38 +00:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-07-10 13:55:04 +08:00
2025-07-13 14:28:41 +08:00
2024-08-10 18:06:10 +00:00
2025-07-08 09:04:21 +00:00
2025-02-26 13:07:12 +01:00
2024-07-17 13:34:18 +00:00
2025-08-06 16:55:50 +01:00
2025-03-11 09:46:34 -07:00
2025-05-16 16:14:24 +02:00
2025-06-10 23:31:05 +08:00
2025-06-27 22:13:01 +02:00
2025-04-12 15:09:06 +08:00
2024-07-29 08:26:52 +10:00
2025-09-22 18:18:36 +02:00
2025-07-28 16:32:11 -06:00
2025-04-16 17:42:39 +00:00
2025-04-12 15:09:06 +08:00
2025-06-12 15:10:12 +02:00
2025-07-31 10:38:34 +02:00
2024-10-09 20:58:27 +11:00
2025-06-16 23:04:33 +00:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-06-28 11:18:46 -04:00
2025-06-12 15:10:12 +02:00
2024-06-15 17:34:43 +02:00
2025-06-12 15:10:12 +02:00
2024-08-15 15:44:29 +02:00
2025-06-12 15:10:12 +02:00
2025-04-07 20:48:40 +00:00
2025-09-26 22:04:00 +02:00
2025-03-24 11:33:56 +01:00
2025-05-22 20:12:50 +02:00
2025-09-26 22:04:00 +02:00
2025-06-19 07:38:01 +08:00
2024-08-29 10:15:17 -04:00
2025-06-17 10:49:28 +08:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-07-17 13:34:18 +00:00
2024-09-27 15:53:26 -07:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2024-06-03 18:35:54 +00:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-08-15 10:26:03 -07:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2024-06-19 16:43:22 -04:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-07-17 14:54:52 +08:00
2025-06-19 07:38:01 +08:00
2024-08-10 18:06:10 +00:00
2025-06-19 07:38:01 +08:00
2024-07-17 13:34:18 +00:00
2025-06-12 15:10:12 +02:00
2024-09-22 19:11:29 -04:00
2024-07-19 12:27:47 -04:00
2024-07-17 15:33:17 -04:00
2025-01-15 01:05:16 +08:00
2024-07-17 15:33:06 -04:00
2024-12-17 14:33:10 -08:00
2024-07-17 15:33:05 -04:00
2024-07-17 15:33:07 -04:00
2024-07-19 12:27:47 -04:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2024-06-18 16:30:26 -04:00
2024-06-07 11:12:24 +02:00
2025-04-12 15:09:06 +08:00
2025-07-31 10:38:34 +02:00
2025-07-15 12:52:39 +02:00
2024-06-20 18:20:12 +02:00
2024-08-02 10:06:20 -04:00
2024-09-22 19:11:29 -04:00
2024-07-26 10:17:39 -04:00
2024-06-30 12:40:24 +02:00
2025-01-22 17:22:39 +01:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2024-12-23 09:40:38 +00:00
2024-07-17 13:34:18 +00:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2024-12-23 20:23:57 +08:00
2025-06-10 23:31:05 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2024-07-17 13:34:18 +00:00
2025-06-12 15:10:12 +02:00
2024-07-25 11:08:07 -04:00
2025-06-19 07:38:01 +08:00
2024-07-29 08:26:52 +10:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-10-22 19:43:22 +08:00
2025-06-12 15:10:12 +02:00
2025-06-19 07:38:01 +08:00
2025-01-22 12:26:50 +08:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-02-07 01:41:54 +08:00
2024-12-23 09:40:25 +00:00
2025-08-06 23:59:47 +00:00
2025-08-06 23:59:47 +00:00
2024-06-26 11:44:23 -04:00
2024-12-23 09:40:25 +00:00
2025-06-12 15:10:12 +02:00
2025-07-10 13:55:04 +08:00
2025-07-30 00:39:43 -04:00
2025-07-10 13:55:04 +08:00
2025-07-18 12:02:43 +00:00
2024-07-22 23:21:59 +02:00
2024-07-23 11:47:11 +02:00
2024-07-18 16:27:34 -04:00
2025-08-11 23:31:07 +02:00
2024-10-27 21:23:28 -04:00
2024-07-19 12:27:47 -04:00
2025-07-31 10:38:34 +02:00
2025-09-26 08:14:25 +02:00
2025-07-22 14:28:48 +02:00
2025-05-15 11:33:28 -07:00
2025-06-12 15:10:12 +02:00
2024-08-06 11:03:18 -04:00
2024-08-06 11:03:18 -04:00
2024-10-10 15:02:47 +01:00
2024-07-29 08:26:52 +10:00
2025-06-12 15:10:12 +02:00
2024-06-26 13:32:30 -04:00
2024-06-25 14:27:43 -04:00
2024-07-22 13:25:39 -04:00
2025-08-10 15:50:40 +02:00
2024-07-19 12:27:47 -04:00
2024-08-10 18:06:10 +00:00
2024-05-28 11:41:53 -04:00
2024-07-18 16:27:34 -04:00
2024-07-18 09:28:30 -04:00
2025-06-12 15:10:12 +02:00
2024-07-18 09:28:30 -04:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-06-19 14:39:09 -04:00
2025-06-12 15:10:12 +02:00
2025-06-16 23:04:33 +00:00
2024-06-25 12:46:58 -04:00
2025-07-30 00:39:43 -04:00
2025-06-12 15:10:12 +02:00
2024-12-23 09:40:25 +00:00
2024-05-31 21:30:08 +10:00
2024-06-11 09:53:31 -04:00
2024-07-17 13:34:18 +00:00
2025-06-19 07:38:01 +08:00
2024-09-22 11:40:00 +02:00
2025-08-04 17:05:20 -07:00
2025-07-08 16:40:10 +08:00
2025-09-06 00:21:04 +02:00
2025-06-12 15:10:12 +02:00
2025-09-21 21:13:15 +02:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2024-07-19 12:27:46 -04:00
2025-07-10 13:55:04 +08:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-07-10 13:55:04 +08:00
2025-06-24 09:06:49 +00:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2024-05-14 17:02:20 -04:00
2024-08-04 02:57:18 +00:00
2024-06-22 17:06:47 +00:00
2024-07-29 08:26:52 +10:00
2024-06-07 11:12:24 +02:00
2025-06-12 15:10:12 +02:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2024-07-29 08:26:52 +10:00
2025-06-12 15:10:12 +02:00
2024-10-09 20:58:27 +11:00
2025-07-16 23:40:49 +08:00
2024-07-17 13:34:18 +00:00
2024-07-17 13:34:18 +00:00
2025-04-12 15:09:06 +08:00
2024-06-17 14:37:49 -04:00
2024-06-29 11:55:30 +02:00
2025-06-12 15:10:12 +02:00
2024-07-25 11:17:34 -04:00
2024-06-07 11:12:24 +02:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-07-05 12:47:19 -04:00
2025-06-12 15:10:12 +02:00
2024-07-19 12:27:47 -04:00
2024-06-25 12:46:58 -04:00
2024-09-10 17:23:05 +02:00
2024-09-10 17:23:05 +02:00
2025-08-13 09:42:46 +08:00
2025-02-23 21:23:36 +08:00
2024-10-09 20:58:27 +11:00
2024-10-09 20:58:27 +11:00
2025-06-12 15:10:12 +02:00
2025-08-13 09:42:46 +08:00
2025-08-13 09:42:46 +08:00
2025-09-22 11:29:54 +01:00
2025-09-22 11:29:54 +01:00
2024-07-17 13:34:18 +00:00
2024-07-17 13:34:18 +00:00
2024-06-26 11:44:11 -04:00
2025-06-12 15:10:12 +02:00
2024-08-26 18:20:32 +08:00
2024-08-10 18:06:10 +00:00
2025-04-01 18:29:39 +02:00
2025-07-10 13:55:04 +08:00
2025-09-12 20:53:28 +02:00
2024-11-02 23:39:29 +11:00
2024-08-26 18:20:32 +08:00
2024-08-26 18:20:32 +08:00
2024-05-31 21:30:08 +10:00
2025-06-19 07:38:01 +08:00
2024-07-16 15:34:04 -04:00
2025-07-10 13:55:04 +08:00
2024-08-05 10:11:53 -04:00
2025-06-16 23:04:33 +00:00
2024-07-12 16:41:28 -04:00
2025-02-26 19:09:51 +01:00
2025-02-26 19:09:51 +01:00
2025-02-26 19:09:51 +01:00
2025-07-10 13:55:04 +08:00
2025-07-10 13:55:04 +08:00
2025-08-08 09:17:35 -07:00
2025-07-10 13:55:04 +08:00
2024-06-25 15:02:19 -04:00
2024-06-08 13:10:55 +02:00
2024-08-10 18:06:10 +00:00
2024-06-20 16:09:39 -04:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2024-12-23 09:40:38 +00:00
2025-06-19 07:38:01 +08:00
2025-04-12 15:09:06 +08:00
2025-06-23 14:34:08 -07:00
2024-08-29 10:15:17 -04:00
2025-06-12 15:10:12 +02:00
2024-06-30 12:40:24 +02:00
2024-07-19 12:27:26 -04:00
2024-07-17 13:34:18 +00:00
2025-07-10 13:55:04 +08:00
2025-06-19 07:38:01 +08:00
2024-08-10 18:06:10 +00:00
2024-07-04 10:05:56 -04:00
2024-08-10 18:06:10 +00:00
2025-07-08 09:21:35 +00:00
2025-08-27 23:44:49 +02:00
2025-07-08 09:04:21 +00:00
2025-07-08 09:04:21 +00:00
2025-07-08 09:04:21 +00:00
2025-09-12 20:53:28 +02:00
2025-06-19 07:38:01 +08:00
2024-07-18 11:14:07 -04:00
2025-08-23 12:15:37 +02:00
2025-09-03 00:06:44 -07:00
2025-08-15 10:26:03 -07:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-09-22 20:36:18 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-02-27 13:22:16 +00:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-09-03 09:09:18 -07:00
2024-09-22 19:11:29 -04:00
2025-07-27 05:21:47 +02:00
2025-08-27 23:44:49 +02:00
2025-05-08 15:24:05 +02:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-06-12 15:10:12 +02:00
2025-08-21 16:08:00 +07:00
2025-08-21 16:08:00 +07:00
2025-08-21 16:08:00 +07:00
2024-09-05 08:43:38 +00:00
2025-09-03 15:24:49 +02:00
2025-06-19 07:38:01 +08:00
2024-05-17 11:49:20 -04:00
2025-06-29 22:29:12 +08:00
2024-07-17 13:34:18 +00:00
2025-06-16 23:04:33 +00:00
2025-07-10 13:55:04 +08:00
2025-06-12 15:10:12 +02:00
2025-07-10 13:55:04 +08:00
2024-07-17 15:33:07 -04:00
2025-06-18 18:39:25 +08:00
2024-05-31 21:30:08 +10:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-07-10 13:55:04 +08:00
2025-06-19 07:38:01 +08:00
2024-07-08 09:59:08 -04:00
2025-06-19 07:38:01 +08:00
2025-06-19 07:38:01 +08:00
2025-04-12 15:09:06 +08:00
2024-07-29 08:26:52 +10:00
2025-01-31 04:30:51 +00:00
2025-07-10 13:55:04 +08:00
2025-08-06 18:01:07 +00:00
2025-06-12 15:10:12 +02:00
2024-09-16 21:31:41 +08:00
2024-09-16 21:31:41 +08:00
2024-09-16 21:31:41 +08:00
2025-06-18 18:39:25 +08:00
2025-03-24 11:33:56 +01:00
2025-08-27 23:44:49 +02:00
2024-08-26 18:20:32 +08:00
2024-07-19 16:48:21 -04:00
2024-07-02 11:37:59 -04:00
2025-06-21 13:16:57 -07:00
2025-06-19 07:38:01 +08:00
2024-10-09 20:58:27 +11:00
2025-01-29 08:11:01 +00:00
2025-06-12 15:10:12 +02:00
2025-06-12 15:10:12 +02:00
2025-06-17 10:53:11 +08:00
2025-06-17 10:53:11 +08:00
2024-07-17 16:11:11 -04:00
2025-06-12 15:10:12 +02:00
2025-07-13 14:28:49 +08:00
2024-06-28 16:28:57 -04:00
2025-06-27 10:03:29 +00:00
2025-05-27 11:12:15 +02:00
2025-03-01 22:11:02 +01:00
2024-07-29 08:26:52 +10:00
2024-07-29 08:26:52 +10:00
2024-07-29 08:26:52 +10:00
2024-07-01 20:29:57 +02:00
2025-02-08 22:12:13 +00:00
2024-07-29 08:26:52 +10:00
2024-08-02 11:54:20 -04:00
2025-08-12 12:07:14 +02:00
2024-07-29 08:26:52 +10:00
2024-07-19 10:27:14 -07:00
2024-07-29 08:26:52 +10:00
2024-07-29 08:26:52 +10:00
2024-07-29 08:26:52 +10:00
2025-08-15 15:48:46 +03:00
2024-07-29 08:26:52 +10:00
2024-07-29 08:26:52 +10:00
2024-06-28 20:32:31 +03:00
2024-06-07 11:12:24 +02:00
2024-07-29 08:26:52 +10:00
2024-08-04 02:57:17 +00:00
2025-03-02 05:56:56 +08:00