Debug tool for tilelang

[DD-DuDa]

with T.Kernel(T.ceildiv(seqlen_q, block_M), heads * batch, num_split, threads=128 * 2) as (bx, by, bz):
      Q_shared = T.alloc_shared([block_M, dim], dtype)
      acc_s = T.alloc_fragment([block_M, block_N], accum_dtype)

      T.copy(Q[bid, 0, hid, :], Q_shared[0, :])

      # block level shared memory tensor, and register level tensor
      if bx == 0 and by == 0 and bz == 0 and threadIdx.x == 0:
          print(Q_shared)
          print(acc_s)

For debug, is it possible to print the intermediate value like that?

[DD-DuDa]

Sorry for not having any experience on TVM ...

We can use @tvm.register_func to achieve it.

@tvm.register_func
def tvm_callback_cuda_postproc(code, target):
    print("Original CUDA Code:")
    print(code)
    
    # remember to return code
    return code 

Here is my code:

@tvm.register_func
def tvm_callback_cuda_postproc(code, target):
    print("Original CUDA Code:")
    print(code)

    # Debug code with proper escaping for newlines
    debug_code = r"""
    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0) {
        for (int i = 0; i < 16; ++i) {
            printf("Q_shared[%d] = %f\\n", i, (float)Q_shared[i]);
        }
    }
    """

    # Regex pattern to locate the shared memory assignment line
    assignment_pattern = r"(\*\(uint4\*\)\(Q_shared \+ \(\(\(int\)threadIdx\.x\) \* 8\)\) = \*\(uint4\*\)\(Q \+ \(\(\(\(int\)blockIdx\.y\) \* 128\) \+ \(\(\(int\)threadIdx\.x\) \* 8\)\)\);)"  # Regex to locate shared memory assignment

    # Insert debug code before and after the assignment
    modified_code = re.sub(
        assignment_pattern,
        debug_code.strip() + r"\n\1\n" + debug_code.strip(),  # Insert debug code before and after the assignment
        code,
        count=1  # Only replace the first occurrence
    )

    print("Modified CUDA Code:")
    print(modified_code)

    # Return the modified CUDA code
    return modified_code


def flashdecoding(batch, heads, seqlen_q, seqlen_kv, dim, is_casual, num_split, tune=False):
    scale = (1.0 / dim) ** 0.5 * 1.44269504  # log2(e)
    shape_q    = [batch, seqlen_q, heads, dim]
    shape_kv   = [batch, seqlen_kv, heads, dim]
    part_shape = [batch, seqlen_q, heads, num_split, dim]
    dtype      = "float16"
    accum_dtype = "float"

    def kernel_func(block_M, block_N):
        
        @T.macro
        def flash_attn_split(
            Q: T.Buffer(shape_q, dtype),
            K: T.Buffer(shape_kv, dtype),
            V: T.Buffer(shape_kv, dtype),
            Output: T.Buffer(shape_q, dtype),
        ):
            with T.Kernel(T.ceildiv(seqlen_q, block_M), heads * batch, num_split, threads=128 * 2) as (bx, by, bz):
                Q_shared = T.alloc_shared([block_M, dim], dtype)

                mid = bx
                hid = by % heads
                bid = by // heads
                sid = bz

                T.copy(Q[bid, 0, hid, :], Q_shared[0, :])

        @T.prim_func
        def main(
                Q: T.Buffer(shape_q, dtype),
                K: T.Buffer(shape_kv, dtype),
                V: T.Buffer(shape_kv, dtype),
                Output: T.Buffer(shape_q, dtype),
        ):
            flash_attn_split(Q, K, V, Output)

        return main

    def kernel(block_M, block_N):
        return kernel_func(block_M, block_N)

    return kernel

I can obtain the log:

Q_shared[0] = 0.000000
Q_shared[1] = 0.000000
Q_shared[2] = 0.000000
Q_shared[3] = 0.000000
Q_shared[4] = 0.000000
Q_shared[5] = 0.000000
Q_shared[6] = 0.000000
Q_shared[7] = 0.000000
Q_shared[8] = 0.000000
Q_shared[9] = 0.000000
Q_shared[10] = 0.000000
Q_shared[11] = 0.000000
Q_shared[12] = 0.000000
Q_shared[13] = 0.000000
Q_shared[14] = 0.000000
Q_shared[15] = 0.000000
Q_shared[0] = -1.178711
Q_shared[1] = -1.483398
Q_shared[2] = 0.819824
Q_shared[3] = 0.268311
Q_shared[4] = -0.950195
Q_shared[5] = 0.772461
Q_shared[6] = -0.782715
Q_shared[7] = 0.696777
Q_shared[8] = -1.754883
Q_shared[9] = -0.028671
Q_shared[10] = -1.160156
Q_shared[11] = 0.000826
Q_shared[12] = -0.229858
Q_shared[13] = -0.006107
Q_shared[14] = 0.492432
Q_shared[15] = 0.971191
Original CUDA Code:
#include <tl_templates/cuda/gemm.h>
#include <tl_templates/cuda/copy.h>
#include <tl_templates/cuda/reduce.h>
#include <tl_templates/cuda/ldsm.h>
#include <tl_templates/cuda/threadblock_swizzle.h>

extern "C" __global__ void __launch_bounds__(256) main_kernel(half_t* __restrict__ Q) {
  extern __shared__ __align__(1024) half_t Q_shared[];
  if (((int)threadIdx.x) < 16) {
    *(uint4*)(Q_shared + (((int)threadIdx.x) * 8)) = *(uint4*)(Q + ((((int)blockIdx.y) * 128) + (((int)threadIdx.x) * 8)));
  }
}


Modified CUDA Code:
#include <tl_templates/cuda/gemm.h>
#include <tl_templates/cuda/copy.h>
#include <tl_templates/cuda/reduce.h>
#include <tl_templates/cuda/ldsm.h>
#include <tl_templates/cuda/threadblock_swizzle.h>

extern "C" __global__ void __launch_bounds__(256) main_kernel(half_t* __restrict__ Q) {
  extern __shared__ __align__(1024) half_t Q_shared[];
  if (((int)threadIdx.x) < 16) {
    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0) {
        for (int i = 0; i < 16; ++i) {
            printf("Q_shared[%d] = %f\n", i, (float)Q_shared[i]);
        }
    }
*(uint4*)(Q_shared + (((int)threadIdx.x) * 8)) = *(uint4*)(Q + ((((int)blockIdx.y) * 128) + (((int)threadIdx.x) * 8)));
if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 && threadIdx.x == 0) {
        for (int i = 0; i < 16; ++i) {
            printf("Q_shared[%d] = %f\n", i, (float)Q_shared[i]);
        }
    }
  }
}

This issue can be closed ~

Or it maybe helpful for reference

[LeiWang1999]

Yes, we’re currently working on documentation for debugging. I also think we should introduce a new statement, T.print, to print tensors.

[DD-DuDa]

Sure thing! It'd be great that it can print sth like CuTe does.

print_layout, print_tensor... with coordinate and even slices.

[LeiWang1999]

PR #45 introduced T.print, which allows print for an entire buffer or expressions, examples are in testing/python/debug/test_tilelang_debug_print.py, we need another pr for debug documentations.

[LeiWang1999]

debug tutorial: https://tilelang.tile-ai.cn/tutorials/debug_tools_for_tilelang.html

[LeiWang1999]

closed as has been resolved :)