TypeTree support in autodiff

[KMJ-007]

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

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

TypeTree(vec![Type {
    offset: 0, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: 0, size: 4, kind: Float,
        child: TypeTree::new()
    }])
}])

Input 1: data: &[f32]

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

Output: f32

TypeTree(vec![Type {
    offset: 0, 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):

; 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

}

With TypeTrees (our goal):

// Enzyme knows:

// - %p1 points to struct with f32 at +0, f32 at +4, i32 at +8

// - Only the f32 fields need derivatives

// - Can generate efficient derivative code directly

TypeTrees - Offset and -1 Explained

Type Structure

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

struct Point {

x: f32, // offset 0, size 4

y: f32, // offset 4, size 4

id: i32, // offset 8, size 4

}

TypeTree for &Point:

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

])

Offset -1 (Special: "Everywhere")

Means "this pattern repeats for ALL elements"

Example 1: Array [f32; 100]

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]

// Pointer to slice data

TypeTree(vec![Type {

	offset: 0, 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

struct Container {

	header: i64, // offset 0

	data: [f32; 1000], // offset 8, but elements use -1

}

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

}])

}

])

[KMJ-007]

Currently, I have implemented only for memcpy

[KMJ-007]

r? @ZuseZ4

[rustbot]

Some changes occurred in compiler/rustc_ast/src/expand/autodiff_attrs.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

[rustbot]

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

[KMJ-007]

CI is failing, fixing them!

[rustbot]

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

[rustbot]

Some changes occurred in src/tools/cargo

cc @ehuss

[rustbot]

This PR was rebased onto a different master commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[KMJ-007]

tidy is failing for typo EnzymeTypeTreeShiftIndiciesEq to EnzymeTypeTreeShiftIndicesEq, but that function on enzyme side also has the same name,

what should i do?

[rust-log-analyzer]

The job tidy failed! Check out the build log: (web) (plain enhanced) (plain)

Click to see the possible cause of the failure (guessed by this bot)

finished building tool typos
error: `Indicies` should be `Indices`
##[error]  --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:113:42
    |
113 |         pub(crate) fn EnzymeTypeTreeShiftIndiciesEq(
    |                                          ^^^^^^^^
    |
error: `Indicies` should be `Indices`
##[error]  --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:226:45
    |
226 |     pub(crate) unsafe fn EnzymeTypeTreeShiftIndiciesEq(
    |                                             ^^^^^^^^
    |
error: `Indicies` should be `Indices`
##[error]  --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:303:32
    |
303 |             EnzymeTypeTreeShiftIndiciesEq(
    |                                ^^^^^^^^
    |
tidy error: checks with external tool 'typos' failed
some tidy checks failed
Command `/checkout/obj/build/x86_64-unknown-linux-gnu/stage1-tools-bin/rust-tidy /checkout /checkout/obj/build/x86_64-unknown-linux-gnu/stage0/bin/cargo /checkout/obj/build 4 /node/bin/npm --extra-checks=py,cpp,js,spellcheck` failed with exit code 1
Created at: src/bootstrap/src/core/build_steps/tool.rs:1583:23
Executed at: src/bootstrap/src/core/build_steps/test.rs:1225:29

Command has failed. Rerun with -v to see more details.
Bootstrap failed while executing `test --stage 0 src/tools/tidy tidyselftest --extra-checks=py,cpp,js,spellcheck`
Build completed unsuccessfully in 0:02:27
  local time: Sun Aug 24 07:48:39 UTC 2025
  network time: Sun, 24 Aug 2025 07:48:39 GMT
##[error]Process completed with exit code 1.