[Feature]:Add device assert

Author: yyttt6

src/tl_templates/cuda/debug.h

@@ -257,3 +257,10 @@ __device__ void debug_print_buffer_value<int16_t>(const char *msg,
          msg, blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, threadIdx.y,
          threadIdx.z, buf_name, index, (int32_t)var);
 }
+
+__device__ void device_assert(bool cond, const char *msg) {
+  if (!cond) {
+    printf("Device assert failed: %s\n", msg);
+    assert(0);
+  }
+}

testing/python/debug/test_device_assert.py

@@ -0,0 +1,27 @@
+# type: ignore
+
+import tilelang
+import tilelang.testing
+import tilelang.language as T
+import os
+
+os.environ["TILELANG_DEBUG"] = "1"
+
+
+def test_device_assert():
+
+    @T.prim_func
+    def program():
+        with T.Kernel(threads=128):
+            tid = T.get_thread_binding()
+            T.device_assert(tid != 0, "Assertion Trigger !")
+
+    jit_kernel = tilelang.compile(program, target="cuda")
+    print(jit_kernel.kernel_source)
+    profiler = jit_kernel.get_profiler()
+    profiler.run_once()
+
+
+if __name__ == "__main__":
+
+    tilelang.testing.main()

tilelang/language/__init__.py

@@ -64,7 +64,7 @@
     cumsum,  # noqa: F401
     finalize_reducer,  # noqa: F401
 )
-from .print import print  # noqa: F401
+from .print import print, device_assert  # noqa: F401
 from .customize import (
     atomic_max,  # noqa: F401
     atomic_min,  # noqa: F401

tilelang/language/print.py

@@ -1,8 +1,9 @@
 """
 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.
+It includes functionality to print variables, print values in buffers, conditionally execute debug prints and assert.
 """
 
+import os
 from tvm import tir
 from typing import Any
 from tilelang.language.kernel import get_thread_bindings
@@ -133,6 +134,18 @@ def print_local_buffer_with_condition(condition: tir.PrimExpr,
                             buffer[coords])
 
 
+@macro
+def device_assert(condition: tir.PrimExpr, msg: str = ""):
+    """
+    Device-side assert emulation.
+    In debug mode (TILELANG_DEBUG=1), performs Python-side conditional check
+    and injects a printf + abort on the device if desired.
+    """
+    debug_mode = os.environ.get("TILELANG_DEBUG", "0") not in ("0", "", "false", "False")
+    if debug_mode:
+        tir.call_extern("void", "device_assert", condition, msg)
+
+
 def print(obj: Any, msg: str = "", warp_group_id: int = 0, warp_id: int = 0) -> tir.PrimExpr:
     """
     A generic print function that handles both TIR buffers and primitive expressions.