saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add helper functions to dump 4d tensors in CPU for debugging (#21043)

Browse source 
Add some helper functions to dump 4D tensors to help debugging.

Example to use it:
(1) Change DUMP_TENSOR_LEVEL from 0 to 2 in
contrib_ops/cpu/utils/debug_macros.h to enable dumping. Without
enabling, the dumping code will not be built into ORT binary.
(2) Add a few lines to dump tensors like
```
DUMP_CPU_TENSOR_INIT();
DUMP_CPU_TENSOR("tensor name", tensor_data, dim0, dim1, dim2, dim3);
```

Changes:
- [x] Add functions to dump 4D int32/int64/float/half tensors in CPU
- [x] Add functions to dump 4D int32/int64 tensors in CUDA
- [x] Change namespace (remove .transformers from namespace, and move
files to utils directory)
This commit is contained in:
Tianlei Wu 2024-06-14 17:32:27 -07:00 committed by GitHub
parent 8f0e896c95
commit f25cf19375
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 185 additions and 92 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h
View file

@ -9,6 +9,7 @@
#include "core/common/common.h"
#include "core/common/safeint.h"
#include "core/framework/op_kernel.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
namespace onnxruntime {
namespace contrib {
@ -204,12 +205,18 @@ class AttentionCPUBase : public AttentionBase {
});
}
DUMP_CPU_TENSOR_INIT();
DUMP_CPU_TENSOR("Q", Q, batch_size, num_heads_, sequence_length, head_size);
DUMP_CPU_TENSOR("QK (scaled)", attention_probs, batch_size, num_heads_, sequence_length, total_sequence_length);
// attention_probs(B, N, S, T) = Softmax(attention_probs)
{
const int N = batch_size * num_heads_ * sequence_length;
const int D = total_sequence_length;
ComputeAttentionSoftmaxInplace(attention_probs, N, D, tp);
}
DUMP_CPU_TENSOR("Softmax(QK)", attention_probs, batch_size, num_heads_, sequence_length, total_sequence_length);
}
template <typename T>

2
onnxruntime/contrib_ops/cpu/transformers/beam_search.cc
View file

@ -30,7 +30,7 @@
#include "contrib_ops/cpu/transformers/beam_search.h"
#include "contrib_ops/cpu/transformers/logits_processor.h"
#include "contrib_ops/cpu/transformers/sequences.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/beam_search_scorer.h"
#include "contrib_ops/cpu/transformers/beam_search_impl_gpt.h"
#include "contrib_ops/cpu/transformers/beam_search_impl_t5.h"

14
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.cc
View file

@ -288,7 +288,7 @@ Status ProcessLogits(const OrtValue& logits, //
const transformers::IGenerationParameters* parameters, // parameters
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper) { // tensor dumper
const IConsoleDumper* dumper) { // tensor dumper
#ifndef DEBUG_GENERATION
ORT_UNUSED_PARAMETER(dumper);
#endif
@ -450,7 +450,7 @@ Status GreedySearchProcessLogits(
bool do_sampling, // whether to do sampling
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper) { // tensor dumper
const IConsoleDumper* dumper) { // tensor dumper
int batch_size = parameters->batch_size;
int vocab_size = parameters->vocab_size;
@ -810,7 +810,7 @@ Status UpdateDecoderFeeds(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper) {
const IConsoleDumper* dumper) {
ORT_UNUSED_PARAMETER(stream);
ORT_UNUSED_PARAMETER(beam_indices_gpu);
ORT_UNUSED_PARAMETER(input_sequence_len);
@ -952,7 +952,7 @@ template Status ProcessLogits<float>(
const transformers::IGenerationParameters* parameters,
int step,
Stream* stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status GreedySearchProcessLogits<float>(
const OrtValue& logits,
@ -966,7 +966,7 @@ template Status GreedySearchProcessLogits<float>(
bool do_sampling,
int step,
Stream* ort_stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status DeviceCopy<float>(
gsl::span<float> target,
@ -1017,7 +1017,7 @@ template Status UpdateDecoderFeeds<float>(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status UpdateDecoderFeeds<MLFloat16>(
AllocatorPtr allocator,
@ -1037,7 +1037,7 @@ template Status UpdateDecoderFeeds<MLFloat16>(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template void ExpandInputs<int32_t>(const OrtValue& input, int num_beams, AllocatorPtr allocator, OrtValue& expanded);

12
onnxruntime/contrib_ops/cpu/transformers/generation_device_helper.h
View file

@ -106,7 +106,7 @@ using ProcessLogitsFunc = std::function<Status(
const transformers::IGenerationParameters* parameters, // parameters
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper)>; // tensor dumper
const IConsoleDumper* dumper)>; // tensor dumper
template <typename T>
using GreedySearchProcessLogitsFunc = std::function<Status(
@ -121,7 +121,7 @@ using GreedySearchProcessLogitsFunc = std::function<Status(
bool do_sampling, // whether to do sampling
int step, // iteration counter
Stream* ort_stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper)>; // tensor dumper
const IConsoleDumper* dumper)>; // tensor dumper
template <typename T>
using DeviceCopyFunc = std::function<Status(
@ -182,7 +182,7 @@ using UpdateDecoderFeedsFunc = std::function<Status(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper)>;
const IConsoleDumper* dumper)>;
//------------------------------------------------
// Modified functions for Whisper Model
@ -277,7 +277,7 @@ Status ProcessLogits(const OrtValue& logits, //
const transformers::IGenerationParameters* parameters, // parameters
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper); // tensor dumper
const IConsoleDumper* dumper); // tensor dumper
template <typename T>
Status GreedySearchProcessLogits(const OrtValue& logits, // logits output of subgraph
@ -291,7 +291,7 @@ Status GreedySearchProcessLogits(const OrtValue& logits,
bool do_sampling, // whether to do sampling
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper); // tensor dumper
const IConsoleDumper* dumper); // tensor dumper
template <typename T>
Status DeviceCopy(gsl::span<T> target,
@ -367,7 +367,7 @@ Status UpdateDecoderFeeds(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
// ---------------------------------------------------------------
// Functions for encoder-decoder model with float input like Whisper

3
onnxruntime/contrib_ops/cpu/transformers/generation_shared.h
View file

@ -7,7 +7,8 @@
#include <random>
#include "core/common/gsl.h"
#include "core/framework/allocator.h"
#include "contrib_ops/cpu/utils/console_dumper.h"
#include "core/framework/ort_value.h"
#include "contrib_ops/cpu/utils/debug_macros.h"
namespace onnxruntime {

2
onnxruntime/contrib_ops/cpu/transformers/greedy_search.cc
View file

@ -30,7 +30,7 @@
#include "contrib_ops/cpu/transformers/greedy_search.h"
#include "contrib_ops/cpu/transformers/logits_processor.h"
#include "contrib_ops/cpu/transformers/sequences.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/greedy_search_impl_gpt.h"
using namespace ONNX_NAMESPACE;

2
onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc
View file

@ -8,7 +8,7 @@
#include "core/common/span_utils.h"
#include "core/providers/cpu/math/softmax_shared.h"
#include "contrib_ops/cpu/transformers/logits_processor.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include <vector>
#include <numeric>
#include <algorithm>

2
onnxruntime/contrib_ops/cpu/transformers/logits_processor.h
View file

@ -6,7 +6,7 @@
#include "core/common/inlined_containers.h"
#include "contrib_ops/cpu/transformers/sequences.h"
#include "contrib_ops/cpu/transformers/beam_search_parameters.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/greedy_search_parameters.h"
#include "contrib_ops/cpu/transformers/sampling_parameters.h"
#include "contrib_ops/cpu/transformers/generation_shared.h"

2
onnxruntime/contrib_ops/cpu/transformers/sampling.cc
View file

@ -14,7 +14,7 @@
#include "contrib_ops/cpu/transformers/sampling.h"
#include "contrib_ops/cpu/transformers/logits_processor.h"
#include "contrib_ops/cpu/transformers/sequences.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/greedy_search_impl_gpt.h"
using namespace ONNX_NAMESPACE;

2
onnxruntime/contrib_ops/cpu/transformers/sampling_cpu_helper.h
View file

@ -61,7 +61,7 @@ Status Sample(AllocatorPtr& allocator,
transformers::ISamplingState<T>* sampling_state,
transformers::IGreedySearchState<T>* greedy_state,
const transformers::IGenerationParameters* parameters,
const transformers::IConsoleDumper* dumper) {
const IConsoleDumper* dumper) {
ORT_UNUSED_PARAMETER(dumper);
gsl::span<T>& sorted_scores = sampling_state->sorted_scores;

2
onnxruntime/contrib_ops/cpu/transformers/subgraph_base.cc
View file

@ -9,7 +9,7 @@
#include "core/providers/cpu/tensor/utils.h"
#include "core/common/gsl.h"
#include "contrib_ops/cpu/transformers/subgraph_base.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
namespace onnxruntime {
namespace contrib {

2
onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc
View file

@ -8,7 +8,7 @@
#include "core/providers/cpu/tensor/utils.h"
#include "core/common/gsl.h"
#include "contrib_ops/cpu/transformers/subgraph_gpt.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
namespace onnxruntime {
namespace contrib {

2
onnxruntime/contrib_ops/cpu/transformers/subgraph_t5_decoder.cc
View file

@ -8,7 +8,7 @@
#include "core/providers/cpu/tensor/utils.h"
#include "core/common/gsl.h"
#include "contrib_ops/cpu/transformers/subgraph_t5_decoder.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/generation_device_helper.h"
#include "contrib_ops/cpu/transformers/sequences.h"

2
onnxruntime/contrib_ops/cpu/transformers/subgraph_whisper_decoder.cc
View file

@ -9,7 +9,7 @@
#include "core/common/gsl.h"
#include "contrib_ops/cpu/transformers/subgraph_t5_decoder.h"
#include "contrib_ops/cpu/transformers/subgraph_whisper_decoder.h"
#include "contrib_ops/cpu/transformers/dump_tensor.h"
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "contrib_ops/cpu/transformers/generation_device_helper.h"
#include "contrib_ops/cpu/transformers/sequences.h"

13
onnxruntime/contrib_ops/cpu/utils/console_dumper.h
View file

@ -4,12 +4,11 @@
#pragma once
#include <string>
#include "core/framework/ort_value.h"
// #define DEBUG_GENERATION 1 // uncomment it for debugging generation (like beam search etc)
#include "core/framework/float16.h"
#include "contrib_ops/cpu/utils/debug_macros.h"
namespace onnxruntime {
namespace contrib {
namespace transformers {
class IConsoleDumper {
public:
@ -22,10 +21,17 @@ class IConsoleDumper {
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 float* tensor, int dim0, int dim1, int dim2, int dim3) const = 0;
virtual void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const = 0;
virtual void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2, int dim3) const = 0;
virtual void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2, int dim3) 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;
@ -35,6 +41,5 @@ class IConsoleDumper {
bool is_enabled_;
};
} // namespace transformers
} // namespace contrib
} // namespace onnxruntime

41
onnxruntime/contrib_ops/cpu/utils/debug_macros.h Normal file
View file

@ -0,0 +1,41 @@
#pragma once
// #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
#define DUMP_CPU_TENSOR_LEVEL DUMP_TENSOR_LEVEL
// For CPU tensor dumping.
#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__)
#else
#define DUMP_CPU_TENSOR_INIT()
#define DUMP_CPU_TENSOR(...)
#endif
#if DUMP_CPU_TENSOR_LEVEL > 1
#define DUMP_CPU_TENSOR_D(...) cpu_dumper.Print(__VA_ARGS__)
#else
#define DUMP_CPU_TENSOR_D(...)
#endif
// For GPU tensor dumping.
#if DUMP_TENSOR_LEVEL > 0
#define DUMP_TENSOR_INIT() onnxruntime::contrib::cuda::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

43
onnxruntime/contrib_ops/cpu/transformers/dump_tensor.cc → onnxruntime/contrib_ops/cpu/utils/dump_tensor.cc
View file

@ -1,15 +1,15 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "dump_tensor.h"
#include <iomanip>
#include "contrib_ops/cpu/utils/dump_tensor.h"
#include "core/framework/print_tensor_utils.h"
#include "contrib_ops/cpu/utils/debug_macros.h"
namespace onnxruntime {
namespace contrib {
namespace transformers {
#ifdef DEBUG_GENERATION
#if DUMP_CPU_TENSOR_LEVEL > 0
template <typename T>
void DumpCpuTensor(const char* name, const T* tensor, int dim0, int dim1) {
@ -147,6 +147,30 @@ void CpuTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int
DumpCpuTensor<int32_t>(name, tensor, dim0, dim1, dim2);
}
void CpuTensorConsoleDumper::Print(const char* name, const float* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (!is_enabled_)
return;
DumpCpuTensor<float>(name, tensor, dim0 * dim1, dim2, dim3);
}
void CpuTensorConsoleDumper::Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (!is_enabled_)
return;
DumpCpuTensor<MLFloat16>(name, tensor, dim0 * dim1, dim2, dim3);
}
void CpuTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (!is_enabled_)
return;
DumpCpuTensor<int64_t>(name, tensor, dim0 * dim1, dim2, dim3);
}
void CpuTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (!is_enabled_)
return;
DumpCpuTensor<int32_t>(name, tensor, dim0 * dim1, dim2, dim3);
}
void CpuTensorConsoleDumper::Print(const char* name, const Tensor& tensor) const {
if (!is_enabled_)
return;
@ -207,6 +231,18 @@ void CpuTensorConsoleDumper::Print(const char*, const int64_t*, int, int, int) c
void CpuTensorConsoleDumper::Print(const char*, const int32_t*, int, int, int) const {
}
void CpuTensorConsoleDumper::Print(const char*, const float*, int, int, int, int) const {
}
void CpuTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int, int, int) const {
}
void CpuTensorConsoleDumper::Print(const char*, const int64_t*, int, int, int, int) const {
}
void CpuTensorConsoleDumper::Print(const char*, const int32_t*, int, int, int, int) const {
}
void CpuTensorConsoleDumper::Print(const char*, const Tensor&) const {
}
@ -221,6 +257,5 @@ void CpuTensorConsoleDumper::Print(const char*, const std::string&, bool) const
#endif
} // namespace transformers
} // namespace contrib
} // namespace onnxruntime

13
onnxruntime/contrib_ops/cpu/transformers/dump_tensor.h → onnxruntime/contrib_ops/cpu/utils/dump_tensor.h
View file

@ -3,12 +3,11 @@
#pragma once
#include <string>
#include "core/framework/tensorprotoutils.h"
#include "contrib_ops/cpu/transformers/generation_shared.h"
#include "core/framework/ort_value.h"
#include "contrib_ops/cpu/utils/console_dumper.h"
namespace onnxruntime {
namespace contrib {
namespace transformers {
class CpuTensorConsoleDumper : public IConsoleDumper {
public:
@ -19,16 +18,22 @@ class CpuTensorConsoleDumper : public IConsoleDumper {
void Print(const char* name, const size_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const int32_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const Tensor& value) const override;
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;
};
} // namespace transformers
} // namespace contrib
} // namespace onnxruntime

2
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu
View file

@ -35,7 +35,7 @@ limitations under the License.
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h"
#include "contrib_ops/cpu/bert/attention_base.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
#include "contrib_ops/cuda/bert/attention_impl.h"

2
onnxruntime/contrib_ops/cuda/bert/attention_prepare_qkv.cu
View file

@ -4,7 +4,7 @@
#include "contrib_ops/cuda/bert/attention_impl.h"
#include "core/providers/cuda/cu_inc/common.cuh"
#include "contrib_ops/cuda/bert/add_bias_transpose.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
using namespace onnxruntime::cuda;

2
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu
View file

@ -36,7 +36,7 @@ limitations under the License.
#include "contrib_ops/cuda/bert/add_bias_transpose.h"
#include "contrib_ops/cpu/bert/attention_base.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/flash_attention/flash_api.h"
#include "contrib_ops/cuda/bert/group_query_attention_impl.h"

2
onnxruntime/contrib_ops/cuda/bert/packed_attention_impl.cu
View file

@ -13,7 +13,7 @@
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/rotary_embedding_util.h"
#include "contrib_ops/cuda/bert/flash_attention/flash_api.h"

2
onnxruntime/contrib_ops/cuda/bert/packed_multihead_attention_impl.cu
View file

@ -14,7 +14,7 @@
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/mha_runner.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cross_attention/fmha_cross_attention.h"
#include "contrib_ops/cuda/bert/bert_padding.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h"
#include "contrib_ops/cuda/bert/rotary_embedding_util.h"
#include "contrib_ops/cuda/bert/flash_attention/flash_api.h"

2
onnxruntime/contrib_ops/cuda/diffusion/group_norm_impl.cu
View file

@ -26,7 +26,7 @@
#include "core/providers/cuda/cuda_common.h"
#include "core/providers/cuda/cu_inc/common.cuh"
#include "contrib_ops/cuda/diffusion/group_norm_impl.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/diffusion/group_norm_common_base.h"
#include "contrib_ops/cuda/diffusion/group_norm_impl_kernel.cuh"

2
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu
View file

@ -2,7 +2,7 @@
// Licensed under the MIT License.
#include "contrib_ops/cuda/sparse/sparse_attention_impl.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/rotary_embedding_impl.h"
#include "contrib_ops/cuda/bert/group_query_attention_impl.h"
#include "contrib_ops/cpu/bert/attention_common.h"

2
onnxruntime/contrib_ops/cuda/sparse/sparse_attention_v1/sparse_attention_common.h
View file

@ -3,7 +3,7 @@
#pragma once
#include "core/providers/cuda/cuda_common.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cuda/bert/tensorrt_fused_multihead_attention/cudaDriverWrapper.h"
#define CU_CHECK(expr, driver) cuErrCheck(expr, *driver)

4
onnxruntime/contrib_ops/cuda/transformers/beam_search.cc
View file

@ -5,7 +5,7 @@
#include "core/providers/cuda/cuda_execution_provider.h"
#include "contrib_ops/cuda/transformers/beam_search.h"
#include "contrib_ops/cuda/transformers/generation_device_helper.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
namespace onnxruntime {
namespace contrib {
@ -56,7 +56,7 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<MLFloat16>()}),
WhisperBeamSearch);
transformers::CudaTensorConsoleDumper g_cuda_dumper;
CudaTensorConsoleDumper g_cuda_dumper;
BeamSearch::BeamSearch(const OpKernelInfo& info)
: onnxruntime::contrib::transformers::BeamSearch(info) {

20
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.cc
View file

@ -12,7 +12,7 @@
#include "contrib_ops/cuda/bert/transformer_cuda_common.h"
#include <cuda_runtime.h>
#include "contrib_ops/cuda/transformers/generation_cuda_impl.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "contrib_ops/cpu/transformers/logits_processor.h"
#include "contrib_ops/cpu/transformers/generation_shared.h"
#include "contrib_ops/cpu/transformers/subgraph_t5_decoder.h"
@ -332,7 +332,7 @@ Status ProcessLogits(const OrtValue& logits, //
const transformers::IGenerationParameters* parameters, // parameters
int step, // iteration counter
Stream* ort_stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper) { // tensor dumper
const IConsoleDumper* dumper) { // tensor dumper
#ifdef ENABLE_NVTX_PROFILE
profile::NvtxNestedRangeCreator processLogitsRange("ProcessLogits", profile::Color::Red);
@ -824,7 +824,7 @@ Status GreedySearchProcessLogits(
bool do_sampling, // whether to do sampling
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper) { // tensor dumper
const IConsoleDumper* dumper) { // tensor dumper
#ifdef ENABLE_NVTX_PROFILE
profile::NvtxNestedRangeCreator processLogitsRange("ProcessLogits", profile::Color::Red);
@ -1242,7 +1242,7 @@ Status UpdateDecoderFeeds(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper) {
const IConsoleDumper* dumper) {
// last_outputs: logits, present_key_self_0, present_value_self_0, ...
// next_inputs: input_ids,
// encoder_attention_mask, encoder_hidden_states,
@ -1448,7 +1448,7 @@ template Status ProcessLogits<float>(
const transformers::IGenerationParameters* parameters,
int step,
Stream* ort_stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status GreedySearchProcessLogits<float>(
const OrtValue& logits,
@ -1462,7 +1462,7 @@ template Status GreedySearchProcessLogits<float>(
bool do_sampling,
int step,
Stream* ort_stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status DeviceCopy<float>(
gsl::span<float> target,
@ -1519,7 +1519,7 @@ template Status ProcessLogits<MLFloat16>(
const transformers::IGenerationParameters* parameters,
int step,
Stream* ort_stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status GreedySearchProcessLogits<MLFloat16>(
const OrtValue& logits,
@ -1533,7 +1533,7 @@ template Status GreedySearchProcessLogits<MLFloat16>(
bool do_sampling,
int step,
Stream* ort_stream,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status UpdateGptFeeds<MLFloat16>(
AllocatorPtr allocator,
@ -1572,7 +1572,7 @@ template Status UpdateDecoderFeeds<float>(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status UpdateDecoderFeeds<MLFloat16>(
AllocatorPtr allocator,
@ -1592,7 +1592,7 @@ template Status UpdateDecoderFeeds<MLFloat16>(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template Status ExpandBuffer<int32_t>(
Stream* ort_stream,

6
onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h
View file

@ -72,7 +72,7 @@ Status ProcessLogits(const OrtValue& logits, //
const transformers::IGenerationParameters* parameters, // parameters
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper); // tensor dumper
const IConsoleDumper* dumper); // tensor dumper
template <typename T>
Status GreedySearchProcessLogits(const OrtValue& logits, // logits output of subgraph
@ -86,7 +86,7 @@ Status GreedySearchProcessLogits(const OrtValue& logits,
bool do_sampling, // whether to do sampling
int step, // iteration counter
Stream* stream, // cuda stream (for CUDA only)
const transformers::IConsoleDumper* dumper); // tensor dumper
const IConsoleDumper* dumper); // tensor dumper
template <typename T>
Status DeviceCopy(gsl::span<T> target,
@ -138,7 +138,7 @@ Status UpdateDecoderFeeds(
bool past_present_share_buffer,
bool need_cache_indir,
transformers::Sequences& sequences,
const transformers::IConsoleDumper* dumper);
const IConsoleDumper* dumper);
template <typename T>
Status ExpandBuffer(

4
onnxruntime/contrib_ops/cuda/transformers/greedy_search.cc
View file

@ -5,7 +5,7 @@
#include "core/providers/cuda/cuda_execution_provider.h"
#include "contrib_ops/cuda/transformers/greedy_search.h"
#include "contrib_ops/cuda/transformers/generation_device_helper.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
namespace onnxruntime {
namespace contrib {
@ -27,7 +27,7 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<MLFloat16>()}),
GreedySearch);
transformers::CudaTensorConsoleDumper g_cuda_dumper_greedysearch;
CudaTensorConsoleDumper g_cuda_dumper_greedysearch;
GreedySearch::GreedySearch(const OpKernelInfo& info)
: onnxruntime::contrib::transformers::GreedySearch(info) {

4
onnxruntime/contrib_ops/cuda/transformers/sampling.cc
View file

@ -5,7 +5,7 @@
#include "core/providers/cuda/cuda_execution_provider.h"
#include "contrib_ops/cuda/transformers/sampling.h"
#include "contrib_ops/cuda/transformers/generation_device_helper.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
namespace onnxruntime {
namespace contrib {
@ -28,7 +28,7 @@ ONNX_OPERATOR_KERNEL_EX(
DataTypeImpl::GetTensorType<MLFloat16>()}),
Sampling);
transformers::CudaTensorConsoleDumper g_cuda_dumper_sampling;
CudaTensorConsoleDumper g_cuda_dumper_sampling;
Sampling::Sampling(const OpKernelInfo& info)
: onnxruntime::contrib::transformers::Sampling(info) {

2
onnxruntime/contrib_ops/cuda/transformers/sampling_cuda_helper.h
View file

@ -26,7 +26,7 @@ Status Sample(AllocatorPtr& allocator,
transformers::IGreedySearchState<T>* greedy_state,
const transformers::IGenerationParameters* parameters,
int step,
const transformers::IConsoleDumper* dumper) {
const IConsoleDumper* dumper) {
ORT_UNUSED_PARAMETER(dumper);
typedef typename ToCudaType<T>::MappedType CudaT;

20
onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc → onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.cc
View file

@ -4,13 +4,12 @@
#include <cuda_runtime_api.h>
#include "core/providers/cuda/cuda_common.h"
#include "core/framework/print_tensor_utils.h"
#include "contrib_ops/cuda/transformers/dump_cuda_tensor.h"
#include "contrib_ops/cuda/utils/dump_cuda_tensor.h"
#include "core/platform/env_var_utils.h"
namespace onnxruntime {
namespace contrib {
namespace cuda {
namespace transformers {
#if DUMP_TENSOR_LEVEL > 0
@ -218,6 +217,11 @@ void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (is_enabled_)
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, dim2, dim3, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, true);
@ -228,6 +232,11 @@ void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (is_enabled_)
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, dim2, dim3, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const float* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<float>(name, tensor, dim0, dim1, true);
@ -325,12 +334,18 @@ void CudaTensorConsoleDumper::Print(const char*, const int32_t*, int, int) const
void CudaTensorConsoleDumper::Print(const char*, const int32_t*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int32_t*, int, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int64_t*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int64_t*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int64_t*, int, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int) const {
}
@ -380,7 +395,6 @@ void CudaTensorConsoleDumper::Print(const char*, const std::string&, bool) const
}
#endif
} // namespace transformers
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

26
onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h → onnxruntime/contrib_ops/cuda/utils/dump_cuda_tensor.h
View file

@ -7,28 +7,11 @@
#include "core/framework/ort_value.h"
#include "contrib_ops/cpu/utils/console_dumper.h"
#define DUMP_TENSOR_LEVEL 0 // change it to 1 or 2 if want to enable dumping for code not in generation.
#if DUMP_TENSOR_LEVEL > 0
#define DUMP_TENSOR_INIT() onnxruntime::contrib::cuda::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
namespace onnxruntime {
namespace contrib {
namespace cuda {
namespace transformers {
class CudaTensorConsoleDumper : public onnxruntime::contrib::transformers::IConsoleDumper {
class CudaTensorConsoleDumper : public onnxruntime::contrib::IConsoleDumper {
public:
CudaTensorConsoleDumper() = default;
virtual ~CudaTensorConsoleDumper() {}
@ -37,13 +20,15 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::transformers::ICons
void Print(const char* name, const int32_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const float* tensor, int dim0, int dim1) const override;
void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2, int dim3) const;
void Print(const char* name, const float* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const half* tensor, int dim0, int dim1) const;
void Print(const char* name, const half* tensor, int dim0, int dim1, int dim2) const;
@ -51,7 +36,7 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::transformers::ICons
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const override;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2) const override;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const;
void Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const override;
void Print(const char* name, const BFloat16* tensor, int dim0, int dim1) const;
void Print(const char* name, const BFloat16* tensor, int dim0, int dim1, int dim2) const;
@ -63,7 +48,6 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::transformers::ICons
void Print(const char* name, const std::string& value, bool end_line) const override;
};
} // namespace transformers
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime

9
onnxruntime/core/framework/print_tensor_utils.h
View file

@ -8,7 +8,6 @@
namespace onnxruntime {
namespace utils {
constexpr int64_t kDefaultSnippetEdgeItems = 3;
constexpr int64_t kDefaultSnippetThreshold = 200;
@ -40,12 +39,12 @@ inline void PrintValue(const T& value) {
// Explicit specialization
template <>
inline void PrintValue(const MLFloat16& value) {
inline void PrintValue(const onnxruntime::MLFloat16& value) {
std::cout << std::setprecision(8) << value.ToFloat();
}
template <>
inline void PrintValue(const BFloat16& value) {
inline void PrintValue(const onnxruntime::BFloat16& value) {
std::cout << std::setprecision(8) << value.ToFloat();
}
@ -220,7 +219,9 @@ DEF_PRINT_CPU_TENSOR_FULL_3D_INT4(Int4x2)
DEF_PRINT_CPU_TENSOR_FULL_3D_INT4(UInt4x2)
template <typename T>
void PrintCpuTensor(const Tensor& tensor, int threshold = kDefaultSnippetThreshold, int edge_items = kDefaultSnippetEdgeItems) {
void PrintCpuTensor(const onnxruntime::Tensor& tensor,
int threshold = kDefaultSnippetThreshold,
int edge_items = kDefaultSnippetEdgeItems) {
const auto& shape = tensor.Shape();
auto num_items = shape.Size();
if (num_items == 0) {