[mlir] Default CPU lowering pipeline

ai_bench/mlir/__init__.py

@@ -1,5 +1,7 @@
 from .compile import cpu_backend
+from .pipeline import cpu_pipeline
 
 __all__ = [
     "cpu_backend",
+    "cpu_pipeline",
 ]

ai_bench/mlir/pipeline.py

@@ -0,0 +1,57 @@
+from mlir import ir
+from mlir.passmanager import PassManager
+
+
+def cpu_pipeline(module: ir.Module) -> ir.Module:
+    """
+    The default lowering pipeline for CPU.
+    Lowers MLIR ops within the module to MLIR LLVM IR dialect.
+
+    The pipeline focuses on enabling end-to-end lowering for various
+    generic kernel modules.
+
+    Performance is currently secondary and not representative.
+
+    Args:
+        module: MLIR module coming from PyTorch importer.
+    Returns:
+        MLIR module with lowered IR.
+    """
+
+    # Use standard C interface wrappers for functions.
+    pm = PassManager("builtin.module", module.context)
+    pm.add("func.func(llvm-request-c-wrappers)")
+
+    # Bufferize.
+    pm.add("eliminate-empty-tensors")
+    pm.add(
+        "one-shot-bufferize{function-boundary-type-conversion=identity-layout-map bufferize-function-boundaries}"
+    )
+    pm.add("drop-equivalent-buffer-results")
+    pm.add("buffer-deallocation-pipeline")
+    pm.add("convert-bufferization-to-memref")
+    pm.add("cse")
+    pm.add("canonicalize")
+
+    # Lower to LLVM.
+    pm.add("convert-linalg-to-loops")
+    pm.add("math-expand-ops")
+    pm.add("expand-strided-metadata")
+    pm.add("canonicalize")
+
+    pm.add("convert-vector-to-scf")
+    pm.add("lower-affine")
+    pm.add("convert-scf-to-cf")
+    pm.add("convert-vector-to-llvm")
+    pm.add("convert-math-to-libm")
+    pm.add("convert-to-llvm")
+    pm.add("reconcile-unrealized-casts")
+
+    # Cleanup
+    pm.add("cse")
+    pm.add("canonicalize")
+
+    # IR is transformed in-place.
+    pm.run(module.operation)
+
+    return module

backends/mlir/cpu/KernelBench/level1/100_HingeLoss.py

@@ -0,0 +1,22 @@
+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):
+    """
+    A model that computes Hinge Loss for binary classification tasks.
+
+    Parameters:
+        None
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, predictions, targets):
+        return torch.mean(torch.clamp(1 - predictions * targets, min=0))

backends/mlir/cpu/KernelBench/level1/10_3D_tensor_matrix_multiplication.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Performs 3D tensor-matrix multiplication.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A, B):
+        """
+        Performs 3D tensor-matrix multiplication.
+
+        Args:
+            A (torch.Tensor): Input 3D tensor of shape (N, M, K).
+            B (torch.Tensor): Input matrix of shape (K, L).
+
+        Returns:
+            torch.Tensor: Output tensor of shape (N, M, L), resulting from the multiplication of A and B along the last dimension of A.
+        """
+        return torch.matmul(A, B)

backends/mlir/cpu/KernelBench/level1/11_4D_tensor_matrix_multiplication.py

@@ -0,0 +1,37 @@
+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):
+    """
+    Performs 4D tensor-matrix multiplication:
+        C[b, i, j, k] = sum_l A[b, i, j, l] * B[l, k]
+
+    Args:
+        A (torch.Tensor): Input 4D tensor of shape (b, i, j, l)
+        B (torch.Tensor): Input matrix of shape (l, k)
+
+    Returns:
+        torch.Tensor: Output 4D tensor of shape (b, i, j, k)
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A, B):
+        """
+        Performs the 4D tensor-matrix multiplication.
+
+        Args:
+            A (torch.Tensor): Input 4D tensor of shape (b, i, j, l)
+            B (torch.Tensor): Input matrix of shape (l, k)
+
+        Returns:
+            torch.Tensor: Output 4D tensor of shape (b, i, j, k)
+        """
+        return torch.einsum("bijl,lk->bijk", A, B)

backends/mlir/cpu/KernelBench/level1/13_Matmul_for_symmetric_matrices.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs a single matrix multiplication (C = A * B) with A and B being symmetric matrices.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A, B):
+        """
+        Performs matrix multiplication of two symmetric matrices.
+
+        Args:
+            A (torch.Tensor): Input matrix A, shape (N, N), symmetric.
+            B (torch.Tensor): Input matrix B, shape (N, N), symmetric.
+
+        Returns:
+            torch.Tensor: Output matrix C, shape (N, N).
+        """
+        return torch.matmul(A, B)

backends/mlir/cpu/KernelBench/level1/14_Matmul_for_upper_triangular_matrices.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs matrix multiplication (C = A * B) for upper triangular matrices.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A, B):
+        """
+        Performs matrix multiplication for upper triangular matrices.
+
+        Args:
+            A (torch.Tensor): Upper triangular matrix of shape (N, N).
+            B (torch.Tensor): Upper triangular matrix of shape (N, N).
+
+        Returns:
+            torch.Tensor: The product of A and B, also an upper triangular matrix of shape (N, N).
+        """
+        return torch.triu(torch.matmul(A, B))

backends/mlir/cpu/KernelBench/level1/15_Matmul_for_lower_triangular_matrices.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs a matrix multiplication (C = A * B) where A and B are lower triangular matrices.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A, B):
+        """
+        Performs matrix multiplication of lower triangular matrices A and B.
+
+        Args:
+            A (torch.Tensor): Lower triangular matrix of shape (N, N).
+            B (torch.Tensor): Lower triangular matrix of shape (N, N).
+
+        Returns:
+            torch.Tensor: The result of matrix multiplication C of shape (N, N).
+        """
+        return torch.tril(torch.matmul(A, B))

backends/mlir/cpu/KernelBench/level1/16_Matmul_with_transposed_A.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs a single matrix multiplication (C = A * B)
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+        """
+        Performs matrix multiplication.
+
+        Args:
+            A: Input tensor of shape (M, K).
+            B: Input tensor of shape (K, N).
+
+        Returns:
+            Output tensor of shape (M, N).
+        """
+        return torch.matmul(A.T, B)

backends/mlir/cpu/KernelBench/level1/17_Matmul_with_transposed_B.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs a single matrix multiplication (C = A * B)
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+        """
+        Performs matrix multiplication.
+
+        Args:
+            A: Input tensor of shape (M, K).
+            B: Input tensor of shape (K, N).
+
+        Returns:
+            Output tensor of shape (M, N).
+        """
+        return torch.matmul(A, B.T)

backends/mlir/cpu/KernelBench/level1/18_Matmul_with_transposed_both.py

@@ -0,0 +1,29 @@
+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):
+    """
+    Simple model that performs a single matrix multiplication (C = A * B)
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
+        """
+        Performs matrix multiplication.
+
+        Args:
+            A: Input tensor of shape (M, K).
+            B: Input tensor of shape (K, N).
+
+        Returns:
+            Output tensor of shape (M, N).
+        """
+        return torch.matmul(A.T, B.T)

backends/mlir/cpu/KernelBench/level1/19_ReLU.py

@@ -0,0 +1,28 @@
+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):
+    """
+    Simple model that performs a ReLU activation.
+    """
+
+    def __init__(self):
+        super(Model, self).__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Applies ReLU activation to the input tensor.
+
+        Args:
+            x (torch.Tensor): Input tensor of any shape.
+
+        Returns:
+            torch.Tensor: Output tensor with ReLU applied, same shape as input.
+        """
+        return torch.relu(x)