[vit][compile]error: unknown: unsupported by backend contract: module initializers,How to lower vit_b_16 to linalg?

[followtheart]

When I compile vit model into LINALG_ON_TENSORS,got this error;any help?

import numpy as np
import torch
import torchvision.models as models
import torch_mlir

model = models.vit_b_16(weights=models.ViT_B_16_Weights.IMAGENET1K_V1)
model.eval()
input_batch = torch.randn(1,3,224,224)
vit_mod = torch_mlir.compile(model, input_batch, 
                            output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, 
                            use_tracing=False)

Error info:

Traceback (most recent call last):
  File "/home/jq/code/demo/torch_mlir-vit.py", line 9, in <module>
    vit_mod = torch_mlir.compile(model, input_batch, 
  File "/home/jq/code/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/__init__.py", line 451, in compile
    run_pipeline_with_repro_report(
  File "/home/jq/code/torch-mlir/build/tools/torch-mlir/python_packages/torch_mlir/torch_mlir/compiler_utils.py", line 69, in run_pipeline_with_repro_report
    raise TorchMlirCompilerError(trimmed_message) from None
torch_mlir.compiler_utils.TorchMlirCompilerError: Lowering TorchScript IR -> Torch Backend IR failed with the following diagnostics:


python exception: Failure while executing pass pipeline:
error: unknown: unsupported by backend contract: module initializers
note: unknown: see current operation: "torch.initialize.global_slots"(%20, %21, %22, %23, %32, %33, %34, %35, %44, %45, %46, %47, %56, %57, %58, %59, %68, %69, %70, %71, %80, %81, %82, %83, %92, %93, %94, %95, %104, %105, %106, %107, %116, %117, %118, %119, %128, %129, %130, %131, %140, %141, %142, %143, %152, %153, %154, %155) <{slotSymNames = [@encoder.layers.encoder_layer_0.self_attention.in_proj_weight, @encoder.layers.encoder_layer_0.self_attention.in_proj_bias, @encoder.layers.encoder_layer_0.self_attention.out_proj.weight, @encoder.layers.encoder_layer_0.self_attention.out_proj.bias, @encoder.layers.encoder_layer_1.self_attention.in_proj_weight, @encoder.layers.encoder_layer_1.self_attention.in_proj_bias, @encoder.layers.encoder_layer_1.self_attention.out_proj.weight, @encoder.layers.encoder_layer_1.self_attention.out_proj.bias, @encoder.layers.encoder_layer_2.self_attention.in_proj_weight, @encoder.layers.encoder_layer_2.self_attention.in_proj_bias, @encoder.layers.encoder_layer_2.self_attention.out_proj.weight, @encoder.layers.encoder_layer_2.self_attention.out_proj.bias, @encoder.layers.encoder_layer_3.self_attention.in_proj_weight, @encoder.layers.encoder_layer_3.self_attention.in_proj_bias, @encoder.layers.encoder_layer_3.self_attention.out_proj.weight, @encoder.layers.encoder_layer_3.self_attention.out_proj.bias, @encoder.layers.encoder_layer_4.self_attention.in_proj_weight, @encoder.layers.encoder_layer_4.self_attention.in_proj_bias, @encoder.layers.encoder_layer_4.self_attention.out_proj.weight, @encoder.layers.encoder_layer_4.self_attention.out_proj.bias, @encoder.layers.encoder_layer_5.self_attention.in_proj_weight, @encoder.layers.encoder_layer_5.self_attention.in_proj_bias, @encoder.layers.encoder_layer_5.self_attention.out_proj.weight, @encoder.layers.encoder_layer_5.self_attention.out_proj.bias, @encoder.layers.encoder_layer_6.self_attention.in_proj_weight, @encoder.layers.encoder_layer_6.self_attention.in_proj_bias, @encoder.layers.encoder_layer_6.self_attention.out_proj.weight, @encoder.layers.encoder_layer_6.self_attention.out_proj.bias, @encoder.layers.encoder_layer_7.self_attention.in_proj_weight, @encoder.layers.encoder_layer_7.self_attention.in_proj_bias, @encoder.layers.encoder_layer_7.self_attention.out_proj.weight, @encoder.layers.encoder_layer_7.self_attention.out_proj.bias, @encoder.layers.encoder_layer_8.self_attention.in_proj_weight, @encoder.layers.encoder_layer_8.self_attention.in_proj_bias, @encoder.layers.encoder_layer_8.self_attention.out_proj.weight, @encoder.layers.encoder_layer_8.self_attention.out_proj.bias, @encoder.layers.encoder_layer_9.self_attention.in_proj_weight, @encoder.layers.encoder_layer_9.self_attention.in_proj_bias, @encoder.layers.encoder_layer_9.self_attention.out_proj.weight, @encoder.layers.encoder_layer_9.self_attention.out_proj.bias, @encoder.layers.encoder_layer_10.self_attention.in_proj_weight, @encoder.layers.encoder_layer_10.self_attention.in_proj_bias, @encoder.layers.encoder_layer_10.self_attention.out_proj.weight, @encoder.layers.encoder_layer_10.self_attention.out_proj.bias, @encoder.layers.encoder_layer_11.self_attention.in_proj_weight, @encoder.layers.encoder_layer_11.self_attention.in_proj_bias, @encoder.layers.encoder_layer_11.self_attention.out_proj.weight, @encoder.layers.encoder_layer_11.self_attention.out_proj.bias]}> : (!torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>, !torch.tensor<[2304,768],f32>, !torch.tensor<[2304],f32>, !torch.tensor<[768,768],f32>, !torch.tensor<[768],f32>) -> ()
note: unknown: this is likely due to InlineGlobalSlots being unable to inline a global slot

For Torch-MLIR developers, the error can be reproduced with:
$ torch-mlir-opt -pass-pipeline='builtin.module(torchscript-module-to-torch-backend-pipeline{backend-legal-ops=aten.flatten.using_ints extra-library=})' /tmp/VisionTransformer.mlir
Add '-mlir-print-ir-after-all -mlir-disable-threading' to get the IR dump for debugging purpose.

[followtheart]

related code:

torch-mlir/lib/Dialect/Torch/Transforms/LowerToBackendContract.cpp

Lines 175 to 195 in ae72eec

	static bool satisfiesBackendContract(ModuleOp module,
	const ConversionTarget &target,
	bool actuallyEmitDiagnostics = false) {
	// We do not permit `torch.global_slot`'s in the backend contract, since
	// support for them is not widespread, and this does not align with PyTorch's
	// more tracing-based direction.
	//
	// We just check for the GlobalSlotModuleInitializerOp since its verifier
	// ensures that the set of global slots matches those initialized by the
	// module initializer.
	auto walkResult0 = module.walk([&](Torch::GlobalSlotModuleInitializerOp op) {
	if (actuallyEmitDiagnostics) {
	// Report the error on the terminator to avoid dumping the whole
	// initializer itself, which can have pages of ops in it.
	op.getBody()
	->getTerminator()
	->emitError("unsupported by backend contract: module initializers")
	.attachNote()
	.append("this is likely due to InlineGlobalSlots being unable to "
	"inline a global slot");

torch-mlir/test/Dialect/Torch/lower-to-backend-contract-error.mlir

Lines 1 to 10 in ae72eec

	// RUN: torch-mlir-opt -torch-lower-to-backend-contract -split-input-file -verify-diagnostics %s
	torch.global_slot.module_initializer {
	%0 = torch.constant.int 1
	// expected-error @+2 {{unsupported by backend contract: module initializers}}
	// expected-note @+1 {{this is likely due to}}
	torch.initialize.global_slots [
	@slot0(%0 : !torch.int)
	]
	}

[maxbartel]

Same behavior with torchvision.models.mobilenet_v2()

[stellaraccident]

I'm not sure that we have project members who know this part very well anymore: it is part of the original TorchScript backend that we are trying to move off of.

@ramiro050 do you have any experience with this?

[ramiro050]

This is one of those cases where there is so much happening that could trip up torch-mlir that it is hard to know if the InlineGlobalSlots error is the issue or just a consequence of another simplification pass no working correctly.

From looking at the IR, there are a few ops that need support in torch-mlir:

torch.operator "aten.__contains__.str_list"
torch.operator "aten._native_multi_head_attention"
torch.operator "aten.is_autocast_enabled"
torch.operator "aten.is_grad_enabled"
torch.operator "aten.split.sizes"
torch.operator "aten.unflatten.int"
torch.operator "prim.is_nested"
torch.operator "prim.requires_grad"
torch.operator "prim.type"

Not having patterns for these ops could be getting in the way of simplifications that would allow inlining global slots.

Additionally, there is quite a bit of control flow in the IR. Torch-MLIR failing to simplify the control flow could be another reason for the error. One way to get around the control flow issue is to set use_tracing=True.

[followtheart]

This script may help.origin from https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py ,and fixed one line

import numpy as np
import torch
import torchvision.models as models
#import torch_mlir

import torch
from torch import nn

from einops import rearrange, repeat
from einops.layers.torch import Rearrange


#from https://github.com/lucidrains/vit-pytorch/blob/main/vit_pytorch/vit.py
# helpers

def pair(t):
    return t if isinstance(t, tuple) else (t, t)

# classes

class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout = 0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.norm = nn.LayerNorm(dim)

        self.attend = nn.Softmax(dim = -1)
        self.dropout = nn.Dropout(dropout)

        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x):
        x = self.norm(x)

        qkv = self.to_qkv(x).chunk(3, dim = -1)
        # q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
        q, k, v = [rearrange(t, 'b n (h d) -> b h n d', h=self.heads) for t in qkv]

        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale

        attn = self.attend(dots)
        attn = self.dropout(attn)

        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout),
                FeedForward(dim, mlp_dim, dropout = dropout)
            ]))

    def forward(self, x):
        for attn, ff in self.layers:
            x = attn(x) + x
            x = ff(x) + x

        return self.norm(x)

class ViT(nn.Module):
    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
        super().__init__()
        image_height, image_width = pair(image_size)
        patch_height, patch_width = pair(patch_size)

        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'

        num_patches = (image_height // patch_height) * (image_width // patch_width)
        patch_dim = channels * patch_height * patch_width
        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_height, p2 = patch_width),
            nn.LayerNorm(patch_dim),
            nn.Linear(patch_dim, dim),
            nn.LayerNorm(dim),
        )

        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

        self.pool = pool
        self.to_latent = nn.Identity()

        self.mlp_head = nn.Linear(dim, num_classes)

    def forward(self, img):
        x = self.to_patch_embedding(img)
        b, n, _ = x.shape

        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = b)
        x = torch.cat((cls_tokens, x), dim=1)
        x += self.pos_embedding[:, :(n + 1)]
        x = self.dropout(x)

        x = self.transformer(x)

        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

        x = self.to_latent(x)
        return self.mlp_head(x)

and

vit = ViT(
    image_size = 224,
    patch_size = 16,
    num_classes = 1000,
    dim = 768,
    depth = 12,
    heads = 12,
    mlp_dim = 3072,
    dropout = 0.1,
    emb_dropout = 0.1
)

img = torch.randn(1, 3, 224, 224)
vit.eval()
#lowering to linalg
vit_mod = torch_mlir.compile(vit, img, 
                            output_type=torch_mlir.OutputType.LINALG_ON_TENSORS, 
                            use_tracing=True)