From 95b5a4fd675bd125086d3878df1d4cd3d47d2485 Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <ntv@google.com>
Date: Thu, 12 Dec 2019 07:32:36 -0800
Subject: [PATCH] Move cpu runner utils templates to .h

This allows reusing the implementation in various places by just including and permits more easily writing test functions without explicit template instantiations.

This also modifies UnrankedMemRefType to take a template type parameter since it cannot be type agnostic atm.

PiperOrigin-RevId: 285187711
---
 .../mlir-cpu-runner/include/mlir_runner_utils.h    | 181 +++++++++++++++++++--
 mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp    | 149 +----------------
 mlir/test/mlir-cpu-runner/utils.mlir               |   4 +-
 3 files changed, 175 insertions(+), 159 deletions(-)

diff --git a/mlir/test/mlir-cpu-runner/include/mlir_runner_utils.h b/mlir/test/mlir-cpu-runner/include/mlir_runner_utils.h
index 8a58493..ba68295 100644
--- a/mlir/test/mlir-cpu-runner/include/mlir_runner_utils.h
+++ b/mlir/test/mlir-cpu-runner/include/mlir_runner_utils.h
@@ -17,7 +17,9 @@
 #ifndef MLIR_CPU_RUNNER_MLIRUTILS_H_
 #define MLIR_CPU_RUNNER_MLIRUTILS_H_
 
+#include <assert.h>
 #include <cstdint>
+#include <iostream>
 
 #ifdef _WIN32
 #ifndef MLIR_RUNNER_UTILS_EXPORT
@@ -79,7 +81,7 @@ template <typename T> struct StridedMemRefType<T, 0> {
 };
 
 // Unranked MemRef
-struct UnrankedMemRefType {
+template <typename T> struct UnrankedMemRefType {
   int64_t rank;
   void *descriptor;
 };
@@ -103,26 +105,12 @@ void printMemRefMetaData(StreamType &os, StridedMemRefType<T, 0> &V) {
      << " offset = " << V.offset;
 }
 
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_f32(UnrankedMemRefType *M);
-
-template <typename StreamType>
-void printUnrankedMemRefMetaData(StreamType &os, UnrankedMemRefType &V) {
+template <typename T, typename StreamType>
+void printUnrankedMemRefMetaData(StreamType &os, UnrankedMemRefType<T> &V) {
   os << "Unranked Memref rank = " << V.rank << " "
      << "descriptor@ = " << reinterpret_cast<float *>(V.descriptor) << " ";
 }
 
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_0d_f32(StridedMemRefType<float, 0> *M);
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_1d_f32(StridedMemRefType<float, 1> *M);
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_2d_f32(StridedMemRefType<float, 2> *M);
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_3d_f32(StridedMemRefType<float, 3> *M);
-extern "C" MLIR_RUNNER_UTILS_EXPORT void
-print_memref_4d_f32(StridedMemRefType<float, 4> *M);
-
 template <typename T, int Dim, int... Dims> struct Vector {
   Vector<T, Dims...> vector[Dim];
 };
@@ -135,6 +123,165 @@ using Vector3D = Vector<T, D1, D2, D3>;
 template <int D1, int D2, int D3, int D4, typename T>
 using Vector4D = Vector<T, D1, D2, D3, D4>;
 
+////////////////////////////////////////////////////////////////////////////////
+// Templated instantiation follows.
+////////////////////////////////////////////////////////////////////////////////
+namespace impl {
+template <typename T, int M, int... Dims>
+std::ostream &operator<<(std::ostream &os, const Vector<T, M, Dims...> &v);
+
+template <int... Dims> struct StaticSizeMult {
+  static constexpr int value = 1;
+};
+
+template <int N, int... Dims> struct StaticSizeMult<N, Dims...> {
+  static constexpr int value = N * StaticSizeMult<Dims...>::value;
+};
+
+static void printSpace(std::ostream &os, int count) {
+  for (int i = 0; i < count; ++i) {
+    os << ' ';
+  }
+}
+
+template <typename T, int M, int... Dims> struct VectorDataPrinter {
+  static void print(std::ostream &os, const Vector<T, M, Dims...> &val);
+};
+
+template <typename T, int M, int... Dims>
+void VectorDataPrinter<T, M, Dims...>::print(std::ostream &os,
+                                             const Vector<T, M, Dims...> &val) {
+  static_assert(M > 0, "0 dimensioned tensor");
+  static_assert(sizeof(val) == M * StaticSizeMult<Dims...>::value * sizeof(T),
+                "Incorrect vector size!");
+  // First
+  os << "(" << val.vector[0];
+  if (M > 1)
+    os << ", ";
+  if (sizeof...(Dims) > 1)
+    os << "\n";
+  // Kernel
+  for (unsigned i = 1; i + 1 < M; ++i) {
+    printSpace(os, 2 * sizeof...(Dims));
+    os << val.vector[i] << ", ";
+    if (sizeof...(Dims) > 1)
+      os << "\n";
+  }
+  // Last
+  if (M > 1) {
+    printSpace(os, sizeof...(Dims));
+    os << val.vector[M - 1];
+  }
+  os << ")";
+}
+
+template <typename T, int M, int... Dims>
+std::ostream &operator<<(std::ostream &os, const Vector<T, M, Dims...> &v) {
+  VectorDataPrinter<T, M, Dims...>::print(os, v);
+  return os;
+}
+
+template <typename T, int N> struct MemRefDataPrinter {
+  static void print(std::ostream &os, T *base, int64_t rank, int64_t offset,
+                    int64_t *sizes, int64_t *strides);
+  static void printFirst(std::ostream &os, T *base, int64_t rank,
+                         int64_t offset, int64_t *sizes, int64_t *strides);
+  static void printLast(std::ostream &os, T *base, int64_t rank, int64_t offset,
+                        int64_t *sizes, int64_t *strides);
+};
+
+template <typename T> struct MemRefDataPrinter<T, 0> {
+  static void print(std::ostream &os, T *base, int64_t rank, int64_t offset,
+                    int64_t *sizes = nullptr, int64_t *strides = nullptr);
+};
+
+template <typename T, int N>
+void MemRefDataPrinter<T, N>::printFirst(std::ostream &os, T *base,
+                                         int64_t rank, int64_t offset,
+                                         int64_t *sizes, int64_t *strides) {
+  os << "[";
+  MemRefDataPrinter<T, N - 1>::print(os, base, rank, offset, sizes + 1,
+                                     strides + 1);
+  // If single element, close square bracket and return early.
+  if (sizes[0] <= 1) {
+    os << "]";
+    return;
+  }
+  os << ", ";
+  if (N > 1)
+    os << "\n";
+}
+
+template <typename T, int N>
+void MemRefDataPrinter<T, N>::print(std::ostream &os, T *base, int64_t rank,
+                                    int64_t offset, int64_t *sizes,
+                                    int64_t *strides) {
+  printFirst(os, base, rank, offset, sizes, strides);
+  for (unsigned i = 1; i + 1 < sizes[0]; ++i) {
+    printSpace(os, rank - N + 1);
+    MemRefDataPrinter<T, N - 1>::print(os, base, rank, offset + i * strides[0],
+                                       sizes + 1, strides + 1);
+    os << ", ";
+    if (N > 1)
+      os << "\n";
+  }
+  if (sizes[0] <= 1)
+    return;
+  printLast(os, base, rank, offset, sizes, strides);
+}
+
+template <typename T, int N>
+void MemRefDataPrinter<T, N>::printLast(std::ostream &os, T *base, int64_t rank,
+                                        int64_t offset, int64_t *sizes,
+                                        int64_t *strides) {
+  printSpace(os, rank - N + 1);
+  MemRefDataPrinter<T, N - 1>::print(os, base, rank,
+                                     offset + (sizes[0] - 1) * (*strides),
+                                     sizes + 1, strides + 1);
+  os << "]";
+}
+
+template <typename T>
+void MemRefDataPrinter<T, 0>::print(std::ostream &os, T *base, int64_t rank,
+                                    int64_t offset, int64_t *sizes,
+                                    int64_t *strides) {
+  os << base[offset];
+}
+
+template <typename T, int N> void printMemRef(StridedMemRefType<T, N> &M) {
+  static_assert(N > 0, "Expected N > 0");
+  printMemRefMetaData(std::cout, M);
+  std::cout << " data = " << std::endl;
+  MemRefDataPrinter<T, N>::print(std::cout, M.data, N, M.offset, M.sizes,
+                                 M.strides);
+  std::cout << std::endl;
+}
+
+template <typename T> void printMemRef(StridedMemRefType<T, 0> &M) {
+  std::cout << "\nMemref base@ = " << M.data << " rank = " << 0
+            << " offset = " << M.offset << " data = " << std::endl;
+  std::cout << "[";
+  MemRefDataPrinter<T, 0>::print(std::cout, M.data, 0, M.offset);
+  std::cout << "]" << std::endl;
+}
+} // namespace impl
+
+////////////////////////////////////////////////////////////////////////////////
+// Currently exposed C API.
+////////////////////////////////////////////////////////////////////////////////
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_f32(UnrankedMemRefType<float> *M);
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_0d_f32(StridedMemRefType<float, 0> *M);
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_1d_f32(StridedMemRefType<float, 1> *M);
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_2d_f32(StridedMemRefType<float, 2> *M);
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_3d_f32(StridedMemRefType<float, 3> *M);
+extern "C" MLIR_RUNNER_UTILS_EXPORT void
+print_memref_4d_f32(StridedMemRefType<float, 4> *M);
+
 extern "C" MLIR_RUNNER_UTILS_EXPORT void
 print_memref_vector_4x4xf32(StridedMemRefType<Vector2D<4, 4, float>, 2> *M);
 
diff --git a/mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp b/mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp
index 056ff65..f8007d7 100644
--- a/mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp
+++ b/mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp
@@ -21,153 +21,20 @@
 //===----------------------------------------------------------------------===//
 
 #include "include/mlir_runner_utils.h"
-#include <assert.h>
-#include <iostream>
-
-template <int... Dims> struct StaticSizeMult {
-  static constexpr int value = 1;
-};
-
-template <int N, int... Dims> struct StaticSizeMult<N, Dims...> {
-  static constexpr int value = N * StaticSizeMult<Dims...>::value;
-};
-
-static void printSpace(std::ostream &os, int count) {
-  for (int i = 0; i < count; ++i) {
-    os << ' ';
-  }
-}
-
-template <typename T, int M, int... Dims> struct VectorDataPrinter {
-  static void print(std::ostream &os, const Vector<T, M, Dims...> &val);
-};
-
-template <typename T, int M, int... Dims>
-void VectorDataPrinter<T, M, Dims...>::print(std::ostream &os,
-                                             const Vector<T, M, Dims...> &val) {
-  static_assert(M > 0, "0 dimensioned tensor");
-  static_assert(sizeof(val) == M * StaticSizeMult<Dims...>::value * sizeof(T),
-                "Incorrect vector size!");
-  // First
-  os << "(" << val.vector[0];
-  if (M > 1)
-    os << ", ";
-  if (sizeof...(Dims) > 1)
-    os << "\n";
-  // Kernel
-  for (unsigned i = 1; i + 1 < M; ++i) {
-    printSpace(os, 2 * sizeof...(Dims));
-    os << val.vector[i] << ", ";
-    if (sizeof...(Dims) > 1)
-      os << "\n";
-  }
-  // Last
-  printSpace(os, sizeof...(Dims));
-  os << val.vector[M - 1] << ")";
-}
-
-template <typename T, int M, int... Dims>
-std::ostream &operator<<(std::ostream &os, const Vector<T, M, Dims...> &v) {
-  VectorDataPrinter<T, M, Dims...>::print(os, v);
-  return os;
-}
-
-template <typename T, int N> struct MemRefDataPrinter {
-  static void print(std::ostream &os, T *base, int64_t rank, int64_t offset,
-                    int64_t *sizes, int64_t *strides);
-  static void printFirst(std::ostream &os, T *base, int64_t rank,
-                         int64_t offset, int64_t *sizes, int64_t *strides);
-  static void printLast(std::ostream &os, T *base, int64_t rank, int64_t offset,
-                        int64_t *sizes, int64_t *strides);
-};
-
-template <typename T> struct MemRefDataPrinter<T, 0> {
-  static void print(std::ostream &os, T *base, int64_t rank, int64_t offset,
-                    int64_t *sizes = nullptr, int64_t *strides = nullptr);
-};
-
-template <typename T, int N>
-void MemRefDataPrinter<T, N>::printFirst(std::ostream &os, T *base,
-                                         int64_t rank, int64_t offset,
-                                         int64_t *sizes, int64_t *strides) {
-  os << "[";
-  MemRefDataPrinter<T, N - 1>::print(os, base, rank, offset, sizes + 1,
-                                     strides + 1);
-  // If single element, close square bracket and return early.
-  if (sizes[0] <= 1) {
-    os << "]";
-    return;
-  }
-  os << ", ";
-  if (N > 1)
-    os << "\n";
-}
-
-template <typename T, int N>
-void MemRefDataPrinter<T, N>::print(std::ostream &os, T *base, int64_t rank,
-                                    int64_t offset, int64_t *sizes,
-                                    int64_t *strides) {
-  printFirst(os, base, rank, offset, sizes, strides);
-  for (unsigned i = 1; i + 1 < sizes[0]; ++i) {
-    printSpace(os, rank - N + 1);
-    MemRefDataPrinter<T, N - 1>::print(os, base, rank, offset + i * strides[0],
-                                       sizes + 1, strides + 1);
-    os << ", ";
-    if (N > 1)
-      os << "\n";
-  }
-  if (sizes[0] <= 1)
-    return;
-  printLast(os, base, rank, offset, sizes, strides);
-}
-
-template <typename T, int N>
-void MemRefDataPrinter<T, N>::printLast(std::ostream &os, T *base, int64_t rank,
-                                        int64_t offset, int64_t *sizes,
-                                        int64_t *strides) {
-  printSpace(os, rank - N + 1);
-  MemRefDataPrinter<T, N - 1>::print(os, base, rank,
-                                     offset + (sizes[0] - 1) * (*strides),
-                                     sizes + 1, strides + 1);
-  os << "]";
-}
-
-template <typename T>
-void MemRefDataPrinter<T, 0>::print(std::ostream &os, T *base, int64_t rank,
-                                    int64_t offset, int64_t *sizes,
-                                    int64_t *strides) {
-  os << base[offset];
-}
-
-template <typename T, int N> void printMemRef(StridedMemRefType<T, N> &M) {
-  static_assert(N > 0, "Expected N > 0");
-  printMemRefMetaData(std::cout, M);
-  std::cout << " data = " << std::endl;
-  MemRefDataPrinter<T, N>::print(std::cout, M.data, N, M.offset, M.sizes,
-                                 M.strides);
-  std::cout << std::endl;
-}
-
-template <typename T> void printMemRef(StridedMemRefType<T, 0> &M) {
-  std::cout << "\nMemref base@ = " << M.data << " rank = " << 0
-            << " offset = " << M.offset << " data = [";
-  MemRefDataPrinter<T, 0>::print(std::cout, M.data, 0, M.offset);
-  std::cout << "]" << std::endl;
-}
 
 extern "C" void
 print_memref_vector_4x4xf32(StridedMemRefType<Vector2D<4, 4, float>, 2> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
 
-extern "C" void print_memref_f32(UnrankedMemRefType *M) {
+extern "C" void print_memref_f32(UnrankedMemRefType<float> *M) {
   printUnrankedMemRefMetaData(std::cout, *M);
   int rank = M->rank;
   void *ptr = M->descriptor;
 
 #define MEMREF_CASE(RANK)                                                      \
   case RANK:                                                                   \
-    printMemRef(*(static_cast<StridedMemRefType<float, RANK> *>(ptr)));        \
+    impl::printMemRef(*(static_cast<StridedMemRefType<float, RANK> *>(ptr)));  \
     break
 
   switch (rank) {
@@ -182,17 +49,17 @@ extern "C" void print_memref_f32(UnrankedMemRefType *M) {
 }
 
 extern "C" void print_memref_0d_f32(StridedMemRefType<float, 0> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
 extern "C" void print_memref_1d_f32(StridedMemRefType<float, 1> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
 extern "C" void print_memref_2d_f32(StridedMemRefType<float, 2> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
 extern "C" void print_memref_3d_f32(StridedMemRefType<float, 3> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
 extern "C" void print_memref_4d_f32(StridedMemRefType<float, 4> *M) {
-  printMemRef(*M);
+  impl::printMemRef(*M);
 }
diff --git a/mlir/test/mlir-cpu-runner/utils.mlir b/mlir/test/mlir-cpu-runner/utils.mlir
index 099b856..2a56920 100644
--- a/mlir/test/mlir-cpu-runner/utils.mlir
+++ b/mlir/test/mlir-cpu-runner/utils.mlir
@@ -12,7 +12,9 @@ func @print_0d() {
   dealloc %A : memref<f32>
   return
 }
-// PRINT-0D: Memref base@ = {{.*}} rank = 0 offset = 0 data = [2]
+// PRINT-0D: Unranked Memref rank = 0 descriptor@ = {{.*}}
+// PRINT-0D: Memref base@ = {{.*}} rank = 0 offset = 0 data =
+// PRINT-0D: [2]
 
 func @print_1d() {
   %f = constant 2.00000e+00 : f32
-- 
2.7.4