update workflows

update worklfow

.github/workflows/build-run-kernel.yml

@@ -8,15 +8,15 @@ on:
 
 jobs:
   build-and-run-kernels:
-    runs-on: ubuntu-latest
-    container: 
-      image: ghcr.io/kuleuven-micas/snax:v0.2.7
+    runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v3
-      - name: Install snax-mlir
-        run: python3 -m pip install '-e.[dev]'
+      - uses: prefix-dev/[email protected]
+        with:
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
       - name: Build and run kernels
-        run: make allrun
+        run: pixi run make allrun
         working-directory: kernels/${{ matrix.kernel }}
     strategy:
       matrix:

.github/workflows/build-run-mlperf-tiny.yml

@@ -9,9 +9,7 @@ on:
 
 jobs:
   build-and-run-networks:
-    runs-on: ubuntu-latest
-    container: 
-      image: ghcr.io/kuleuven-micas/snax:v0.2.7
+    runs-on: ubuntu-24.04
     strategy:
       matrix:
         model:
@@ -25,15 +23,16 @@ jobs:
             dir: kernels/mlperf_tiny_kws
     steps:
       - uses: actions/checkout@v3
-      - name: Install snax-mlir
-        run: |
-          python3 -m pip install '-e.[dev,nn]'
+      - uses: prefix-dev/[email protected]
+        with:
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
       - name: snax-opt network compilation
         run: |
-          make ${{ matrix.model.name }}.o
+          pixi run make ${{ matrix.model.name }}.o
         working-directory: ${{ matrix.model.dir }}
       - name: vanilla network compilation
         run: |
-          make ${{ matrix.model.name }}.no-snax-opt.o
+          pixi run make ${{ matrix.model.name }}.no-snax-opt.o
         working-directory: ${{ matrix.model.dir }}
 

.github/workflows/code-formatting.yml

@@ -6,7 +6,7 @@ on:
 
 jobs:
   code-formatting:
-    runs-on: ubuntu-22.04
+    runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v3
 
@@ -17,4 +17,4 @@ jobs:
 
     - name: Run code formatting checks with pre-commit
       uses: pre-commit/[email protected]
-
+

.github/workflows/lit-tests.yml

@@ -12,18 +12,16 @@ permissions:
 jobs:
   python-tests:
 
-    runs-on: ubuntu-latest
-    container: 
-      image: ghcr.io/kuleuven-micas/snax:v0.2.7
+    runs-on: ubuntu-24.04
 
     steps:
     - uses: actions/checkout@v3
-    - name: Install snax-mlir
-      shell: bash
-      run: |
-        python3 -m pip install '-e.[dev]'
+    - uses: prefix-dev/[email protected]
+      with:
+        cache: true
+        cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
     - name: Test with lit
       shell: bash
       run: |
-        lit tests/filecheck -v
+        pixi run lit tests/filecheck -v
 

.github/workflows/python-tests.yml

@@ -12,18 +12,16 @@ permissions:
 jobs:
   python-tests:
 
-    runs-on: ubuntu-latest
-    container: 
-      image: ghcr.io/kuleuven-micas/snax:v0.2.7
+    runs-on: ubuntu-24.04
 
     steps:
     - uses: actions/checkout@v3
-    - name: Install snax-mlir
-      shell: bash
-      run: |
-        python3 -m pip install '-e.[dev]'
+    - uses: prefix-dev/[email protected]
+      with:
+        cache: true
+        cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
     - name: Test with pytest
       shell: bash
       run: |
-        python3 -m pytest .
+        pixi run python -m pytest .
 

.github/workflows/run-benchmarks.yml

@@ -10,14 +10,13 @@ on:
 
 jobs:
   run-benchmarks:
-    runs-on: ubuntu-latest
-    container:
-      image: ghcr.io/kuleuven-micas/snax:v0.2.7
+    runs-on: ubuntu-24.04
     steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
-      - name: Install snax-mlir
-        run: python3 -m pip install '-e.[dev,viz]'
+      - uses: actions/checkout@v4
+      - uses: actions/prefix-dev/[email protected]
+        with:
+          cache: true
+          cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
       - name: Run benchmarks
         run: python3 genbenchmark.py
         working-directory: benchmarks/${{ matrix.kernel }}

compiler/tools/snax_opt_main.py

@@ -45,6 +45,9 @@
 from compiler.transforms.test.debug_to_func import DebugToFuncPass
 from compiler.transforms.test.insert_debugs import InsertDebugPass
 from compiler.transforms.test.test_add_mcycle_around_launch import AddMcycleAroundLaunch
+from compiler.transforms.test.test_transform_dialect_erase_schedule import (
+    TestTransformDialectEraseSchedule,
+)
 from compiler.transforms.test_add_mcycle_around_loop import AddMcycleAroundLoopPass
 from compiler.transforms.test_remove_memref_copy import RemoveMemrefCopyPass
 
@@ -124,6 +127,10 @@ def __init__(
         super().register_pass(SnaxBufferize.name, lambda: SnaxBufferize)
         super().register_pass(FuseStreamingRegions.name, lambda: FuseStreamingRegions)
         super().register_pass(AllocToGlobalPass.name, lambda: AllocToGlobalPass)
+        super().register_pass(
+            TestTransformDialectEraseSchedule.name,
+            lambda: TestTransformDialectEraseSchedule,
+        )
 
         # arg handling
         arg_parser = argparse.ArgumentParser(description=description)

compiler/transforms/test/test_transform_dialect_erase_schedule.py

@@ -0,0 +1,31 @@
+from xdsl.context import MLContext
+from xdsl.dialects import builtin, transform
+from xdsl.passes import ModulePass
+from xdsl.pattern_rewriter import (
+    PatternRewriter,
+    PatternRewriteWalker,
+    RewritePattern,
+    op_type_rewrite_pattern,
+)
+
+
+class RemoveSchedule(RewritePattern):
+    @op_type_rewrite_pattern
+    def match_and_rewrite(
+        self,
+        op: transform.NamedSequenceOp | transform.SequenceOp,
+        rewriter: PatternRewriter,
+        /,
+    ):
+        rewriter.erase_matched_op()
+
+
+class TestTransformDialectEraseSchedule(ModulePass):
+    """
+    Copy of the test-transform-dialect-erase-schedule pass in MLIR
+    """
+
+    name = "test-transform-dialect-erase-schedule"
+
+    def apply(self, ctx: MLContext, op: builtin.ModuleOp) -> None:
+        PatternRewriteWalker(RemoveSchedule()).rewrite_module(op)

kernels/streamer_matmul/Makefile

@@ -9,7 +9,7 @@ TESTS =
 TESTS += quantized_matmul.x
 TESTS += tiled_quantized_matmul.x
 
-SNAXOPTFLAGS = -p convert-linalg-to-kernel,insert-accfg-op{accelerator=snax_gemmx},dispatch-kernels,convert-linalg-to-stream,fuse-streaming-regions,snax-bufferize,alloc-to-global,set-memory-space,set-memory-layout,realize-memref-casts,insert-sync-barrier,dispatch-regions{nb_cores=2},convert-stream-to-snax-stream,convert-linalg-to-accfg,convert-accfg-to-csr,snax-copy-to-dma,memref-to-snax,snax-to-func,clear-memory-space
+SNAXOPTFLAGS = -p test-transform-dialect-erase-schedule,convert-linalg-to-kernel,insert-accfg-op{accelerator=snax_gemmx},dispatch-kernels,convert-linalg-to-stream,fuse-streaming-regions,snax-bufferize,alloc-to-global,set-memory-space,set-memory-layout,realize-memref-casts,insert-sync-barrier,dispatch-regions{nb_cores=2},convert-stream-to-snax-stream,convert-linalg-to-accfg,convert-accfg-to-csr,snax-copy-to-dma,memref-to-snax,snax-to-func,clear-memory-space
 
 CFLAGS += -std=gnu11
 CFLAGS += -Wall -Wextra

pyproject.toml

@@ -17,6 +17,7 @@ dev = [
     "lit",
     "numpy<2.0.0",
     "pytest",
+    "pyright",
 ]
 
 # All optional dependencies for graphical plotting

runtime/Makefile.rules

@@ -75,7 +75,7 @@ LDFLAGS += -lsnRuntime
 
 # TRANSFORM Pre-Preprocessing
 
-MLIRTRANSFORM = --pass-pipeline='builtin.module(transform-interpreter{debug-bind-trailing-args=linalg.quantized_matmul}, test-transform-dialect-erase-schedule)'
+MLIRTRANSFORM = --pass-pipeline='builtin.module(transform-interpreter{debug-bind-trailing-args=linalg.quantized_matmul})'
 
 %.mlir: %.transform.mlir
 	$(MLIROPT) $(MLIRTRANSFORM) -o $@ $<