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?