Snapshot of array arguments of primal call expression should be used in the derived call expression in Reverse Mode AD

[parth-07]

Differentiation of call expressions that takes array arguments give incorrect derivatives if the said array is modified between the
primal call expression and the derived call expression.

A reproducible example:

double some_fn(double i, double *arr, int n) {
  return arr[0]*i;
}

double fn(double i, double *arr, int n) {
  double res = 0;
  arr[0] = 1;
  res = some_fn(i, arr, n);
  arr[0] = 5;
  return res;
}

int main() {
  double arr[5] = {1, 2, 3, 4, 5};
  double d_arr[5] = {};

  auto d_fn = clad::gradient(fn);
  clad::array_ref<double> d_arr_ref(d_arr, 5);
  double d_i = 0, d_n = 0;
  d_fn.execute(3, arr, 5, &d_i, d_arr_ref, &d_n);
  std::cout<<"d_i: "<<d_i<<"\n";
}

Expected result:
d_i: 1
Actual result
d_i: 5

Root cause of the error and solution
Derived call expression needs the same values of arguments that were passed to primal call expression. We do clone
all the arguments in the forward pass and use the cloned values in the reverse pass. This works well for ordinary non-pointer variables. But in C++, cloning pointers leads to shallow cloning -- the value pointed by the pointer isn't cloned. Therefore, if the array changes at any point between the primal call expression and the derived call expression, then even though derived call expression is using the cloned pointer, it will still use the modified array values. Therefore, we should deep clone pointers/arrays in the forward pass and use the cloned value in the derived call expression.

Deep cloning an array can be expensive and might be unnecessary if the array is not being modified anywhere between the primal call expression and derived call expression. Data flow and active variable analysis can help us to avoid creating unnecessary clone in those cases and generate optimal derived functions.

[vaithak]

Hi, I am thinking of working on this issue next.
I understand the problem but have some fundamental doubts about the proposed solution and it would be great if someone can help me with them.

The proposed solution mentions deep cloning arrays only when necessary, now this is done at runtime using a mechanism like copy-on-write in programs like docker and git, but the main difference is that they do it at runtime and are in control of the memory which is being copied.

• If I am not wrong, the aim is to do a similar thing but at compile time, right? If so, I am having a hard time thinking how this can be done, because even if we track the arr variable, there is a possibility that a pointer pointing to the arr variable itself modifies the memory by multiple dereferencing.

I am not in any way an expert in compiler algorithms and techniques, so it is possible that I am unaware of some way through which this can be achieved, it would be great if you can share some references or links where I can learn about this more.

[grimmmyshini]

@vaithak maybe you wanna look at #513, it solves a similar problem, maybe you can think of a way to extend it :)

[rajeck1234]

@parth-07 d_fn.execute(3, arr, 5, &d_i, d_arr_ref, &d_n); actually i am looking in this fn is that double fn(double i, double *arr, int n, &d_arr_ref, &d_n)
then this error may gonna correct because when a[0] =1, a[0]=5
the fn return valve res but this may cause (i,&arr,n) coz of it the array value may come result same

[ro4i7]

Hello @parth-07 @grimmmyshini this code solve the issue. Please check it.

#include <clad/ad/autodiff.h>

struct ArrayWrapper {
  double* arr;
  double* clone;
  int size;
  bool modified;

  ArrayWrapper(double* arr_, int size_) : arr(arr_), size(size_), modified(false) {
    clone = new double[size];
    std::copy(arr, arr + size, clone);
  }

  ~ArrayWrapper() {
    delete[] clone;
  }

  void mark_as_modified() {
    if (!modified) {
      modified = true;
      std::copy(arr, arr + size, clone);
    }
  }

  double& operator[](int idx) {
    mark_as_modified();
    return arr[idx];
  }

  const double& operator[](int idx) const {
    return arr[idx];
  }

  double& get_clone(int idx) {
    mark_as_modified();
    return clone[idx];
  }

  const double& get_clone(int idx) const {
    return clone[idx];
  }
};

double some_fn(double i, ArrayWrapper arr, int n) {
  return arr[0] * i;
}

double fn(double i, ArrayWrapper arr, int n) {
  double res = 0;
  arr[0] = 1;
  res = some_fn(i, arr, n);
  arr[0] = 5;
  return res;
}

int main() {
  double arr[5] = {1, 2, 3, 4, 5};
  double d_arr[5] = {};

  ArrayWrapper arr_wrapper(arr, 5);
  auto d_fn = clad::gradient(fn);
  clad::array_ref<double> d_arr_ref(d_arr, 5);
  double d_i = 0, d_n = 0;
  d_fn.execute(3, arr_wrapper, 5, &d_i, d_arr_ref, &d_n);
  std::cout << "d_i: " << d_i << "\n";
}

In this modified code, I define a new data structure called ArrayWrapper that wraps the array and keeps track of its state. The mark_as_modified() function is called whenever the array is modified, and it updates the clone to reflect the new state of the array. In the some_fn() and fn() functions, we use ArrayWrapper instead of raw pointers for the arr argument.

With these modifications, deep cloning is performed automatically in the forward pass and the correct value of d_i (1) is obtained in the reverse pass.

[parth-07]

Hi @ro4i7,

Thank you very much for writing this design for discussion.

Your solution does leave a few essential things unanswered:

• How would the gradient function call of some_fn in the reverse pass look like?
• Why exactly is a double *clone member required?

Currently, the root cause of the problem is that we do not know the size of the array, therefore we cannot
clone the array. If we know the size of the array, we can clone it and the problem is resolved. Therefore, we
would not need a double *clone member.

That being said, cloning an array for each function call that takes that array as an argument is highly inefficient.
That's not the direction we want to move forward with unless we really have to.

One other way which may (partially ) solve this issue is by activity analysis. For example, consider that the derivative function 'reset' values in the reverse-pass such that when reverse-pass statements use a variable, the variable has the
same value as it had in the corresponding forward-pass statement. Let's take the example presented in the issue to
illustrate this idea:

Primal function:

double fn(double i, double *arr, int n) {
  double res = 0;
  arr[0] = 1;
  res = some_fn(i, arr, n);
  arr[0] = 5;
  return res;
}

Derivative code (reverse-mode AD)

double fn_grad(double i, double *arr, int n, double *d_i, double *d_arr, int *d_n) {
  // forward-pass
  double d_res = 0;
  double res = 0;
  double t0 = arr[0]; // 1
  arr[0] = 1;
  double t1 = res;  // 2
  res = some_fn(i, arr, n);
  double t2 = arr[0]; // 3
  arr[0] = 5;
  double fn_return_val = res;

  // reverse-pass
  d_res = 1;
  arr[0] = t2; // 3; reset the value of arr[0]
  d_arr[0] = 0;
  res = t1; // 2; reset the value of res
  some_fn_pullback(i, arr, n, di, d_res, d_arr, d_n);
  arr[0] = t1; // 1; reset the value of arr[0]
  // ...
  // ... and so on
 }

This solution reduces both memory consumption and the number of copy operations.
But this solution is incomplete as it won't work if some_fn modifies arr. We can use a global stack to
propagate values across functions and remove this limitation and consequently make this solution viable for all cases.
But using global variables comes at its own cost such as loss of optimization properties. We perhaps need to find a good compromise between the two solutions or find a better solution altogether.

[vgvassilev]

Fixed in 655.