New compiler macro design

[bytesnake]

Requirements

As extern block design

The current syntax relies heavily on blackbox and does not extend easily to custom rules, this PR adds a new autodiff ABI with the following syntax (copied from #31)

extern "autodiff" {
    fn primal(e: &KM, nh: &NH) -> KM {
        let mut dpsi_de = KM::zero();
            d_psi(&e, &mut dpsi_de, &nh, 1.0);
            let b = e.cauchy_green();
            dpsi_de * b
    }

    #[autodiff(Forward, DuplicatedNoNeed, Duplicated, Const)]
    fn d_stress(e: &KM, tangent_e: &KM, nh: &NH) -> KM;

    #[autodiff(Forward, DuplicatedNoNeed, Const, Duplicated)]
    fn d2_stress(e: &KM, nh: &NH, tangent_nh: &NH) -> KM;
}

or shorthand (actually is a bit confusing and would move this part to external proc-macro)

#[autodiff(Forward, DuplicatedNoNeed, Duplicated, Const)]
extern "autodiff" fn d_stress(e: &KM, nh: &NH) -> KM {
    ...
}

Some observations:

• the autodiff ABI allows us to ensure that no inlining happens
• extern block groups functions of the same primal with different rules

Custom tapes

When specifying a custom tape struct, the user is responsible for defining both augmented primal and custom derivatives:

#[autodiff_custom_tape=Tape]
extern "autodiff" {
    fn primal(e: &KM, nh: &NH) -> (Tape, KM) {
        // calculate primal value + fill up tape with intermediate results
        // ...
    }

    #[autodiff(Forward, DuplicatedNoNeed, Duplicated, Const)]
    fn d_stress_custom(e: &KM, tangent_e: &KM, nh: &NH, tape: Tape) -> KM {
        // calculate gradients with custom tape
        // ...
    }
}

As trait design

[ZuseZ4]

For visualization, can you please show how the self example in case of a function associated to a struct would look like? I think we recently had one example from Jed?

Also, we need to check with someone from rustc/llvm how to disable inlining. "Really" disabling it seems much harder than I expected it in the past. I agree that we really should have it though.

[bytesnake]

oh well, I did not consider this yet. In the best case we allow extern blocks for impls, which maps to groups

struct Data {
   ...
}

extern "autodiff" impl Data {
    fn primal(&self, ..) {
       ..
    }

    fn d_stress(..);
}

if this contradicts with extern block rules, we could also allow impl and mod to carry the compiler macro. This adds then implicit ABI rules to all functions in a block

[ZuseZ4]

Yeah sorry, just came to my mind first as a challenge for this design.

I think places allowing extern are quite limited and I am not sure if bjorn had in mind that we would start allowing it in more and more places when he suggested a new ABI for us. Having an own ABI does however generally look good to me, Enzyme gains easily outperform the small abi gains which the rust default abi will bring (probably excluding some wild corner cases) and it makes things easier and more stable for us.

However, some blocking issues for now are:

https://doc.rust-lang.org/reference/items/external-blocks.html#functions A function declared in an extern block is implicitly unsafe. When coerced to a function pointer, a function declared in an extern block has type unsafe extern "abi" for<'l1, ..., 'lm> fn(A1, ..., An) -> R, where 'l1, ... 'lm are its lifetime parameters, A1, ..., An are the declared types of its parameters and R is the declared return type.

We can't except the whole rust hpc/ml/.. ecosystem to move to unsafe