ABI/Layout handling for the automatic differentiation feature

Basic Information

• Name: Marcelo Dominguez
• Github: https://github.com/sa4dus
• Email: dmmarcelo27@gmail.com
• LinkedIn: https://www.linkedin.com/in/dmmarcelo
• Location: Madrid, Spain (UTC+1:00)
• Mentor: Manuel Drehwald (@ZuseZ4), Oli Scherer (@oli-obk)

About me

I am currently pursuing a Dual Degree in Mathematics and Software Engineering at King Juan Carlos University. I started programming at 15 with Python, then moved into web development using JavaScript, React, Express.js, and MongoDB. Upon entering university, I developed a strong interest in systems programming with C and Rust, as well as neural networks. My main passion lies in Rust, particularly in systems programming and compiler design.

Contributions

Prevent ICE in Autodiff Validation by Emitting User-Friendly Errors: #138231

Project Details

Size: Medium (~175 hours)

Abstract

The integration of automatic differentiation (autodiff) into the Rust compiler using Enzyme has encountered issues due to conflicting ABI transformations in Rust. These adjustments (sometimes optimizations) generally improve performance but interfere with autodiff functionality. The goal of this project is to prevent undesired ABI modifications for functions marked with the #[rustc_autodiff] attribute without impacting the performance of other functions.

[!NOTE] #[autodiff] attribute is a macro frontend that generates a placeholder function and applies two #[rustc_autodiff] attributes: one to the original function and one to the new dummy function (which will later become the derivative of the original function).

Primary Goals

• Ensure that rustc does not apply ABI adjustments to functions with the #[rustc_autodiff] attribute.
• Develop a general solution that enables the #[autodiff(..)] attribute to handle arbitrary function headers correctly.

Secondary Goals

• Fix bugs in the autodiff frontend.
• Contribute to the EnzymeAD Rust documentation.

Constraints

• The changes must not introduce any performance overhead for functions that do not have the #[rustc_autodiff] attribute.

Technical Approach

The goal is to identify and disable undesired adjustments for functions marked with #[autodiff]. The following steps outline the technical process:

1. Determine and Categorize ABI Adjustments:

Goal: Establish a clear scope by identifying where known optimizations introduce ABI adjustments that affect autodiff. Additionally, detect any new potential changes that may not have been previously considered.

Implied actions:

• Map out function headers and call sites where ABI modifications occur.
• Analyze how Rust’s layout and calling conventions interact with these adjustments.
• Categorize optimization-driven changes in function signatures to define key areas for intervention.
• Investigate potential ABI-related modifications that may have gone unnoticed and assess their impact.

2. Lower #[autodiff(..)] Attribute Information to ABI Level

Goal: Make the compiler explicitly recognize, at the ABI level, that a function will be differentiated, ensuring that ABI adjustments are prevented only for these functions.

Implied Actions:

• Modify FnAbi layout to encode differentiation-related information.
• Introduce a mechanism that allows the compiler to lower differentiation status, preventing ABI adjustments specifically for autodiff functions.
• Ensure this information is correctly propagated through Rust’s compilation pipeline.

3. Prevent ABI Adjustments That Break Autodiff

Goal: Prevent adjustments that interfere with autodiff, while allowing necessary compiler optimizations.

Implied Actions:

• Modify layout handling to avoid breaking autodiff functions.
• Use AttributeKind::OptimizeNone where necessary.
• Validate that these changes do not negatively impact performance for non-autodiff functions.

4. Testing and Refinement

Goal: Validate correctness and robustness through testing, ensuring both functional accuracy and proper code generation.

Implied Actions:

• Develop functional tests covering all identified problematic cases.
• Implement codegen tests to verify that the expected ABI adjustments (or lack thereof) occur at the LLVM IR level.
• Refine the approach based on common function headers and performance benchmarks.

Examples of Incorrectly Handled Function Headers

First, we will examine those transformations that occur at OptLevel::Aggressive, but not at OptLevel::None. These will be considered optimizations.

• fn foo(x: &[T;n]), where T is an Scalar type and n is less or equal than 2. See playground. In this family of function headers, the individual elements of the slice are passed as value.

  Optimization Level	OptLevel::Aggressive
  Expected	ptr align 4 %x
  Actual	float %x.0.val, float %x.4.val

• fn foo(f: fn(f32) -> f32, x: f32). See playground. This optimization is triggered even if applying the #[inline(never)] attribute to the callback function.

  Optimization Level	OptLevel::Aggressive
  Expected	ptr %f, float %x
  Actual	float noundef %x

• fn foo<'a>(x: &'a f32, /* static */ y: &'a f32). See playground. The static variable is elided from the function signature.

  Optimization Level	OptLevel::Aggressive
  Expected	ptr align 4 %x, ptr align 4 %y
  Actual	float %x.0.val

• fn foo(x: &dyn MyTrait). See playground. // Note(Sa4dUs): i really don't know what's the expected behaviour here, but changing from two params to zero doesn't seem good for Enzyme backend. (UPDATE: pending the response from the creator of Enzyme on whether to include it in the proposal, but if something extra were to be added, this would possibly be the most interesting candidate.)

  Optimization Level	OptLevel::Aggressive
  Debug	ptr align 1 %x.0, ptr align 8 %x.1
  Release	()

Now, we have two cases that also cause issues with the Enzyme backend, but these transformations occur even at OptLevel::None. This is likely due to how Rust handles functions at a low level.

• fn foo(x: (T, ...)), where T is an Scalar type and n is less or equal than 2. See playground. Same as the previous case.

  Optimization Level	OptLevel::No
  Expected	ptr align 4 %x
  Actual	float %x.0.val, float %x.4.val

• fn foo(i: T), where T is a struct with 2 or less scalar fields See playground. This is also triggered even with OptLevel::No, so it may be related to the previous case.

  Optimization Level	OptLevel::No
  Expected	ptr align 4 %x
  Actual	float %i.0, float %i.1

Deliverables

• Midterm Evaluation: The system should pass basic test cases, though still in early stages (not production-ready).
• Final Evaluation: Autodiff will handle most function headers correctly, with a robust test suite and comprehensive documentation.

Project Timeline

Community Bonding Period (May 8 - June 1)

• Engage with the mentor and Rust contributors to discuss approaches for lowering some trace of #[autodiff] attribute into ABI level. (Step 1)
• Refine the project scope and establish priorities.(Step 1)
• Review the relevant codebase and documentation
• Identify and analyze problematic optimizations that may cause issues. (Step 1)

Week 1 (June 2 - June 8)

• Define a set of tests that must pass by the end of the project to establish clear development objectives. (Step 4)
• Design formal layout modifications to enable LLVM codegen to determine whether a function should be differentiated. (Step 2)

Week 2-4 (June 9 - June 29)

• Lower function #[autodiff(..)] attribute information to ABI level by modifying the FnAbi layout. (Step 2)

Week 5-6 (June 30 - July 13)

• Disable architecture-agnostic optimizations, particularly those related to argument handling in LLVM codegen. (Step 3)
• Ensure most common cases are working properly for midterm evaluation. (Step 4)
• Buffer time for debugging and adjustments.

Week 7 (July 14 - July 20)

• Submit midterm evaluation.
• Investigate potential architecture-specific optimizations.

Week 8-9 (July 21 - August 3)

• Extend the solution to support more complex function headers that may introduce additional challenges. (Step 4)

Week 10 (August 4 - August 10)

• Conduct performance benchmarks and resolve any potential performance issues affecting non-autodiff functions. (Step 4)
• Expand test coverage to ensure comprehensive validation. (Step 4)

Week 11 (August 11 - August 17)

• If needed, refactor code for clarity, readability, and maintainability. (Step 4)

Week 12 (August 18 - August 24)

• Address bugs or regressions identified during testing and ensure all issues are resolved. (Step 4)
• Finalize documentation and ensure all tests are passing. (Step 4)
• Buffer time for final adjustments.

Final Week (August 25 - September 1)

• Submit the final evaluation, completing the project with a robust solution, full documentation, and extensive test coverage.

What’s next?

After the GSoC project concludes, I plan to continue contributing to the Rust compiler and maintain the improvements made to the autodiff functionality, ensuring that it remains robust and efficient for future development.