chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,291 @@
|
||||
# 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.
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
import numpy as np
|
||||
|
||||
from paddle.base import core
|
||||
|
||||
|
||||
def _print_tensor_in_gpu(tensor):
|
||||
"""
|
||||
Print a GPU tensor's dtype, shape, and all data values directly from
|
||||
the device using a single-thread CUDA kernel (device-side printf).
|
||||
|
||||
This function is **CUDA Graph safe**: no host/device memory transfer
|
||||
is performed (shape is passed via kernel-argument registers), so it
|
||||
can be called inside a CUDA Graph capture region.
|
||||
|
||||
Args:
|
||||
tensor (paddle.Tensor): A GPU DenseTensor to print. Must already
|
||||
reside on a CUDA device (call ``tensor.cuda()`` first if needed).
|
||||
|
||||
Raises:
|
||||
ValueError: If PaddlePaddle is not compiled with CUDA support.
|
||||
InvalidArgument: If the tensor is not a DenseTensor or not on GPU.
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> # doctest: +REQUIRES(env:GPU)
|
||||
>>> import paddle
|
||||
>>> paddle.device.set_device('gpu')
|
||||
>>> x = paddle.to_tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
>>> paddle.utils.gpu_utils._print_tensor_in_gpu(x)
|
||||
|
||||
"""
|
||||
if not core.is_compiled_with_cuda():
|
||||
raise ValueError(
|
||||
"paddle.utils._print_tensor_in_gpu is not supported in "
|
||||
"CPU-only PaddlePaddle. Please reinstall PaddlePaddle with GPU "
|
||||
"support to call this API."
|
||||
)
|
||||
core.eager._print_tensor_in_gpu(tensor)
|
||||
|
||||
|
||||
# Mapping from the dtype name printed by DebugPrintGPUTensor to a numpy dtype.
|
||||
_DTYPE_STR_TO_NUMPY = {
|
||||
'FLOAT32': np.float32,
|
||||
'FLOAT64': np.float64,
|
||||
'FLOAT16': np.float16,
|
||||
'BFLOAT16': np.float32, # NumPy has no bfloat16; promote to float32
|
||||
'INT32': np.int32,
|
||||
'INT64': np.int64,
|
||||
'INT16': np.int16,
|
||||
'INT8': np.int8,
|
||||
'UINT8': np.uint8,
|
||||
'BOOL': np.bool_,
|
||||
}
|
||||
|
||||
# Mapping from the dtype name to the paddle dtype string.
|
||||
_DTYPE_STR_TO_PADDLE = {
|
||||
'FLOAT32': 'float32',
|
||||
'FLOAT64': 'float64',
|
||||
'FLOAT16': 'float16',
|
||||
'BFLOAT16': 'bfloat16',
|
||||
'INT32': 'int32',
|
||||
'INT64': 'int64',
|
||||
'INT16': 'int16',
|
||||
'INT8': 'int8',
|
||||
'UINT8': 'uint8',
|
||||
'BOOL': 'bool',
|
||||
}
|
||||
|
||||
|
||||
def _parse_tensor_from_gpu_print(text):
|
||||
"""
|
||||
Reconstruct a ``paddle.Tensor`` from the text output produced by
|
||||
:func:`paddle.utils.gpu_utils._print_tensor_in_gpu`.
|
||||
|
||||
The expected input format is the output written to stdout by the
|
||||
``DebugPrintGPUTensor`` CUDA kernel::
|
||||
|
||||
[TensorDebug] dtype : FLOAT32
|
||||
[TensorDebug] shape : [2, 3]
|
||||
[TensorDebug] numel : 6
|
||||
[TensorDebug] data :
|
||||
[[1, 2, 3],
|
||||
[4, 5, 6]]
|
||||
|
||||
Special cases handled:
|
||||
|
||||
* **Scalar (0-D tensor)** – the data line is ``[TensorDebug] data : <value>``
|
||||
(no newline before the value).
|
||||
* **Empty tensor (numel == 0)** – the data line ends with ``[]``.
|
||||
|
||||
Args:
|
||||
text (str): The captured stdout string from a call to
|
||||
``paddle.utils.gpu_utils._print_tensor_in_gpu``.
|
||||
|
||||
Returns:
|
||||
paddle.Tensor: A GPU tensor with the dtype and values recovered from
|
||||
*text*. Call ``.cpu()`` on the result if you need a CPU tensor.
|
||||
|
||||
Raises:
|
||||
ValueError: If *text* cannot be parsed (missing header lines, unknown
|
||||
dtype, mismatched numel, …).
|
||||
|
||||
Examples:
|
||||
.. code-block:: pycon
|
||||
|
||||
>>> import paddle
|
||||
>>> import os, sys, tempfile
|
||||
>>> # Capture the C-level stdout written by the CUDA printf kernel.
|
||||
>>> def capture(tensor):
|
||||
... with tempfile.NamedTemporaryFile(mode='w+', suffix='.txt', delete=False) as f:
|
||||
... path = f.name
|
||||
... sys.stdout.flush()
|
||||
... old = os.dup(1)
|
||||
... fd = os.open(path, os.O_WRONLY | os.O_TRUNC)
|
||||
... os.dup2(fd, 1)
|
||||
... os.close(fd)
|
||||
... paddle.utils.gpu_utils._print_tensor_in_gpu(tensor)
|
||||
... paddle.device.synchronize()
|
||||
... sys.stdout.flush()
|
||||
... os.dup2(old, 1)
|
||||
... os.close(old)
|
||||
... text = open(path).read()
|
||||
... os.remove(path)
|
||||
... return text
|
||||
>>> x = paddle.to_tensor([[1.0, 2.0], [3.0, 4.0]])
|
||||
>>> text = capture(x)
|
||||
>>> y = paddle.utils.gpu_utils._parse_tensor_from_gpu_print(text)
|
||||
>>> print(y)
|
||||
"""
|
||||
import paddle # local import to avoid circular imports at module load time
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 1. Extract header fields #
|
||||
# ------------------------------------------------------------------ #
|
||||
dtype_m = re.search(r'\[TensorDebug\] dtype\s*:\s*(\w+)', text)
|
||||
shape_m = re.search(r'\[TensorDebug\] shape\s*:\s*\[([^\]]*)\]', text)
|
||||
numel_m = re.search(r'\[TensorDebug\] numel\s*:\s*(\d+)', text)
|
||||
data_m = re.search(r'\[TensorDebug\] data\s*:(.*)', text, re.DOTALL)
|
||||
|
||||
if not dtype_m:
|
||||
raise ValueError(
|
||||
"_parse_tensor_from_gpu_print: could not find "
|
||||
"'[TensorDebug] dtype' line in the provided text."
|
||||
)
|
||||
if not shape_m:
|
||||
raise ValueError(
|
||||
"_parse_tensor_from_gpu_print: could not find "
|
||||
"'[TensorDebug] shape' line in the provided text."
|
||||
)
|
||||
if not numel_m:
|
||||
raise ValueError(
|
||||
"_parse_tensor_from_gpu_print: could not find "
|
||||
"'[TensorDebug] numel' line in the provided text."
|
||||
)
|
||||
if not data_m:
|
||||
raise ValueError(
|
||||
"_parse_tensor_from_gpu_print: could not find "
|
||||
"'[TensorDebug] data' line in the provided text."
|
||||
)
|
||||
|
||||
dtype_str = dtype_m.group(1).strip().upper()
|
||||
shape_raw = shape_m.group(1).strip()
|
||||
numel = int(numel_m.group(1).strip())
|
||||
data_raw = data_m.group(1).strip()
|
||||
|
||||
if dtype_str not in _DTYPE_STR_TO_NUMPY:
|
||||
raise ValueError(
|
||||
f"_parse_tensor_from_gpu_print: unknown dtype '{dtype_str}'. "
|
||||
f"Supported: {list(_DTYPE_STR_TO_NUMPY.keys())}"
|
||||
)
|
||||
|
||||
np_dtype = _DTYPE_STR_TO_NUMPY[dtype_str]
|
||||
paddle_dtype = _DTYPE_STR_TO_PADDLE[dtype_str]
|
||||
|
||||
# Parse shape: "" means scalar (0-D), otherwise comma-separated ints.
|
||||
if shape_raw == '':
|
||||
shape = []
|
||||
else:
|
||||
shape = [int(s.strip()) for s in shape_raw.split(',') if s.strip()]
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# 2. Parse data section #
|
||||
# ------------------------------------------------------------------ #
|
||||
gpu_place = paddle.CUDAPlace(0)
|
||||
|
||||
if numel == 0:
|
||||
# Empty tensor – no data values to parse.
|
||||
arr = np.empty(shape, dtype=np_dtype)
|
||||
t = paddle.to_tensor(arr, dtype=paddle_dtype, place=gpu_place)
|
||||
return t
|
||||
|
||||
if len(shape) == 0:
|
||||
# Scalar (0-D tensor): data_raw is the single value directly.
|
||||
value_str = data_raw
|
||||
arr = np.array(_parse_value(value_str, dtype_str), dtype=np_dtype)
|
||||
t = paddle.to_tensor(arr, dtype=paddle_dtype, place=gpu_place)
|
||||
return t
|
||||
|
||||
# General N-D case: strip all brackets, whitespace, and split by commas.
|
||||
# The nested bracket notation produced by PrintTensorKernel uses standard
|
||||
# Python list syntax, so ast.literal_eval can reconstruct the nested list
|
||||
# directly – but only for numeric dtypes. For 'BOOL', PrintValue emits
|
||||
# "True"/"False" which are valid Python literals too.
|
||||
flat_values = _extract_flat_values(data_raw, dtype_str, numel)
|
||||
|
||||
arr = np.array(flat_values, dtype=np_dtype).reshape(shape)
|
||||
t = paddle.to_tensor(arr, dtype=paddle_dtype, place=gpu_place)
|
||||
return t
|
||||
|
||||
|
||||
def _parse_value(value_str, dtype_str):
|
||||
"""Parse a single scalar value string from the debug output."""
|
||||
value_str = value_str.strip()
|
||||
if dtype_str == 'BOOL':
|
||||
if value_str == 'True':
|
||||
return True
|
||||
elif value_str == 'False':
|
||||
return False
|
||||
else:
|
||||
raise ValueError(
|
||||
f"_parse_tensor_from_gpu_print: cannot parse bool value "
|
||||
f"'{value_str}'."
|
||||
)
|
||||
return float(value_str)
|
||||
|
||||
|
||||
def _extract_flat_values(data_raw, dtype_str, numel):
|
||||
"""
|
||||
Extract a flat list of Python scalars from the nested-bracket string
|
||||
produced by PrintTensorKernel.
|
||||
|
||||
Strategy: remove all bracket characters and split the remaining string
|
||||
on commas / whitespace to get individual token strings, then convert
|
||||
each token to the appropriate scalar type.
|
||||
"""
|
||||
# Remove all '[' and ']' characters.
|
||||
flat_str = data_raw.replace('[', '').replace(']', '')
|
||||
|
||||
# Split on commas (values are separated by ", " or ",\n indent").
|
||||
tokens = re.split(r'[,\s]+', flat_str)
|
||||
tokens = [t.strip() for t in tokens if t.strip()]
|
||||
|
||||
if len(tokens) != numel:
|
||||
raise ValueError(
|
||||
f"_parse_tensor_from_gpu_print: expected {numel} values but "
|
||||
f"found {len(tokens)} tokens in the data section. "
|
||||
f"Parsed tokens: {tokens[:20]}{'...' if len(tokens) > 20 else ''}"
|
||||
)
|
||||
|
||||
if dtype_str == 'BOOL':
|
||||
result = []
|
||||
for tok in tokens:
|
||||
if tok == 'True':
|
||||
result.append(True)
|
||||
elif tok == 'False':
|
||||
result.append(False)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"_parse_tensor_from_gpu_print: cannot parse bool token "
|
||||
f"'{tok}'."
|
||||
)
|
||||
return result
|
||||
|
||||
# Integer dtypes must preserve exact values.
|
||||
# GPU printf may emit scientific notation (e.g. "1e+06"), so use
|
||||
# Decimal for lossless string-to-integer conversion.
|
||||
from decimal import Decimal
|
||||
|
||||
_INT_DTYPES = {'INT8', 'INT16', 'INT32', 'INT64', 'UINT8'}
|
||||
if dtype_str in _INT_DTYPES:
|
||||
return [int(Decimal(t)) for t in tokens]
|
||||
|
||||
return [float(t) for t in tokens]
|
||||
Reference in New Issue
Block a user