diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index ba47ae7bc4f1..53d6df48934d 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -144,7 +144,11 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): if groups == 1: if layout == "NCHW": assert kernel_layout == "OIHW" - if data.dtype in ("int8", "uint8") and kernel.dtype in ("int8", "uint8"): + if ( + target.kind.name == "cuda" + and data.dtype in ("int8", "uint8") + and kernel.dtype in ("int8", "uint8") + ): assert data.dtype == kernel.dtype strategy.add_implementation( wrap_compute_conv2d(topi.cuda.conv2d_nchw_int8), @@ -293,7 +297,7 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): "Unsupported shape for conv2d HWNC.\ Need to satisfy tensor core schedule." ) - elif layout == "NCHW4c" and data.dtype in ["int8", "uint8"]: + elif target.kind.name == "cuda" and layout == "NCHW4c" and data.dtype in ["int8", "uint8"]: assert kernel_layout == "OIHW4o4i" strategy.add_implementation( wrap_compute_conv2d(topi.cuda.conv2d_NCHWc_int8, True), @@ -353,7 +357,8 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): ic_chunk = in_channels // 4 if ( - data.dtype in ["int8", "uint8"] + target.kind.name == "cuda" + and data.dtype in ["int8", "uint8"] and kernel.dtype in ["int8", "uint8"] and channels % groups == 0 and out_channels % groups == 0 diff --git a/tests/python/relay/test_op_level2.py b/tests/python/relay/test_op_level2.py index 8117539e611a..44f211dd9f8a 100644 --- a/tests/python/relay/test_op_level2.py +++ b/tests/python/relay/test_op_level2.py @@ -325,10 +325,6 @@ def test_run( kernel_size, ): target = tvm.target.Target(target) - if target.kind.name == "vulkan" and dtype == "int8": - # The schedule selection incorrectly picks an - # implementation that requires NCHWc packed input. - pytest.xfail("Known failing test for vulkan") x = relay.var("x", shape=dshape, dtype=dtype) w = relay.var("w", shape=kshape, dtype=dtype)