16x8x16 trans working

Author: masahi

tests/python/unittest/test_mma_16x8x16_4k_tune_simple.py

@@ -411,20 +411,20 @@ def tile_wmma_fragment(block_read, height):
 
 f = tvm.build(sch.mod["main"], target="cuda", name="dense")
 
-# dev = tvm.device("cuda", 0)
-# a_np = np.random.uniform(size=(N, K)).astype("float16")
-# b_np = np.random.uniform(size=(K, M)).astype("float16")
-# c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
-# a = tvm.nd.array(a_np, dev)
-# b = tvm.nd.array(b_np, dev)
-# c = tvm.nd.array(np.zeros((M, N), dtype="float32"), dev)
-
-# print(f.imported_modules[0].get_source())
-# f(a, b, c)
+dev = tvm.device("cuda", 0)
+a_np = np.random.uniform(size=(N, K)).astype("float16")
+b_np = np.random.uniform(size=(K, M)).astype("float16")
+c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
+a = tvm.nd.array(a_np, dev)
+b = tvm.nd.array(b_np, dev)
+c = tvm.nd.array(np.zeros((M, N), dtype="float32"), dev)
+
+print(f.imported_modules[0].get_source())
+f(a, b, c)
 # tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
 # print("ok")
 
-# evaluator = f.time_evaluator(f.entry_name, dev, number=1000)
-# gflops = (N * M * K) * 2 / 1e9
-# time_ms = evaluator(a, b, c).mean * 1e3
-# print("matmul with tensor core: %f ms, %f GFLOPS" % (time_ms, gflops / (time_ms / 1e3)))
+evaluator = f.time_evaluator(f.entry_name, dev, number=10)
+gflops = (N * M * K) * 2 / 1e9
+time_ms = evaluator(a, b, c).mean * 1e3
+print("matmul with tensor core: %f ms, %f GFLOPS" % (time_ms, gflops / (time_ms / 1e3)))

tests/python/unittest/test_mma_16x8x16_simple.py

@@ -317,16 +317,17 @@ def fetch_to_shared(block, idx):
 target = "cuda"
 
 f = tvm.build(sch.mod["main"], target=target, name="dense")
-# dev = tvm.device(target, 0)
+dev = tvm.device(target, 0)
 
-# a_np = np.random.uniform(size=(16, K)).astype("float16")
-# b_np = np.random.uniform(size=(K, K)).astype("float16")
-# c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
+a_np = np.random.uniform(size=(16, K)).astype("float16")
+b_np = np.random.uniform(size=(K, K)).astype("float16")
+c_np = np.dot(a_np.astype("float32"), b_np.astype("float32"))
 
-# a = tvm.nd.array(a_np, dev)
-# b = tvm.nd.array(b_np, dev)
-# c = tvm.nd.array(np.zeros((16, K), dtype="float32"), dev)
+a = tvm.nd.array(a_np, dev)
+b = tvm.nd.array(b_np, dev)
+c = tvm.nd.array(np.zeros((16, K), dtype="float32"), dev)
 
-# # print(f.imported_modules[0].get_source())
-# f(a, b, c)
-# tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
+# print(f.imported_modules[0].get_source())
+f(a, b, c)
+tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-3)
+print("ok")