Support printing tensors with custom names or titles ; support delimiter of multi dim tensor ; simplify code

[fzyzcjy]

Close #1

API is suboptimal but I have no time to improve it currently, anyway it seems useable now.

output

demo without cuda graph...

[[[0.0321, 0.9337, 0.2087, 0.6069, 0.7234, ], [0.3619, 0.8504, 0.5817, 0.5964, 0.5141, ], [0.7377, 0.5986, 0.0033, 0.6992, 0.9506, ], [0.7565, 0.6974, 0.4397, 0.2722, 0.0544, ], ], [[0.5629, 0.0225, 0.0032, 0.3086, 0.6959, ], [0.3602, 0.4634, 0.4939, 0.2671, 0.4978, ], [0.2733, 0.8136, 0.5757, 0.0800, 0.9290, ], [0.5862, 0.9824, 0.5481, 0.0201, 0.5500, ], ], [[0.6986, 0.0111, 0.3925, 0.7489, 0.4052, ], [0.7369, 0.8827, 0.6377, 0.9266, 0.3367, ], [0.5987, 0.4838, 0.3857, 0.7841, 0.5796, ], [0.6315, 0.3088, 0.9585, 0.2998, 0.3314, ], ], ]

[[[0.0321, 0.9337, 0.2087, ], [0.3619, 0.8504, 0.5817, ], [0.7377, 0.5986, 0.0033, ], [0.7565, 0.6974, 0.4397, ], ], [[0.5629, 0.0225, 0.0032, ], [0.3602, 0.4634, 0.4939, ], [0.2733, 0.8136, 0.5757, ], [0.5862, 0.9824, 0.5481, ], ], [[0.6986, 0.0111, 0.3925, ], [0.7369, 0.8827, 0.6377, ], [0.5987, 0.4838, 0.3857, ], [0.6315, 0.3088, 0.9585, ], ], ]

[[0, 5, 10, 15, ], [20, 25, 30, 35, ], [40, 45, 50, 55, ], ]

[[[5, 6, 7, 8, ], [10, 11, 12, 13, ], ], ]
demo for all types...
name=for int32, addr=140590144552960, shape=(3), stride=(1)
[3, 4, 5, ]
name=for int64, addr=140590144552960, shape=(3), stride=(1)
[3, 4, 5, ]
name=for float, addr=140590144552960, shape=(3), stride=(1)
[1.5000, 2.5000, 3.5000, ]
demo for all dims...
name=for 1D, addr=140590144552960, shape=(3), stride=(1)
[3, 4, 5, ]
name=for 2D, addr=140590144552960, shape=(2, 3), stride=(3, 1)
[[1, 2, 3, ], [3, 4, 5, ], ]
name=for 3D, addr=140590144552960, shape=(2, 2, 3), stride=(6, 3, 1)
[[[1, 2, 3, ], [3, 4, 5, ], ], [[10, 20, 30, ], [30, 40, 50, ], ], ]
start warmup...
start graph capture...
start replay...

[[-0.8896, -0.4331, ], [-1.0596, 1.5400, ], ]
addr=140590144553472, 
[[1.0000, 1.0000, ], [1.0000, 1.0000, ], ]

[[-1.3223, -1.3223, ], [0.4805, 0.4805, ], ]
name=This is name for part of z1, addr=140590329102848, shape=(2), stride=(2)
[-2.6445, 0.9609, ]
name=This is name for z2, 
[[-5.2891, -5.2891, ], [1.9219, 1.9219, ], ]

[gemini-code-assist]

Summary of Changes

Hello @fzyzcjy, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly improves the tensor debugging utility by introducing the capability to print tensors with custom, user-defined names or titles. This feature provides clearer identification of tensors in debug logs, making it easier to track specific data during complex computations. The changes span both the CUDA C++ kernels, which now accept and print these names, and the Python interface, which handles the encoding and efficient transfer of string data to the GPU.

Highlights

• Feature: Custom Tensor Naming: Implemented the ability to associate and print custom string names or titles alongside tensor data during debugging. This enhances the clarity and traceability of debug output by providing more context for printed tensors.
• C++ Kernel Refactoring: Introduced a new __device__ helper function PrintCommon in csrc/debug_print.cu to centralize the logic for printing tensor addresses and the new custom names. All existing Print*Tensor*D CUDA kernels were updated to accept an additional name_ptr argument and utilize this new common helper, reducing code duplication.
• Python API and Memory Management: Refactored the Python debug_print.print_tensor function to accept an optional name argument. This involved creating _Buffer and _DebugPrinter classes to manage pre-allocated GPU memory for storing string names, ensuring efficient allocation and transfer of these names to the CUDA kernels. An initialize() function was added to set up the _DebugPrinter singleton.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request adds a valuable feature for debugging by allowing custom names to be printed with tensors. The C++ changes are well-structured, refactoring common code into a helper function. However, the Python implementation has a few significant issues that need to be addressed:

• A critical bug in tensor creation that prevents the feature from working.
• A resource management issue where the name buffer is never reset, which will cause crashes in long-running applications.
• An API design that requires explicit initialization, making it less user-friendly.

I've provided suggestions to fix these issues. Once they are addressed, this will be a great addition to the library.

[gemini-code-assist]

There is a critical bug in how the temporary tensor for the name is created. torch.empty() expects a shape (a tuple of integers), but it's being passed list(name_bytes) + [0], which is a list of byte values. This will create a tensor with incorrect dimensions (e.g., for name="a", it tries to create a tensor of shape (97, 0)), causing a RuntimeError during the copy_ operation.

You should use torch.tensor() to create a tensor from data.

Suggested change

	tmp = torch.empty(list(name_bytes) + [0], dtype=torch.uint8, device="cpu")
	tmp = torch.tensor(list(name_bytes) + [0], dtype=torch.uint8, device="cpu")

[gemini-code-assist]

The _Buffer class implements a simple bump allocator by incrementing _used_len but never resets it. In applications that capture CUDA graphs multiple times or in other long-running scenarios, this will lead to buffer exhaustion and an AssertionError.

Consider adding a mechanism to reset the buffer. For example, you could add a reset() method to _DebugPrinter that users can call before a new graph capture session to reset the _used_len of each buffer.

[gemini-code-assist]

The current implementation requires the user to explicitly call initialize() before using print_tensor. If not called, print_tensor will raise an error because _printer is None. This is not ideal for usability, and the provided example.py does not call initialize(), meaning it would fail.

A better approach is to use lazy initialization, where the _DebugPrinter is created on the first call to print_tensor. This simplifies the API for the user by removing the need for an explicit initialization step.

Suggested change

	def initialize():
	global _printer
	assert _printer is None
	_printer = _DebugPrinter()
	def print_tensor(x: torch.Tensor, print_ptr: bool = False):
	print_tensor_kernel(x, print_ptr)
	def print_tensor(x: torch.Tensor, name: str = "", print_ptr: bool = False):
	_printer(x=x, name=name, print_ptr=print_ptr)
	def print_tensor(x: torch.Tensor, name: str = "", print_ptr: bool = False):
	global _printer
	if _printer is None:
	_printer = _DebugPrinter()
	_printer(x=x, name=name, print_ptr=print_ptr)

[gemini-code-assist]

Using %lld to print a pointer address is not portable and can lead to undefined behavior. The correct format specifier for a void* pointer is %p.

    printf("addr: %p\n", x);