rename vnni.py to x86.py

masahi · Apr 6, 2022 · 0f0682d · 0f0682d
1 parent f88c31e
commit 0f0682d
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 7 deletions.
diff --git a/python/tvm/tir/tensor_intrin/__init__.py b/python/tvm/tir/tensor_intrin/__init__.py
@@ -16,4 +16,4 @@
 # under the License.
 # pylint: disable=unused-import
 """Intrinsics for tensorization."""
-from . import vnni
+from . import x86
diff --git a/python/tvm/tir/tensor_intrin/vnni.py → python/tvm/tir/tensor_intrin/x86.py b/python/tvm/tir/tensor_intrin/vnni.py → python/tvm/tir/tensor_intrin/x86.py
@@ -23,7 +23,7 @@
 
 
 @T.prim_func
-def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
+def dot_product_16x4_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
     A = T.match_buffer(a, (4,), "uint8", offset_factor=1)
     B = T.match_buffer(b, (16, 4), "int8", offset_factor=1)
     C = T.match_buffer(c, (16,), "int32", offset_factor=1)
@@ -41,7 +41,7 @@ def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
 
 
 @T.prim_func
-def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
+def dot_product_16x4_vnni_impl(a: T.handle, b: T.handle, c: T.handle) -> None:
     A = T.match_buffer(a, (4,), "uint8", offset_factor=1)
     B = T.match_buffer(b, (16, 4), "int8", offset_factor=1)
     C = T.match_buffer(c, (16,), "int32", offset_factor=1)
@@ -66,6 +66,6 @@ def dot_product_intrin(a: T.handle, b: T.handle, c: T.handle) -> None:
         )
 
 
-INTRIN_NAME = "dot_16x1x16_uint8_int8_int32_cascadelake"
+VNNI_INTRIN = "dot_16x4_vnni"
 
-TensorIntrin.register(INTRIN_NAME, dot_product_desc, dot_product_intrin)
+TensorIntrin.register(VNNI_INTRIN, dot_product_16x4_desc, dot_product_16x4_vnni_impl)
diff --git a/tests/python/unittest/test_meta_schedule_tune_relay.py b/tests/python/unittest/test_meta_schedule_tune_relay.py
@@ -42,7 +42,7 @@
 from tvm.target.target import Target
 from tvm.tir.schedule import BlockRV, Schedule
 from tvm.tir.schedule.trace import Trace
-from tvm.tir.tensor_intrin.vnni import INTRIN_NAME as VNNI_INTRIN
+from tvm.tir.tensor_intrin.x86 import VNNI_INTRIN
 
 
 logging.basicConfig()

diff --git a/tests/python/unittest/test_tir_schedule_tensorize.py b/tests/python/unittest/test_tir_schedule_tensorize.py
@@ -22,7 +22,7 @@
 from tvm import tir, te
 from tvm.script import tir as T
 from tvm.tir.schedule.testing import verify_trace_roundtrip
-from tvm.tir.tensor_intrin.vnni import INTRIN_NAME as VNNI_INTRIN
+from tvm.tir.tensor_intrin.x86 import VNNI_INTRIN
 
 # fmt: off
 # pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks