[CI][Examples] Enable testing of examples through lit, on both x86 and Arm (#22)

Adds lit as a test dependency via pyproject.toml. Adds FileCheck as an
external dependency -- for CI we use Ubuntu's llvm-dev package. Makes CI
run for both x86_64 and AArch64.

Author: rolfmorel

.github/workflows/examples.yml

@@ -3,11 +3,17 @@ name: Lighthouse Examples
 on:
   workflow_dispatch:
   push:
+    branches:
+      - 'main'
   pull_request:
 
 jobs:
   Examples:
-    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        arch: [x86_64, AArch64]
+
+    runs-on: ${{ matrix.arch == 'AArch64' && 'ubuntu-24.04-arm' || 'ubuntu-24.04' }}
     steps:
       - uses: actions/checkout@v5
 
@@ -16,21 +22,10 @@ jobs:
 
       - name: Install the project and its dependencies
         run: |
-          uv sync
+          sudo apt-get install -y llvm-dev # Obtain FileCheck, used in testing.
           uv sync --extra ingress-torch-cpu
 
-      - name: Run MLP From File
-        run: |-
-          uv run python/examples/ingress/torch/mlp_from_file.py
-
-      - name: Run MLP From Module
-        run: |-
-          uv run python/examples/ingress/torch/mlp_from_model.py
-
-      - name: Run Compile And Run
-        run: |-
-          uv run python/examples/mlir/compile_and_run.py
-
-      - name: Run apply basic schedule to basic payload
-        run: |-
-          uv run python/examples/schedule/transform_a_payload_according_to_a_schedule.py
+      - name: Run lit-enabled examples as tests
+        run: |
+          export FILECHECK=FileCheck-18 # Ubuntu's llvm-dev appends a version number.
+          uv run lit python/examples  # Makes sure to substitute FileCheck for $FILECHECK

README.md

@@ -117,3 +117,17 @@ pip install .[ingress_torch_cpu] \
   --extra-index-url https://download.pytorch.org/whl \
   --only-binary :all:
 ```
+
+## Running tests
+
+Running the tests is as simple as `lit .` in the root of the project.
+
+We assume that the [`FileCheck`](https://llvm.org/docs/CommandGuide/FileCheck.html) and [`lit`](https://llvm.org/docs/CommandGuide/lit.html) executables are available on the `PATH`.
+
+<details>
+<summary>
+Obtaining <code>FileCheck</code> and <code>lit</code>.
+</summary>
+To obtain the <a href="https://pypi.org/project/lit">Python package for <code>lit</code></a>, simply run <code>uv sync</code> (<code>lit</code> is included in the "dev" dependency group).
+In case the <code>FileCheck</code> executable happens to be available under a different name/location, e.g. as <code>FileCheck-18</code> from Ubuntu's <code>llvm-dev</code> package, set the <code>FILECHECK</code> environment variable when invoking <code>lit</code>.
+</details>

lit.cfg.py

@@ -0,0 +1,16 @@
+import os
+
+import lit.formats
+from lit.TestingConfig import TestingConfig
+
+# Imagine that, all your variables defined and with types!
+assert isinstance(config := eval("config"), TestingConfig)
+
+config.name = "Lighthouse test suite"
+config.test_format = lit.formats.ShTest(True)
+config.test_source_root = os.path.dirname(__file__)
+config.test_exec_root = os.path.dirname(__file__) + "/lit.out"
+
+config.substitutions.append(("%PYTHON", "uv run"))
+if filecheck_path := os.environ.get("FILECHECK"):
+    config.substitutions.append(("FileCheck", filecheck_path))

pyproject.toml

@@ -6,6 +6,11 @@ dependencies = [
     "mlir-python-bindings==20251118+99630eb1b"
 ]
 
+[dependency-groups]
+dev = [
+  "lit==18.1.8"  # Tool to configure, discover and run tests
+]
+
 [project.optional-dependencies]
 ingress_torch_mlir = [
     "torch-mlir==20251122.639"
@@ -66,6 +71,7 @@ name = "torch_mlir"
 url = "https://github.com/llvm/torch-mlir-release/releases/expanded_assets/dev-wheels"
 explicit = true
 
+# Derive package version from version in code
 [build-system]
 requires = ["setuptools>=68", "wheel"]
 build-backend = "setuptools.build_meta"

python/examples/ingress/torch/lit.local.cfg

@@ -0,0 +1 @@
+config.excludes = ["MLPModel"]

python/examples/ingress/torch/mlp_from_file.py

@@ -1,3 +1,5 @@
+# RUN: %PYTHON %s
+
 """
 Example demonstrating how to load a PyTorch model to MLIR using Lighthouse
 without initializing the model class on the user's side.

python/examples/ingress/torch/mlp_from_model.py

@@ -1,3 +1,5 @@
+# RUN: %PYTHON %s
+
 """
 Example demonstrating how to load an already initialized PyTorch model
 to MLIR using Lighthouse.

python/examples/lit.local.cfg

@@ -0,0 +1 @@
+config.suffixes = {'.py'}

python/examples/mlir/compile_and_run.py

@@ -1,3 +1,5 @@
+# RUN: %PYTHON %s
+
 import torch
 import argparse
 

python/examples/schedule/transform_a_payload_according_to_a_schedule.py

@@ -1,3 +1,5 @@
+# RUN: %PYTHON %s | FileCheck %s
+
 # Simply demonstrates applying a schedule to a payload.
 # To do so generates a basic payload and a basic schedule, purely as an example.
 
@@ -24,16 +26,23 @@ def example_payload() -> Module:
     with InsertionPoint(payload.body):
         matrixType = RankedTensorType.get([16, 16], F32Type.get())
 
+        # NB: Do the CHECKing on the transformed output:
+        # CHECK-LABEL: result of applying schedule to payload
+        # CHECK: func.func @fold_add_on_two_matmuls
+        # CHECK-SAME:      (%[[MATRIX_A:.*]]: {{.*}}, %[[MATRIX_B:.*]]: {{.*}}, %[[WEIGHTS:.*]]: {{.*}})
         @func.func(matrixType, matrixType, matrixType)
         def fold_add_on_two_matmuls(matrixA, matrixB, weights):
             empty = tensor.empty(matrixType.shape, matrixType.element_type)
             c0 = arith.constant(F32Type.get(), 0.0)
+            # CHECK: %[[ZERO_INIT:.*]] = linalg.fill
             zero_init = linalg.fill(c0, outs=[empty])
+            # CHECK: %[[A_X_WEIGHTS:.*]] = linalg.matmul ins(%[[MATRIX_A]], %[[WEIGHTS]]{{.*}}) outs(%[[ZERO_INIT]]
             A_x_weights = linalg.matmul(matrixA, weights, outs=[zero_init])
-            empty2 = tensor.empty(matrixType.shape, matrixType.element_type)
-            zero_init2 = linalg.fill(c0, outs=[empty2])
-            B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init2])
+            # CHECK: %[[RES:.*]] = linalg.matmul ins(%[[MATRIX_B]], %[[WEIGHTS]]{{.*}}) outs(%[[A_X_WEIGHTS]]
+            B_x_weights = linalg.matmul(matrixB, weights, outs=[zero_init])
+            # CHECK-NOT: linalg.add
             added = linalg.add(A_x_weights, B_x_weights, outs=[empty])
+            # CHECK: return %[[RES]]
             return added
 
     print(payload)