[mlir] KernelBench level 3 kernels

backends/mlir/cpu/KernelBench/level3/13_DenseNet121TransitionLayer.py

@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(self, num_input_features: int, num_output_features: int):
+        """
+        :param num_input_features: The number of input feature maps
+        :param num_output_features: The number of output feature maps
+        """
+        super(Model, self).__init__()
+        self.transition = nn.Sequential(
+            nn.BatchNorm2d(num_input_features),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(
+                num_input_features, num_output_features, kernel_size=1, bias=False
+            ),
+            nn.AvgPool2d(kernel_size=2, stride=2),
+        )
+
+    def forward(self, x):
+        """
+        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
+        :return: Downsampled tensor with reduced number of feature maps
+        """
+        return self.transition(x)

backends/mlir/cpu/KernelBench/level3/14_DenseNet121DenseBlock.py

@@ -0,0 +1,46 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(self, num_layers: int, num_input_features: int, growth_rate: int):
+        """
+        :param num_layers: The number of layers in the dense block
+        :param num_input_features: The number of input feature maps
+        :param growth_rate: The growth rate for the dense block (new features added per layer)
+        """
+        super(Model, self).__init__()
+        layers = []
+        for i in range(num_layers):
+            layers.append(
+                self._make_layer(num_input_features + i * growth_rate, growth_rate)
+            )
+        self.layers = nn.ModuleList(layers)
+
+    def _make_layer(self, in_features: int, growth_rate: int):
+        """
+        Creates a single layer with BatchNorm, ReLU, Conv2D, and Dropout.
+        """
+        return nn.Sequential(
+            nn.BatchNorm2d(in_features),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(in_features, growth_rate, kernel_size=3, padding=1, bias=False),
+            nn.Dropout(0.0),
+        )
+
+    def forward(self, x):
+        """
+        :param x: Input tensor of shape (batch_size, num_input_features, height, width)
+        :return: Concatenated output tensor with shape (batch_size, num_output_features, height, width)
+        """
+        features = [x]
+        for layer in self.layers:
+            new_feature = layer(x)
+            features.append(new_feature)
+            x = torch.cat(features, 1)  # Concatenate along channel axis
+        return x

backends/mlir/cpu/KernelBench/level3/17_SqueezeNetFireModule.py

@@ -0,0 +1,45 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(
+        self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels
+    ):
+        """
+        :param in_channels: Number of input channels
+        :param squeeze_channels: Number of output channels for the squeeze layer
+        :param expand1x1_channels: Number of output channels for the 1x1 expand layer
+        :param expand3x3_channels: Number of output channels for the 3x3 expand layer
+        """
+        super(Model, self).__init__()
+
+        self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
+        self.squeeze_activation = nn.ReLU(inplace=True)
+
+        self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
+        self.expand1x1_activation = nn.ReLU(inplace=True)
+
+        self.expand3x3 = nn.Conv2d(
+            squeeze_channels, expand3x3_channels, kernel_size=3, padding=1
+        )
+        self.expand3x3_activation = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """
+        :param x: Input tensor, shape (batch_size, in_channels, height, width)
+        :return: Output tensor, shape (batch_size, expand1x1_channels + expand3x3_channels, height, width)
+        """
+        x = self.squeeze_activation(self.squeeze(x))
+        return torch.cat(
+            [
+                self.expand1x1_activation(self.expand1x1(x)),
+                self.expand3x3_activation(self.expand3x3(x)),
+            ],
+            1,
+        )

backends/mlir/cpu/KernelBench/level3/18_SqueezeNet.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+class FireModule(nn.Module):
+    def __init__(
+        self, in_channels, squeeze_channels, expand1x1_channels, expand3x3_channels
+    ):
+        """
+        :param in_channels: Number of input channels
+        :param squeeze_channels: Number of output channels for the squeeze layer
+        :param expand1x1_channels: Number of output channels for the 1x1 expand layer
+        :param expand3x3_channels: Number of output channels for the 3x3 expand layer
+        """
+        super(FireModule, self).__init__()
+
+        self.squeeze = nn.Conv2d(in_channels, squeeze_channels, kernel_size=1)
+        self.squeeze_activation = nn.ReLU(inplace=True)
+
+        self.expand1x1 = nn.Conv2d(squeeze_channels, expand1x1_channels, kernel_size=1)
+        self.expand1x1_activation = nn.ReLU(inplace=True)
+
+        self.expand3x3 = nn.Conv2d(
+            squeeze_channels, expand3x3_channels, kernel_size=3, padding=1
+        )
+        self.expand3x3_activation = nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        """
+        :param x: Input tensor, shape (batch_size, in_channels, height, width)
+        :return: Output tensor, shape (batch_size, expand1x1_channels + expand3x3_channels, height, width)
+        """
+        x = self.squeeze_activation(self.squeeze(x))
+        return torch.cat(
+            [
+                self.expand1x1_activation(self.expand1x1(x)),
+                self.expand3x3_activation(self.expand3x3(x)),
+            ],
+            1,
+        )
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(self, num_classes=1000):
+        """
+        :param num_classes: Number of output classes
+        """
+        super(Model, self).__init__()
+
+        self.features = nn.Sequential(
+            nn.Conv2d(3, 96, kernel_size=7, stride=2),
+            nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+            FireModule(96, 16, 64, 64),
+            FireModule(128, 16, 64, 64),
+            FireModule(128, 32, 128, 128),
+            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+            FireModule(256, 32, 128, 128),
+            FireModule(256, 48, 192, 192),
+            FireModule(384, 48, 192, 192),
+            FireModule(384, 64, 256, 256),
+            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True),
+            FireModule(512, 64, 256, 256),
+        )
+
+        self.classifier = nn.Sequential(
+            nn.Dropout(p=0.0),
+            nn.Conv2d(512, num_classes, kernel_size=1),
+            nn.ReLU(inplace=True),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+
+    def forward(self, x):
+        """
+        :param x: Input tensor, shape (batch_size, 3, height, width)
+        :return: Output tensor, shape (batch_size, num_classes)
+        """
+        x = self.features(x)
+        x = self.classifier(x)
+        return torch.flatten(x, 1)

backends/mlir/cpu/KernelBench/level3/1_MLP.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(self, input_size, layer_sizes, output_size):
+        """
+        :param input_size: The number of input features
+        :param layer_sizes: A list of ints containing the sizes of each hidden layer
+        :param output_size: The number of output features
+        """
+        super(Model, self).__init__()
+
+        layers = []
+        current_input_size = input_size
+
+        for layer_size in layer_sizes:
+            layers.append(nn.Linear(current_input_size, layer_size))
+            layers.append(nn.ReLU())
+            current_input_size = layer_size
+
+        layers.append(nn.Linear(current_input_size, output_size))
+
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, x):
+        """
+        :param x: The input tensor, shape (batch_size, input_size)
+        :return: The output tensor, shape (batch_size, output_size)
+        """
+        return self.network(x)

backends/mlir/cpu/KernelBench/level3/21_EfficientNetMBConv.py

@@ -0,0 +1,79 @@
+import torch
+import torch.nn as nn
+
+import ai_bench.mlir
+
+
+@torch.compile(
+    dynamic=False, backend=ai_bench.mlir.cpu_backend(ai_bench.mlir.cpu_pipeline)
+)
+class Model(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, expand_ratio):
+        """
+        MBConv block implementation.
+
+        :param in_channels: Number of input channels.
+        :param out_channels: Number of output channels.
+        :param kernel_size: Kernel size for the depthwise convolution.
+        :param stride: Stride for the depthwise convolution.
+        :param expand_ratio: Expansion ratio for the intermediate channels.
+        """
+        super(Model, self).__init__()
+
+        self.use_residual = stride == 1 and in_channels == out_channels
+        hidden_dim = in_channels * expand_ratio
+
+        if expand_ratio != 1:
+            self.expand_conv = nn.Sequential(
+                nn.Conv2d(
+                    in_channels,
+                    hidden_dim,
+                    kernel_size=1,
+                    stride=1,
+                    padding=0,
+                    bias=False,
+                ),
+                nn.BatchNorm2d(hidden_dim),
+                nn.ReLU6(inplace=True),
+            )
+
+        self.depthwise_conv = nn.Sequential(
+            nn.Conv2d(
+                hidden_dim,
+                hidden_dim,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=(kernel_size - 1) // 2,
+                groups=hidden_dim,
+                bias=False,
+            ),
+            nn.BatchNorm2d(hidden_dim),
+            nn.ReLU6(inplace=True),
+        )
+
+        self.project_conv = nn.Sequential(
+            nn.Conv2d(
+                hidden_dim, out_channels, kernel_size=1, stride=1, padding=0, bias=False
+            ),
+            nn.BatchNorm2d(out_channels),
+        )
+
+    def forward(self, x):
+        """
+        Forward pass of the MBConv block.
+
+        :param x: The input tensor, shape (batch_size, in_channels, H, W)
+        :return: The output tensor, shape (batch_size, out_channels, H', W')
+        """
+        identity = x
+
+        if hasattr(self, "expand_conv"):
+            x = self.expand_conv(x)
+
+        x = self.depthwise_conv(x)
+        x = self.project_conv(x)
+
+        if self.use_residual:
+            x += identity
+
+        return x