The following code crashes at compilation
import torch
from torch.nn import functional as F
from tvm import relay
from tvm.ir.transform import PassContext
class A(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x, y, z):
return F.linear(x, F.linear(y, z))
model = A()
input_data = tuple(torch.randn(100,100) for _ in range(3))
scripted_model = torch.jit.trace(model, input_data).eval()
shape_list = [('x', (100, 100)), ('y', (100, 100)), ('z', (100, 100))]
mod, params = relay.frontend.from_pytorch(scripted_model, shape_list)
print(mod)
with PassContext(opt_level=3):
relay.build(mod, target='llvm')
The error message: Check failed: (tir::BijectiveLayout(new_src_layout, dst_layout).defined()) is false: Cannot insert layout transform because there are inconvertible layouts: NC v.s. NK10n.
It seems that it’s because F.linear is mapped to nn.dense, which assumes a specific layout of its inputs and output. The whole model is converted to something like nn.dense(x, nn.dense(y, z)) (modulo some transpositions) where the output of nn.dense is used as the second input for the other nn.dense. The former requries NC while the latter requires NK10n.
Maybe we should we nn.matmul in this case? Or we can also change the layout annotation such that it’s not using letters but the indexes to record the relation between the old and new layouts? (For example, NC → I0_I1, NK10n → I0_I1_10i0)