autodiff fails in impl blocks

[ZuseZ4]

I tried this code:
https://github.com/EnzymeAD/rustbook/blob/main/samples/tests/traits/mod.rs

#![feature(autodiff)]
use std::autodiff::autodiff;
trait Volumetric {
    /// Strain energy density
    fn psi(&self, j: f64) -> f64;
    /// Derivative of strain energy with respect to $J$
    fn d_psi(&self, j: f64, b_psi: f64) -> (f64, f64);
    /// The volumetric contribution to the second Piola-Kirchhoff stress is
    /// a scalar multiplied by $C^{-1}$ where $C = I + 2E$ in terms of the
    /// Green-Lagrange strain $E$. The derivative of $J$ with respect to $E$
    /// is $J C^{-1}$. We'll call the volumetric contribution that scalar
    /// multiple of $C^{-1}$.
    fn stress(&self, j: f64) -> f64 {
        let (_, d_psi) = self.d_psi(j, 1.0);
        d_psi * j
    }
}

struct Ogden {
    k: f64,
}
impl Ogden {
    pub fn stress_analytic(&self, j: f64) -> f64 {
        self.k * 0.5 * (j * j - 1.0)
    }
}
impl Volumetric for Ogden {
    #[autodiff(d_psi, Reverse, Const, Active, Active)]
    fn psi(&self, j: f64) -> f64 {
        self.k * 0.25 * (j * j - 1.0 - 2.0 * j.ln())
    }
}

fn main() {
    let j = 0.8;
    let vol = Ogden { k: 1.0 };
    let s = vol.stress(j);
    let s_ref = vol.stress_analytic(j);
    assert!((s - s_ref).abs() < 1e-15, "{}", s - s_ref);
}

I expected to see this happen: Passes the assertion

Instead, this happened:

➜  ad git:(main) RUSTFLAGS="-Zautodiff=Enable" cargo +enzyme run  --release
   Compiling ad v0.1.0 (/home/manuel/prog/ad)
error: autodiff must be applied to function
  --> src/main.rs:29:5
   |
29 | /     fn psi(&self, j: f64) -> f64 {
30 | |         self.k * 0.25 * (j * j - 1.0 - 2.0 * j.ln())
31 | |     }
   | |_____^

error[E0046]: not all trait items implemented, missing: `d_psi`
  --> src/main.rs:27:1
   |
7  |     fn d_psi(&self, j: f64, b_psi: f64) -> (f64, f64);
   |     -------------------------------------------------- `d_psi` from trait
...
27 | impl Volumetric for Ogden {
   | ^^^^^^^^^^^^^^^^^^^^^^^^^ missing `d_psi` in implementation

For more information about this error, try `rustc --explain E0046`.

Meta

rustc --version --verbose:

build from source

Backtrace

<backtrace>

This is an easy bug and a got way to get started. It used to work a few weeks ago, so it's likely that just one of the other improvement PRs accidentally regressed this test, since it's not upstream yet. Supporting autodiff in impl blocks is of course important, so it would be nice to fix this regression.

To prevent further regressions, this autodiff invocation (under impl Volumentric for Ogden) should be added in tests/pretty/autodiff. We emit the error as dcx.emit_err(errors::AutoDiffInvalidApplication { span: item.span() }); in compiler/rustc_builtin_macros/src/autodiff.rs so the bugfix likely goes there, this PR could be an inspiration: https://github.com/rust-lang/rust/pull/138314/files

Tracking:

• Tracking Issue for autodiff #124509

[SpencerWightman]

@rustbot claim

[ZuseZ4]

@SpencerWightman Did you have a chance to look into the issue, or do you have any questions?

[SpencerWightman]

@ZuseZ4 Yes, I've made some progress. Sorry for holding it up so long. I'll make a PR today or drop the claim.

[Shourya742]

Hi @SpencerWightman, I'm currently exploring autodiff and would be happy to collaborate on this if you're interested. Otherwise, please let me know if it's okay for me to take it on—I'm excited to contribute either way!

[ZuseZ4]

@SpencerWightman no rush if you don't get to it today, but I hope to have a PR for this ready around the 22nd, so it has a bit of time to land before I enable autodiff on the 26th. Feel free to experiment with it a bit more, otherwise I can probably come up with a fix on the 22nd. @Shourya742 I pinged you in another issue with something you could look at if you have time.

[SpencerWightman]

@rustbot release-assignment
@ZuseZ4 @Shourya742 I don't want to hold this up any longer. Thanks for letting me have a shot at it.

[Shourya742]

@ZuseZ4, Is it ok now, if I take this up now?

[Shourya742]

@rustbot claim