test fix

apache · Dec 23, 2022 · 243225e · 243225e
1 parent 6d550ca
commit 243225e
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/python/tvm/tir/tensor_intrin/x86.py b/python/tvm/tir/tensor_intrin/x86.py
@@ -67,12 +67,12 @@ def dot_product_16x4_u8i8i32_vnni(
         )
 
 
+mem_scope="global"
 @T.prim_func
-def dot_product_16x4_u8i8i32_avx512(
-    A: T.Buffer((4,), "uint8", offset_factor=1),
-    B: T.Buffer((16, 4), "int8", offset_factor=1),
-    C: T.Buffer((16,), "int32", offset_factor=1),
-) -> None:
+def dot_product_16x4_u8i8i32_avx512(a: T.handle, b: T.handle, c: T.handle,) -> None:
+    A = T.match_buffer(a, (4,), "uint8", offset_factor=1, scope=mem_scope)
+    B = T.match_buffer(b, (16, 4), "int8", offset_factor=1, scope=mem_scope)
+    C = T.match_buffer(c, (16,), "int32", offset_factor=1, scope=mem_scope)
     with T.block("root"):
         T.reads(C[0:16], A[0:4], B[0:16, 0:4])
         T.writes(C[0:16])