Skip to content

Commit

Permalink
Move cpu runner utils templates to .h
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@tensorflower-gardener
Nicolas Vasilache authored and tensorflower-gardener committed Dec 12, 2019
1 parent 9b85582 commit 95b5a4f
Show file tree
Hide file tree
Showing 3 changed files with 175 additions and 159 deletions.
181 changes: 164 additions & 17 deletions mlir/test/mlir-cpu-runner/include/mlir_runner_utils.h
View file
Expand Up @@ -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
Expand Down Expand Up @@ -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;
};
Expand All @@ -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];
};
Expand All @@ -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);

Expand Down
149 changes: 8 additions & 141 deletions mlir/test/mlir-cpu-runner/mlir_runner_utils.cpp
View file
Expand Up @@ -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) {
Expand All @@ -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);
}

0 comments on commit 95b5a4f

Please sign in to comment.