diff --git a/python/tvm/contrib/nvcc.py b/python/tvm/contrib/nvcc.py index 89548b74866b..bc11e4a867e4 100644 --- a/python/tvm/contrib/nvcc.py +++ b/python/tvm/contrib/nvcc.py @@ -269,15 +269,34 @@ def have_int8(compute_version): return False -def have_tensorcore(compute_version): +def have_tensorcore(compute_version=None, target=None): """Either TensorCore support is provided in the compute capability or not Parameters ---------- - compute_version : str - compute capability of a GPU (e.g. "7.0") + compute_version : str, optional + compute capability of a GPU (e.g. "7.0"). + + target : tvm.target.Target, optional + The compilation target, will be used to determine arch if compute_version + isn't specified. """ + if compute_version is None: + if tvm.gpu(0).exist: + compute_version = tvm.gpu(0).compute_version + else: + if target is None or "arch" not in target.attrs: + warnings.warn( + "Tensorcore will be disabled due to no CUDA architecture specified." + "Try specifying it by adding '-arch=sm_xx' to your target." + ) + return False + compute_version = target.attrs["arch"] + # Compute version will be in the form "sm_{major}{minor}" + major, minor = compute_version.split("_")[1] + compute_version = major + "." + minor major, _ = parse_compute_version(compute_version) + if major == 7: return True diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index fc80c9ed6171..09ed475ad4e6 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -197,7 +197,7 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): if judge_winograd_autotvm: if ( target.kind.name == "cuda" - and nvcc.have_tensorcore(tvm.gpu(0).compute_version) + and nvcc.have_tensorcore(target=target) and judge_winograd_tensorcore ): strategy.add_implementation( @@ -215,7 +215,7 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): ) if ( target.kind.name == "cuda" - and nvcc.have_tensorcore(tvm.gpu(0).compute_version) + and nvcc.have_tensorcore(target=target) and ( (N % 16 == 0 and CI % 16 == 0 and CO % 16 == 0) or (N % 8 == 0 and CI % 16 == 0 and CO % 32 == 0) @@ -436,7 +436,7 @@ def conv2d_winograd_without_weight_transfrom_strategy_cuda(attrs, inputs, out_ty ) if ( target.kind.name == "cuda" - and nvcc.have_tensorcore(tvm.gpu(0).compute_version) + and nvcc.have_tensorcore(target=target) and judge_winograd_tensorcore ): strategy.add_implementation( @@ -563,7 +563,7 @@ def conv3d_strategy_cuda(attrs, inputs, out_type, target): N, _, _, _, _ = get_const_tuple(data.shape) _, _, _, CI, CO = get_const_tuple(kernel.shape) if target.kind.name == "cuda": - if nvcc.have_tensorcore(tvm.gpu(0).compute_version): + if nvcc.have_tensorcore(target=target): if ( (N % 16 == 0 and CI % 16 == 0 and CO % 16 == 0) or (N % 8 == 0 and CI % 16 == 0 and CO % 32 == 0) @@ -679,7 +679,7 @@ def dense_strategy_cuda(attrs, inputs, out_type, target): plevel=5, ) if target.kind.name == "cuda": - if nvcc.have_tensorcore(tvm.gpu(0).compute_version): + if nvcc.have_tensorcore(target=target): if ( (i % 16 == 0 and b % 16 == 0 and o % 16 == 0) or (i % 16 == 0 and b % 8 == 0 and o % 32 == 0)