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

Autodiff Upstreaming - enzyme frontend #129458

Open
wants to merge 21 commits into
base: master
Choose a base branch
from

Conversation

ZuseZ4
Copy link
Contributor

@ZuseZ4 ZuseZ4 commented Aug 23, 2024

This is an upstream PR for the autodiff rustc_builtin_macro that is part of the autodiff feature.

For the full implementation, see: #129175

Content:
It contains a new #[autodiff(<args>)] rustc_builtin_macro, as well as a #[rustc_autodiff] builtin attribute.
The autodiff macro is applied on function f and will expand to a second function df (name given by user).
It will add a dummy body to df to make sure it type-checks. The body will later be replaced by enzyme on llvm-ir level,
we therefore don't really care about the content. Most of the changes (700 from 1.2k) are in compiler/rustc_builtin_macros/src/autodiff.rs, which expand the macro. Nothing except expansion is implemented for now.
I have a fallback implementation for relevant functions in case that rustc should be build without autodiff support. The default for now will be off, although we want to flip it later (once everything landed) to on for nightly. For the sake of CI, I have flipped the defaults, I'll revert this before merging.

Dummy function Body:
The first line is an inline_asm nop to make inlining less likely (I have additional checks to prevent this in the middle end of rustc. If f gets inlined too early, we can't pass it to enzyme and thus can't differentiate it.
If df gets inlined too early, the call site will just compute this dummy code instead of the derivatives, a correctness issue. The following black_box lines make sure that none of the input arguments is getting optimized away before we replace the body.

Motivation:
The user facing autodiff macro can verify the user input. Then I write it as args to the rustc_attribute, so from here on I can know that these values should be sensible. A rustc_attribute also turned out to be quite nice to attach this information to the corresponding function and carry it till the backend.
This is also just an experiment, I expect to adjust the user facing autodiff macro based on user feedback, to improve usability.

As a simple example of what this will do, we can see this expansion:
From:

#[autodiff(df, Reverse, Duplicated, Const, Active)]
pub fn f1(x: &[f64], y: f64) -> f64 {
    unimplemented!()
}

to

#[rustc_autodiff]
#[inline(never)]
pub fn f1(x: &[f64], y: f64) -> f64 {
    ::core::panicking::panic("not implemented")
}
#[rustc_autodiff(Reverse, Duplicated, Const, Active,)]
#[inline(never)]
pub fn df(x: &[f64], dx: &mut [f64], y: f64, dret: f64) -> f64 {
    unsafe { asm!("NOP"); };
    ::core::hint::black_box(f1(x, y));
    ::core::hint::black_box((dx, dret));
    ::core::hint::black_box(f1(x, y))
}

I will add a few more tests once I figured out why rustc rebuilds every time I touch a test.

Tracking:

@rustbot
Copy link
Collaborator

rustbot commented Aug 23, 2024

r? @pnkfelix

rustbot has assigned @pnkfelix.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

@rustbot rustbot added S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. T-libs Relevant to the library team, which will review and decide on the PR/issue. labels Aug 23, 2024
@rust-log-analyzer

This comment has been minimized.

@traviscross traviscross added the F-autodiff `#![feature(autodiff)]` label Aug 23, 2024
@ZuseZ4 ZuseZ4 changed the title implement a working autodiff frontend Autodiff Upstreaming - enzyme frontend Aug 23, 2024
@bjorn3 bjorn3 mentioned this pull request Aug 24, 2024
9 tasks
@ZuseZ4 ZuseZ4 marked this pull request as ready for review August 24, 2024 20:25
@rust-log-analyzer

This comment has been minimized.

@bors
Copy link
Contributor

bors commented Aug 25, 2024

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

let mut attr: ast::Attribute = ast::Attribute {
kind: ast::AttrKind::Normal(rustc_ad_attr.clone()),
//id: ast::DUMMY_TR_ID,
id: ast::AttrId::from_u32(12341), // TODO: fix
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would be happy for suggestions on what ID to use or create here.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion (for other reviewers as well) [REV-ATTR-ID]: I feel like this needs to be a fresh AttrId, like via ecx.sess.psess.attr_id_generator.mk_attr_id(), otherwise we might get very kaboom behavior if the AttrIds collide.

pub fn mk_attr_id(&self) -> AttrId {
let id = self.0.fetch_add(1, Ordering::Relaxed);
assert!(id != u32::MAX);
AttrId::from_u32(id)
}
}

But I'm not too sure if this is right.

@jieyouxu jieyouxu self-assigned this Aug 30, 2024
Copy link
Member

@jieyouxu jieyouxu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the work on autodiff! I have some feedback and questions. As you may have gathered, I am not very knowledgeable about autodiff. If I ask some questions for more context / explanation, it might be good to encode them as comments in the impl itself for more context. So that if someone else (or even yourself) comes back later to try to change this impl, they are better equipped to figure out what this is doing.

EDIT: please ignore panic! -> bug! suggestions as that might not be available yet in macro expansion here.

@@ -2733,6 +2733,13 @@ impl FnRetTy {
FnRetTy::Ty(ty) => ty.span,
}
}

pub fn has_ret(&self) -> bool {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Discussion [REV-1]: I find this somewhat misleading, because unspecified return type can still have non-unit () return types (e.g. closure type inference). Maybe has_specified_ret_ty or something... What's the significance of distinguishing between specified and non-specified return types?

pub enum FnRetTy {
/// Returns type is not specified.
///
/// Functions default to `()` and closures default to inference.
/// Span points to where return type would be inserted.
Default(Span),

I need to check what the call-site of this actually wants, seems a bit strange to want this.

Copy link
Member

@compiler-errors compiler-errors Aug 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is very much sketchy 👍

Ideally don't distinguish between -> () and no return, because in Rust these are precisely equal, but in any case, this should not be a helper on the AST that other users may accidentally use since it's 99.9% wrong.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think closures are relevant enough to spend time on adding support for it, my macro rejects everything but functions.
I don't care about the difference between -> () and the implicit () return. It looks like it worked by coincidence before correctly, but I've added a check for explicit -> () return and also moved the function out of ast.rs

Copy link
Member

@jieyouxu jieyouxu Sep 7, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question: does it matter to autodiff then if the return type is a type-alias-to-unit or effectively-unit?

// e.g.
type Unit = ();
// or 
#[repr(transparent)]
pub struct UnitButFancy { inner: () }
// or
pub struct Sandwich(());

Then the user may have written

pub fn foo() -> Unit /* or `UnitButFancy` or `Sandwich` */ {}

Would this be a problem?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't really know how the implementation would behave here.

The current mvp strongly relies on the user not doing "strange" things, I don't see this experiment as anywhere close of handling type aliases or malicious users correctly. Sicne there are still bugs left that users in good faith will run into, I'd just leave this as an unhandled case. This type doesn't have any floats in it so const would be the only reasonable annotation, but I don't think spending time on catching this would be a good investment.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That seems fine, probably worth remarking this as a comment.

compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
}
}

pub fn valid_ret_activity(mode: DiffMode, activity: DiffActivity) -> bool {
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Question [REV-CTXT]: I can of course read the implementation below, but for someone in the future who's trying to change this code, what's the high-level description of what constitutes a "valid" ret activity? What's a DiffActivity? Why is e.g. DiffMode::Forward + activity == DiffActivity::Dual a valid ret activity but DiffMode::Forward + activity == DiffActivity::Const not?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion: could we hoist definition of DiffActivity and perhaps add some explanation for its purpose before its usage here? I think that would make it easier to follow.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both of your examples are allowed, did you mean something else? I also have added some more docs and bundled all enum/struct definitions at the top, does that help?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both of your examples are allowed, did you mean something else?

Err probably, substitute whatever I wrote with a combination that's not allowed (if there's any). In any case, could some invalid combination cause the enzyme backend to kaboom if we don't catch them here?

I also have added some more docs and bundled all enum/struct definitions at the top, does that help?

Yes that seems fine.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, every single invalid combination that we don't catch will make Enzyme go kaboom. If we're lucky we just run into an enzyme assertion with a more or less helpful error, in the worstcase we get a segfault.

compiler/rustc_builtin_macros/src/autodiff.rs Outdated Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Outdated Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Show resolved Hide resolved

// Test that reverse mode ad macros are expanded correctly.

#[autodiff(df, Reverse, Duplicated, Const, Active)]
Copy link
Member

@jieyouxu jieyouxu Aug 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Problem && Suggestion: I think we want at least test cases for:

  • Positive test cases for what autodiff should be applied to.
  • Negative test cases for when autodiff is misapplied to AST nodes and their diagnostics. In particular, closures, statements and expressions.
  • #[autodiff] malformed attribute (where args?) error
  • #[autodiff = ""] invalid attribute syntax
  • #[autodiff()] where arg
  • #[autodiff(df)] + fn df() <- what if I already have a df in value namespace?
  • #[autodiff(df, Reverse)] + enum Foo { Reverse } + use Foo::Reverse; <- what if I already have a Reverse in type (enum variant decl) and value (enum variant constructor) namespace?
  • #[autodiff(df)] <- is this minimally acceptable?
  • #[autodiff(df, Reverse)] <- valid mode
  • #[autodiff(df, Debug)] <- invalid mode
  • #[autodiff(df, Forward, Reverse)] <- is this valid
  • target fn has specified valid return type e.g. -> f64
  • target fn has unspecified return type fn foo() {}
  • target fn has specified but invalid return type e.g. -> Owo
  • target fn has aliased f32/f64 return types (currently unsupported) -> F64Alias
  • target fn has #[repr(transparent)] struct F64Trans { inner: f64 } return type (currently unsupported) -> F64Trans.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All of these cases now give error messages instead of ICEs.
I also made sure that we don't return after the first error, but continue to do something sufficiently okish that we can first expand all autodiff macros before we abort compilation.
We probably could include even better errors in the future, but I left those as fixme's for now.
Do they look good to you?

@rustbot rustbot added S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. and removed S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. labels Aug 30, 2024
@ZuseZ4 ZuseZ4 force-pushed the enzyme-frontend branch 2 times, most recently from fb5fa71 to 9d8ec28 Compare September 3, 2024 01:58
@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
compiler/rustc_ast/src/expand/autodiff_attrs.rs Outdated Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Outdated Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Outdated Show resolved Hide resolved
compiler/rustc_builtin_macros/src/autodiff.rs Outdated Show resolved Hide resolved
@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@ZuseZ4
Copy link
Contributor Author

ZuseZ4 commented Sep 22, 2024

working my way through it. Already lead to some nice improvement, because I now return incorrect function signatures for df if I encounter a usage error, knowing that compilation will abort later anyway. By not erroring immediately we can catch more errors in one go.

  • Positive test cases for what autodiff should be applied to.
  • Negative test cases for when autodiff is misapplied to AST nodes and their diagnostics. In particular, closures, statements and expressions.
  • #[autodiff] malformed attribute (where args?) error
  • #[autodiff = ""] invalid attribute syntax
  • #[autodiff()] where arg
  • #[autodiff(df)] + fn df() <- what if I already have a df in value namespace?
  • #[autodiff(df, Reverse)] + enum Foo { Reverse } + use Foo::Reverse; <- what if I already have a Reverse in type (enum variant decl) and value (enum variant constructor) namespace?
  • #[autodiff(df)] <- is this minimally acceptable?
  • #[autodiff(df, Reverse)] <- valid mode
  • #[autodiff(df, Debug)] <- invalid mode
  • #[autodiff(df, Forward, Reverse)] <- is this valid
  • target fn has specified valid return type e.g. -> f64
  • target fn has unspecified return type fn foo() {}
  • target fn has specified but invalid return type e.g. -> Owo
  • target fn has aliased f32/f64 return types (currently unsupported) -> F64Alias
  • target fn has #[repr(transparent)] struct F64Trans { inner: f64 } return type (currently unsupported) -> F64Trans.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer

This comment has been minimized.

@rust-log-analyzer
Copy link
Collaborator

The job mingw-check-tidy failed! Check out the build log: (web) (plain)

Click to see the possible cause of the failure (guessed by this bot)

COPY host-x86_64/mingw-check/validate-toolstate.sh /scripts/
COPY host-x86_64/mingw-check/validate-error-codes.sh /scripts/

# NOTE: intentionally uses python2 for x.py so we can test it still works.
# validate-toolstate only runs in our CI, so it's ok for it to only support python3.
ENV SCRIPT TIDY_PRINT_DIFF=1 python2.7 ../x.py test \
           --stage 0 src/tools/tidy tidyselftest --extra-checks=py:lint,cpp:fmt
# This file is autogenerated by pip-compile with Python 3.10
# by the following command:
#
#    pip-compile --allow-unsafe --generate-hashes reuse-requirements.in
---
#13 2.784 Building wheels for collected packages: reuse
#13 2.785   Building wheel for reuse (pyproject.toml): started
#13 3.031   Building wheel for reuse (pyproject.toml): finished with status 'done'
#13 3.032   Created wheel for reuse: filename=reuse-4.0.3-cp310-cp310-manylinux_2_35_x86_64.whl size=132715 sha256=dfa09868353292d98f811d3efdb0d54d07389e808efc71d68e3b93c514bf8bec
#13 3.032   Stored in directory: /tmp/pip-ephem-wheel-cache-gduy_hnj/wheels/3d/8d/0a/e0fc6aba4494b28a967ab5eaf951c121d9c677958714e34532
#13 3.035 Installing collected packages: boolean-py, binaryornot, tomlkit, reuse, python-debian, markupsafe, license-expression, jinja2, chardet, attrs
#13 3.439 Successfully installed attrs-23.2.0 binaryornot-0.4.4 boolean-py-4.0 chardet-5.2.0 jinja2-3.1.4 license-expression-30.3.0 markupsafe-2.1.5 python-debian-0.1.49 reuse-4.0.3 tomlkit-0.13.0
#13 3.439 WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv
#13 3.978 Collecting virtualenv
#13 3.978 Collecting virtualenv
#13 4.016   Downloading virtualenv-20.26.5-py3-none-any.whl (6.0 MB)
#13 4.102      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 6.0/6.0 MB 71.5 MB/s eta 0:00:00
#13 4.147 Collecting distlib<1,>=0.3.7
#13 4.150   Downloading distlib-0.3.8-py2.py3-none-any.whl (468 kB)
#13 4.158      ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 468.9/468.9 KB 88.1 MB/s eta 0:00:00
#13 4.198 Collecting filelock<4,>=3.12.2
#13 4.202   Downloading filelock-3.16.1-py3-none-any.whl (16 kB)
#13 4.241 Collecting platformdirs<5,>=3.9.1
#13 4.245   Downloading platformdirs-4.3.6-py3-none-any.whl (18 kB)
#13 4.331 Installing collected packages: distlib, platformdirs, filelock, virtualenv
#13 4.530 Successfully installed distlib-0.3.8 filelock-3.16.1 platformdirs-4.3.6 virtualenv-20.26.5
#13 DONE 4.6s

#14 [7/8] COPY host-x86_64/mingw-check/validate-toolstate.sh /scripts/
#14 DONE 0.0s
---
DirectMap4k:      200640 kB
DirectMap2M:     8187904 kB
DirectMap1G:    10485760 kB
##[endgroup]
Executing TIDY_PRINT_DIFF=1 python2.7 ../x.py test            --stage 0 src/tools/tidy tidyselftest --extra-checks=py:lint,cpp:fmt
+ TIDY_PRINT_DIFF=1 python2.7 ../x.py test --stage 0 src/tools/tidy tidyselftest --extra-checks=py:lint,cpp:fmt
    Finished `dev` profile [unoptimized] target(s) in 0.04s
##[endgroup]
downloading https://ci-artifacts.rust-lang.org/rustc-builds-alt/6ce376774c0bc46ac8be247bca93ff5a1287a8fc/rust-dev-nightly-x86_64-unknown-linux-gnu.tar.xz
extracting /checkout/obj/build/cache/llvm-6ce376774c0bc46ac8be247bca93ff5a1287a8fc-true/rust-dev-nightly-x86_64-unknown-linux-gnu.tar.xz to /checkout/obj/build/x86_64-unknown-linux-gnu/ci-llvm
---
##[endgroup]
fmt check
fmt: checked 5573 files
tidy check
tidy error: /checkout/compiler/rustc_passes/messages.ftl: message `passes_autodiff_attr` appears before `passes_allow_incoherent_impl`, but is alphabetically later than it
run `./x.py test tidy --bless` to sort the file correctly
tidy error: /checkout/compiler/rustc_builtin_macros/messages.ftl: message `builtin_macros_autodiff_unknown_activity` appears before `builtin_macros_autodiff`, but is alphabetically later than it
run `./x.py test tidy --bless` to sort the file correctly
tidy error: /checkout/compiler/rustc_builtin_macros/messages.ftl: message `builtin_macros_autodiff_not_build` appears before `builtin_macros_autodiff_mode_activity`, but is alphabetically later than it
run `./x.py test tidy --bless` to sort the file correctly
tidy error: /checkout/compiler/rustc_builtin_macros/messages.ftl: message `builtin_macros_autodiff_number_activities` appears before `builtin_macros_autodiff_mode`, but is alphabetically later than it
run `./x.py test tidy --bless` to sort the file correctly
tidy error: /checkout/compiler/rustc_builtin_macros/messages.ftl: message `builtin_macros_autodiff_ty_activity` appears before `builtin_macros_asm_clobber_abi`, but is alphabetically later than it
run `./x.py test tidy --bless` to sort the file correctly
tidy: Skipping binary file check, read-only filesystem
##[error]tidy error: /checkout/compiler/rustc_builtin_macros/src/autodiff.rs:246: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/compiler/rustc_builtin_macros/src/autodiff.rs:253: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/compiler/rustc_builtin_macros/src/autodiff.rs:611: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/compiler/rustc_builtin_macros/src/autodiff.rs:631: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/compiler/rustc_feature/src/builtin_attrs.rs:755: TODO is used for tasks that should be done before merging a PR; If you want to leave a message in the codebase use FIXME
##[error]tidy error: /checkout/tests/pretty/autodiff_illegal.rs:45: trailing whitespace
##[error]tidy error: /checkout/compiler/rustc_builtin_macros/src/lib.rs:18: line not in alphabetical order
##[error]tidy error: /checkout/library/core/src/macros/mod.rs:1556: feature gate autodiff has inconsistent `issue`: "none" mismatches the previous `issue` of "12345"
##[error]tidy error: /checkout/library/std/src/lib.rs:621: feature gate autodiff has inconsistent `issue`: "12345" mismatches the previous `issue` of "none"
removing old virtual environment
creating virtual environment at '/checkout/obj/build/venv' using 'python3.10'
Requirement already satisfied: pip in ./build/venv/lib/python3.10/site-packages (24.2)
All checks passed!
checking C++ file formatting
some tidy checks failed
some tidy checks failed
Command has failed. Rerun with -v to see more details.
  local time: Sun Sep 22 07:15:26 UTC 2024
  network time: Sun, 22 Sep 2024 07:15:26 GMT
##[error]Process completed with exit code 1.
Post job cleanup.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
F-autodiff `#![feature(autodiff)]` S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. T-libs Relevant to the library team, which will review and decide on the PR/issue.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

8 participants