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!

[antoyo]

Could you please move this comment right above the call to new_call()?

[ZuseZ4]

@KMJ-007 The TODO might be a bit too ambitious, I don't think anyone is currently working on autodiff for GCC,so I would just remove it and the code change. Also, even if they want to add it one day, they should rather do it from scratch to be able to learn some lessons from Enzyme, so I'm not sure if TypeTrees would even be a thing for them.

[antoyo]

Oh, I was talking about the existing TODO:

// TODO(antoyo): handle aligns and is_volatile.

I would like it to stay where it was before this change.

[KMJ-007]

@KMJ-007 The TODO might be a bit too ambitious, I don't think anyone is currently working on autodiff for GCC,so I would just remove it and the code change. Also, even if they want to add it one day, they should rather do it from scratch to be able to learn some lessons from Enzyme, so I'm not sure if TypeTrees would even be a thing for them.

I added it for temp, because CI was failing for gcc, i think i shouldn't have changed memcpy for gcc, but somewhere they are sharing the same trait or struct, I don't remember exactly, but will look into it again, and handle the changes so we don't break gcc or change it

[rustbot]

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

[rustbot]

⚠️ Warning ⚠️

• Some commits in this PR modify submodules.

[rustbot]

Some changes occurred in compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

cc @ZuseZ4

This PR modifies run-make tests.

cc @jieyouxu

[rust-log-analyzer]

The job pr-check-2 failed! Check out the build log: (web) (plain enhanced) (plain)

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

[RUSTC-TIMING] rustc_mir_build test:false 3.396
    Checking rustc_mir_transform v0.0.0 (/checkout/compiler/rustc_mir_transform)
[RUSTC-TIMING] rustc_passes test:false 1.854
    Checking rustc_borrowck v0.0.0 (/checkout/compiler/rustc_borrowck)
error: unused imports: `c_char` and `c_uint`
 --> compiler/rustc_codegen_llvm/src/builder/autodiff.rs:1:20
  |
1 | use std::os::raw::{c_char, c_uint};
  |                    ^^^^^^  ^^^^^^
  |
  = note: `-D unused-imports` implied by `-D warnings`
  = help: to override `-D warnings` add `#[allow(unused_imports)]`

---

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:275:21
    |
275 |         let inner = unsafe { EnzymeNewTypeTree() };
    |                     ^^^^^^ unnecessary `unsafe` block
    |
    = note: `-D unused-unsafe` implied by `-D warnings`
    = help: to override `-D warnings` add `#[allow(unused_unsafe)]`

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:280:21
    |
280 |         let inner = unsafe { EnzymeNewTypeTreeCT(t, ctx) };
    |                     ^^^^^^ unnecessary `unsafe` block

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:285:9
    |
---

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:318:21
    |
318 |         let inner = unsafe { EnzymeNewTypeTreeTR(self.inner) };
    |                     ^^^^^^ unnecessary `unsafe` block

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:325:19
    |
325 |         let ptr = unsafe { EnzymeTypeTreeToString(self.inner) };
    |                   ^^^^^^ unnecessary `unsafe` block

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:333:9
    |
333 |         unsafe {
    |         ^^^^^^ unnecessary `unsafe` block

error: unnecessary `unsafe` block
   --> compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs:349:9
    |
349 |         unsafe { EnzymeFreeTypeTree(self.inner) }
    |         ^^^^^^ unnecessary `unsafe` block

[RUSTC-TIMING] rustc_codegen_llvm test:false 3.424
error: could not compile `rustc_codegen_llvm` (lib) due to 10 previous errors
warning: build failed, waiting for other jobs to finish...