rustc_ast/expand/
typetree.rs

Help
1//! This module contains the definition of the `TypeTree` and `Type` structs.
2//! They are thin Rust wrappers around the TypeTrees used by Enzyme as the LLVM based autodiff
3//! backend. The Enzyme TypeTrees currently have various limitations and should be rewritten, so the
4//! Rust frontend obviously has the same limitations. The main motivation of TypeTrees is to
5//! represent how a type looks like "in memory". Enzyme can deduce this based on usage patterns in
6//! the user code, but this is extremely slow and not even always sufficient. As such we lower some
7//! information from rustc to help Enzyme. For a full explanation of their design it is necessary to
8//! analyze the implementation in Enzyme core itself. As a rough summary, `-1` in Enzyme speech means
9//! everywhere. That is `{0:-1: Float}` means at index 0 you have a ptr, if you dereference it it
10//! will be floats everywhere. Thus `* f32`. If you have `{-1:int}` it means int's everywhere,
11//! e.g. [i32; N]. `{0:-1:-1 float}` then means one pointer at offset 0, if you dereference it there
12//! will be only pointers, if you dereference these new pointers they will point to array of floats.
13//! Generally, it allows byte-specific descriptions.
14//! FIXME: This description might be partly inaccurate and should be extended, along with
15//! adding documentation to the corresponding Enzyme core code.
16//! FIXME: Rewrite the TypeTree logic in Enzyme core to reduce the need for the rustc frontend to
17//! provide typetree information.
18//! FIXME: We should also re-evaluate where we create TypeTrees from Rust types, since MIR
19//! representations of some types might not be accurate. For example a vector of floats might be
20//! represented as a vector of u8s in MIR in some cases.
21
22use std::fmt;
23
24use crate::expand::{Decodable, Encodable, HashStable_Generic};
25
26#[derive(Clone, Copy, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
27pub enum Kind {
28    Anything,
29    Integer,
30    Pointer,
31    Half,
32    Float,
33    Double,
34    Unknown,
35}
36
37#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
38pub struct TypeTree(pub Vec<Type>);
39
40impl TypeTree {
41    pub fn new() -> Self {
42        Self(Vec::new())
43    }
44    pub fn all_ints() -> Self {
45        Self(vec![Type { offset: -1, size: 1, kind: Kind::Integer, child: TypeTree::new() }])
46    }
47    pub fn int(size: usize) -> Self {
48        let mut ints = Vec::with_capacity(size);
49        for i in 0..size {
50            ints.push(Type {
51                offset: i as isize,
52                size: 1,
53                kind: Kind::Integer,
54                child: TypeTree::new(),
55            });
56        }
57        Self(ints)
58    }
59}
60
61#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
62pub struct FncTree {
63    pub args: Vec<TypeTree>,
64    pub ret: TypeTree,
65}
66
67#[derive(Clone, Eq, PartialEq, Encodable, Decodable, Debug, HashStable_Generic)]
68pub struct Type {
69    pub offset: isize,
70    pub size: usize,
71    pub kind: Kind,
72    pub child: TypeTree,
73}
74
75impl Type {
76    pub fn add_offset(self, add: isize) -> Self {
77        let offset = match self.offset {
78            -1 => add,
79            x => add + x,
80        };
81
82        Self { size: self.size, kind: self.kind, child: self.child, offset }
83    }
84}
85
86impl fmt::Display for Type {
87    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
88        <Self as fmt::Debug>::fmt(self, f)
89    }
90}
rustc_ast/expand/typetree.rs

rustc_ast/expand/
typetree.rs
