fix usage of autodiff macro with inner functions

[haenoe]

This PR adds additional handling into the expansion step of the std::autodiff macro (in compiler/rustc_builtin_macros/src/autodiff.rs), which allows the macro to be applied to inner functions.

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

fn main() {
    #[autodiff(d_inner, Forward, Dual, DualOnly)]
    fn inner(x: f32) -> f32 {
        x * x
    }
}

Previously, the compiler didn't allow this due to only handling Annotatable::Item and Annotatable::AssocItem and missing the handling of Annotatable::Stmt. This resulted in the rather generic error

error: autodiff must be applied to function
 --> src/main.rs:6:5
  |
6 | /     fn inner(x: f32) -> f32 {
7 | |         x * x
8 | |     }
  | |_____^

error: could not compile `enzyme-test` (bin "enzyme-test") due to 1 previous error

This issue was originally reported here.

Quick question: would it make sense to add a ui test to ensure there is no regression on this?
This is my first contribution, so I'm extra grateful for any piece of feedback!! :D

r? @oli-obk

Tracking issue for autodiff: #124509

[ZuseZ4]

Thank you for looking into this! Yes, adding a test to avoid regressions would be great.
UI tests are usually more towards stdout/stderr (or runtime) testing. Since this one now works, I would add it under tests/pretty/autodiff_reverse.{rs|pp}. You can just add an example by renaming main to f6, and use --bless when running the test.
With git diff you can then check if the change is what you expect it to be.
The whole test suite just got reworked recently, so it has quite nice docs these days: https://rustc-dev-guide.rust-lang.org/tests/compiletest.html
https://rustc-dev-guide.rust-lang.org/tests/intro.html
and https://rustc-dev-guide.rust-lang.org/tests/best-practices.html are probably the relevant ones.

[oli-obk]

@bors delegate=ZuseZ4

[bors]

✌️ @ZuseZ4, you can now approve this pull request!

If @oli-obk told you to "r=me" after making some further change, please make that change, then do @bors r=@oli-obk

[haenoe]

Thank you for looking into this! Yes, adding a test to avoid regressions would be great. UI tests are usually more towards stdout/stderr (or runtime) testing. Since this one now works, I would add it under tests/pretty/autodiff_reverse.{rs|pp}. You can just add an example by renaming main to f6, and use --bless when running the test. With git diff you can then check if the change is what you expect it to be. The whole test suite just got reworked recently, so it has quite nice docs these days: https://rustc-dev-guide.rust-lang.org/tests/compiletest.html https://rustc-dev-guide.rust-lang.org/tests/intro.html and https://rustc-dev-guide.rust-lang.org/tests/best-practices.html are probably the relevant ones.

@ZuseZ4 , thank you for providing more insights into this! The pretty printer tests do indeed sound like a good place for this. Also I've got to express how cool and insightful the docs are -- this is really not something to be taken as granted :D

However, while looking into this I noticed that both tests/pretty/autodiff_{forward,reverse}.rs are currently failing due to the fact that they still include functions which use the autodiff macro in conjuction with DiffMode::{Forward,Reverse}First. These two were removed with commit 061abbc36928cce784c54463c266f4d43d14d419.

It makes sense to open a separate PR fixing these issues before continuing with the inner function test, right?
Have a nice evening ^^'

[ZuseZ4]

Yes I agree, I think one reason why I ended up being a rustc developer rather than working on other projects is the documentation. I don't have a pc this weekend, but my batching PR contains fixes for those. You can cherry pick them if you want, or I'll make a PR for that on Monday. I had hoped to get it done for the weekend, but unfortunately ran out of time.

[haenoe]

Okidoki! Don't stress, I'll just rebase when the changes are in master.
Enjoy your weekend!

[bors]

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

[haenoe]

Hey @ZuseZ4
Rebased this to master and also tried to integrate your fixes regarding the duplicated attributes via your same_attribute function.
Quick question: Do you think this simple test suffices to test the expansion logic? Or would you add one/two more advanced usages (e.g. multiple derivatives).

[ZuseZ4]

looks good, but I will check again later when I'm on my computer. For the test I would just add a second autodiff macro with an in any way different config to the same function/test, to verify that even after future refactors we don't add inline or autodiff attributes multiple time to the source functions.
Also, we usually use bug! instead of unreachable if we have previously tested assumptions, can you move them over?

[bors]

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