Add custom grad multisource example

7_multisource_custom/Makefile

@@ -0,0 +1,13 @@
+all: lib.exe
+
+clean:
+	rm -f *.o *.ll *.exe
+
+%.o: %.c
+	../dockerscript.sh clang-12 -c /host/$^ -O2 -Xclang -load -Xclang /Enzyme/enzyme/build/Enzyme/ClangEnzyme-12.so -ffast-math -o /host/$@
+
+lib.exe: myblas.o multisource.o
+	../dockerscript.sh clang-12 /host/myblas.o /host/multisource.o -O2 -Xclang -load -Xclang /Enzyme/enzyme/build/Enzyme/ClangEnzyme-12.so -ffast-math -o /host/$@
+
+run-%: %.exe
+	../dockerscript.sh /host/$^ 3

7_multisource_custom/multisource.c

@@ -0,0 +1,50 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+#include <assert.h>
+#include "myblas.h"
+
+void __enzyme_double(void*, size_t);
+
+double dotabs(struct complex* alpha, struct complex* beta, int n) {
+  struct complex prod = myblas_cdot(alpha, beta, n);
+  __enzyme_double(&prod, sizeof(struct complex));
+  return myblas_cabs(prod);
+}
+
+void __enzyme_autodiff(void*, ...);
+
+int main(int argc, char *argv[]) {
+  int n = 3;
+  if (argc > 1) {
+    n = atoi(argv[1]);
+  }
+
+
+  struct complex *A = (struct complex*)malloc(sizeof(struct complex) * n);
+  assert(A != 0);
+  for(int i=0; i<n; i++)
+   A[i] = (struct complex){(i+1), (i+2)};
+
+  struct complex *grad_A = (struct complex*)malloc(sizeof(struct complex) * n);
+  assert(grad_A != 0);
+  for(int i=0; i<n; i++)
+   grad_A[i] = (struct complex){0,0};
+
+
+
+  struct complex *B = (struct complex*)malloc(sizeof(struct complex) * n);
+  assert(B != 0);
+  for(int i=0; i<n; i++)
+   B[i] = (struct complex){-3-i, 2*i};
+
+  struct complex *grad_B = (struct complex*)malloc(sizeof(struct complex) * n);
+  assert(grad_B != 0);
+  for(int i=0; i<n; i++)
+   grad_B[i] = (struct complex){0,0};
+
+  __enzyme_autodiff((void*)dotabs, A, grad_A, B, grad_B, n);
+  printf("Gradient dotabs(A)[0] = %f\n", grad_A[0].r);
+
+  return 0;
+}

7_multisource_custom/myblas.c

@@ -0,0 +1,33 @@
+#include "myblas.h"
+
+struct complex myblas_cdot(struct complex* x, struct complex *y, int n) {
+  struct complex sum = { 0, 0 };
+  for (int i=0; i<n; i++) {
+    sum.r += x[i].r * y[i].r - x[i].i * y[i].i;
+    sum.i += x[i].r * y[i].i + x[i].i * y[i].r;
+  }
+  return sum;
+}
+
+double myblas_cabs(struct complex x) {
+    return x.r * x.r + x.i * x.i;
+}
+
+struct TapeAndComplex __enzyme_augmentfwd(void*, ...);
+struct TapeAndDouble __enzyme_augmentfwd2(void*, ...);
+void __enzyme_reverse(void*, ...);
+struct complex __enzyme_reverse2(void*, ...);
+struct TapeAndComplex myblas_cdot_fwd(struct complex* x, struct complex *dx, struct complex *y, struct complex* dy, int n, int dn) {
+   return __enzyme_augmentfwd((void*)myblas_cdot, x, dx, y, dy, n);
+}
+
+void myblas_cdot_rev(struct complex* x, struct complex* dx, struct complex *y, struct complex* dy, int n, int dn, struct complex* diffret, void* tape) {
+   return __enzyme_reverse((void*)myblas_cdot, x, dx, y, dy, enzyme_dup, n, dn, enzyme_byref, sizeof(struct complex), diffret, tape);
+}
+
+struct TapeAndDouble myblas_cabs_fwd(struct complex x) {
+   return __enzyme_augmentfwd2((void*)myblas_cabs, x);
+}
+struct complex myblas_cabs_rev(struct complex x, double diffret, void* tape) {
+   return __enzyme_reverse2((void*)myblas_cabs, x, diffret, tape);
+}

7_multisource_custom/myblas.h

@@ -0,0 +1,29 @@
+extern int enzyme_const, enzyme_dup, enzyme_out;
+extern int enzyme_byref;
+
+struct complex {
+    double r;
+    double i;
+};
+
+struct complex myblas_cdot(struct complex* x, struct complex *y, int n);
+
+double myblas_cabs(struct complex x);
+
+struct TapeAndComplex {
+    void* tape;
+    struct complex res;
+};
+
+struct TapeAndComplex myblas_cdot_fwd(struct complex* x, struct complex* dx, struct complex *y, struct complex* dy, int n, int dn);
+void myblas_cdot_rev(struct complex* x, struct complex* dx, struct complex *y, struct complex* dy, int n, int dn, struct complex* diffret, void* tape);
+
+struct TapeAndDouble {
+    void* tape;
+    double res;
+};
+
+struct TapeAndDouble myblas_cabs_fwd(struct complex x);
+struct complex myblas_cabs_rev(struct complex x, double diffret, void* tape);
+void* __enzyme_register_gradient1[] = { (void*)myblas_cdot, (void*)myblas_cdot_fwd, (void*)myblas_cdot_rev, "byref_6" };
+void* __enzyme_register_gradient2[] = { (void*)myblas_cabs, (void*)myblas_cabs_fwd, (void*)myblas_cabs_rev };

docker/Dockerfile

@@ -5,13 +5,13 @@ ENV DEBIAN_FRONTEND=noninteractive
 RUN apt-get -y update && apt-get install -y --no-install-recommends curl gnupg lsb-core software-properties-common \
     && curl -fsSL https://apt.llvm.org/llvm-snapshot.gpg.key|apt-key add - \
     && apt-add-repository "deb http://apt.llvm.org/`lsb_release -c | cut -f2`/ llvm-toolchain-`lsb_release -c | cut -f2`-12 main" \
-    && apt-get install -y --no-install-recommends autoconf cmake ninja-build gcc g++ libtool gfortran llvm-12-dev lld-12 clang-12 libopenmpi-dev openmpi-bin git \
+    && apt-get install -y --no-install-recommends autoconf cmake ninja-build gcc g++ libtool gfortran llvm-12-dev lld-12 clang-12 libclang-12-dev libopenmpi-dev openmpi-bin git \
     && apt-get autoremove -y --purge \
     && apt-get clean -y \
     && rm -rf /var/lib/apt/lists/*
 
 # Get & install Enzyme
-RUN git clone https://github.com/wsmoses/Enzyme.git \
+RUN git clone https://github.com/wsmoses/Enzyme.git -b byref \
     && cd Enzyme/enzyme \
     && mkdir build && cd build \
     && cmake -G Ninja .. -DCMAKE_BUILD_TYPE=Debug \