support autogenerating transformer block in vision transformer

Makefile

@@ -65,6 +65,19 @@ build-docker:
 		docker pull $(img); \
 	fi
 
+build-docker-python13:
+	docker build --build-arg VHLS_PATH=$(vhls) --build-arg VHLS_VERSION=$(vhls_version) -f Docker/Dockerfile-$(PLATFORM)-python13 --tag mase-ubuntu2204-docker-python13 Docker; \
+
+shell-python13:
+	docker run -it --shm-size 256m \
+        --hostname mase-ubuntu2204-docker-python13 \
+        -w /workspace \
+        -v /$(USER_PREFIX)/$(shell whoami)/.gitconfig:/root/.gitconfig \
+        -v /$(USER_PREFIX)/$(shell whoami)/.ssh:/root/.ssh \
+        -v /$(USER_PREFIX)/$(shell whoami)/.mase:/root/.mase:z \
+        -v $(shell pwd):/workspace:z \
+        $(DOCKER_RUN_EXTRA_ARGS) \
+        $(img) /bin/bash
 shell:
 	docker run -it --shm-size 256m \
         --hostname mase-ubuntu2204 \

a_cx_mxint_quant/__init__.py

@@ -0,0 +1,86 @@
+from .module_level_tranform import vit_module_level_quantize
+from .quantizers import mxint_hardware, mxint_quant_block
+
+from .linear import MXIntLinear
+from .attention import MXIntAttention
+from .module_level_tranform import MXIntLayerNorm, MXIntGELU
+from .modules import MXIntPatchEmbed, MXIntAddition
+from mase_components import get_module_dependencies
+VIT_CUSTOM_OPS = {
+    "modules": {
+        MXIntPatchEmbed: {
+            "args": {
+                "data_in": "data_in",
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_patch_embed",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_patch_embed"
+            ),
+        },
+        MXIntAttention: {
+            "args": {
+                "data_in": "data_in",
+                "dim": "config",
+                "num_heads": "config",
+                "qkv_bias": "config",
+                "qk_norm": None,
+                "attn_drop": None,
+                "proj_drop": None,
+                "norm_layer": None,
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_vit_attention_wrap",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_vit_attention_wrap"
+            ),
+        },
+        MXIntLayerNorm: {
+            "args": {
+                "data_in": "data_in",
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_layernorm",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_layernorm"
+            ),
+        },
+        MXIntGELU: {
+            "args": {
+                "data_in": "data_in",
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_gelu",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_gelu"
+            ),
+        },
+        MXIntLinear: {
+            "args": {
+                "data_in": "data_in",
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_linear",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_linear"
+            ),
+        },
+        MXIntAddition: {
+            "args": {
+                "input_0": "data_in",
+                "input_1": "data_in",
+                "q_config": "config",
+            },
+            "toolchain": "INTERNAL_RTL",
+            "module": "mxint_addition",
+            "dependence_files": get_module_dependencies(
+                "linear_layers/mxint_operators/mxint_addition"
+            ),
+        },
+    },
+}

a_cx_mxint_quant/attention.py

@@ -0,0 +1,192 @@
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from torch.nn import functional as F
+
+from .attention_head import _ViTSelfAttentionHeadBase, ViTSelfAttentionHeadInteger
+
+from chop.nn.quantized.modules.linear import (
+    LinearInteger,
+)
+from chop.nn.quantized.functional import fixed_softermax
+from chop.nn.quantizers import integer_quantizer
+from chop.nn.quantized.functional import matmul_integer
+
+from typing import Optional, Tuple, Union
+
+from .linear import MXIntLinear
+from .attention_head import MXIntViTAttentionHead
+
+class _ViTAttentionBase(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.query = nn.Linear(dim, dim, bias=qkv_bias)
+        self.key = nn.Linear(dim, dim, bias=qkv_bias)
+        self.value = nn.Linear(dim, dim, bias=qkv_bias)
+        self.self_attention = _ViTSelfAttentionHeadBase(
+            dim=self.head_dim, num_heads=num_heads, attn_drop=attn_drop
+        )
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+
+        def _tensor_reshape(x):
+            return x.reshape(B, -1, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+
+        q, k, v = (
+            _tensor_reshape(self.query(x)),
+            _tensor_reshape(self.key(x)),
+            _tensor_reshape(self.value(x)),
+        )
+        x = self.self_attention(q, k, v)
+        x = x.transpose(1, 2).reshape(B, N, C)
+
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+class ViTAttentionInteger(_ViTAttentionBase):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        norm_layer: nn.Module = nn.LayerNorm,
+        q_config: dict = None,
+        floor=True,
+    ) -> None:
+        super().__init__(dim, num_heads, qkv_bias, qk_norm, attn_drop, proj_drop)
+        self.q_config = q_config
+        self.query = LinearInteger(
+            dim,
+            dim,
+            bias=qkv_bias,
+            config={
+                "data_in_width": q_config["data_in_width"],
+                "data_in_frac_width": q_config["data_in_frac_width"],
+                "weight_width": q_config["qkv_weight_width"],
+                "weight_frac_width": q_config["qkv_weight_frac_width"],
+                "bias_width": q_config["qkv_bias_width"],
+                "bias_frac_width": q_config["qkv_bias_frac_width"],
+            },
+            out_config={
+                "data_out_width": q_config["qkv_width"],
+                "data_out_frac_width": q_config["qkv_frac_width"],
+            },
+            floor=floor,
+        )
+        self.key = LinearInteger(
+            dim,
+            dim,
+            bias=qkv_bias,
+            config={
+                "data_in_width": q_config["data_in_width"],
+                "data_in_frac_width": q_config["data_in_frac_width"],
+                "weight_width": q_config["qkv_weight_width"],
+                "weight_frac_width": q_config["qkv_weight_frac_width"],
+                "bias_width": q_config["qkv_bias_width"],
+                "bias_frac_width": q_config["qkv_bias_frac_width"],
+            },
+            out_config={
+                "data_out_width": q_config["qkv_width"],
+                "data_out_frac_width": q_config["qkv_frac_width"],
+            },
+            floor=floor,
+        )
+        self.value = LinearInteger(
+            dim,
+            dim,
+            bias=qkv_bias,
+            config={
+                "data_in_width": q_config["data_in_width"],
+                "data_in_frac_width": q_config["data_in_frac_width"],
+                "weight_width": q_config["qkv_weight_width"],
+                "weight_frac_width": q_config["qkv_weight_frac_width"],
+                "bias_width": q_config["qkv_bias_width"],
+                "bias_frac_width": q_config["qkv_bias_frac_width"],
+            },
+            out_config={
+                "data_out_width": q_config["qkv_width"],
+                "data_out_frac_width": q_config["qkv_frac_width"],
+            },
+            floor=floor,
+        )
+        self.self_attention = ViTSelfAttentionHeadInteger(
+            dim=self.head_dim,
+            num_heads=num_heads,
+            attn_drop=attn_drop,
+            q_config={
+                "query_width": q_config["qkv_width"],
+                "query_frac_width": q_config["qkv_frac_width"],
+                "key_width": q_config["qkv_width"],
+                "key_frac_width": q_config["qkv_frac_width"],
+                "value_width": q_config["qkv_width"],
+                "value_frac_width": q_config["qkv_frac_width"],
+                "qkmm_out_width": q_config["qkmm_out_width"],
+                "qkmm_out_frac_width": q_config["qkmm_out_frac_width"],
+                "softmax_exp_width": q_config["softmax_exp_width"],
+                "softmax_exp_frac_width": q_config["softmax_exp_frac_width"],
+                "softmax_out_frac_width": q_config["softmax_out_frac_width"],
+                "svmm_out_width": q_config["svmm_out_width"],
+                "svmm_out_frac_width": q_config["svmm_out_frac_width"],
+            },
+            floor=floor,
+        )
+        self.proj = LinearInteger(
+            dim,
+            dim,
+            config={
+                "data_in_width": q_config["svmm_out_width"],
+                "data_in_frac_width": q_config["svmm_out_frac_width"],
+                "weight_width": q_config["proj_weight_width"],
+                "weight_frac_width": q_config["proj_weight_frac_width"],
+                "bias_width": q_config["proj_bias_width"],
+                "bias_frac_width": q_config["proj_bias_frac_width"],
+            },
+            out_config={
+                "data_out_width": q_config["data_out_width"],
+                "data_out_frac_width": q_config["data_out_frac_width"],
+            },
+            floor=floor,
+        )
+
+class MXIntAttention(_ViTAttentionBase):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_norm: bool = False,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+        q_config: dict = None,
+    ) -> None:
+        super().__init__(dim, num_heads, qkv_bias, qk_norm, attn_drop, proj_drop)
+        self.q_config = q_config
+
+        # Replace attention with MXIntViTAttentionHead
+        # self.self_attention = MXIntViTAttentionHead(
+        #     dim=self.head_dim,
+        #     num_heads=num_heads,
+        #     attn_drop=attn_drop,
+        #     q_config=q_config
+        # )