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.

[rustbot]

Some changes occurred in compiler/rustc_builtin_macros/src/autodiff.rs

cc @ZuseZ4

[haenoe]

Hey @ZuseZ4!
Sorry for the long silence, I didn't have access to my computer :/
But now I am able to be more responsive, so let me know if you have any more feedback!

[ZuseZ4]

No worries, and the test looks good, thanks!
However something else seem to have broken in the meantime, maybe you can try to rebase locally onto the latest master and see if that builds?
Also, #138740 also modifies a lot of code, including our frontend, so you'll likely need to rebase once it lands anyway.

[rustbot]

These commits modify the Cargo.lock file. Unintentional changes to Cargo.lock can be introduced when switching branches and rebasing PRs.

If this was unintentional then you should revert the changes before this PR is merged.
Otherwise, you can ignore this comment.

[haenoe]

Thanks for the quick reply!
I didn't try to compile after the rebase. Should be fixed now ^-^'
Quick question: it felt a bit weird adjusting the Cargo.toml to get access to the rustc_middle::bug macro. It seems the macro wasn't used in the rustc_bultin_macros before. Did I overlook something here?

EDIT: forgot to install the tidy git hook again 🤦

[bors]

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

[oli-obk]

@bors r+ rollup

[bors]

📌 Commit 284672a has been approved by oli-obk

It is now in the queue for this repository.

[ZuseZ4]

Thank you two! But @oli-obk I'd be cautious with a rollup here, it will probably get a merge conflict in ~1hr, once my batching PR lands. I think github's merge algorithm likes to err on the site of caution.

[oli-obk]

@bors r-

[ZuseZ4]

Unfortunate, but as predicted. @haenoe the diff should hopefully be easy to fix, can you rebase one more time, so we can add it back to the CI?

[bors]

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

[haenoe]

Yess of course! Will look at it when I'm at my computer in ~2 hours.
Glad to see things moving so quickly :D