TBR analysis doesn't include the for loop's step in the forward pass when there's a statement inside with a body of its own

[kchristin22]

double f(double val) {
  for (int i=1; i<5; ++i) {
    val += val;
  }
  return val;
}

Gradient:

• Without TBR:

void f_grad(double val, double *_d_val) {
    unsigned long _t0;
    int _d_i = 0;
    int i = 0;
    clad::tape<double> _t1 = {};
    _t0 = 0;
    for (i = 1; i < 5; ++i) {
        _t0++;
        clad::push(_t1, val);
        val += val;
    }
    goto _label0;
  _label0:
    *_d_val += 1;
    for (; _t0; _t0--) {
        --i;
        {
            val = clad::pop(_t1);
            double _r_d0 = *_d_val;
            *_d_val += _r_d0;
        }
    }
}

• With TBR:

void f_grad(double val, double *_d_val) {
    unsigned long _t0;
    int _d_i = 0;
    int i = 0;
    _t0 = 0;
    for (i = 1; i < 5; ++i) {
        _t0++;
        val += val;
    }
    goto _label0;
  _label0:
    *_d_val += 1;
    for (; _t0; _t0--) {
        {
            double _r_d0 = *_d_val;
            *_d_val += _r_d0;
        }
    }
}

[PetroZarytskyi]

The code looks good to me. Could you explain what concerns you in particular? How should the for loop's step be included?

[kchristin22]

Yes, here the result is not compromised, but if it depended on the value of i (e.g. val += i) then it would be wrong. Hence, the increment of i should be included similarly to the case where TBR is not enabled.

[PetroZarytskyi]

Do you mean --i in the reverse pass? It's not included because we don't need the value of i. It is the purpose of TBR analysis to detect what values we don't need to restore to compute the derivatives. Adding val += i would not change anything because the transformation is linear. However, if we differentiate a function where the value of i is necessary, we will still have --i. e.g.

double f(double val) {
    for (int i=1; i<5; ++i) {
        val += val * i;
    }
    return val;
}

gives us

void f_grad(double val, double *_d_val) {
    unsigned long _t0;
    int _d_i = 0;
    int i = 0;
    clad::tape<double> _t1 = {};
    _t0 = 0;
    for (i = 1; i < 5; ++i) {
        _t0++;
        clad::push(_t1, val);
        val += val * i;
    }
    goto _label0;
  _label0:
    *_d_val += 1;
    for (; _t0; _t0--) {
        --i;
        {
            val = clad::pop(_t1);
            double _r_d0 = *_d_val;
            *_d_val += _r_d0 * i;
            _d_i += val * _r_d0;
        }
    }
}

where we see both i and val be stored and restored because we need them to compute the derivatives.

[kchristin22]

Agreed, but when there's an if or a switch statement or generally a statement with a body, its body isn't traversed I suppose and hence the increment is not included.

double fn21(double val) {
  double res = 0;
  for (int i=1; i<5; ++i) {
    if (i == 3)
      continue;
    res += i * val;
  }
  return res;
}

• With TBR:

void fn21_grad(double val, double *_d_val) {
    double _d_res = 0;
    unsigned long _t0;
    int _d_i = 0;
    int i = 0;
    clad::tape<bool> _t2 = {};
    clad::tape<unsigned long> _t3 = {};
    double res = 0;
    _t0 = 0;
    for (i = 1; i < 5; ++i) {
        _t0++;
        bool _t1 = i == 3;
        {
            if (_t1) {
                clad::push(_t3, 1UL);
                continue;
            }
            clad::push(_t2, _t1);
        }
        res += i * val;
        clad::push(_t3, 2UL);
    }
    goto _label0;
  _label0:
    _d_res += 1;
    for (; _t0; _t0--)
        switch (clad::pop(_t3)) {
          case 2UL:
            ;
            {
                double _r_d0 = _d_res;
                _d_i += _r_d0 * val;
                *_d_val += i * _r_d0;
            }
            if (clad::pop(_t2))
              case 1UL:
                ;
        }
}

Without TBR:

void fn21_grad(double val, double *_d_val) {
    double _d_res = 0;
    unsigned long _t0;
    int _d_i = 0;
    int i = 0;
    clad::tape<bool> _t2 = {};
    clad::tape<unsigned long> _t3 = {};
    clad::tape<double> _t4 = {};
    double res = 0;
    _t0 = 0;
    for (i = 1; i < 5; ++i) {
        _t0++;
        bool _t1 = i == 3;
        {
            if (_t1) {
                clad::push(_t3, 1UL);
                continue;
            }
            clad::push(_t2, _t1);
        }
        clad::push(_t4, res);
        res += i * val;
        clad::push(_t3, 2UL);
    }
    goto _label0;
  _label0:
    _d_res += 1;
    for (; _t0; _t0--) {
        --i;
        switch (clad::pop(_t3)) {
          case 2UL:
            ;
            {
                res = clad::pop(_t4);
                double _r_d0 = _d_res;
                _d_i += _r_d0 * val;
                *_d_val += i * _r_d0;
            }
            if (clad::pop(_t2))
              case 1UL:
                ;
        }
    }
}

I will update the issue title as this case was the one that made me open the issue in the first place.

[PetroZarytskyi]

Okay, this is definitely a bug. I will need more time to find out why this happens. Thank you for pointing this out.

[vgvassilev]

@kchristin22, do you mind looking in the tbr analysis and try to pinpoint the underlying problem? Petro is overloaded with other things in the next couple of weeks.

[kchristin22]

Unfortunately my schedule is also pretty packed these weeks. I'll try to have a look if possible.

[vgvassilev]

Fixed by #1051.