tilelang.language.print¶

This module provides macros and utilities for debugging TileLang (tl) programs. It includes functionality to print variables, print values in buffers, and conditionally execute debug prints.

Functions¶

print_var(var[, msg])

Prints the value of a TIR primitive expression (PrimExpr) for debugging purposes.

print_var_with_condition(condition, var[, msg])

Conditionally prints a TIR primitive expression (PrimExpr) if a given condition is True.

print_global_buffer_with_condition(condition, buffer, ...)

Conditionally prints the values of a flattened TIR buffer if the condition is True.

print_shared_buffer_with_condition(condition, buffer, ...)

Conditionally prints the values of a flattened TIR buffer if the condition is True.

print_fragment_buffer_with_condition(condition, ...[, msg])

Conditionally prints the values of a flattened TIR buffer if the condition is True.

print_local_buffer_with_condition(condition, buffer, elems)

Conditionally prints the values of a flattened TIR buffer if the condition is True.

print(obj[, msg, warp_group_id, warp_id])

A generic print function that handles both TIR buffers and primitive expressions.

Module Contents¶

tilelang.language.print.print_var(var, msg='')¶

Prints the value of a TIR primitive expression (PrimExpr) for debugging purposes.

Parameters:

  • var (tir.PrimExpr) – The variable or expression to be printed.

  • msg (str)

Returns:

The TIR expression for the debug print operation.

Return type:

tir.PrimExpr

tilelang.language.print.print_var_with_condition(condition, var, msg='')¶

Conditionally prints a TIR primitive expression (PrimExpr) if a given condition is True.

Parameters:

  • condition (tir.PrimExpr) – A TIR expression representing the condition to check.

  • var (tir.PrimExpr) – The variable or expression to be printed.

  • msg (str)

Returns:

The TIR expression for the debug print operation, if the condition is True.

Return type:

tir.PrimExpr

tilelang.language.print.print_global_buffer_with_condition(condition, buffer, elems, msg='')¶

Conditionally prints the values of a flattened TIR buffer if the condition is True.

Parameters:

  • condition (tvm.tir.PrimExpr)

  • buffer (tvm.tir.Buffer)

  • elems (int)

  • msg (str)

Return type:

tvm.tir.PrimExpr

tilelang.language.print.print_shared_buffer_with_condition(condition, buffer, elems, msg='')¶

Conditionally prints the values of a flattened TIR buffer if the condition is True.

Parameters:

  • condition (tir.PrimExpr) – A TIR expression representing the condition to check.

  • buffer (tir.Buffer) – The buffer whose values need to be printed.

  • elems (int) – The number of elements in the buffer to print.

  • msg (str)

Returns:

The TIR expression for the debug print operation.

Return type:

tir.PrimExpr

tilelang.language.print.print_fragment_buffer_with_condition(condition, buffer, elems, msg='')¶

Conditionally prints the values of a flattened TIR buffer if the condition is True.

Parameters:

  • condition (tir.PrimExpr) – A TIR expression representing the condition to check.

  • buffer (tir.Buffer) – The buffer whose values need to be printed.

  • elems (int) – The number of elements in the buffer to print.

  • msg (str)

Returns:

The TIR expression for the debug print operation.

Return type:

tir.PrimExpr

tilelang.language.print.print_local_buffer_with_condition(condition, buffer, elems, msg='')¶

Conditionally prints the values of a flattened TIR buffer if the condition is True.

Parameters:

  • condition (tir.PrimExpr) – A TIR expression representing the condition to check.

  • buffer (tir.Buffer) – The buffer whose values need to be printed.

  • elems (int) – The number of elements in the buffer to print.

  • msg (str)

Returns:

The TIR expression for the debug print operation.

Return type:

tir.PrimExpr

tilelang.language.print.print(obj, msg='', warp_group_id=0, warp_id=0)¶

A generic print function that handles both TIR buffers and primitive expressions.

  • If the input is a TIR buffer, it prints its values, but only on the first thread (tx=0, ty=0, tz=0).

  • If the input is a TIR primitive expression, it prints its value directly.

Parameters:

  • obj (Any) – The object to print. It can be either a tir.Buffer or tir.PrimExpr.

  • msg (str) – An optional message to include in the print statement.

  • warp_group_id (int) – The warp group id to print.

  • warp_id (int) – The warp id to print.

  • warp_id. (print thread will be warp_group_id * warp_group_size +)

Returns:

The TIR expression for the debug print operation.

Return type:

tir.PrimExpr

Raises:

ValueError – If the input object type is unsupported.