diff --git a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h index 630c533c47..4f5873cc9f 100644 --- a/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h +++ b/onnxruntime/contrib_ops/cpu/transformers/generation_shared.h @@ -7,7 +7,7 @@ #include #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 diff --git a/onnxruntime/contrib_ops/cpu/utils/console_dumper.h b/onnxruntime/contrib_ops/cpu/utils/console_dumper.h new file mode 100644 index 0000000000..97e0a6597d --- /dev/null +++ b/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 +#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 diff --git a/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc b/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc index 3046a58040..b31f5d243e 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc +++ b/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(kTensorSnippetThreshold, + onnxruntime::utils::kDefaultSnippetThreshold); + snippet_edge_items = ParseEnvironmentVariableWithDefault(kTensorSnippetEdgeItems, + onnxruntime::utils::kDefaultSnippetEdgeItems); + } + ~DumpTensorConfig() {} +}; + template class PinnedHostBuffer { public: @@ -38,8 +72,6 @@ class PinnedHostBuffer { T* buffer_; }; -constexpr int64_t kGpuSnippetThreshold = 0; - template 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>(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(num_items)) { - onnxruntime::utils::PrintCpuTensorSnippet(*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(*data, dim0, dim1, snippet_edge_items); } else { onnxruntime::utils::PrintCpuTensorFull(*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>(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(num_items)) { - onnxruntime::utils::PrintCpuTensorSnippet(*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(*data, dim0, dim1, dim2, snippet_edge_items); } else { onnxruntime::utils::PrintCpuTensorFull(*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>(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(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((*data) + i * dim1 * dim2 * dim3, dim1, dim2, dim3, onnxruntime::utils::kDefaultSnippetEdgeItems); + onnxruntime::utils::PrintCpuTensorSnippet((*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((*data) + i * dim1 * dim2 * dim3, dim1, dim2, dim3); } diff --git a/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h b/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h index 341c0c134a..264ecd7cfe 100644 --- a/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h +++ b/onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h @@ -3,10 +3,9 @@ #pragma once -#include #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 { diff --git a/onnxruntime/core/framework/print_tensor_utils.h b/onnxruntime/core/framework/print_tensor_utils.h index 2cfbd75bcb..9509ca2646 100644 --- a/onnxruntime/core/framework/print_tensor_utils.h +++ b/onnxruntime/core/framework/print_tensor_utils.h @@ -4,6 +4,7 @@ #include #include +#include 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(value); } -template <> inline void PrintValue(const int8_t& value) { +template <> +inline void PrintValue(const int8_t& value) { std::cout << static_cast(value); }