TypeTree support in autodiff #144197

KMJ-007 · 2025-07-19T22:18:34Z

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 · 2025-07-19T23:50:35Z

Currently, I have implemented only for memcpy

KMJ-007 · 2025-07-19T23:50:53Z

r? @ZuseZ4

rustbot · 2025-07-19T23:51:03Z

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 · 2025-07-20T06:10:39Z

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

KMJ-007 · 2025-07-21T12:01:35Z

CI is failing, fixing them!

antoyo · 2025-07-21T18:11:48Z

compiler/rustc_codegen_gcc/src/builder.rs

+            // For now, we just ignore the TypeTree since gcc backend doesn't support autodiff yet
+            // When autodiff support is added to gcc backend, this should attach TypeTree information
+            // as function attributes similar to how LLVM backend does it.
+        }
        // TODO(antoyo): handle aligns and is_volatile.


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

@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.

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 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 · 2025-07-23T17:33:55Z

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

Signed-off-by: Karan Janthe <[email protected]>

rustbot · 2025-07-23T17:38:02Z

⚠️ Warning ⚠️

Some commits in this PR modify submodules.

rustbot · 2025-07-23T18:24:58Z

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

cc @ZuseZ4

This PR modifies run-make tests.

cc @jieyouxu

Signed-off-by: Karan Janthe <[email protected]>

rust-log-analyzer · 2025-07-24T05:10:01Z

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...

ZuseZ4

I haven't reviewed compiler/rustc_middle/src/ty/mod.rs yet, but left a few notes to all other pieces.

ZuseZ4 · 2025-07-28T22:30:25Z

compiler/rustc_ast/src/expand/autodiff_attrs.rs

@@ -85,6 +86,9 @@ pub struct AutoDiffItem {
    /// The name of the function being generated
    pub target: String,
    pub attrs: AutoDiffAttrs,
+    // Type Tree support


redundant comment.

ZuseZ4 · 2025-07-28T22:34:48Z

compiler/rustc_codegen_gcc/src/builder.rs

@@ -1368,6 +1369,7 @@ impl<'a, 'gcc, 'tcx> BuilderMethods<'a, 'tcx> for Builder<'a, 'gcc, 'tcx> {
        _src_align: Align,
        size: RValue<'gcc>,
        flags: MemFlags,
+        _tt: Option<FncTree>,


The gcc people would probably appreciate a commend clarifying, that these tt are an LLVM only feature and thus shouldn't be set for gcc

ZuseZ4 · 2025-07-28T22:35:34Z

compiler/rustc_codegen_llvm/src/back/lto.rs

-                );
+                // Remove 'noinline' if present (it should be there in most cases)
+                if attributes::has_attr(function, Function, llvm::AttributeKind::NoInline) {
+                    attributes::remove_from_llfn(function, Function, llvm::AttributeKind::NoInline);


I'd probably merge the frontend rework first, then you can drop these lines.

ZuseZ4 · 2025-07-28T22:37:06Z

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

+}
+
+#[cfg(not(llvm_enzyme))]
+#[allow(dead_code)]


Why does this need a dead_code? Are you also gating the caller side? Then you shouldn't need this dummy.

ZuseZ4 · 2025-07-28T22:39:31Z

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2673,4 +2673,5 @@ unsafe extern "C" {

    pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
    pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);
+


ZuseZ4 · 2025-07-28T22:40:10Z

compiler/rustc_middle/src/ty/mod.rs

@@ -10,6 +10,7 @@
 //! ["The `ty` module: representing types"]: https://rustc-dev-guide.rust-lang.org/ty.html

 #![allow(rustc::usage_of_ty_tykind)]
+#![allow(unused_imports)]


needs to be cleaned up before landing.

ZuseZ4 · 2025-07-28T22:42:57Z

tests/run-make/autodiff/type-trees/type-analysis/memcpy/memcpy.check

@@ -0,0 +1,9 @@
+CHECK: test_memcpy - {[-1]:Float@double} |{[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Float@double, [-1,16]:Float@double, [-1,24]:Float@double, [-1,32]:Float@double, [-1,40]:Float@double, [-1,48]:Float@double, [-1,56]:Float@double}:{}
+


Shouldn't this test also check for the existance of an llvm.memcpy call?

ZuseZ4 · 2025-07-29T18:19:31Z

compiler/rustc_codegen_llvm/src/llvm/enzyme_ffi.rs

+    use super::{CConcreteType, CTypeTreeRef, Context};
+
+    // TypeTree function fallbacks
+    pub(crate) fn EnzymeNewTypeTree() -> CTypeTreeRef {


mark this (and all other new ones) as unsafe, and your CI failures should be gone.

ZuseZ4 · 2025-07-29T18:22:01Z

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

+            .expect("got a non-UTF8 data-layout from LLVM");
+
+    // Attribute name that Enzyme recognizes for TypeTree information
+    let attr_name = "enzyme_type";


@KMJ-007 Here you mark your attribute as enzyme_type, but in the test below I can't spot any metadata with that name. Are you sure it is getting lowered and not just dropped? Maybe add a few dbg statements to see till where exists.

rustbot added F-autodiff `#![feature(autodiff)]` S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jul 19, 2025

KMJ-007 mentioned this pull request Jul 19, 2025

[WIP] TypeTree support in autodiff #143490

Closed

This comment has been minimized.

Sign in to view

rustbot added the A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. label Jul 19, 2025

This comment has been minimized.

Sign in to view

rustbot assigned ZuseZ4 Jul 19, 2025

KMJ-007 marked this pull request as ready for review July 19, 2025 23:50

rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. and removed S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 19, 2025

This comment has been minimized.

Sign in to view

antoyo reviewed Jul 21, 2025

View reviewed changes

This comment has been minimized.

Sign in to view

rustbot added has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 23, 2025

KMJ-007 added 3 commits July 23, 2025 17:37

basic typetree compiling going to llvm level next

0a56a89

Signed-off-by: Karan Janthe <[email protected]>

memcpy compiling

a03bfe5

Signed-off-by: Karan Janthe <[email protected]>

fmt

aad9b6e

Signed-off-by: Karan Janthe <[email protected]>

KMJ-007 added 6 commits July 23, 2025 17:37

typoes fixing

07685f3

Signed-off-by: Karan Janthe <[email protected]>

typo

d584119

Signed-off-by: Karan Janthe <[email protected]>

fixing todo issue

cd9753d

Signed-off-by: Karan Janthe <[email protected]>

trailing whitespace

46e89c8

Signed-off-by: Karan Janthe <[email protected]>

fixing white space issue because of docs comments

2544376

Signed-off-by: Karan Janthe <[email protected]>

fixing gcc memcpy

94aa8bc

Signed-off-by: Karan Janthe <[email protected]>

rust-cloud-vms bot force-pushed the type-tree branch from d3f78d7 to 94aa8bc Compare July 23, 2025 17:37

rustbot removed has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Jul 23, 2025

This comment has been minimized.

Sign in to view

wip, fixing ci

96ac876

rustbot added the A-run-make Area: port run-make Makefiles to rmake.rs label Jul 23, 2025

This comment has been minimized.

Sign in to view

formatting

340aa40

Signed-off-by: Karan Janthe <[email protected]>

This comment has been minimized.

Sign in to view

KMJ-007 added 3 commits July 23, 2025 20:05

fixing ffi issue

54573be

noinline wait for sometimes

9ffdfe8

fixing ir

e93b3e3

Signed-off-by: Karan Janthe <[email protected]>

This comment has been minimized.

Sign in to view

fixing ci

dbc871a

Signed-off-by: Karan Janthe <[email protected]>

This comment has been minimized.

Sign in to view

KMJ-007 requested a review from ZuseZ4 July 24, 2025 04:57

fixing ci

6ff80be

Signed-off-by: Karan Janthe <[email protected]>

ZuseZ4 reviewed Jul 28, 2025

View reviewed changes

ZuseZ4 reviewed Jul 29, 2025

View reviewed changes

		@@ -2673,4 +2673,5 @@ unsafe extern "C" {

		pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
		pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);

		@@ -0,0 +1,9 @@
		CHECK: test_memcpy - {[-1]:Float@double} \|{[-1]:Pointer, [-1,0]:Float@double, [-1,8]:Float@double, [-1,16]:Float@double, [-1,24]:Float@double, [-1,32]:Float@double, [-1,40]:Float@double, [-1,48]:Float@double, [-1,56]:Float@double}:{}

TypeTree support in autodiff #144197

Are you sure you want to change the base?

TypeTree support in autodiff #144197

Conversation

KMJ-007 commented Jul 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TypeTrees for Autodiff

What are TypeTrees?

Structure

Example: fn compute(x: &f32, data: &[f32]) -> f32

Why Needed?

What Enzyme Does With This Information:

TypeTrees - Offset and -1 Explained

Type Structure

Offset Values

Regular Offset (0, 4, 8, etc.)

Offset -1 (Special: "Everywhere")

Example 1: Array [f32; 100]

Example 2: Slice &[i32]

Example 3: Mixed Structure

Uh oh!

This comment has been minimized.

This comment has been minimized.

KMJ-007 commented Jul 19, 2025

Uh oh!

KMJ-007 commented Jul 19, 2025

Uh oh!

rustbot commented Jul 19, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

rustbot commented Jul 20, 2025

Uh oh!

This comment has been minimized.

KMJ-007 commented Jul 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 Jul 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rustbot commented Jul 23, 2025

Uh oh!

This comment has been minimized.

rustbot commented Jul 23, 2025

Uh oh!

This comment has been minimized.

rustbot commented Jul 23, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

rust-log-analyzer commented Jul 24, 2025

Uh oh!

ZuseZ4 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KMJ-007 commented Jul 19, 2025 •

edited

Loading

Example: `fn compute(x: &f32, data: &[f32]) -> f32`

Example 1: Array `[f32; 100]`

Example 2: Slice `&[i32]`

ZuseZ4 Jul 21, 2025 •

edited

Loading