bench gmm: correct gradients so long as we skip log_gamma_distrib

enzyme/benchmarks/ReverseMode/adbench/ba.h

@@ -127,6 +127,19 @@ extern "C" {
         double* reproj_err,
         double* w_err
     );
+
+    void rust2_ba_objective(
+        int n,
+        int m,
+        int p,
+        double const* cams,
+        double const* X,
+        double const* w,
+        int const* obs,
+        double const* feats,
+        double* reproj_err,
+        double* w_err
+    );
 
     void dcompute_reproj_error(
         double const* cam,
@@ -169,6 +182,20 @@ extern "C" {
     );
 
     void adept_compute_zach_weight_error(double const* w, double* dw, double* err, double* derr);
+
+    void rust_dcompute_reproj_error(
+        double const* cam,
+        double * dcam,
+        double const* X,
+        double * dX,
+        double const* w,
+        double * wb,
+        double const* feat,
+        double *err,
+        double *derr
+    );
+
+    void rust_dcompute_zach_weight_error(double const* w, double* dw, double* err, double* derr);
 }
 
 void read_ba_instance(const string& fn,
@@ -486,9 +513,9 @@ int main(const int argc, const char* argv[]) {
       gettimeofday(&start, NULL);
       calculate_jacobian<dcompute_reproj_error, dcompute_zach_weight_error>(input, result);
       gettimeofday(&end, NULL);
-      printf("Enzyme combined %0.6f\n", tdiff(&start, &end));
+      printf("Enzyme c++ combined %0.6f\n", tdiff(&start, &end));
       json enzyme;
-      enzyme["name"] = "Enzyme combined";
+      enzyme["name"] = "Enzyme c++ combined";
       enzyme["runtime"] = tdiff(&start, &end);
       for(unsigned i=0; i<5; i++) {
         printf("%f ", result.J.vals[i]);
@@ -499,6 +526,125 @@ int main(const int argc, const char* argv[]) {
     }
 
     }
+
+    {
+    struct BAInput input;
+    read_ba_instance("data/" + path, input.n, input.m, input.p, input.cams, input.X, input.w, input.obs, input.feats);
+
+    struct BAOutput result = {
+        std::vector<double>(2 * input.p),
+        std::vector<double>(input.p),
+        BASparseMat(input.n, input.m, input.p)
+    };
+
+
+    {
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+    ba_objective(
+        input.n,
+        input.m,
+        input.p,
+        input.cams.data(),
+        input.X.data(),
+        input.w.data(),
+        input.obs.data(),
+        input.feats.data(),
+        result.reproj_err.data(),
+        result.w_err.data()
+    );
+      gettimeofday(&end, NULL);
+      printf("primal c++ t=%0.6f\n", tdiff(&start, &end));
+      json enzyme;
+      enzyme["name"] = "primal c++";
+      enzyme["runtime"] = tdiff(&start, &end);
+      for(unsigned i=0; i<5; i++) {
+        printf("%f ", result.reproj_err[i]);
+        enzyme["result"].push_back(result.reproj_err[i]);
+      }
+      for(unsigned i=0; i<5; i++) {
+        printf("%f ", result.w_err[i]);
+        enzyme["result"].push_back(result.w_err[i]);
+      }
+      printf("\n");
+      test_suite["tools"].push_back(enzyme);
+    }
+    }
+
+
+    {
+    struct BAInput input;
+    read_ba_instance("data/" + path, input.n, input.m, input.p, input.cams, input.X, input.w, input.obs, input.feats);
+
+    struct BAOutput result = {
+        std::vector<double>(2 * input.p),
+        std::vector<double>(input.p),
+        BASparseMat(input.n, input.m, input.p)
+    };
+    {
+
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+    rust2_ba_objective(
+        input.n,
+        input.m,
+        input.p,
+        input.cams.data(),
+        input.X.data(),
+        input.w.data(),
+        input.obs.data(),
+        input.feats.data(),
+        result.reproj_err.data(),
+        result.w_err.data()
+    );
+      gettimeofday(&end, NULL);
+      printf("primal rust t=%0.6f\n", tdiff(&start, &end));
+      json enzyme;
+      enzyme["name"] = "primal rust";
+      enzyme["runtime"] = tdiff(&start, &end);
+      for(unsigned i=0; i<5; i++) {
+        printf("%f ", result.reproj_err[i]);
+        enzyme["result"].push_back(result.reproj_err[i]);
+      }
+      for(unsigned i=0; i<5; i++) {
+        printf("%f ", result.w_err[i]);
+        enzyme["result"].push_back(result.w_err[i]);
+      }
+      printf("\n");
+      test_suite["tools"].push_back(enzyme);
+    }
+    }
+
+    {
+
+    struct BAInput input;
+    read_ba_instance("data/" + path, input.n, input.m, input.p, input.cams, input.X, input.w, input.obs, input.feats);
+
+    struct BAOutput result = {
+        std::vector<double>(2 * input.p),
+        std::vector<double>(input.p),
+        BASparseMat(input.n, input.m, input.p)
+    };
+
+    {
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+      calculate_jacobian<rust_dcompute_reproj_error, rust_dcompute_zach_weight_error>(input, result);
+      gettimeofday(&end, NULL);
+      printf("Enzyme rust combined %0.6f\n", tdiff(&start, &end));
+      json enzyme;
+      enzyme["name"] = "Enzyme rust combined";
+      enzyme["runtime"] = tdiff(&start, &end);
+      for(unsigned i=0; i<5; i++) {
+        printf("%f ", result.J.vals[i]);
+        enzyme["result"].push_back(result.J.vals[i]);
+      }
+      printf("\n");
+      test_suite["tools"].push_back(enzyme);
+    }
+
+    }
+
     test_suite["llvm-version"] = __clang_version__;
     test_suite["mode"] = "ReverseMode";
     test_suite["batch-size"] = 1;

enzyme/benchmarks/ReverseMode/adbench/gmm.h

@@ -33,6 +33,17 @@ struct GMMParameters {
 };
 
 extern "C" {
+void gmm_objective(
+    int d,
+    int k,
+    int n,
+    double const* alphas,
+    double const* means,
+    double const* icf,
+    double const* x,
+    Wishart wishart,
+    double* err
+);
     void dgmm_objective(int d, int k, int n, const double *alphas, double *
             alphasb, const double *means, double *meansb, const double *icf,
             double *icfb, const double *x, Wishart wishart, double *err, double *
@@ -47,6 +58,15 @@ extern "C" {
         alphasb, const double *means, double *meansb, const double *icf,
         double *icfb, const double *x, Wishart wishart, double *err, double *
         errb);
+
+    void rust_dgmm_objective(int d, int k, int n, const double *alphas, double *
+            alphasb, const double *means, double *meansb, const double *icf,
+            double *icfb, const double *x, Wishart &wishart, double *err, double *
+            errb);
+
+    void rust_gmm_objective(int d, int k, int n, const double *alphas, 
+            const double *means, const double *icf,
+             const double *x, Wishart &wishart, double *err);
 }
 
 void read_gmm_instance(const string& fn,
@@ -123,10 +143,7 @@ void read_gmm_instance(const string& fn,
     fclose(fid);
 }
 
-typedef void(*deriv_t)(int d, int k, int n, const double *alphas, double *alphasb, const double *means, double *meansb, const double *icf,
-            double *icfb, const double *x, Wishart wishart, double *err, double *errb);
-
-template<deriv_t deriv>
+template<auto  deriv>
 void calculate_jacobian(struct GMMInput &input, struct GMMOutput &result)
 {
     double* alphas_gradient_part = result.gradient.data();
@@ -159,6 +176,25 @@ void calculate_jacobian(struct GMMInput &input, struct GMMOutput &result)
     );
 }
 
+template<auto  deriv>
+double primal(struct GMMInput &input)
+{
+    double tmp = 0.0;       // stores fictive result
+                            // (Tapenade doesn't calculate an original function in reverse mode)
+    deriv(
+        input.d,
+        input.k,
+        input.n,
+        input.alphas.data(),
+        input.means.data(),
+        input.icf.data(),
+        input.x.data(),
+        input.wishart,
+        &tmp
+    );
+    return tmp;
+}
+
 int main(const int argc, const char* argv[]) {
     printf("starting main\n");
 
@@ -167,9 +203,11 @@ int main(const int argc, const char* argv[]) {
 
     std::vector<std::string> paths;// = { "1k/gmm_d10_K100.txt" };
 
-    getTests(paths, "data/1k", "1k/");
-    getTests(paths, "data/2.5k", "2.5k/");
-    getTests(paths, "data/10k", "10k/");
+    //getTests(paths, "data/1k", "1k/");
+    //getTests(paths, "data/2.5k", "2.5k/");
+    //getTests(paths, "data/10k", "10k/");
+    //paths.push_back("1k/gmm_d128_K100.txt");
+    paths.push_back("1k/gmm_d2_K5.txt");
 
     std::ofstream jsonfile("results.json", std::ofstream::trunc);
     json test_results;
@@ -257,6 +295,7 @@ int main(const int argc, const char* argv[]) {
       gettimeofday(&start, NULL);
       calculate_jacobian<dgmm_objective>(input, result);
       gettimeofday(&end, NULL);
+      printf("Enzyme c++ combined %0.6f\n", tdiff(&start, &end));
       json enzyme;
       enzyme["name"] = "Enzyme combined";
       enzyme["runtime"] = tdiff(&start, &end);
@@ -269,6 +308,61 @@ int main(const int argc, const char* argv[]) {
       test_suite["tools"].push_back(enzyme);
     }
 
+    }
+
+    {
+
+    struct GMMInput input;
+    read_gmm_instance("data/" + path, &input.d, &input.k, &input.n,
+        input.alphas, input.means, input.icf, input.x, input.wishart, params.replicate_point);
+
+    int Jcols = (input.k * (input.d + 1) * (input.d + 2)) / 2;
+
+    struct GMMOutput result = { 0, std::vector<double>(Jcols) };
+
+    {
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+      auto res = primal<gmm_objective>(input);
+      gettimeofday(&end, NULL);
+      printf("c++ primal combined t=%0.6f, err=%f\n", tdiff(&start, &end), res);
+
+      json primal;
+      primal["name"] = "C++ primal";
+      primal["runtime"] = tdiff(&start, &end);
+      primal["result"].push_back(res);
+      test_suite["tools"].push_back(primal);
+    }
+    {
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+      auto res = primal<rust_gmm_objective>(input);
+      gettimeofday(&end, NULL);
+      printf("rust primal combined t=%0.6f, err=%f\n", tdiff(&start, &end), res);
+      json primal;
+      primal["name"] = "Rust primal";
+      primal["runtime"] = tdiff(&start, &end);
+      primal["result"].push_back(res);
+      test_suite["tools"].push_back(primal);
+    }
+    {
+      struct timeval start, end;
+      gettimeofday(&start, NULL);
+      calculate_jacobian<rust_dgmm_objective>(input, result);
+      gettimeofday(&end, NULL);
+      printf("Enzyme rust combined %0.6f\n", tdiff(&start, &end));
+      json enzyme;
+      enzyme["name"] = "Rust Enzyme combined";
+      enzyme["runtime"] = tdiff(&start, &end);
+      for (unsigned i = result.gradient.size() - 5;
+           i < result.gradient.size(); i++) {
+        printf("%f ", result.gradient[i]);
+        enzyme["result"].push_back(result.gradient[i]);
+      }
+      printf("\n");
+      test_suite["tools"].push_back(enzyme);
+    }
+
     }
     test_suite["llvm-version"] = __clang_version__;
     test_suite["mode"] = "ReverseMode";

enzyme/benchmarks/ReverseMode/ba/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "bars"
+version = "0.1.0"
+edition = "2021"
+
+
+[lib]
+crate-type = ["cdylib"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[profile.release]
+lto = "fat"
+
+[profile.dev]
+lto = "fat"
+
+[dependencies]
+libm = { version = "0.2.8", optional = true }