From 7d6524daffaf2197ddfbd67db2fd6ec5b5d691a7 Mon Sep 17 00:00:00 2001
From: Sergiu Deitsch <sergiu.deitsch@gmail.com>
Date: Tue, 1 Feb 2022 20:52:27 +0100
Subject: [PATCH] Support fma Jet

Change-Id: I9b1ad49611e3e6f117190d56512ef2ec4bdbb1c1
---
 include/ceres/jet.h        | 12 ++++++++++++
 internal/ceres/jet_test.cc |  6 ++++++
 2 files changed, 18 insertions(+)

diff --git a/include/ceres/jet.h b/include/ceres/jet.h
index 9a6631752..3048cdce2 100644
--- a/include/ceres/jet.h
+++ b/include/ceres/jet.h
@@ -452,6 +452,7 @@ using std::exp;
 using std::exp2;
 using std::expm1;
 using std::floor;
+using std::fma;
 using std::fmax;
 using std::fmin;
 using std::hypot;
@@ -713,6 +714,17 @@ inline Jet<T, N> hypot(const Jet<T, N>& x,
 }
 #endif  // defined(CERES_HAS_CPP17)
 
+// Like x * y + z but rounded only once.
+template <typename T, int N>
+inline Jet<T, N> fma(const Jet<T, N>& x,
+                     const Jet<T, N>& y,
+                     const Jet<T, N>& z) {
+  // d/dx fma(x, y, z) = y
+  // d/dy fma(x, y, z) = x
+  // d/dz fma(x, y, z) = 1
+  return Jet<T, N>(fma(x.a, y.a, z.a), y.a * x.v + x.a * y.v + z.v);
+}
+
 template <typename T, int N>
 inline Jet<T, N> fmax(const Jet<T, N>& x, const Jet<T, N>& y) {
   using std::isgreater;
diff --git a/internal/ceres/jet_test.cc b/internal/ceres/jet_test.cc
index e613b3f0f..fb7ee81e8 100644
--- a/internal/ceres/jet_test.cc
+++ b/internal/ceres/jet_test.cc
@@ -645,6 +645,12 @@ TEST(Jet, Midpoint) {
 
 #endif  // defined(CERES_HAS_CPP20)
 
+TEST(Jet, Fma) {
+  J v = fma(x, y, z);
+  J w = x * y + z;
+  EXPECT_THAT(v, IsAlmostEqualTo(w));
+}
+
 TEST(Jet, Fmax) {
   EXPECT_THAT(fmax(x, y), IsAlmostEqualTo(x));
   EXPECT_THAT(fmax(y, x), IsAlmostEqualTo(x));