move intrin decl for vector type

apache · Mar 28, 2022 · c19ecc1 · c19ecc1
1 parent 3dd7f04
commit c19ecc1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/python/tvm/script/tir/intrin.py b/python/tvm/script/tir/intrin.py
@@ -102,6 +102,16 @@ def float64(imm, span):
     return imm.astype("float64", span)
 
 
+@register
+def int32x16(imm, span):
+    return imm.astype("int32x16", span)
+
+
+@register
+def int32x4(imm, span):
+    return imm.astype("int32x4", span)
+
+
 @register
 def min_value(dtype, span):
     return tvm.tir.min_value(dtype, span)

diff --git a/tests/python/unittest/test_meta_schedule_tune_relay.py b/tests/python/unittest/test_meta_schedule_tune_relay.py
@@ -34,7 +34,6 @@
 from tvm.meta_schedule.utils import derived_object
 from tvm.target.target import Target
 from tvm.script import tir as T
-from tvm.script.registry import register
 from tvm._ffi import register_func
 import tempfile
 
@@ -326,11 +325,6 @@ def get_output(data, lib):
         assert np.allclose(actual_output, expected_output, rtol=1e-4, atol=2e-4)
 
 
-@register
-def int32x16(imm, span):
-    return imm.astype("int32x16", span)
-
-
 @T.prim_func
 def dot_product_desc(a: T.handle, b: T.handle, c: T.handle) -> None:
     A = T.match_buffer(a, (4,), "uint8", offset_factor=1)