Merge branch 'main' into kim/stim/stim-printer-parser

.github/workflows/ci-core.yml

@@ -22,12 +22,6 @@ jobs:
     steps:
     - uses: actions/checkout@v4
 
-    - name: Install native dependencies
-      uses: awalsh128/cache-apt-pkgs-action@latest
-      with:
-        packages: mesa-vulkan-drivers
-        version: 1.0
-
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
@@ -41,15 +35,9 @@ jobs:
       run: |
         pip install --upgrade pip
 
-    - name: Install the package locally
-      run: pip install -r requirements.txt
-
-    - name: Test with pytest
-      run: |
-        pytest -W error
+    - name: Install the package and dev dependencies locally
+      run: pip install -e ".[dev]"
 
-    - name: Execute lit tests
+    - name: Run all tests
       run: |
-        export PYTHONPATH=$(pwd)
-        lit -v tests/filecheck/
-        lit -v docs/Toy/examples/
+        make tests-functional

.github/workflows/ci-mlir.yml

@@ -20,7 +20,7 @@ jobs:
 
     env:
       LLVM_SYMBOLIZER_PATH: /usr/lib/llvm-11/bin/llvm-symbolizer
-      MLIR-Version: 98e674c9f16d677d95c67bc130e267fae331e43c
+      MLIR-Version: d401987fe349a87c53fe25829215b080b70c0c1a
     steps:
     - uses: actions/checkout@v4
       with:
@@ -116,12 +116,12 @@ jobs:
         export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/
         pytest --nbval docs/mlir_interoperation.ipynb --maxfail 1 -vv
 
-    - name: Test MLIR dependent marimo notebooks
+    - name: Test ONNX-dependent marimo notebooks
       run: |
         cd xdsl
         # Add mlir-opt to the path
         export PATH=$PATH:${GITHUB_WORKSPACE}/llvm-project/build/bin/
-        make tests-marimo-mlir
+        make tests-marimo-onnx
 
     - name: Combine coverage data
       run: |

Makefile

@@ -16,7 +16,7 @@ TESTS_COVERAGE_FILE = ${COVERAGE_FILE}.tests
 # these targets don't produce files:
 .PHONY: ${VENV_DIR}/ venv clean clean-caches filecheck pytest pytest-nb tests-toy tests
 .PHONY: rerun-notebooks precommit-install precommit pyright tests-marimo
-.PHONY: tests-marimo-mlir coverage coverage-tests coverage-filecheck-tests
+.PHONY: coverage coverage-tests coverage-filecheck-tests
 .PHONY: coverage-report-html coverage-report-md
 
 # set up the venv with all dependencies for development
@@ -49,6 +49,7 @@ pytest:
 pytest-nb:
 	pytest -W error --nbval -vv docs \
 		--ignore=docs/mlir_interoperation.ipynb \
+		--ignore=docs/Toy \
 		--nbval-current-env
 
 # run tests for Toy tutorial
@@ -58,7 +59,15 @@ filecheck-toy:
 pytest-toy:
 	pytest docs/Toy/toy/tests
 
-tests-toy: filecheck-toy pytest-toy
+.PHONY: pytest-toy-nb
+pytest-toy-nb:
+	@if python -c "import riscemu" > /dev/null 2>&1; then \
+		pytest -W error --nbval -vv docs/Toy --nbval-current-env; \
+	else \
+		echo "riscemu is not installed, skipping tests."; \
+	fi
+
+tests-toy: filecheck-toy pytest-toy pytest-toy-nb
 
 tests-marimo:
 	@for file in docs/marimo/*.py; do \
@@ -71,24 +80,33 @@ tests-marimo:
 	done
 	@echo "All marimo tests passed successfully."
 
-tests-marimo-mlir:
-	@if ! command -v mlir-opt > /dev/null 2>&1; then \
-		echo "MLIR is not installed, skipping tests."; \
-		exit 0; \
+.PHONY: tests-marimo-onnx
+tests-marimo-onnx:
+	@if python -c "import onnx" > /dev/null 2>&1; then \
+		echo "onnx is installed, running tests."; \
+		if ! command -v mlir-opt > /dev/null 2>&1; then \
+			echo "MLIR is not installed, skipping tests."; \
+			exit 0; \
+		fi; \
+		for file in docs/marimo/onnx/*.py; do \
+			echo "Running $$file"; \
+			error_message=$$(python3 "$$file" 2>&1) || { \
+				echo "Error running $$file"; \
+				echo "$$error_message"; \
+				exit 1; \
+			}; \
+		done; \
+		echo "All marimo onnx tests passed successfully."; \
+	else \
+		echo "onnx is not installed, skipping tests."; \
 	fi
-	@echo "MLIR is installed, running tests."
-	@for file in docs/marimo/mlir/*.py; do \
-		echo "Running $$file"; \
-		error_message=$$(python3 "$$file" 2>&1) || { \
-			echo "Error running $$file"; \
-			echo "$$error_message"; \
-			exit 1; \
-		}; \
-	done
-	@echo "All marimo mlir tests passed successfully."
 
 # run all tests
-tests: pytest tests-toy filecheck pytest-nb tests-marimo tests-marimo-mlir pyright
+tests-functional: pytest tests-toy filecheck pytest-nb tests-marimo tests-marimo-onnx
+	@echo All functional tests done.
+
+# run all tests
+tests: tests-functional pyright
 	@echo All tests done.
 
 # re-generate the output from all jupyter notebooks in the docs directory

README.md

@@ -44,7 +44,7 @@ pip install xdsl
 
 *Note:* This version of xDSL is validated against a specific MLIR version,
 interoperability with other versions may result in problems. The supported
-MLIR version is commit `98e674c9f16d677d95c67bc130e267fae331e43c`.
+MLIR version is commit `d401987fe349a87c53fe25829215b080b70c0c1a`.
 
 ### Subprojects With Extra Dependencies
 

docs/Toy/toy/dialects/toy.py

@@ -17,7 +17,16 @@
     UnrankedTensorType,
     f64,
 )
-from xdsl.ir import Attribute, Block, Dialect, Operation, OpResult, Region, SSAValue
+from xdsl.ir import (
+    Attribute,
+    Block,
+    Dialect,
+    Operation,
+    OpResult,
+    OpTraits,
+    Region,
+    SSAValue,
+)
 from xdsl.irdl import (
     IRDLOperation,
     attr_def,
@@ -28,6 +37,7 @@
     opt_operand_def,
     region_def,
     result_def,
+    traits_def,
     var_operand_def,
     var_result_def,
 )
@@ -77,7 +87,7 @@ class ConstantOp(IRDLOperation):
     value = attr_def(DenseIntOrFPElementsAttr)
     res = result_def(TensorTypeF64)
 
-    traits = frozenset((Pure(),))
+    traits = traits_def(Pure())
 
     def __init__(self, value: DenseIntOrFPElementsAttr):
         super().__init__(result_types=[value.type], attributes={"value": value})
@@ -138,7 +148,7 @@ class AddOp(IRDLOperation):
     rhs = operand_def(AnyTensorTypeF64Constr)
     res = result_def(AnyTensorTypeF64Constr)
 
-    traits = frozenset((Pure(), InferAddOpShapeTrait()))
+    traits = traits_def(Pure(), InferAddOpShapeTrait())
 
     def __init__(self, lhs: SSAValue, rhs: SSAValue):
         if isa(lhs.type, TensorTypeF64):
@@ -204,7 +214,7 @@ class FuncOp(IRDLOperation):
     function_type = attr_def(FunctionType)
     sym_visibility = opt_attr_def(StringAttr)
 
-    traits = frozenset((SymbolOpInterface(), FuncOpCallableInterface()))
+    traits = traits_def(SymbolOpInterface(), FuncOpCallableInterface())
 
     def __init__(
         self,
@@ -317,7 +327,7 @@ class MulOp(IRDLOperation):
     rhs = operand_def(AnyTensorTypeF64Constr)
     res = result_def(AnyTensorTypeF64Constr)
 
-    traits = frozenset((Pure(), InferMulOpShapeTrait()))
+    traits = traits_def(Pure(), InferMulOpShapeTrait())
 
     def __init__(self, lhs: SSAValue, rhs: SSAValue):
         if isa(lhs.type, TensorTypeF64):
@@ -375,7 +385,7 @@ class ReturnOp(IRDLOperation):
     name = "toy.return"
     input = opt_operand_def(AnyTensorTypeF64Constr)
 
-    traits = frozenset([IsTerminator()])
+    traits = traits_def(IsTerminator())
 
     def __init__(self, input: SSAValue | None = None):
         return super().__init__(operands=[input])
@@ -408,7 +418,7 @@ class ReshapeOp(IRDLOperation):
     # We expect that the reshape operation returns a statically shaped tensor.
     res = result_def(TensorTypeF64)
 
-    traits = frozenset((Pure(), ReshapeOpHasCanonicalizationPatternsTrait()))
+    traits = traits_def(Pure(), ReshapeOpHasCanonicalizationPatternsTrait())
 
     def __init__(self, arg: SSAValue, shape: list[int]):
         if not isattr(arg.type, AnyTensorTypeF64Constr):
@@ -468,8 +478,8 @@ class TransposeOp(IRDLOperation):
     arg = operand_def(AnyTensorTypeF64Constr)
     res = result_def(AnyTensorTypeF64Constr)
 
-    traits = frozenset(
-        (
+    traits = OpTraits(
+        lambda: (
             Pure(),
             InferTransposeOpShapeTrait(),
             TransposeOpHasCanonicalizationPatternsTrait(),
@@ -515,7 +525,7 @@ class CastOp(IRDLOperation):
     arg = operand_def(AnyTensorTypeF64Constr)
     res = result_def(AnyTensorTypeF64Constr)
 
-    traits = frozenset((Pure(), InferCastOpShapeTrait()))
+    traits = traits_def(Pure(), InferCastOpShapeTrait())
 
     def __init__(self, arg: SSAValue, res: AnyTensorTypeF64 | None = None):
         if res is None:

docs/Toy/toy/frontend/ir_gen.py

@@ -1,12 +1,13 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import NoReturn
 
 from xdsl.builder import Builder
 from xdsl.dialects.builtin import ModuleOp, TensorType, UnrankedTensorType, f64
 from xdsl.ir import Block, Region, SSAValue
+from xdsl.utils.scoped_dict import ScopedDict
 
 from ..dialects.toy import (
     AddOp,
@@ -43,24 +44,6 @@ class IRGenError(Exception):
     pass
 
 
-@dataclass
-class ScopedSymbolTable:
-    "A mapping from variable names to SSAValues, append-only"
-
-    table: dict[str, SSAValue] = field(default_factory=dict)
-
-    def __contains__(self, __o: object) -> bool:
-        return __o in self.table
-
-    def __getitem__(self, __key: str) -> SSAValue:
-        return self.table[__key]
-
-    def __setitem__(self, __key: str, __value: SSAValue) -> None:
-        if __key in self:
-            raise AssertionError(f"Cannot add value for key {__key} in scope {self}")
-        self.table[__key] = __value
-
-
 @dataclass(init=False)
 class IRGen:
     """
@@ -80,7 +63,7 @@ class IRGen:
     is stateful, in particular it keeps an "insertion point": this is where
     the next operations will be introduced."""
 
-    symbol_table: ScopedSymbolTable | None = None
+    symbol_table: ScopedDict[str, SSAValue] | None = None
     """
     The symbol table maps a variable name to a value in the current scope.
     Entering a function creates a new scope, and the function arguments are
@@ -156,7 +139,7 @@ def ir_gen_function(self, function_ast: FunctionAST) -> FuncOp:
         parent_builder = self.builder
 
         # Create a scope in the symbol table to hold variable declarations.
-        self.symbol_table = ScopedSymbolTable()
+        self.symbol_table = ScopedDict[str, SSAValue]()
 
         proto_args = function_ast.proto.args
 

docs/irdl.ipynb

@@ -465,7 +465,7 @@
    "source": [
     "# The constraint sets the constraint variable if the attribute verifies\n",
     "var_constraint.verify(i32, constraint_context)\n",
-    "print(constraint_context.variables[\"T\"])"
+    "print(constraint_context.get_variable(\"T\"))"
    ]
   },
   {
@@ -1058,14 +1058,15 @@
     }
    ],
    "source": [
-    "from xdsl.irdl import ConstraintVar\n",
+    "from typing import ClassVar\n",
+    "from xdsl.irdl import base, VarConstraint\n",
     "\n",
     "\n",
     "@irdl_op_definition\n",
     "class BinaryOp(IRDLOperation):\n",
     "    name = \"binary_op\"\n",
     "\n",
-    "    T = Annotated[IntegerType, ConstraintVar(\"T\")]\n",
+    "    T: ClassVar = VarConstraint(\"T\", base(IntegerType))\n",
     "\n",
     "    lhs = operand_def(T)\n",
     "    rhs = operand_def(T)\n",
@@ -1427,15 +1428,15 @@
     }
    ],
    "source": [
-    "from xdsl.irdl import Block, Region, region_def\n",
+    "from xdsl.irdl import Block, Region, region_def, traits_def\n",
     "from xdsl.traits import NoTerminator\n",
     "\n",
     "\n",
     "@irdl_op_definition\n",
     "class WhileOp(IRDLOperation):\n",
     "    name = \"while_op\"\n",
     "    value = region_def()\n",
-    "    traits = frozenset([NoTerminator()])\n",
+    "    traits = traits_def(NoTerminator())\n",
     "\n",
     "\n",
     "region = Region(Block(arg_types=[i32]))\n",
@@ -1475,11 +1476,11 @@
       "Operation does not verify: result at position 0 does not verify:\n",
       "attribute i32 expected from variable 'T', but got i64\n",
       "\n",
-      "%0 = \"arith.addi\"(%1, %1) : (i32, i32) -> i64\n",
+      "%0 = \"arith.addi\"(%1, %1) <{\"overflowFlags\" = #arith.overflow<none>}> : (i32, i32) -> i64\n",
       "\n",
       "\n",
       "\n",
-      "%0 = \"arith.addi\"(%1, %1) : (i32, i32) -> i64\n",
+      "%0 = \"arith.addi\"(%1, %1) <{\"overflowFlags\" = #arith.overflow<none>}> : (i32, i32) -> i64\n",
       "\n",
       "\n"
      ]

docs/marimo/linalg_snitch.py

@@ -554,7 +554,7 @@ def __(
 
     snitch_c_shaped = ShapedArray(TypedPtr.new_float64([0.0] * c_len), c_shape)
 
-    register_implementations(snitch_interpreter, ctx, include_wgpu=False)
+    register_implementations(snitch_interpreter, ctx, include_wgpu=False, include_onnx=False)
 
     snitch_interpreter.call_op(
         "matmul",

docs/marimo/onnx/README.md

@@ -0,0 +1,4 @@
+# Marimo notebooks that depend on ONNX and mlir-opt
+
+For these notebooks to run as intended, the `onnx` package needs to be installed and `mlir-opt` needs to be in the path.
+Please see the [MLIR Interoperation](../../mlir_interoperation.md) document for more information.

docs/mlir_interoperation.ipynb

@@ -167,7 +167,7 @@
       "\"builtin.module\"() ({\r\n",
       "  %0 = \"arith.constant\"() <{value = 1 : i32}> : () -> i32\r\n",
       "  %1 = \"arith.constant\"() <{value = 2 : i32}> : () -> i32\r\n",
-      "  %2 = \"arith.addi\"(%0, %1) : (i32, i32) -> i32\r\n",
+      "  %2 = \"arith.addi\"(%0, %1) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32\r\n",
       "}) : () -> ()\r\n",
       "\r\n"
      ]

docs/mlir_interoperation.md

@@ -25,4 +25,4 @@ mlir-opt --convert-scf-to-cf --convert-cf-to-llvm --convert-func-to-llvm \
 The generated `tmp.ll` file contains LLVM IR, so it can be directly passed to
 the clang compiler. Notice that a `main` function is required for clang to
 build. The functionality is tested with the MLIR git commit hash:
-98e674c9f16d677d95c67bc130e267fae331e43c
+d401987fe349a87c53fe25829215b080b70c0c1a