Skip to content

Commit

Permalink
Add support for C-style memory alloc and free in forward mode AD
Browse files Browse the repository at this point in the history
  • Loading branch information
@vaithak
vaithak committed Feb 23, 2024
1 parent ce677a8 commit 7e85226
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 4 deletions.
25 changes: 25 additions & 0 deletions include/clad/Differentiator/BuiltinDerivatives.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,31 @@ CUDA_HOST_DEVICE void clamp_pullback(const T& v, const T& lo, const T& hi,
#endif

} // namespace std

// NOLINTBEGIN(cppcoreguidelines-no-malloc)
// NOLINTBEGIN(cppcoreguidelines-owning-memory)
inline ValueAndPushforward<void*, void*> malloc_pushforward(size_t sz,
size_t d_sz) {
return {malloc(sz), malloc(sz)};
}

inline ValueAndPushforward<void*, void*>
calloc_pushforward(size_t n, size_t sz, size_t d_n, size_t d_sz) {
return {calloc(n, sz), calloc(n, sz)};
}

inline ValueAndPushforward<void*, void*>
realloc_pushforward(void* ptr, size_t sz, void* d_ptr, size_t d_sz) {
return {realloc(ptr, sz), realloc(d_ptr, sz)};
}

inline void free_pushforward(void* ptr, void* d_ptr) {
free(ptr);
free(d_ptr);
}
// NOLINTEND(cppcoreguidelines-owning-memory)
// NOLINTEND(cppcoreguidelines-no-malloc)

// These are required because C variants of mathematical functions are
// defined in global namespace.
using std::abs_pushforward;
Expand Down
43 changes: 43 additions & 0 deletions test/ForwardMode/Pointer.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -133,18 +133,61 @@ double fn6 (double i) {
// CHECK-NEXT: return _d_res;
// CHECK-NEXT: }

double fn7(double i) {
double *p = (double*)malloc(8UL /*sizeof(double)*/);
*p = i;
T *t = (T*)calloc(1, sizeof(T));
t->i = i;
double res = *p + t->i;
p = (double*)realloc(p, 2*sizeof(double));
p[1] = 2*i;
res += p[1];
free(t);
free(p);
return res;
}

// CHECK: double fn7_darg0(double i) {
// CHECK-NEXT: double _d_i = 1;
// CHECK-NEXT: {{(clad::)?}}ValueAndPushforward<void *, void *> _t0 = clad::custom_derivatives::malloc_pushforward(8UL, 0UL);
// CHECK-NEXT: double *_d_p = (double *)_t0.pushforward;
// CHECK-NEXT: double *p = (double *)_t0.value;
// CHECK-NEXT: *_d_p = _d_i;
// CHECK-NEXT: *p = i;
// CHECK-NEXT: {{(clad::)?}}ValueAndPushforward<void *, void *> _t1 = clad::custom_derivatives::calloc_pushforward(1, sizeof(T), 0, sizeof(T));
// CHECK-NEXT: T *_d_t = (T *)_t1.pushforward;
// CHECK-NEXT: T *t = (T *)_t1.value;
// CHECK-NEXT: _d_t->i = _d_i;
// CHECK-NEXT: t->i = i;
// CHECK-NEXT: double _d_res = *_d_p + _d_t->i;
// CHECK-NEXT: double res = *p + t->i;
// CHECK-NEXT: unsigned long _t2 = sizeof(double);
// CHECK-NEXT: {{(clad::)?}}ValueAndPushforward<void *, void *> _t3 = clad::custom_derivatives::realloc_pushforward(p, 2 * _t2, _d_p, 0 * _t2 + 2 * sizeof(double));
// CHECK-NEXT: _d_p = (double *)_t3.pushforward;
// CHECK-NEXT: p = (double *)_t3.value;
// CHECK-NEXT: _d_p[1] = 0 * i + 2 * _d_i;
// CHECK-NEXT: p[1] = 2 * i;
// CHECK-NEXT: _d_res += _d_p[1];
// CHECK-NEXT: res += p[1];
// CHECK-NEXT: clad::custom_derivatives::free_pushforward(t, _d_t);
// CHECK-NEXT: clad::custom_derivatives::free_pushforward(p, _d_p);
// CHECK-NEXT: return _d_res;
// CHECK-NEXT: }

int main() {
INIT_DIFFERENTIATE(fn1, "i");
INIT_DIFFERENTIATE(fn2, "i");
INIT_DIFFERENTIATE(fn3, "i");
INIT_DIFFERENTIATE(fn4, "i");
INIT_DIFFERENTIATE(fn5, "i");
INIT_DIFFERENTIATE(fn6, "i");
INIT_DIFFERENTIATE(fn7, "i");

TEST_DIFFERENTIATE(fn1, 3, 5); // CHECK-EXEC: {5.00}
TEST_DIFFERENTIATE(fn2, 3, 5); // CHECK-EXEC: {5.00}
TEST_DIFFERENTIATE(fn3, 3, 5); // CHECK-EXEC: {6.00}
TEST_DIFFERENTIATE(fn4, 3, 5); // CHECK-EXEC: {16.00}
TEST_DIFFERENTIATE(fn5, 3, 5); // CHECK-EXEC: {57.00}
TEST_DIFFERENTIATE(fn6, 3); // CHECK-EXEC: {1.00}
TEST_DIFFERENTIATE(fn7, 3); // CHECK-EXEC: {4.00}
}
58 changes: 54 additions & 4 deletions test/Gradient/Pointers.C
View file
Original file line number Diff line number Diff line change
Expand Up @@ -433,7 +433,13 @@ double structPointer (double x) {
double cStyleMemoryAlloc(double x, size_t n) {
T* t = (T*)malloc(n * sizeof(T));
t->x = x;
double res = t->x;
double* p = (double*)calloc(1, sizeof(double));
*p = x;
double res = t->x + *p;
p = (double*)realloc(p, 2*sizeof(double));
p[1] = 2*x;
res += p[1];
free(p);
free(t);
return res;
}
Expand All @@ -442,22 +448,66 @@ double cStyleMemoryAlloc(double x, size_t n) {
// CHECK-NEXT: size_t _d_n = 0;
// CHECK-NEXT: T *_d_t = 0;
// CHECK-NEXT: double _t0;
// CHECK-NEXT: double *_d_p = 0;
// CHECK-NEXT: double _t1;
// CHECK-NEXT: double _d_res = 0;
// CHECK-NEXT: double *_t2;
// CHECK-NEXT: double *_t3;
// CHECK-NEXT: double _t4;
// CHECK-NEXT: double _t5;
// CHECK-NEXT: _d_t = (T *)malloc(n * sizeof(T));
// CHECK-NEXT: T *t = (T *)malloc(n * sizeof(T));
// CHECK-NEXT: _t0 = t->x;
// CHECK-NEXT: t->x = x;
// CHECK-NEXT: double res = t->x;
// CHECK-NEXT: _d_p = (double *)calloc(1, sizeof(double));
// CHECK-NEXT: double *p = (double *)calloc(1, sizeof(double));
// CHECK-NEXT: _t1 = *p;
// CHECK-NEXT: *p = x;
// CHECK-NEXT: double res = t->x + *p;
// CHECK-NEXT: _t2 = p;
// CHECK-NEXT: _t3 = _d_p;
// CHECK-NEXT: _d_p = (double *)realloc(_d_p, 2 * sizeof(double));
// CHECK-NEXT: p = (double *)realloc(p, 2 * sizeof(double));
// CHECK-NEXT: _t4 = p[1];
// CHECK-NEXT: p[1] = 2 * x;
// CHECK-NEXT: _t5 = res;
// CHECK-NEXT: res += p[1];
// CHECK-NEXT: goto _label0;
// CHECK-NEXT: _label0:
// CHECK-NEXT: _d_res += 1;
// CHECK-NEXT: _d_t->x += _d_res;
// CHECK-NEXT: {
// CHECK-NEXT: res = _t5;
// CHECK-NEXT: double _r_d3 = _d_res;
// CHECK-NEXT: _d_p[1] += _r_d3;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: p[1] = _t4;
// CHECK-NEXT: double _r_d2 = _d_p[1];
// CHECK-NEXT: _d_p[1] -= _r_d2;
// CHECK-NEXT: * _d_x += 2 * _r_d2;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: p = _t2;
// CHECK-NEXT: _d_p = _t3;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: _d_t->x += _d_res;
// CHECK-NEXT: *_d_p += _d_res;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: *p = _t1;
// CHECK-NEXT: double _r_d1 = *_d_p;
// CHECK-NEXT: *_d_p -= _r_d1;
// CHECK-NEXT: * _d_x += _r_d1;
// CHECK-NEXT: }
// CHECK-NEXT: {
// CHECK-NEXT: t->x = _t0;
// CHECK-NEXT: double _r_d0 = _d_t->x;
// CHECK-NEXT: _d_t->x -= _r_d0;
// CHECK-NEXT: * _d_x += _r_d0;
// CHECK-NEXT: }
// CHECK-NEXT: free(p);
// CHECK-NEXT: free(_d_p);
// CHECK-NEXT: free(t);
// CHECK-NEXT: free(_d_t);
// CHECK-NEXT: }
Expand Down Expand Up @@ -570,5 +620,5 @@ int main() {
auto d_cStyleMemoryAlloc = clad::gradient(cStyleMemoryAlloc, "x");
d_x = 0;
d_cStyleMemoryAlloc.execute(5, 7, &d_x);
printf("%.2f\n", d_x); // CHECK-EXEC: 1.00
printf("%.2f\n", d_x); // CHECK-EXEC: 4.00
}

0 comments on commit 7e85226

Please sign in to comment.