[tvm] branch main updated: [FIX, AUTOTVM] Add flop counts to cublas (#7297)

[co...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

comaniac pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 218048e  [FIX,AUTOTVM] Add flop counts to cublas (#7297)
218048e is described below

commit 218048ec8306d4f3dc2e9a8dcca187804ad8a254
Author: Tristan Konolige <tr...@gmail.com>
AuthorDate: Sat Jan 23 00:48:40 2021 -0800

    [FIX,AUTOTVM] Add flop counts to cublas (#7297)
---
 python/tvm/autotvm/task/task.py | 1 +
 python/tvm/contrib/cublas.py    | 4 ++--
 python/tvm/topi/cuda/dense.py   | 5 ++++-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/tvm/autotvm/task/task.py b/python/tvm/autotvm/task/task.py
index c8b50ad..52f0996 100644
--- a/python/tvm/autotvm/task/task.py
+++ b/python/tvm/autotvm/task/task.py
@@ -580,6 +580,7 @@ def compute_flop(sch):
                 pass
             else:
                 raise FlopCalculationError(
+                    f"{op.name} is not supported by autotvm. "
                     "Only support te.compute currently. "
                     "Other ops like tvm.te.scan/te.extern is not supported"
                 )
diff --git a/python/tvm/contrib/cublas.py b/python/tvm/contrib/cublas.py
index 9a36fa5..e01b09c 100644
--- a/python/tvm/contrib/cublas.py
+++ b/python/tvm/contrib/cublas.py
@@ -48,7 +48,7 @@ def matmul(lhs, rhs, transa=False, transb=False, dtype=None):
             "tvm.contrib.cublas.matmul", ins[0], ins[1], outs[0], transa, transb
         ),
         dtype=dtype,
-        name="C",
+        name="matmul_cublas",
     )
 
 
@@ -82,5 +82,5 @@ def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
             "tvm.contrib.cublas.batch_matmul", ins[0], ins[1], outs[0], transa, transb
         ),
         dtype=dtype,
-        name="C",
+        name="batch_matmul_cublas",
     )
diff --git a/python/tvm/topi/cuda/dense.py b/python/tvm/topi/cuda/dense.py
index 85b9b19..f8abe4d 100644
--- a/python/tvm/topi/cuda/dense.py
+++ b/python/tvm/topi/cuda/dense.py
@@ -17,7 +17,7 @@
 # pylint: disable=invalid-name, unused-argument
 """Schedule for dense operator"""
 import logging
-from tvm import te
+from tvm import te, tir
 import tvm.autotvm as autotvm
 from tvm.autotvm.task.space import SplitEntity
 from tvm.contrib import cublas
@@ -44,6 +44,9 @@ def dense_cublas(cfg, data, weight, bias=None, out_dtype=None):
     matmul = cublas.matmul(data, weight, False, True)
     if isinstance(batch, int):
         cfg.add_flop(batch * in_dim * out_dim * 2)
+    elif isinstance(batch, tir.IntImm):
+        cfg.add_flop(batch.value * in_dim * out_dim * 2)
+    # if we get a te.Var, we cannot add flop counts
     if bias is not None:
         matmul = te.compute(
             (batch, out_dim), lambda i, j: matmul[i, j] + bias[j], tag=tag.BROADCAST