Skip to content

Commit 1d7a470

Browse files
author
GitHub Actions
committed
Update docs
1 parent 603044c commit 1d7a470

File tree

3 files changed

+29
-5
lines changed

3 files changed

+29
-5
lines changed

_sources/tutorials/debug_tools_for_tilelang.md.txt

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,10 +108,22 @@ Hence, by registering a Python function named `tilelang_callback_cuda_postproc`,
108108
import tilelang
109109
import tilelang.language as T
110110
from tilelang import tvm
111+
from tilelang.engine.callback import register_cuda_postproc_callback
111112

112-
@tvm.register_func
113+
@register_cuda_postproc_callback
113114
def tilelang_callback_cuda_postproc(code, _):
114-
# ...existing code...
115+
print(code) # print the final CUDA code
116+
code = "// modified by tilelang_callback_cuda_postproc\n" + code
117+
return code
118+
119+
kernel = tilelang.compile(matmul, target="cuda")
120+
kernel_source = kernel.get_kernel_source()
121+
print(kernel_source)
122+
'''
123+
// modified by tilelang_callback_cuda_postproc
124+
#include "cuda_runtime.h"
125+
...
126+
'''
115127
```
116128

117129
### Runtime Debug Prints with `T.print`

0 commit comments

Comments
 (0)