add test

Author: LeiWang1999

testing/python/components/test_storage_rewrite_detect_inplace.py

@@ -0,0 +1,62 @@
+import tilelang
+import tilelang.testing
+from tilelang import language as T
+
+
+@tilelang.jit
+def _compile_kernel_without_inplace():
+    num_tokens = T.symbolic("num_tokens")
+
+    @T.prim_func
+    def buggy_kernel(x: T.Tensor[(num_tokens,), "float"]):
+        with T.Kernel(num_tokens, threads=32) as pid:
+            read = T.alloc_var("int")
+            read = x[pid]
+
+            write = T.alloc_var("int")
+            write = read * 2
+            x[pid] = write
+
+    return buggy_kernel
+
+
+@tilelang.jit(
+    pass_configs={
+        tilelang.PassConfigKey.TL_STORAGE_REWRITE_DETECT_INPLACE: True,
+    },
+)
+def _compile_kernel_with_inplace():
+    num_tokens = T.symbolic("num_tokens")
+
+    @T.prim_func
+    def buggy_kernel(x: T.Tensor[(num_tokens,), "float"]):
+        with T.Kernel(num_tokens, threads=32) as pid:
+            read = T.alloc_var("int")
+            read = x[pid]
+
+            write = T.alloc_var("int")
+            write = read * 2
+            x[pid] = write
+
+    return buggy_kernel
+
+
+def _get_device_kernel_script(detect_inplace: bool) -> str:
+    if detect_inplace:
+        kernel = _compile_kernel_with_inplace()
+    else:
+        kernel = _compile_kernel_without_inplace()
+    source = kernel.get_kernel_source()
+    return source
+
+def test_storage_rewrite_detect_inplace_toggle():
+    script_off = _get_device_kernel_script(detect_inplace=False)
+    script_on = _get_device_kernel_script(detect_inplace=True)
+    
+    assert script_off.count("read = (read * 2);") == 0
+    assert script_on.count("read = (read * 2);") > 0
+    
+
+
+if __name__ == "__main__":
+    tilelang.testing.main()