[Operator] Resolve matmul to batch_matmul when lower than sm80 (#146)

* Resolve matmul to batch_matmul when lower than sm80 * Resolve matmul to batch_matmul when lower than sm80
hidet-org · Mar 24, 2023 · 3a87bed · 3a87bed
1 parent 5b3ab00
commit 3a87bed
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/python/hidet/graph/ops/definitions/matmul/resolve.py b/python/hidet/graph/ops/definitions/matmul/resolve.py
@@ -173,6 +173,9 @@ def resolve_f16(self, op: Operator) -> Optional[List[Tensor]]:
         if not (a.dtype == dtypes.float16 and b.dtype == dtypes.float16 and a.shape[-1] % 8 == b.shape[-1] % 8 == 0):
             return None
 
+        if hidet.cuda.compute_capability() < (8, 0):
+            return None
+
         parallel_k = self.get_config('parallel_k', default='default')  # 'default', 'search', 2, 4, ...
         if isinstance(parallel_k, str):
             if parallel_k == 'default':