326 lines
11 KiB
Plaintext
326 lines
11 KiB
Plaintext
// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "paddle/phi/kernels/legacy/gpu/tensor_debug.h"
|
|
|
|
#include <cuda_bf16.h>
|
|
#include <cuda_fp16.h>
|
|
#include <cuda_runtime.h>
|
|
#include <stdint.h>
|
|
|
|
#include "paddle/common/enforce.h"
|
|
#include "paddle/phi/common/bfloat16.h"
|
|
#include "paddle/phi/common/data_type.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
|
|
namespace phi {
|
|
|
|
// Maximum tensor rank supported. Shapes are passed as plain int64_t values
|
|
// in registers to avoid any host/device memory transfers (which are forbidden
|
|
// inside a CUDA Graph capture region).
|
|
static constexpr int kMaxDims = 9;
|
|
|
|
// dtype tag passed as an integer so we need no device-side char* pointer.
|
|
enum class DebugDtype : int {
|
|
FLOAT32 = 0,
|
|
FLOAT64,
|
|
FLOAT16,
|
|
BFLOAT16,
|
|
INT32,
|
|
INT64,
|
|
INT16,
|
|
INT8,
|
|
UINT8,
|
|
BOOL,
|
|
};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helper: dtype name from enum tag (device-side string literal, no pointer)
|
|
// ---------------------------------------------------------------------------
|
|
__device__ static const char* DtypeName(DebugDtype tag) {
|
|
switch (tag) {
|
|
case DebugDtype::FLOAT32:
|
|
return "FLOAT32";
|
|
case DebugDtype::FLOAT64:
|
|
return "FLOAT64";
|
|
case DebugDtype::FLOAT16:
|
|
return "FLOAT16";
|
|
case DebugDtype::BFLOAT16:
|
|
return "BFLOAT16";
|
|
case DebugDtype::INT32:
|
|
return "INT32";
|
|
case DebugDtype::INT64:
|
|
return "INT64";
|
|
case DebugDtype::INT16:
|
|
return "INT16";
|
|
case DebugDtype::INT8:
|
|
return "INT8";
|
|
case DebugDtype::UINT8:
|
|
return "UINT8";
|
|
case DebugDtype::BOOL:
|
|
return "BOOL";
|
|
default:
|
|
return "UNKNOWN";
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Shape is passed as a fixed-size struct so all values live in registers.
|
|
// No device-memory pointer is needed.
|
|
// ---------------------------------------------------------------------------
|
|
struct ShapeArgs {
|
|
int64_t dims[kMaxDims];
|
|
int ndim;
|
|
};
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Device helpers: compute multi-dim indices, count leading/trailing bracket
|
|
// events for shape-aware nested printing.
|
|
//
|
|
// Algorithm (single thread, linear scan):
|
|
// For linear index i:
|
|
// * compute per-dimension indices idx[d] = (i / stride[d]) % dims[d]
|
|
// * "open_dims" = number of rightmost dims whose idx is 0
|
|
// -> print that many '[' before the value
|
|
// * "close_dims" = number of rightmost dims at their last position
|
|
// -> print that many ']' after the value, then
|
|
// newline+indent
|
|
// Special cases: scalar (ndim==0) and 1-D tensors.
|
|
// ---------------------------------------------------------------------------
|
|
|
|
// Compute strides for C-contiguous layout.
|
|
__device__ static void ComputeStrides(const ShapeArgs& shape,
|
|
int64_t* strides) {
|
|
strides[shape.ndim - 1] = 1;
|
|
for (int d = shape.ndim - 2; d >= 0; --d) {
|
|
strides[d] = strides[d + 1] * shape.dims[d + 1];
|
|
}
|
|
}
|
|
|
|
// For element at linear index i, how many trailing dimensions have idx==0
|
|
// (i.e. we are at the start of a new sub-array in those dims).
|
|
__device__ static int CountOpenDims(int64_t i,
|
|
const ShapeArgs& shape,
|
|
const int64_t* strides) {
|
|
int count = 0;
|
|
for (int d = shape.ndim - 1; d >= 0; --d) {
|
|
if ((i / strides[d]) % shape.dims[d] == 0) {
|
|
++count;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
return count;
|
|
}
|
|
|
|
// For element at linear index i, how many trailing dimensions are at their
|
|
// last position (i.e. we just finished a sub-array in those dims).
|
|
__device__ static int CountCloseDims(int64_t i,
|
|
const ShapeArgs& shape,
|
|
const int64_t* strides) {
|
|
int count = 0;
|
|
for (int d = shape.ndim - 1; d >= 0; --d) {
|
|
int64_t idx_d = (i / strides[d]) % shape.dims[d];
|
|
if (idx_d == shape.dims[d] - 1) {
|
|
++count;
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
return count;
|
|
}
|
|
|
|
// Print the prefix (opening brackets / comma / newline+indent) before element
|
|
// i.
|
|
__device__ static void PrintPrefix(int64_t i,
|
|
const ShapeArgs& shape,
|
|
const int64_t* strides,
|
|
int close_prev) {
|
|
if (i == 0) {
|
|
// Opening brackets for all dimensions.
|
|
for (int d = 0; d < shape.ndim; ++d) printf("[");
|
|
return;
|
|
}
|
|
|
|
if (close_prev > 0) {
|
|
// After closing brackets we start a new row: comma, newline, indent.
|
|
printf(",\n");
|
|
// Indent = (ndim - close_prev) spaces so the opening '[' aligns.
|
|
int indent = shape.ndim - close_prev;
|
|
for (int d = 0; d < indent; ++d) printf(" ");
|
|
for (int d = 0; d < close_prev; ++d) printf("[");
|
|
} else {
|
|
printf(", ");
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Typed print helpers (device-side, called from the kernel body)
|
|
// ---------------------------------------------------------------------------
|
|
template <typename T>
|
|
__device__ static void PrintValue(T v) {
|
|
printf("%.6g", static_cast<double>(v));
|
|
}
|
|
|
|
template <>
|
|
__device__ void PrintValue<__half>(__half v) {
|
|
printf("%.6g", static_cast<double>(__half2float(v)));
|
|
}
|
|
|
|
template <>
|
|
__device__ void PrintValue<__nv_bfloat16>(__nv_bfloat16 v) {
|
|
printf("%.6g", static_cast<double>(__bfloat162float(v)));
|
|
}
|
|
|
|
template <>
|
|
__device__ void PrintValue<bool>(bool v) {
|
|
printf("%s", v ? "True" : "False");
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// CUDA kernel: single thread, shape-aware nested printing.
|
|
// All state (shape, strides) lives in registers -- CUDA Graph safe.
|
|
// ---------------------------------------------------------------------------
|
|
template <typename T>
|
|
__global__ void PrintTensorKernel(const T* data,
|
|
int64_t numel,
|
|
ShapeArgs shape,
|
|
DebugDtype dtag) {
|
|
// -- header ----------------------------------------------------------------
|
|
printf("[TensorDebug] dtype : %s\n", DtypeName(dtag));
|
|
|
|
printf("[TensorDebug] shape : [");
|
|
for (int i = 0; i < shape.ndim; ++i) {
|
|
if (i > 0) printf(", ");
|
|
printf("%lld", static_cast<long long>(shape.dims[i])); // NOLINT
|
|
}
|
|
printf("]\n");
|
|
|
|
printf("[TensorDebug] numel : %lld\n",
|
|
static_cast<long long>(numel)); // NOLINT
|
|
|
|
// -- data: shape-aware nested printing -------------------------------------
|
|
printf("[TensorDebug] data :");
|
|
|
|
if (numel == 0) {
|
|
printf(" []\n");
|
|
return;
|
|
}
|
|
|
|
// Scalar (0-D tensor)
|
|
if (shape.ndim == 0) {
|
|
printf(" ");
|
|
PrintValue(data[0]);
|
|
printf("\n");
|
|
return;
|
|
}
|
|
|
|
printf("\n");
|
|
|
|
// Compute strides in registers (no device malloc).
|
|
int64_t strides[kMaxDims] = {};
|
|
ComputeStrides(shape, strides);
|
|
|
|
int close_prev = 0; // closing brackets printed after the previous element
|
|
|
|
for (int64_t i = 0; i < numel; ++i) {
|
|
PrintPrefix(i, shape, strides, close_prev);
|
|
PrintValue(data[i]);
|
|
|
|
close_prev = CountCloseDims(i, shape, strides);
|
|
// Print closing brackets after the value (before the next prefix).
|
|
for (int d = 0; d < close_prev; ++d) printf("]");
|
|
}
|
|
|
|
printf("\n");
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helper: build ShapeArgs on the host and launch the kernel.
|
|
// No dynamic allocation, no memcpy, no stream sync -> CUDA Graph safe.
|
|
// ---------------------------------------------------------------------------
|
|
template <typename T>
|
|
static void LaunchPrint(const T* d_data,
|
|
int64_t numel,
|
|
const int64_t* host_shape,
|
|
int ndim,
|
|
DebugDtype dtag,
|
|
cudaStream_t stream) {
|
|
ShapeArgs shape{};
|
|
shape.ndim = ndim;
|
|
for (int i = 0; i < ndim && i < kMaxDims; ++i) {
|
|
shape.dims[i] = host_shape[i];
|
|
}
|
|
|
|
// Single-thread kernel, all arguments in registers.
|
|
// No malloc/free/memcpy -> safe inside a CUDA Graph capture region.
|
|
PrintTensorKernel<T><<<1, 1, 0, stream>>>(d_data, numel, shape, dtag);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Public entry: DebugPrintGPUTensor
|
|
// ---------------------------------------------------------------------------
|
|
void DebugPrintGPUTensor(const phi::DenseTensor& tensor, cudaStream_t stream) {
|
|
PADDLE_ENFORCE_EQ(tensor.place().GetType() == phi::AllocationType::GPU,
|
|
true,
|
|
phi::errors::InvalidArgument(
|
|
"DebugPrintGPUTensor only supports GPU DenseTensors. "
|
|
"Please call tensor.cuda() first."));
|
|
|
|
PADDLE_ENFORCE_LE(
|
|
tensor.dims().size(),
|
|
kMaxDims,
|
|
phi::errors::InvalidArgument(
|
|
"DebugPrintGPUTensor: tensor rank %d exceeds kMaxDims (%d).",
|
|
tensor.dims().size(),
|
|
kMaxDims));
|
|
|
|
auto dtype = tensor.dtype();
|
|
int64_t numel = tensor.numel();
|
|
int ndim = tensor.dims().size();
|
|
const int64_t* host_shape = tensor.dims().Get();
|
|
|
|
#define DISPATCH(cpp_type, enum_val) \
|
|
case phi::DataType::enum_val: { \
|
|
LaunchPrint<cpp_type>(tensor.data<cpp_type>(), \
|
|
numel, \
|
|
host_shape, \
|
|
ndim, \
|
|
DebugDtype::enum_val, \
|
|
stream); \
|
|
break; \
|
|
}
|
|
|
|
switch (dtype) {
|
|
DISPATCH(float, FLOAT32)
|
|
DISPATCH(double, FLOAT64)
|
|
DISPATCH(phi::dtype::float16, FLOAT16)
|
|
DISPATCH(phi::dtype::bfloat16, BFLOAT16)
|
|
DISPATCH(int32_t, INT32)
|
|
DISPATCH(int64_t, INT64)
|
|
DISPATCH(int16_t, INT16)
|
|
DISPATCH(int8_t, INT8)
|
|
DISPATCH(uint8_t, UINT8)
|
|
DISPATCH(bool, BOOL)
|
|
default:
|
|
PADDLE_THROW(phi::errors::Unimplemented(
|
|
"DebugPrintGPUTensor: unsupported dtype %s",
|
|
phi::DataTypeToString(dtype)));
|
|
}
|
|
|
|
#undef DISPATCH
|
|
}
|
|
|
|
} // namespace phi
|