[tensor] Allow specifying number of elements in dumpImpl (#2601)

pytorch · Mar 29, 2019 · 0e92a7d · 0e92a7d
1 parent fd7d9a6
commit 0e92a7d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 16 deletions.
diff --git a/include/glow/Base/Tensor.h b/include/glow/Base/Tensor.h
@@ -507,10 +507,14 @@ class Tensor final {
 //                    Tensor Handle
 //===----------------------------------------------------------------------===//
 
+constexpr unsigned MAX_DUMP_ELEMS = 100;
+
 void dumpAsciiImpl(const Tensor *T, llvm::raw_ostream &os);
 void dumpAsciiImpl(const Tensor *T);
 
-void dumpImpl(const Tensor *T, llvm::raw_ostream &os);
+void dumpImpl(const Tensor *T, llvm::raw_ostream &os,
+              unsigned maxNumElem = MAX_DUMP_ELEMS);
+void dumpImpl(const Tensor *T, unsigned maxNumElem);
 void dumpImpl(const Tensor *T);
 
 /// A class that provides indexed access to a tensor. This class has value
@@ -761,8 +765,11 @@ template <class ElemTy> class Handle final {
     return std::all_of(begin(), end(), [=](ElemTy e) { return e == trueZero; });
   }
 
-  void dump(llvm::raw_ostream &os) const { dumpImpl(tensor_, os); }
-  void dump() const { dumpImpl(tensor_); }
+  void dump(llvm::raw_ostream &os, unsigned maxNumElem = MAX_DUMP_ELEMS) const {
+    dumpImpl(tensor_, os, maxNumElem);
+  }
+  void dump(unsigned maxNumElem) const { dumpImpl(tensor_, maxNumElem); }
+  void dump() const { dumpImpl(tensor_, MAX_DUMP_ELEMS); }
 
   /// Fill the array with random data that's close to zero using the
   /// Xavier method, based on the paper [Bengio and Glorot 2010].

diff --git a/lib/Base/Tensor.cpp b/lib/Base/Tensor.cpp
@@ -64,7 +64,8 @@ template <class ElemTy> static char valueToChar(ElemTy input) {
 }
 
 template <class ElemTy>
-static void dumpGenericImpl(Handle<ElemTy> handle, llvm::raw_ostream &os) {
+static void dumpGenericImpl(Handle<ElemTy> handle, llvm::raw_ostream &os,
+                            unsigned maxNumElem) {
   auto shape = handle.dims();
   size_t numDims = shape.size();
   auto &Ty = handle.getType();
@@ -105,8 +106,6 @@ static void dumpGenericImpl(Handle<ElemTy> handle, llvm::raw_ostream &os) {
   llvm::write_double(os, mn, llvm::FloatStyle::Fixed, 3);
   os << "\n";
 
-  const unsigned maxNumElem = 100;
-
   os << "[";
 
   for (size_t i = 0, e = std::min<size_t>(maxNumElem, handle.size()); i < e;
@@ -292,29 +291,34 @@ void glow::dumpAsciiImpl(const Tensor *T, llvm::raw_ostream &os) {
 
 void glow::dumpAsciiImpl(const Tensor *T) { dumpAsciiImpl(T, llvm::outs()); }
 
-void glow::dumpImpl(const Tensor *T, llvm::raw_ostream &os) {
+void glow::dumpImpl(const Tensor *T, llvm::raw_ostream &os,
+                    unsigned maxNumElem) {
   switch (T->getElementType()) {
   case ElemKind::FloatTy:
-    return dumpGenericImpl(T->getHandle<float>(), os);
+    return dumpGenericImpl(T->getHandle<float>(), os, maxNumElem);
   case ElemKind::Float16Ty:
-    return dumpGenericImpl(T->getHandle<float16_t>(), os);
+    return dumpGenericImpl(T->getHandle<float16_t>(), os, maxNumElem);
   case ElemKind::Int8QTy:
-    return dumpGenericImpl(T->getHandle<int8_t>(), os);
+    return dumpGenericImpl(T->getHandle<int8_t>(), os, maxNumElem);
   case ElemKind::Int16QTy:
-    return dumpGenericImpl(T->getHandle<int16_t>(), os);
+    return dumpGenericImpl(T->getHandle<int16_t>(), os, maxNumElem);
   case ElemKind::Int32QTy:
-    return dumpGenericImpl(T->getHandle<int32_t>(), os);
+    return dumpGenericImpl(T->getHandle<int32_t>(), os, maxNumElem);
   case ElemKind::Int32ITy:
-    return dumpGenericImpl(T->getHandle<int32_t>(), os);
+    return dumpGenericImpl(T->getHandle<int32_t>(), os, maxNumElem);
   case ElemKind::Int64ITy:
-    return dumpGenericImpl(T->getHandle<int64_t>(), os);
+    return dumpGenericImpl(T->getHandle<int64_t>(), os, maxNumElem);
   case ElemKind::UInt8FusedQTy:
-    return dumpGenericImpl(T->getHandle<uint8_t>(), os);
+    return dumpGenericImpl(T->getHandle<uint8_t>(), os, maxNumElem);
   case ElemKind::BoolTy:
-    return dumpGenericImpl(T->getHandle<bool>(), os);
+    return dumpGenericImpl(T->getHandle<bool>(), os, maxNumElem);
   }
 }
 
+void glow::dumpImpl(const Tensor *T, unsigned maxNumElem) {
+  dumpImpl(T, llvm::outs(), maxNumElem);
+}
+
 void glow::dumpImpl(const Tensor *T) { dumpImpl(T, llvm::outs()); }
 
 void glow::genericTranspose(const Tensor *src, Tensor *dest,