check dot prod availablity in batch matmul schedule

Author: masahi

python/tvm/topi/cuda/batch_matmul.py

@@ -368,7 +368,16 @@ def _schedule_batch_matmul_int8(cfg, s, output):
     ko, ki = s[batch_matmul_cache].split(ko, factor=4)
     ko, kt = cfg["tile_k"].apply(s, batch_matmul_cache, ko)
     # dp4a tensorize
-    s[batch_matmul_cache].tensorize(ki, _dp4a)
+
+    target = tvm.target.Target.current(allow_none=False)
+    do_tensorize = True
+
+    if "vulkan" in target.keys:
+        do_tensorize = "+dotprod" in target.mattr or target.supports_integer_dot_product
+
+    if do_tensorize:
+        dtypes = (input_x.dtype, input_y.dtype)
+        s[batch_matmul_cache].tensorize(ki, dp4a("shared", "shared", "local", dtypes))
 
     # tile axis
     f, m, n = batch_matmul_op.axis

tests/python/topi/python/test_topi_batch_matmul.py

@@ -128,7 +128,7 @@ def check_device(device):
         f(a, b, c)
         tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
 
-    for device in ["cuda"]:
+    for device in ["cuda", "vulkan -from_device=0"]:
         check_device(device)