diff --git a/.github/workflows/examples.yml b/.github/workflows/examples.yml index 991d947..136b848 100644 --- a/.github/workflows/examples.yml +++ b/.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 - 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 + sudo apt-get install -y llvm-dev # Obtain FileCheck, used in testing. + uv sync --extra ingress-torch-cpu --group test - - 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 diff --git a/README.md b/README.md index d09a4f8..41bfe18 100644 --- a/README.md +++ b/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`. + +
+ +Obtaining FileCheck and lit. + +To obtain the Python package for lit, simply run uv sync --group test. +In case the FileCheck executable happens to be available under a different name/location, e.g. as FileCheck-18 from Ubuntu's llvm-dev package, set the FILECHECK environment variable when invoking lit. +
diff --git a/lit.cfg.py b/lit.cfg.py new file mode 100644 index 0000000..e8b7547 --- /dev/null +++ b/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)) diff --git a/pyproject.toml b/pyproject.toml index e559632..912a34d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,11 @@ dependencies = [ "mlir-python-bindings==20251118+99630eb1b" ] +[dependency-groups] +test = [ + "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" diff --git a/python/examples/ingress/torch/lit.local.cfg b/python/examples/ingress/torch/lit.local.cfg new file mode 100644 index 0000000..2813a1b --- /dev/null +++ b/python/examples/ingress/torch/lit.local.cfg @@ -0,0 +1 @@ +config.excludes = ["MLPModel"] diff --git a/python/examples/ingress/torch/mlp_from_file.py b/python/examples/ingress/torch/mlp_from_file.py index 01edeec..b8afc36 100644 --- a/python/examples/ingress/torch/mlp_from_file.py +++ b/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. diff --git a/python/examples/ingress/torch/mlp_from_model.py b/python/examples/ingress/torch/mlp_from_model.py index b51af80..6621e90 100644 --- a/python/examples/ingress/torch/mlp_from_model.py +++ b/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. diff --git a/python/examples/lit.local.cfg b/python/examples/lit.local.cfg new file mode 100644 index 0000000..73171b0 --- /dev/null +++ b/python/examples/lit.local.cfg @@ -0,0 +1 @@ +config.suffixes = {'.py'} diff --git a/python/examples/mlir/compile_and_run.py b/python/examples/mlir/compile_and_run.py index 009d9d5..0d8b8f1 100644 --- a/python/examples/mlir/compile_and_run.py +++ b/python/examples/mlir/compile_and_run.py @@ -1,3 +1,5 @@ +# RUN: %PYTHON %s + import torch import argparse diff --git a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py b/python/examples/schedule/transform_a_payload_according_to_a_schedule.py index 12702ab..fcbbe7b 100644 --- a/python/examples/schedule/transform_a_payload_according_to_a_schedule.py +++ b/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)