Topi matmul crashed(core dumped) on CPU


#1

Error occurred when performing big matrix multiplication inside my Docker (with TVM 0.6).

Error:
Total size for allocation tensor is constant but exceeds 2^31 - 1.
Aborted (core dumped)

Same error came out outside docker.

import tvm
import topi
import numpy as np
import time


Dimension = 16
Batch_size = 128 * 100 * 100
use_gpu = False

X = tvm.placeholder((Batch_size, Dimension), name="X", dtype="float32")

W = tvm.placeholder((Dimension, Dimension), name="W", dtype="float32")

R = (topi.sum(topi.matmul(topi.matmul(X, W), X, transp_a=False, transp_b=True), 1)
     / topi.sum(topi.matmul(X, X, transp_b=True), 1))

# print(R.shape)

s = tvm.create_schedule(R.op)

print("23333")

# Core dumped happened this line.
f = tvm.build(s, [X, W, R], target_host='llvm', name="my_func")

x_np = np.random.uniform(size=(Batch_size, Dimension)).astype(X.dtype)
w_np = np.random.uniform(size=(Dimension, Dimension)).astype(W.dtype)

if use_gpu:
    ctx = tvm.gpu(0)
else:
    ctx = tvm.cpu(0)


x_nd = tvm.nd.array(
    x_np,
    ctx
)

w_nd = tvm.nd.array(
    w_np,
    ctx
)

c_nd = tvm.nd.array(
    np.zeros((Batch_size,), dtype=R.dtype),
    ctx
)

f(x_nd, w_nd, c_nd)

c_nd.asnumpy()


#2

Result of (XW) X^T is a of size Batchsize x Batchzise which is more than 10^12 elements.