[Testing] Move TMA 1D and test for its functionality

Author: tzj-fxz

examples/elementwise/test_example_elementwise.py

@@ -1,15 +1,10 @@
 import tilelang.testing
 import example_elementwise_add
-import example_elementwise_add_tma_1d
 
 
 def test_example_elementwise_add():
     example_elementwise_add.main()
 
 
-def test_example_elementwise_add_tma_1d():
-    example_elementwise_add_tma_1d.main()
-
-
 if __name__ == "__main__":
     tilelang.testing.main()

examples/elementwise/example_elementwise_add_tma_1d.py renamed to testing/python/language/test_tilelang_language_tma_1d.py

@@ -1,7 +1,7 @@
 import argparse
+import torch
 import tilelang
 import tilelang.language as T
-import torch
 
 
 def ref_program(x, y):
@@ -30,23 +30,29 @@ def elem_add(A: T.Tensor((M, N), in_dtype), B: T.Tensor((M, N), in_dtype), C: T.
     return elem_add
 
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--m", type=int, default=128)
-    parser.add_argument("--n", type=int, default=128)
-    args, _ = parser.parse_known_args()
-    M, N = args.m, args.n
-
+def run_elementwise_add(M, N):
     a = torch.randn(M, N, dtype=torch.float32, device="cuda")
     b = torch.randn(M, N, dtype=torch.float32, device="cuda")
 
     # Default config
-    config = {"block_M": 128, "block_N": 128, "threads": 128}
+    block_M, block_N = 128, 128
+    config = {"block_M": block_M, "block_N": block_N, "threads": 128}
     kernel = elementwise_add(M, N, **config, in_dtype="float32", out_dtype="float32")
 
     out = kernel(a, b)
     torch.testing.assert_close(out, ref_program(a, b), rtol=1e-2, atol=1e-2)
-    print("All passed!")
+
+    code = kernel.get_kernel_source()
+    if N == block_N:
+        assert "tma_load" in code and "CUtensorMap" not in code
+    else:
+        assert "tma_load" in code and "CUtensorMap" in code
+
+
+def main():
+    run_elementwise_add(128, 128)
+    run_elementwise_add(256, 128)
+    run_elementwise_add(256, 256)
 
 
 if __name__ == "__main__":