Skip to content

Conversation

KMJ-007
Copy link
Contributor

@KMJ-007 KMJ-007 commented Jul 19, 2025

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: -1, size: 8, kind: Pointer,
    child: TypeTree(vec![Type {
        offset: -1, size: 4, kind: Float,
        child: TypeTree::new()
    }])
}])

Input 1: data: &[f32]

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

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):

; 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):

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

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 (internal representation):

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:

"enzyme_type"="{[]:Float@float}"

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

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
        }])
    }
])

@rustbot 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
@rust-log-analyzer

This comment has been minimized.

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

This comment has been minimized.

@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 19, 2025

Currently, I have implemented only for memcpy

@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 19, 2025

r? @ZuseZ4

@KMJ-007 KMJ-007 marked this pull request as ready for review July 19, 2025 23:50
@rustbot 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
@rustbot
Copy link
Collaborator

rustbot commented Jul 19, 2025

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

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rustbot
Copy link
Collaborator

rustbot commented Jul 20, 2025

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

@rust-log-analyzer

This comment has been minimized.

@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Jul 21, 2025

CI is failing, fixing them!

@rustbot
Copy link
Collaborator

rustbot commented Jul 23, 2025

Some changes occurred in src/tools/enzyme

cc @ZuseZ4

@rustbot

This comment has been minimized.

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

This comment has been minimized.

@rustbot rustbot removed the has-merge-commits PR has merge commits, merge with caution. label Jul 23, 2025
@KMJ-007
Copy link
Contributor Author

KMJ-007 commented Sep 11, 2025

@ZuseZ4 I have updated the tests and also fixed recursion, need to update the enzyme submodule in order to fix it, but you can review it

@rust-log-analyzer

This comment has been minimized.

@bors
Copy link
Collaborator

bors commented Sep 18, 2025

☔ The latest upstream changes (presumably #146671) made this pull request unmergeable. Please resolve the merge conflicts.

  - Add F128 support to TypeTree Kind enum
  - Implement TypeTree FFI bindings and conversion functions
  - Add typetree.rs module for metadata attachment to LLVM functions
  - Integrate TypeTree generation with autodiff intrinsic pipeline
  - Support scalar types: f32, f64, integers, f16, f128
  - Attach enzyme_type attributes as LLVM string metadata for Enzyme

Signed-off-by: Karan Janthe <[email protected]>
  - Fix nott-flag test to emit LLVM IR and check enzyme_type attributes
  - Replace TODO comments with actual TypeTree metadata verification
  - Test that NoTT flag properly disables TypeTree generation
  - Test that TypeTree enabled generates proper enzyme_type attributes

Signed-off-by: Karan Janthe <[email protected]>
  - Add specific tests for f32, f64, i32, f16, f128 TypeTree generation
  - Verify correct enzyme_type metadata for each scalar type
  - Ensure TypeTree metadata matches expected Enzyme format

Signed-off-by: Karan Janthe <[email protected]>
@rustbot
Copy link
Collaborator

rustbot commented Sep 19, 2025

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.

@rust-log-analyzer

This comment has been minimized.

@rust-cloud-vms rust-cloud-vms bot force-pushed the type-tree branch 2 times, most recently from cc0b6c4 to 2f742ea Compare September 19, 2025 04:21
@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. A-run-make Area: port run-make Makefiles to rmake.rs F-autodiff `#![feature(autodiff)]` S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

8 participants