simplify import

masahi · Apr 7, 2022 · 9a3e508 · 9a3e508
1 parent d8e43ec
commit 9a3e508
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/python/tvm/tir/tensor_intrin/__init__.py b/python/tvm/tir/tensor_intrin/__init__.py
@@ -16,4 +16,5 @@
 # under the License.
 # pylint: disable=unused-import
 """Intrinsics for tensorization."""
-from . import x86
+from .x86 import *
+from .arm_cpu import *
diff --git a/tests/python/unittest/test_tir_schedule_tensorize.py b/tests/python/unittest/test_tir_schedule_tensorize.py
@@ -22,8 +22,11 @@
 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.x86 import VNNI_DOT_16x4_INTRIN as VNNI_INTRIN
-from tvm.tir.tensor_intrin.arm_cpu import ARM_DOT_4x4_i8_NEON_INTRIN, ARM_DOT_4x4_i8_SDOT_INTRIN
+from tvm.tir.tensor_intrin import (
+    VNNI_DOT_16x4_INTRIN,
+    ARM_DOT_4x4_i8_NEON_INTRIN,
+    ARM_DOT_4x4_i8_SDOT_INTRIN,
+)
 
 # fmt: off
 # pylint: disable=no-member,invalid-name,unused-variable,line-too-long,redefined-outer-name,unexpected-keyword-arg,too-many-nested-blocks
@@ -567,7 +570,7 @@ def test_tensorize_vnni():
     sch.reorder(ko, ji, ki)
 
     sch.decompose_reduction(block, ko)
-    sch.tensorize(ji, VNNI_INTRIN)
+    sch.tensorize(ji, VNNI_DOT_16x4_INTRIN)
 
     verify_trace_roundtrip(sch=sch, mod=func)