Using BLAS library with TVM

Questions
1

Hi,

I’m trying to get TVM to use the BLAS library.

So far I’ve: A. Built TVM with BLAS support (openblas) B. Ran the test_cblas.py unit test successfully.

When running the unit test and inspecting the generated LLVM IR, I can see function calls to cblas functions.

When I run the the baseline sample from: https://docs.tvm.ai/tutorials/optimize/opt_gemm.html#preparation-and-baseline

With target set to ‘llvm -libs=cblas’

I don’t see any calls to cblas functions, nor any difference in the measured time vs running without the cblas flag.

What am I doing wrong? :slight_smile:

2

Did you use the compute function (the following one) in the example to do GEMM?

C = tvm.compute(
           (M, N),
           lambda x, y: tvm.sum(A[x, k] * B[k, y], axis=k),
           name='C')

If so, then TVM won’t call CBLAS functions because it doesn’t know the compute function is actually doing GEMM. You need to call the dense operator in Relay/TOPI:

from tvm import relay
a = relay.var('a', shape=(10, 5))
b = relay.var('b', shape=(20, 5))
c = relay.nn.dense(a, b) # shape of c is (10, 20)

Note that the dense op was implemented to perform A(M, K) x B(N, K) instead of normal A(M, K) x B(K, N) computation.