16x8x16 worked

masahi · May 17, 2022 · ced5d8d · ced5d8d
1 parent 3d2c90d
commit ced5d8d
Showing 1 changed file with 13 additions and 16 deletions.
diff --git a/tests/python/unittest/test_mma_16x8x16.py b/tests/python/unittest/test_mma_16x8x16.py
@@ -183,7 +183,7 @@ def mma_sync_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
                 "fp16",
                 "fp32",
                 A.data,
-                A.elem_offset + tx * 8 + 4,
+                A.elem_offset + tx * 8,
                 B.data,
                 B.elem_offset + tx * 8 + 4,
                 C.data,
@@ -369,22 +369,19 @@ def shared_16x16_to_ldmatrix_32x8_layout(i, j):
 
 # lowered = tvm.lower(sch.mod["main"])
 
-# if use_gpu:
-#     target = "vulkan -from_device=0"
-# else:
-#     target = "llvm"
+target = "cuda"
 
-# f = tvm.build(sch.mod["main"], target=target, name="dense")
-# dev = tvm.device(target, 0)
+f = tvm.build(sch.mod["main"], target=target, name="dense")
+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)