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

[CUDA] Add functions to dump bfloat16 tensors (#19266)

Browse source 
### Description
GroupQueryAttention add BFloat16 in
https://github.com/microsoft/onnxruntime/pull/19095, and there is build
error when enable dumping. This supports print bfloat16 tensor to
console.
This commit is contained in:
Tianlei Wu 2024-01-25 09:30:15 -08:00 committed by GitHub
parent 5b06505073
commit 2b285cd78a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 85 additions and 50 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

108
onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.cc
View file

@ -203,23 +203,19 @@ 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 char* name, const float* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<float>(name, tensor, dim0, dim1, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<MLFloat16>(name, tensor, dim0, dim1, true);
}
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);
}
void CudaTensorConsoleDumper::Print(const char* name, const half* tensor, int dim0, int dim1) const {
Print(name, reinterpret_cast<const MLFloat16*>(tensor), dim0, dim1);
void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1) const {
@ -227,9 +223,14 @@ void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1) const {
void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, true);
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, dim2, 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);
}
void CudaTensorConsoleDumper::Print(const char* name, const float* tensor, int dim0, int dim1, int dim2) const {
@ -242,6 +243,11 @@ void CudaTensorConsoleDumper::Print(const char* name, const float* tensor, int d
DumpGpuTensor<float>(name, tensor, dim0, dim1, dim2, dim3, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const MLFloat16* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<MLFloat16>(name, tensor, dim0, dim1, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const MLFloat16* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<MLFloat16>(name, tensor, dim0, dim1, dim2, true);
@ -252,6 +258,25 @@ void CudaTensorConsoleDumper::Print(const char* name, const MLFloat16* tensor, i
DumpGpuTensor<MLFloat16>(name, tensor, dim0, dim1, dim2, dim3, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const BFloat16* tensor, int dim0, int dim1) const {
if (is_enabled_)
DumpGpuTensor<BFloat16>(name, tensor, dim0, dim1, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const BFloat16* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<BFloat16>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const BFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const {
if (is_enabled_)
DumpGpuTensor<BFloat16>(name, tensor, dim0, dim1, dim2, dim3, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const half* tensor, int dim0, int dim1) const {
Print(name, reinterpret_cast<const MLFloat16*>(tensor), dim0, dim1);
}
void CudaTensorConsoleDumper::Print(const char* name, const half* tensor, int dim0, int dim1, int dim2) const {
Print(name, reinterpret_cast<const MLFloat16*>(tensor), dim0, dim1, dim2);
}
@ -260,16 +285,6 @@ void CudaTensorConsoleDumper::Print(const char* name, const half* tensor, int di
Print(name, reinterpret_cast<const MLFloat16*>(tensor), dim0, dim1, dim2, dim3);
}
void CudaTensorConsoleDumper::Print(const char* name, const int64_t* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<int64_t>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const int32_t* tensor, int dim0, int dim1, int dim2) const {
if (is_enabled_)
DumpGpuTensor<int32_t>(name, tensor, dim0, dim1, dim2, true);
}
void CudaTensorConsoleDumper::Print(const char* name, const Tensor& tensor) const {
if (is_enabled_)
DumpGpuTensor(name, tensor);
@ -301,45 +316,54 @@ void CudaTensorConsoleDumper::Print(const char* name, const std::string& value,
}
#else
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const size_t*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const half*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int64_t*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const int32_t*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int, int) const {
void CudaTensorConsoleDumper::Print(const char*, const int32_t*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const half*, 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 int32_t*, int, int, int) const {
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const float*, int, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const MLFloat16*, int, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const BFloat16*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const BFloat16*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const BFloat16*, int, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const half*, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const half*, int, int, int) const {
}
void CudaTensorConsoleDumper::Print(const char*, const half*, int, int, int, int) const {
}

27
onnxruntime/contrib_ops/cuda/transformers/dump_cuda_tensor.h
View file

@ -16,20 +16,31 @@ class CudaTensorConsoleDumper : public onnxruntime::contrib::transformers::ICons
public:
CudaTensorConsoleDumper() = default;
virtual ~CudaTensorConsoleDumper() {}
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;
void Print(const char* name, const size_t* tensor, int dim0, int dim1) const override;
void Print(const char* name, const half* tensor, int dim0, int dim1) const;
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 int32_t* tensor, int dim0, int dim1, int dim2) 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 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 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 half* tensor, int dim0, int dim1) const;
void Print(const char* name, const half* tensor, int dim0, int dim1, int dim2) const;
void Print(const char* name, const half* tensor, int dim0, int dim1, int dim2, int dim3) const;
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 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 BFloat16* tensor, int dim0, int dim1) const;
void Print(const char* name, const BFloat16* tensor, int dim0, int dim1, int dim2) const;
void Print(const char* name, const BFloat16* tensor, int dim0, int dim1, int dim2, int dim3) const;
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;