Skip to content

Commit

Permalink
Add support for 'std::atan2' and 'std::acos'
Browse files Browse the repository at this point in the history
  • Loading branch information
ZeptoStarling authored and vgvassilev committed Sep 18, 2024
1 parent fd2ac31 commit 7567420
Show file tree
Hide file tree
Showing 2 changed files with 75 additions and 0 deletions.
22 changes: 22 additions & 0 deletions include/clad/Differentiator/BuiltinDerivatives.h
Original file line number Diff line number Diff line change
Expand Up @@ -183,6 +183,25 @@ CUDA_HOST_DEVICE ValueAndPushforward<T, T> floor_pushforward(T x, T /*d_x*/) {
return {::std::floor(x), (T)0};
}

template <typename T>
CUDA_HOST_DEVICE ValueAndPushforward<T, T> atan2_pushforward(T y, T x, T d_y,
T d_x) {
return {::std::atan2(y, x),
-(y / ((x * x) + (y * y))) * d_x + x / ((x * x) + (y * y)) * d_y};
}

template <typename T, typename U>
CUDA_HOST_DEVICE void atan2_pullback(T y, T x, U d_z, T* d_y, T* d_x) {
*d_y += x / ((x * x) + (y * y)) * d_z;

*d_x += -(y / ((x * x) + (y * y))) * d_z;
}

template <typename T>
CUDA_HOST_DEVICE ValueAndPushforward<T, T> acos_pushforward(T x, T d_x) {
return {::std::acos(x), ((-1) / (::std::sqrt(1 - x * x))) * d_x};
}

template <typename T>
CUDA_HOST_DEVICE ValueAndPushforward<T, T> ceil_pushforward(T x, T /*d_x*/) {
return {::std::ceil(x), (T)0};
Expand Down Expand Up @@ -321,6 +340,9 @@ inline void free_pushforward(void* ptr, void* d_ptr) {
// These are required because C variants of mathematical functions are
// defined in global namespace.
using std::abs_pushforward;
using std::acos_pushforward;
using std::atan2_pullback;
using std::atan2_pushforward;
using std::ceil_pushforward;
using std::cos_pushforward;
using std::exp_pushforward;
Expand Down
53 changes: 53 additions & 0 deletions test/FirstDerivative/BuiltinDerivatives.C
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,45 @@ double f14(double x) {
return __builtin_pow(x, 3);
}

double f15(double y, double x) {
return std::atan2(y, x);
}

//CHECK: float f15_darg0(float y, float x) {
//CHECK-NEXT: float _d_y = 1;
//CHECK-NEXT: float _d_x = 0;
//CHECK-NEXT: ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::atan2_pushforward(y, x, _d_y, _d_x);
//CHECK-NEXT: return _t0.pushforward;
//CHECK-NEXT: }

//CHECK: float f15_darg1(float y, float x) {
//CHECK-NEXT: float _d_y = 0;
//CHECK-NEXT: float _d_x = 1;
//CHECK-NEXT: ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::atan2_pushforward(y, x, _d_y, _d_x);
//CHECK-NEXT: return _t0.pushforward;
//CHECK-NEXT: }

void f15_grad(double y, double x, double *_d_y, double *_d_x);
//CHECK: void f15_grad(double y, double x, double *_d_y, double *_d_x) {
//CHECK-NEXT: double _t0 = std::atan2(y, x);
//CHECK-NEXT: {
//CHECK-NEXT: double _r0 = 0;
//CHECK-NEXT: double _r1 = 0;
//CHECK-NEXT: clad::custom_derivatives::atan2_pullback(y, x, 2 * 1, &_r0, &_r1);
//CHECK-NEXT: *_d_y += _r0;
//CHECK-NEXT: *_d_x += _r1;
//CHECK-NEXT: }
//CHECK-NEXT: }

float f16(float x) {
return std::acos(x);
}
// CHECK: float f16_darg0(float x) {
//CHECK-NEXT: float _d_x = 1;
//CHECK-NEXT: ValueAndPushforward<float, float> _t0 = clad::custom_derivatives::std::acos_pushforward(x, _d_x);
//CHECK-NEXT: return _t0.pushforward;
//CHECK-NEXT: }

int main () { //expected-no-diagnostics
float f_result[2];
double d_result[2];
Expand Down Expand Up @@ -326,5 +365,19 @@ int main () { //expected-no-diagnostics
auto f14_ddarg0 = clad::differentiate<2>(f14, 0);
printf("Result is = %f\n", f14_ddarg0.execute(1)); //CHECK-EXEC: Result is = 6.000000

auto f15_darg0 = clad::differentiate(f15, 0);
printf("Result is = %f\n", f15_darg0.execute(4, 3)); //CHECK-EXEC: Result is = 0.120000

auto f15_darg1 = clad::differentiate(f15, 1);
printf("Result is = %f\n", f15_darg1.execute(4, 3)); //CHECK-EXEC: Result is = -0.160000

d_result[0] = d_result[1] = 0;
clad::gradient(f15);
f15_grad(4, 3, &d_result[0], &d_result[1]);
printf("Result is = {%f, %f}\n", d_result[0], d_result[1]); //CHECK-EXEC: Result is = {0.240000, -0.320000}

auto f16_darg0 = clad::differentiate(f16, 0);
printf("Result is = %f\n", f16_darg0.execute(0.9)); //CHECK-EXEC: Result is = -2.294157

return 0;
}

0 comments on commit 7567420

Please sign in to comment.