From 25f45cb2db8d8cb27041dfa490ef78ac37addfea Mon Sep 17 00:00:00 2001
From: Sherlock <SherlockNoMad@users.noreply.github.com>
Date: Fri, 22 Mar 2019 10:12:58 -0700
Subject: [PATCH] Introduce Rowwise/Colwise Sum to math util (#656)

---
 onnxruntime/core/util/math.h      | 12 ++++++++++++
 onnxruntime/core/util/math_cpu.cc | 20 ++++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/onnxruntime/core/util/math.h b/onnxruntime/core/util/math.h
index 4c55e6a58f..86edc4884a 100644
--- a/onnxruntime/core/util/math.h
+++ b/onnxruntime/core/util/math.h
@@ -150,6 +150,18 @@ void AddStripedBatch(
     int batch,
     Provider* provider);
 
+// Compute the row-wise sum of a N*D matrix X, and write it to a N
+// dimensional vector y.
+template <typename T, class Provider>
+void RowwiseSum(int N, int D, const T* x, T* y,
+                Provider* provider);
+
+// Compute the column-wise sum of a N*D matrix X, and write it to a D
+// dimensional vector y.
+template <typename T, class Provider>
+void ColwiseSum(int N, int D, const T* x, T* y,
+                Provider* provider);
+
 // Compute the row-wise max of a N*D matrix X, and write it to a N
 // dimensional vector y.
 template <typename T, class Provider>
diff --git a/onnxruntime/core/util/math_cpu.cc b/onnxruntime/core/util/math_cpu.cc
index 8459eb019a..9ba7af5e2d 100644
--- a/onnxruntime/core/util/math_cpu.cc
+++ b/onnxruntime/core/util/math_cpu.cc
@@ -780,6 +780,26 @@ SPECIALIZED_REDUCEMAX(int64_t)
 
 #undef SPECIALIZED_REDUCEMAX
 
+#define SPECIALIZED_ROWWISESUM(T)                                 \
+  template <>                                                     \
+  void RowwiseSum<T, CPUMathUtil>(                                \
+      const int N, const int D, const T* x, T* y, CPUMathUtil*) { \
+    EigenVectorMap<T>(y, N) =                                     \
+        ConstEigenMatrixMap<T>(x, D, N).colwise().sum();          \
+  }
+SPECIALIZED_ROWWISESUM(float)
+#undef SPECIALIZED_ROWWISESUM
+
+#define SPECIALIZED_COLWISESUM(T)                                 \
+  template <>                                                     \
+  void ColwiseSum<T, CPUMathUtil>(                                \
+      const int N, const int D, const T* x, T* y, CPUMathUtil*) { \
+    EigenVectorMap<T>(y, D) =                                     \
+        ConstEigenMatrixMap<T>(x, D, N).rowwise().sum();          \
+  }
+SPECIALIZED_COLWISESUM(float)
+#undef SPECIALIZED_COLWISESUM
+
 #define SPECIALIZED_ROWWISEMAX(T)                                 \
   template <>                                                     \
   void RowwiseMax<T, CPUMathUtil>(                                \