Skip to content

Commit

Permalink
Add debugging helper to dump string, vector and thread id (#21224)
Browse files Browse the repository at this point in the history
### Description

Add some macro to help print data to console for debugging purpose.

Example usage:
```
int input_id;
vector<int> some_vector;

DUMP_CPU_TENSOR_INIT()
DUMP_CPU_TENSOR("some vector", some_vector);
DUMP_STRING("input_id=", input_id);
```

- To enable dump thread id, set environment variable
`ORT_DUMP_THREAD_ID=0`.
- User can disable dumping by environment variable
`ORT_ENABLE_CPU_DUMP=0`.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
  • Loading branch information
tianleiwu authored Jul 2, 2024
1 parent 7be1d4a commit 7df97f1
Show file tree
Hide file tree
Showing 6 changed files with 75 additions and 3 deletions.
2 changes: 2 additions & 0 deletions onnxruntime/contrib_ops/cpu/utils/console_dumper.h
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@ class IConsoleDumper {
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;

virtual void Print(const std::string& value) const = 0;

protected:
bool is_enabled_;
};
Expand Down
3 changes: 3 additions & 0 deletions onnxruntime/contrib_ops/cpu/utils/debug_macros.h
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#pragma once
#include "core/common/make_string.h"

// #define DEBUG_GENERATION 1 // uncomment it for debugging generation (like beam search etc)

Expand All @@ -14,9 +15,11 @@
#if DUMP_CPU_TENSOR_LEVEL > 0
#define DUMP_CPU_TENSOR_INIT() onnxruntime::contrib::CpuTensorConsoleDumper cpu_dumper
#define DUMP_CPU_TENSOR(...) cpu_dumper.Print(__VA_ARGS__)
#define DUMP_STRING(...) cpu_dumper.Print(::onnxruntime::MakeString(__VA_ARGS__))
#else
#define DUMP_CPU_TENSOR_INIT()
#define DUMP_CPU_TENSOR(...)
#define DUMP_STRING(...)
#endif

#if DUMP_CPU_TENSOR_LEVEL > 1
Expand Down
52 changes: 50 additions & 2 deletions onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc
Original file line number Diff line number Diff line change
@@ -1,18 +1,38 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#include <iomanip>
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include <iomanip>
#include <mutex>
#include <thread>
#include <iostream>
#include "core/framework/print_tensor_utils.h"
#include "contrib_ops/cpu/utils/debug_macros.h"
#include "core/platform/env_var_utils.h"

namespace onnxruntime {
namespace contrib {

#if DUMP_CPU_TENSOR_LEVEL > 0

// Environment variable to enable/disable dumping
constexpr const char* kEnableCpuTensorDumper = "ORT_ENABLE_CPU_DUMP";

// Environment variable to enable/disable dumping thread id
constexpr const char* kDumpThreadId = "ORT_DUMP_THREAD_ID";

// To avoid dumping at the same time from multiple threads
static std::mutex s_mutex;

static bool s_output_thread_id = false;

template <typename T>
void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1) {
std::unique_lock<std::mutex> lock(s_mutex);

if (s_output_thread_id)
std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl;

if (nullptr != name) {
std::cout << std::string(name) << std::endl;
}
Expand All @@ -26,6 +46,11 @@ void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1) {

template <typename T>
void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1, int dim2) {
std::unique_lock<std::mutex> lock(s_mutex);

if (s_output_thread_id)
std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl;

if (nullptr != name) {
std::cout << std::string(name) << std::endl;
}
Expand Down Expand Up @@ -93,6 +118,21 @@ void DumpCpuTensor(const char* name, const Tensor& tensor) {
DumpCpuTensor(nullptr, tensor, static_cast<int>(num_rows), static_cast<int>(row_size));
}

CpuTensorConsoleDumper::CpuTensorConsoleDumper() {
is_enabled_ = ParseEnvironmentVariableWithDefault<int>(kEnableCpuTensorDumper, 1) != 0;
s_output_thread_id = ParseEnvironmentVariableWithDefault<int>(kDumpThreadId, 0) != 0;
}

void CpuTensorConsoleDumper::Print(const std::string& value) const {
if (!is_enabled_)
return;

std::unique_lock<std::mutex> lock(s_mutex);
if (s_output_thread_id)
std::cout << "Thread ID:" << std::this_thread::get_id() << std::endl;
std::cout << value << std::endl;
}

void CpuTensorConsoleDumper::Print(const char* name, const float* tensor, int dim0, int dim1) const {
if (!is_enabled_)
return;
Expand Down Expand Up @@ -185,6 +225,8 @@ void CpuTensorConsoleDumper::Print(const char* name, const OrtValue& value) cons
void CpuTensorConsoleDumper::Print(const char* name, int index, bool end_line) const {
if (!is_enabled_)
return;

std::unique_lock<std::mutex> lock(s_mutex);
std::cout << std::string(name) << "[" << index << "]";

if (end_line) {
Expand All @@ -196,6 +238,7 @@ void CpuTensorConsoleDumper::Print(const char* name, const std::string& value, b
if (!is_enabled_)
return;

std::unique_lock<std::mutex> lock(s_mutex);
std::cout << std::string(name) << "=" << value;

if (end_line) {
Expand All @@ -204,6 +247,12 @@ void CpuTensorConsoleDumper::Print(const char* name, const std::string& value, b
}

#else
CpuTensorConsoleDumper::CpuTensorConsoleDumper() {
}

void CpuTensorConsoleDumper::Print(const std::string&) const {
}

void CpuTensorConsoleDumper::Print(const char*, const float*, int, int) const {
}

Expand Down Expand Up @@ -254,7 +303,6 @@ void CpuTensorConsoleDumper::Print(const char*, int, bool) const {

void CpuTensorConsoleDumper::Print(const char*, const std::string&, bool) const {
}

#endif

} // namespace contrib
Expand Down
11 changes: 10 additions & 1 deletion onnxruntime/contrib_ops/cpu/utils/dump_tensor.h
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// Licensed under the MIT License.

#pragma once
#include <vector>
#include <string>
#include "core/framework/ort_value.h"
#include "contrib_ops/cpu/utils/console_dumper.h"
Expand All @@ -11,7 +12,7 @@ namespace contrib {

class CpuTensorConsoleDumper : public IConsoleDumper {
public:
CpuTensorConsoleDumper() = default;
CpuTensorConsoleDumper();
virtual ~CpuTensorConsoleDumper() {}
void Print(const char* name, const float* tensor, int dim0, int dim1) const override;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const override;
Expand All @@ -33,6 +34,14 @@ class CpuTensorConsoleDumper : public IConsoleDumper {
void Print(const char* name, const OrtValue& value) const override;
void Print(const char* name, int index, bool end_line) const override;
void Print(const char* name, const std::string& value, bool end_line) const override;

void Print(const std::string& value) const override;

// Output a vector with a threshold for max number of elements to output. Default threshold 0 means no limit.
template <typename T>
void Print(const char* name, const std::vector<T>& vec, size_t max_count = 0) const {
this->Print(name, vec.data(), 1, static_cast<int>(std::min(max_count, vec.size())));
}
};

} // namespace contrib
Expand Down
8 changes: 8 additions & 0 deletions onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,10 @@ void DumpGpuTensor(const char* name, const Tensor& tensor) {
DumpGpuTensor(nullptr, tensor, static_cast<int>(num_rows), static_cast<int>(row_size));
}

void CudaTensorConsoleDumper::Print(const std::string& value) const {
std::cout << value << std::endl;
}

void CudaTensorConsoleDumper::Print(const char* name, const size_t* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<size_t>(name, tensor, dim0, dim1, true);
Expand Down Expand Up @@ -325,6 +329,10 @@ void CudaTensorConsoleDumper::Print(const char* name, const std::string& value,
}

#else

void CudaTensorConsoleDumper::Print(const std::string&) const {
}

void CudaTensorConsoleDumper::Print(const char*, const size_t*, int, int) const {
}

Expand Down
2 changes: 2 additions & 0 deletions onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,8 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::IConsoleDumper {
void Print(const char* name, const OrtValue& value) const override;
void Print(const char* name, int index, bool end_line) const override;
void Print(const char* name, const std::string& value, bool end_line) const override;

void Print(const std::string& value) const override;
};

} // namespace cuda
Expand Down

0 comments on commit 7df97f1

Please sign in to comment.