Add Environment Variables for cuda tensor dumper (#14780)

(1) Add two environment variables to configure the cuda dumper:
`ORT_TENSOR_SNIPPET_THRESHOLD` and `ORT_TENSOR_SNIPPET_EDGE_ITEMS`
(2) Move IConsoleDumper definition to a separated file console_dumper.h.

onnxruntime/contrib_ops/cpu/transformers/generation_shared.h

@@ -7,7 +7,7 @@
 #include <random>
 #include "core/common/gsl.h"
 #include "core/framework/allocator.h"
-#include "core/framework/ort_value.h"
+#include "contrib_ops/cpu/utils/console_dumper.h"
 
 namespace onnxruntime {
 
@@ -163,50 +163,6 @@ struct IGenerationParameters {
   bool custom_sampling = false;
 };
 
-// #define DEBUG_GENERATION 1  // uncomment it for debugging generation (like beam search etc)
-#ifdef DEBUG_GENERATION
-#define DUMP_TENSOR_LEVEL 2
-#else
-#define DUMP_TENSOR_LEVEL 0  // change it to 1 or 2 if want to enable dumping for code not in generation.
-#endif
-
-#if DUMP_TENSOR_LEVEL > 0
-#define DUMP_TENSOR_INIT() transformers::CudaTensorConsoleDumper dumper
-#define DUMP_TENSOR(...) dumper.Print(__VA_ARGS__)
-#else
-#define DUMP_TENSOR_INIT()
-#define DUMP_TENSOR(...)
-#endif
-#if DUMP_TENSOR_LEVEL > 1
-#define DUMP_TENSOR_D(...) dumper.Print(__VA_ARGS__)
-#else
-#define DUMP_TENSOR_D(...)
-#endif
-
-class IConsoleDumper {
- public:
-  IConsoleDumper() : is_enabled_(true) {}
-  virtual ~IConsoleDumper() {}
-  void Disable() { is_enabled_ = false; }
-  bool IsEnabled() const { return is_enabled_; }
-  virtual void Print(const char* name, const float* tensor, int dim0, int dim1) const = 0;
-  virtual void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const = 0;
-  virtual void Print(const char* name, const size_t* tensor, int dim0, int dim1) const = 0;
-  virtual void Print(const char* name, const int64_t* tensor, int dim0, int dim1) const = 0;
-  virtual void Print(const char* name, const int32_t* tensor, int dim0, int dim1) const = 0;
-  virtual void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2) const = 0;
-  virtual void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2) const = 0;
-  virtual void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const = 0;
-  virtual void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const = 0;
-  virtual void Print(const char* name, const Tensor& value) const = 0;
-  virtual void Print(const char* name, const OrtValue& value) const = 0;
-  virtual void Print(const char* name, int index, bool end_line) const = 0;
-  virtual void Print(const char* name, const std::string& value, bool end_line) const = 0;
-
- protected:
-  bool is_enabled_;
-};
-
 }  // namespace transformers
 }  // namespace contrib
 }  // namespace onnxruntime

onnxruntime/contrib_ops/cpu/utils/console_dumper.h

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include <string>
+#include "core/framework/ort_value.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace transformers {
+
+// #define DEBUG_GENERATION 1  // uncomment it for debugging generation (like beam search etc)
+#ifdef DEBUG_GENERATION
+#define DUMP_TENSOR_LEVEL 2
+#else
+#define DUMP_TENSOR_LEVEL 0  // change it to 1 or 2 if want to enable dumping for code not in generation.
+#endif
+
+#if DUMP_TENSOR_LEVEL > 0
+#define DUMP_TENSOR_INIT() transformers::CudaTensorConsoleDumper dumper
+#define DUMP_TENSOR(...) dumper.Print(__VA_ARGS__)
+#else
+#define DUMP_TENSOR_INIT()
+#define DUMP_TENSOR(...)
+#endif
+#if DUMP_TENSOR_LEVEL > 1
+#define DUMP_TENSOR_D(...) dumper.Print(__VA_ARGS__)
+#else
+#define DUMP_TENSOR_D(...)
+#endif
+
+class IConsoleDumper {
+ public:
+  IConsoleDumper() : is_enabled_(true) {}
+  virtual ~IConsoleDumper() {}
+  void Disable() { is_enabled_ = false; }
+  bool IsEnabled() const { return is_enabled_; }
+  virtual void Print(const char* name, const float* tensor, int dim0, int dim1) const = 0;
+  virtual void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const = 0;
+  virtual void Print(const char* name, const size_t* tensor, int dim0, int dim1) const = 0;
+  virtual void Print(const char* name, const int64_t* tensor, int dim0, int dim1) const = 0;
+  virtual void Print(const char* name, const int32_t* tensor, int dim0, int dim1) const = 0;
+  virtual void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2) const = 0;
+  virtual void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2) const = 0;
+  virtual void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const = 0;
+  virtual void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const = 0;
+  virtual void Print(const char* name, const Tensor& value) const = 0;
+  virtual void Print(const char* name, const OrtValue& value) const = 0;
+  virtual void Print(const char* name, int index, bool end_line) const = 0;
+  virtual void Print(const char* name, const std::string& value, bool end_line) const = 0;
+
+ protected:
+  bool is_enabled_;
+};
+
+}  // namespace transformers
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc

@@ -5,6 +5,7 @@
 #include "core/providers/cuda/cuda_common.h"
 #include "core/framework/print_tensor_utils.h"
 #include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
+#include "core/platform/env_var_utils.h"
 
 namespace onnxruntime {
 namespace contrib {
@@ -12,6 +13,39 @@ namespace cuda {
 namespace transformers {
 
 #if DUMP_TENSOR_LEVEL > 0
+
+// Total number of elements which trigger snippet rather than full dump (default 200). Value 0 disables snippet.
+constexpr const char* kTensorSnippetThreshold = "ORT_TENSOR_SNIPPET_THRESHOLD";
+
+// Number of array items in snippet at beginning and end of each dimension (default 3)
+constexpr const char* kTensorSnippetEdgeItems = "ORT_TENSOR_SNIPPET_EDGE_ITEMS";
+
+class DumpTensorConfig {
+ public:
+  static DumpTensorConfig& instance() {
+    static DumpTensorConfig instance;
+    return instance;
+  }
+
+  DumpTensorConfig(const DumpTensorConfig&) = delete;
+  DumpTensorConfig& operator=(const DumpTensorConfig&) = delete;
+
+  int get_snippet_threshold() const { return snippet_threshold; }
+  int get_snippet_edge_items() const { return snippet_edge_items; }
+
+ private:
+  int snippet_threshold;
+  int snippet_edge_items;
+
+  DumpTensorConfig() {
+    snippet_threshold = ParseEnvironmentVariableWithDefault<int>(kTensorSnippetThreshold,
+                                                                 onnxruntime::utils::kDefaultSnippetThreshold);
+    snippet_edge_items = ParseEnvironmentVariableWithDefault<int>(kTensorSnippetEdgeItems,
+                                                                  onnxruntime::utils::kDefaultSnippetEdgeItems);
+  }
+  ~DumpTensorConfig() {}
+};
+
 template <typename T>
 class PinnedHostBuffer {
  public:
@@ -38,8 +72,6 @@ class PinnedHostBuffer {
   T* buffer_;
 };
 
-constexpr int64_t kGpuSnippetThreshold = 0;
-
 template <typename T>
 void DumpGpuTensor(const char* name, const T* tensor, int dim0, int dim1, bool is_gpu_tensor) {
   // Occasionally, user will need dump CPU tensor in CUDA EP.
@@ -47,15 +79,17 @@ void DumpGpuTensor(const char* name, const T* tensor, int dim0, int dim1, bool i
   int num_items = dim0 * dim1;
   auto data = std::make_shared<PinnedHostBuffer<T>>(num_items);
   CUDA_CALL_THROW(cudaDeviceSynchronize());
-  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T), is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
-
+  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T),
+                             is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
 
   if (nullptr != name) {
     std::cout << std::string(name) << std::endl;
   }
 
-  if (kGpuSnippetThreshold > 0 && kGpuSnippetThreshold < static_cast<int64_t>(num_items)) {
-    onnxruntime::utils::PrintCpuTensorSnippet<T>(*data, dim0, dim1, onnxruntime::utils::kDefaultSnippetEdgeItems);
+  int snippet_threshold = DumpTensorConfig::instance().get_snippet_threshold();
+  int snippet_edge_items = DumpTensorConfig::instance().get_snippet_edge_items();
+  if (snippet_threshold > 0 && snippet_threshold < num_items) {
+    onnxruntime::utils::PrintCpuTensorSnippet<T>(*data, dim0, dim1, snippet_edge_items);
   } else {
     onnxruntime::utils::PrintCpuTensorFull<T>(*data, dim0, dim1);
   }
@@ -66,14 +100,17 @@ void DumpGpuTensor(const char* name, const T* tensor, int dim0, int dim1, int di
   int num_items = dim0 * dim1 * dim2;
   auto data = std::make_shared<PinnedHostBuffer<T>>(num_items);
   CUDA_CALL_THROW(cudaDeviceSynchronize());
-  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T), is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
+  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T),
+                             is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
 
   if (nullptr != name) {
     std::cout << std::string(name) << std::endl;
   }
 
-  if (kGpuSnippetThreshold > 0 && kGpuSnippetThreshold < static_cast<int64_t>(num_items)) {
-    onnxruntime::utils::PrintCpuTensorSnippet<T>(*data, dim0, dim1, dim2, onnxruntime::utils::kDefaultSnippetEdgeItems);
+  int snippet_threshold = DumpTensorConfig::instance().get_snippet_threshold();
+  int snippet_edge_items = DumpTensorConfig::instance().get_snippet_edge_items();
+  if (snippet_threshold > 0 && snippet_threshold < num_items) {
+    onnxruntime::utils::PrintCpuTensorSnippet<T>(*data, dim0, dim1, dim2, snippet_edge_items);
   } else {
     onnxruntime::utils::PrintCpuTensorFull<T>(*data, dim0, dim1, dim2);
   }
@@ -84,19 +121,23 @@ void DumpGpuTensor(const char* name, const T* tensor, int dim0, int dim1, int di
   int num_items = dim0 * dim1 * dim2 * dim3;
   auto data = std::make_shared<PinnedHostBuffer<T>>(num_items);
   CUDA_CALL_THROW(cudaDeviceSynchronize());
-  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T), is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
+  CUDA_CALL_THROW(cudaMemcpy(*data, tensor, num_items * sizeof(T),
+                             is_gpu_tensor ? cudaMemcpyDeviceToHost : cudaMemcpyHostToHost));
 
   if (nullptr != name) {
     std::cout << std::string(name) << std::endl;
   }
 
-  if (kGpuSnippetThreshold > 0 && kGpuSnippetThreshold < static_cast<int64_t>(num_items)) {
-    for(int i = 0; i < dim0; i++) {
+  int snippet_threshold = DumpTensorConfig::instance().get_snippet_threshold();
+  int snippet_edge_items = DumpTensorConfig::instance().get_snippet_edge_items();
+  if (snippet_threshold > 0 && snippet_threshold < num_items) {
+    for (int i = 0; i < dim0; i++) {
       std::cout << "[" << i << "]:" << std::endl;
-      onnxruntime::utils::PrintCpuTensorSnippet<T>((*data) + i * dim1 * dim2 * dim3, dim1, dim2, dim3, onnxruntime::utils::kDefaultSnippetEdgeItems);
+      onnxruntime::utils::PrintCpuTensorSnippet<T>((*data) + i * dim1 * dim2 * dim3, dim1, dim2, dim3,
+                                                   snippet_edge_items);
     }
   } else {
-    for(int i = 0; i < dim0; i++) {
+    for (int i = 0; i < dim0; i++) {
       std::cout << "[" << i << "]:" << std::endl;
       onnxruntime::utils::PrintCpuTensorFull<T>((*data) + i * dim1 * dim2 * dim3, dim1, dim2, dim3);
     }

onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h

@@ -3,10 +3,9 @@
 
 #pragma once
 
-#include <string>
 #include "core/framework/tensorprotoutils.h"
 #include "core/framework/ort_value.h"
-#include "contrib_ops/cpu/transformers/generation_shared.h"
+#include "contrib_ops/cpu/utils/console_dumper.h"
 
 namespace onnxruntime {
 namespace contrib {

onnxruntime/core/framework/print_tensor_utils.h

@@ -4,6 +4,7 @@
 
 #include <functional>
 #include <iomanip>
+#include <iostream>
 
 namespace onnxruntime {
 namespace utils {
@@ -38,19 +39,23 @@ inline void PrintValue(const T& value) {
 }
 
 // Explicit specialization
-template <> inline void PrintValue(const MLFloat16& value) {
+template <>
+inline void PrintValue(const MLFloat16& value) {
   std::cout << std::setprecision(8) << value.ToFloat();
 }
 
-template <> inline void PrintValue(const BFloat16& value) {
+template <>
+inline void PrintValue(const BFloat16& value) {
   std::cout << std::setprecision(8) << value.ToFloat();
 }
 
-template <> inline void PrintValue(const uint8_t& value) {
+template <>
+inline void PrintValue(const uint8_t& value) {
   std::cout << static_cast<uint32_t>(value);
 }
 
-template <> inline void PrintValue(const int8_t& value) {
+template <>
+inline void PrintValue(const int8_t& value) {
   std::cout << static_cast<int32_t>(value);
 }