From eae387b7c349a26fb6c83ed88f8fb9915b055ad4 Mon Sep 17 00:00:00 2001 From: Josh Fromm Date: Sat, 18 Apr 2020 21:09:14 -0700 Subject: [PATCH] Add cuda target check to dense tensorcore schedule. (#5376) --- python/tvm/relay/op/strategy/cuda.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 4e5088f0d85e..81a7feb882de 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -373,15 +373,16 @@ def dense_strategy_cuda(attrs, inputs, out_type, target): wrap_topi_schedule(topi.cuda.schedule_dense_large_batch), name="dense_large_batch.cuda", plevel=5) - if nvcc.have_tensorcore(tvm.gpu(0).compute_version): - if(i % 16 == 0 and b % 16 == 0 and o % 16 == 0) \ - or (i % 16 == 0 and b % 8 == 0 and o % 32 == 0) \ - or (i % 16 == 0 and b % 32 == 0 and o % 8 == 0): - strategy.add_implementation( - wrap_compute_dense(topi.cuda.dense_tensorcore), - wrap_topi_schedule(topi.cuda.schedule_dense_tensorcore), - name="dense_tensorcore.cuda", - plevel=20) + if target.target_name == "cuda": + if nvcc.have_tensorcore(tvm.gpu(0).compute_version): + if(i % 16 == 0 and b % 16 == 0 and o % 16 == 0) \ + or (i % 16 == 0 and b % 8 == 0 and o % 32 == 0) \ + or (i % 16 == 0 and b % 32 == 0 and o % 8 == 0): + strategy.add_implementation( + wrap_compute_dense(topi.cuda.dense_tensorcore), + wrap_topi_schedule(topi.cuda.schedule_dense_tensorcore), + name="dense_tensorcore.cuda", + plevel=20) if target.target_name == "cuda" and "cublas" in target.libs: strategy.add_implementation( wrap_compute_dense(topi.cuda.dense_cublas),