Skip to content

T
tic
Unverified Commit 4672dc76 by Josh Fromm Committed by GitHub
Options

Add cuda target check to dense tensorcore schedule. (#5376)

parent a2d6fe65
Showing with 1 additions and 0 deletions
python/tvm/relay/op/strategy/cuda.py
...@@ -373,6 +373,7 @@ def dense_strategy_cuda(attrs, inputs, out_type, target): ...@@ -373,6 +373,7 @@ def dense_strategy_cuda(attrs, inputs, out_type, target):
wrap_topi_schedule(topi.cuda.schedule_dense_large_batch), wrap_topi_schedule(topi.cuda.schedule_dense_large_batch),
name="dense_large_batch.cuda", name="dense_large_batch.cuda",
plevel=5) plevel=5)
if target.target_name == "cuda":
if nvcc.have_tensorcore(tvm.gpu(0).compute_version): if nvcc.have_tensorcore(tvm.gpu(0).compute_version):
if(i % 16 == 0 and b % 16 == 0 and o % 16 == 0) \ 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 % 8 == 0 and o % 32 == 0) \
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment