huochaitiantang commented on a change in pull request #7814:
URL: https://github.com/apache/tvm/pull/7814#discussion_r611298716
##########
File path: python/tvm/relay/op/strategy/cuda.py
##########
@@ -722,12 +722,20 @@ def dense_strategy_cuda(attrs, inputs, out_type, target):
def batch_matmul_strategy_cuda(attrs, inputs, out_type, target):
"""batch_matmul cuda strategy"""
strategy = _op.OpStrategy()
- strategy.add_implementation(
- wrap_compute_batch_matmul(topi.cuda.batch_matmul),
- wrap_topi_schedule(topi.cuda.schedule_batch_matmul),
- name="batch_matmul.cuda",
- plevel=10,
- )
+ x, y = inputs
+ if x.dtype == "int8" and y.dtype == "int8" and out_type.dtype == "int32":
+ strategy.add_implementation(
+ wrap_compute_batch_matmul(topi.cuda.batch_matmul_int8,
need_out_dtype=True),
+ wrap_topi_schedule(topi.cuda.schedule_batch_matmul_int8),
+ name="batch_matmul_int8.cuda",
Review comment:
We add plevel=10 for the batch_matmul_int8.cuda, which is the same as
dense_int8.cuda
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
