updating the docs

Author: KMJ-007

src/autodiff/type-trees.md

@@ -1,118 +1,163 @@
-# Type Trees in Enzyme
+# TypeTrees for Autodiff
 
-This document describes type trees as used by Enzyme for automatic differentiation.
+## What are TypeTrees?
+Memory layout descriptors for Enzyme. Tell Enzyme exactly how types are structured in memory so it can compute derivatives efficiently.
 
-## What are Type Trees?
-
-Type trees in Enzyme are a way to represent the types of variables, including their activity (e.g., whether they are active, duplicated, or contain duplicated data) for automatic differentiation. They provide a structured way for Enzyme to understand how to handle different data types during the differentiation process.
+## Structure
+```rust
+TypeTree(Vec<Type>)
 
-## Representing Rust Types as Type Trees
+Type {
+    offset: isize,  // byte offset (-1 = everywhere)
+    size: usize,    // size in bytes
+    kind: Kind,     // Float, Integer, Pointer, etc.
+    child: TypeTree // nested structure
+}
+```
 
-Enzyme needs to understand the structure and properties of Rust types to perform automatic differentiation correctly. This is where type trees come in. They provide a detailed map of a type, including pointer indirections and the underlying concrete data types.
+## Example: `fn compute(x: &f32, data: &[f32]) -> f32`
 
-The `-enzyme-rust-type` flag in Enzyme helps in interpreting types more accurately in the context of Rust's memory layout and type system.
+**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()
+    }])
+}])
+```
 
-### Primitive Types
+**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()
+    }])
+}])
+```
 
-#### Floating-Point Types (`f32`, `f64`)
+**Output: `f32`**
+```rust
+TypeTree(vec![Type {
+    offset: -1, size: 4, kind: Float,
+    child: TypeTree::new()
+}])
+```
 
-Consider a Rust reference to a 32-bit floating-point number, `&f32`.
+## 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
 
-In LLVM IR, this might be represented, for instance, as an `i8*` (a generic byte pointer) that is then `bitcast` to a `float*`. Consider the following LLVM IR function:
+## What Enzyme Does With This Information:
 
+Without TypeTrees:
 ```llvm
-define internal void @callee(i8* %x) {
-start:
-  %x.dbg.spill = bitcast i8* %x to float*
-  ; ...
-  ret void
+; Enzyme sees generic LLVM IR:
+define float @distance(i8* %p1, i8* %p2) {
+; Has to guess what these pointers point to
+; Slow analysis of all memory operations
+; May miss optimization opportunities
 }
 ```
 
-When Enzyme analyzes this function (with appropriate flags like `-enzyme-rust-type`), it might produce the following type information for the argument `%x` and the result of the bitcast:
-
+With TypeTrees:
 ```llvm
-i8* %x: {[-1]:Pointer, [-1,0]:Float@float}
-%x.dbg.spill = bitcast i8* %x to float*: {[-1]:Pointer, [-1,0]:Float@float}
+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
+}
 ```
 
-**Understanding the Type Tree: `{[-1]:Pointer, [-1,0]:Float@float}`**
-
-This string is the type tree representation. Let's break it down:
+# TypeTrees - Offset and -1 Explained
 
-*   **`{ ... }`**: This encloses the set of type information for the variable.
-*   **`[-1]:Pointer`**:
-    *   `[-1]` is an index or path. In this context, `-1` often refers to the base memory location or the immediate value pointed to.
-    *   `Pointer` indicates that the variable `%x` itself is treated as a pointer.
-*   **`[-1,0]:Float@float`**:
-    *   `[-1,0]` is a path. It means: start with the base item `[-1]` (the pointer), and then look at offset `0` from the memory location it points to.
-    *   `Float` is the `CConcreteType` (from `enzyme_ffi.rs`, corresponding to `DT_Float`). It signifies that the data at this location is a floating-point number.
-    *   `@float` is a subtype or specific variant of `Float`. In this case, it specifies a single-precision float (like Rust's `f32`).
+## Type Structure
 
-A reference to an `f64` (e.g., `&f64`) is handled very similarly. The LLVM IR might cast to `double*`:
-```llvm
-define internal void @callee(i8* %x) {
-start:
-  %x.dbg.spill = bitcast i8* %x to double*
-  ; ...
-  ret void
+```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)
 }
 ```
 
-And the type tree would be:
+## Offset Values
 
-```llvm
-i8* %x: {[-1]:Pointer, [-1,0]:Float@double}
-```
-The key difference is `@double`, indicating a double-precision float.
-
-This level of detail allows Enzyme to know, for example, that if `x` is an active variable in differentiation, the floating-point value it points to needs to be handled according to AD rules for its specific precision.
-
-### Compound Types
-
-#### Structs
-
-Consider a Rust struct `T` with two `f32` fields (e.g., a reference `&T`):
+### Regular Offset (0, 4, 8, etc.)
+**Specific byte position within a structure**
 
 ```rust
-struct T {
-    x: f32,
-    y: f32,
+struct Point {
+    x: f32, // offset 0, size 4
+    y: f32, // offset 4, size 4
+    id: i32, // offset 8, size 4
 }
-
-// And a function taking a reference to it:
-// fn callee(t: &T) { /* ... */ }
 ```
 
-In LLVM IR, a pointer to this struct might be initially represented as `i8*` and then cast to the specific struct type, like `{ float, float }*`:
+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
-define internal void @callee(i8* %t) {
-start:
-  %t.dbg.spill = bitcast i8* %t to { float, float }*
-  ; ...
-  ret void
-}
+"enzyme_type"="{[]:Float@float}"
 ```
 
-The Enzyme type analysis output for `%t` would be:
+### Offset -1 (Special: "Everywhere")
+**Means "this pattern repeats for ALL elements"**
 
-```llvm
-i8* %t: {[-1]:Pointer, [-1,0]:Float@float, [-1,4]:Float@float}
+#### 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
+}])
 ```
 
-**Understanding the Struct Type Tree: `{[-1]:Pointer, [-1,0]:Float@float, [-1,4]:Float@float}`**
+Instead of listing 100 separate Types with offsets `0,4,8,12...396`
 
-*   **`[-1]:Pointer`**: As before, this indicates that `%t` is a pointer.
-*   **`[-1,0]:Float@float`**:
-    *   This describes the first field of the struct (`x`).
-    *   `[-1,0]` means: from the memory location pointed to by `%t` (`-1`), at offset `0` bytes.
-    *   `Float@float` indicates this field is an `f32`.
-*   **`[-1,4]:Float@float`**:
-    *   This describes the second field of the struct (`y`).
-    *   `[-1,4]` means: from the memory location pointed to by `%t` (`-1`), at offset `4` bytes.
-    *   `Float@float` indicates this field is also an `f32`.
+#### 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
+    }])
+}])
+```
 
-The offset `4` comes from the size of the first field (`f32` is 4 bytes). If the first field were, for example, an `f64` (8 bytes), the second field might be at offset `[-1,8]`. Enzyme uses these offsets to pinpoint the exact memory location of each field within the struct.
+#### Example 3: Mixed Structure
+```rust
+struct Container {
+    header: i64,        // offset 0
+    data: [f32; 1000],  // offset 8, but elements use -1
+}
+```
 
-This detailed mapping is crucial for Enzyme to correctly track the activity of individual struct fields during automatic differentiation.
+```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
+        }])
+    }
+])
+```