AD Cleanup: Avoid dealing with invalid IR by representing out-of-scope accesses through Push and Pop instructions

[saipraveenb25]

The unzipping step creates two copies of the control-flow graph, one holding primal instructions and the other holding differential instructions. The differential blocks run immediately after all the primal ones, so if there is any non-trivial branching, there will be invalid accesses of instructions generated in conditional regions.

Example:

Consider a function f with control-flow:

float f(float x)
{
    float val = 0.f;
    if (x < 0.5)
    {
        float k = 2 * x;
        val = k * x;
    }
    else
    {
        float k = 3 * x;
        val = k * x;
    }

    return val;
}

This will result in the following code after the unzipping step:

DifferentialPair<float> unzipped_f(DifferentialPair<float> dpx)
{
    float val = 0.f;
    float val_d = 0.f;

    if (dpx.p < 0.5)
    {
        float k = 2 * dpx.p;
        val = k * dpx.p;
    }
    else
    {
        float k = 3 * dpx.p;
        val = k * dpx.p;
    }

    if (dpx.p < 0.5)
    {
        float d_k = 2 * dpx.d;

        // Invalid use of 'k' which is defined in an inaccessible scope.
        val_d = d_k * dpx.p + k * d_x; 
    }
    else
    {
        float d_k = 3 * dpx.d;
        
        // Invalid use of 'k' which is defined in an inaccessible scope.
        val_d = d_k * dpx.p + k * d_x;
    }

    return DifferentialPair<float>(val, val_d);
}

This contains invalid out-of-scope accesses.
For now, we allow this until the last step when the checkpointing pass legalizes all these accesses by inserting loads and stores as necessary.

This can be cleaned up in a more principled way by using ideas from "Tape-based AD". This approach assumes a dynamic infinite 'stack' that the primal code can push intermediate values to. The differential blocks then 'pop' this information out and use it.

We can implement this approach by introducing instructions like IRTapePush and IRTapePop

This allows us to have different implementations depending on how we lower these instructions. The current approach would simply be the static method where we introduce a struct with fields for all unique IRTapePush instructions (and arrays wherever IRTapePush occurs within a loop region).
This could also enable user written derivative code to make use of the intermediate context. Right now, reverse-mode user-written derivatives need to use workarounds such as non-differentiable auxiliary parameters, or write a fully self-contained derivative method.