Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

new realloc bug reproducer #190

Open
ZuseZ4 opened this issue Dec 4, 2024 · 1 comment
Open

new realloc bug reproducer #190

ZuseZ4 opened this issue Dec 4, 2024 · 1 comment
Assignees

Comments

@ZuseZ4
Copy link
Member

ZuseZ4 commented Dec 4, 2024

Reproducer:

#![feature(autodiff)]
use std::autodiff::autodiff;

#[autodiff(sin_vec, Reverse, Duplicated, Const, Active)]
fn cos_vec(x: &mut Vec<f32>, y: usize) -> f32 {
    // uses enum internally and breaks
    for i in 0..y {
        x.push(i as f32);
    }
    let res = x.into_iter().collect::<Vec<&mut f32>>();

    *res[0]
}

fn main() {
    let mut x = vec![1.0, 1.0, 1.0];
    let mut d_x = vec![0.0; 3];

    sin_vec(&mut x, &mut d_x, 4, 1.0);

    dbg!(&d_x, &x);
}

RUSTFLAGS="-Z autodiff=LooseTypes,Print" cargo +enzyme run

mod.ll.txt

Demangled function names:
https://play.rust-lang.org/?version=stable&mode=debug&edition=2021&gist=2fdf75b03825f5f63f353ff1122f3e0f

Demangled function names are:

alloc::raw_vec::finish_grow::hce1ff89b620ba31c
alloc::raw_vec::RawVecInner<A>::grow_amortized::h6d11486032b72786
alloc::raw_vec::RawVecInner<A>::reserve::do_reserve_and_handle::hf69592d80ba560b8
alloc::vec::Vec<T,A>::reserve::hfbf4bb7fa5035d73
alloc::vec::Vec<T,A>::extend_trusted::h9fccf647fedf83c0

I'd say reserve is a reasonable function to provide a custom-derivative for @wsmoses , what would we need to register?
Augmented primal is just a reserve for the same additional on the shadow, and do nothing for the reverse?
For context, that's the definition:

pub struct Vec<T, A: Allocator = Global> {
    buf: RawVec<T, A>,
    len: usize,
}

    pub fn reserve(&mut self, additional: usize) {
        self.buf.reserve(self.len, additional);
    }

This example and some playing around in release mode indicates we should also have a custom rule for alloc::raw_vec::RawVecInner<A>::reserve::do_reserve_and_handle
https://play.rust-lang.org/?version=nightly&mode=release&edition=2021&gist=3d81e4ea9dcd7abd8acace3d9e40685c

@ZuseZ4
Copy link
Member Author

ZuseZ4 commented Dec 6, 2024

A best guess on what to do here:

  1. Gett all names for relevant vector reserve calls (here e.g. _ZN5alloc3vec16Vec$LT$T$C$A$GT$7reserve17hfbf4bb7fa5035d73E).
    2A) Call EnzymeRegisterCallHandler (?), using the name from 1)
    2B) Pass a nullptr as CustomFunctionReverse since it shouldn't need anything in the reverse pass
    2C) Pass a Builder and gutils as first and third argument to the CustomAugmentedFunctionForward
    2D) Pass callinst to the reserve() call from 1) as the second arg to CustomAugmentedFunctionForward
    2E) do something (??) with the last three arguments of CustomAugmentedFunctionForward

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants