WSE-3 support, improved color reuse, multi-hop routing, and additional kernels

.github/workflows/python-app.yml

@@ -15,6 +15,9 @@ permissions:
 jobs:
   test-pytest:
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        wse-arch: ["wse2", "wse3"]
 
     steps:
     - uses: actions/checkout@v4
@@ -32,11 +35,19 @@ jobs:
         pip install -e ".[dev]"
 
     - name: Test with pytest
+      env:
+        WSE_ARCH: ${{ matrix.wse-arch }}
       run: |
         pytest
 
+  # WSE-2 and WSE-3 share one runner so Singularity and the SDK are installed
+  # once.
   test-csl:
     runs-on: ubuntu-latest
+    timeout-minutes: 180
+    strategy:
+      matrix:
+        wse-arch: ["wse2", "wse3"]
 
     steps:
     - uses: actions/checkout@v4
@@ -54,29 +65,35 @@ jobs:
         pip install -e ".[dev]"
 
     - name: Install CSL dependencies
+      env:
+        DEBIAN_FRONTEND: noninteractive
       run: |
+        set -euxo pipefail
         . /etc/os-release
         echo "Using Ubuntu version $UBUNTU_CODENAME"
 
-        # Make dependency setup faster
+        # Avoid man-db postinst work during package installs.
         echo 'set man-db/auto-update false' | sudo debconf-communicate >/dev/null
-        sudo dpkg-reconfigure man-db
 
-        sudo apt-get update
-        sudo apt-get install -y build-essential libssl-dev uuid-dev libgpgme11-dev squashfs-tools
-        wget -q -O singularity.deb https://github.com/sylabs/singularity/releases/download/v4.2.1/singularity-ce_4.2.1-${UBUNTU_CODENAME}_amd64.deb
-        sudo apt install ./singularity.deb
+        sudo apt-get update -y -o Acquire::Retries=3
+        sudo apt-get install -y --no-install-recommends \
+          build-essential libssl-dev uuid-dev libgpgme11-dev squashfs-tools
+
+        wget -q -O singularity.deb \
+          "https://github.com/sylabs/singularity/releases/download/v4.2.1/singularity-ce_4.2.1-${UBUNTU_CODENAME}_amd64.deb"
+        # Must pass -y: bare "apt install ./singularity.deb" waits for confirmation and hangs CI.
+        sudo apt-get install -y --no-install-recommends ./singularity.deb
         singularity --version
 
-    - name: Cache dependencies
-      id: cache-deps
+    - name: Cache Cerebras SDK
+      id: cache-sdk
       uses: actions/cache@v4
       with:
         path: cerebras-sdk
-        key: ${{ runner.os }}-deps
+        key: ${{ runner.os }}-cerebras-sdk-1.4.0
 
     - name: Install Cerebras SDK v1.4.0
-      if: steps.cache-deps.outputs.cache-hit != 'true'
+      if: steps.cache-sdk.outputs.cache-hit != 'true'
       run: |
         mkdir cerebras-sdk
         cd cerebras-sdk
@@ -92,7 +109,7 @@ jobs:
     - name: Test CSL with simulator
       run: |
         export PATH=$PATH:`pwd`/cerebras-sdk
-        ./tests/csl_runtime/run_tests.sh
+        WSE_ARCH=${{ matrix.wse-arch }} ./tests/csl_runtime/run_tests.sh
 
   build-package:
     runs-on: ubuntu-latest

README.md

@@ -106,17 +106,18 @@ Sample SpaDA programs are in `samples/`:
 | `samples/stencils.py` | GT4Py stencil definitions (Laplacian, vertical advection, UVBKE, …) |
 | `samples/advanced_stencils.py` | GT4Py definitions for horizontal diffusion kernels |
 | `samples/benchmarks/` | Pre-compiled `.spst`/`.sptl` pairs for five kernels at five domain sizes |
-| `samples/spatial/simple/` | Basic single-PE and streaming operations: `add`, `copy`, `forward_sum`, `backward_sum`, `mult_scalar`, `streaming_copy` |
-| `samples/spatial/blas/` | Dense linear algebra: `axpy`, `matvec`, `gemv`, `gemv_twophase` |
+| `samples/spatial/simple/` | Basic single-PE and streaming operations: `add`, `copy`, `forward_sum`, `backward_sum`, `mult_scalar`, `streaming_copy`, and `exchange_bundle_1D` |
+| `samples/spatial/blas/` | Linear algebra: `axpy`, `matvec`, `gemv`, `gemv_twophase`, `spmv` |
 | `samples/spatial/collectives/` | Reductions (`scalar`, `chain`, `tree`, `twophase` in 1D/2D) and broadcasts (`broadcast_1D`, `broadcast_2D`, and multicast variants) |
+| `samples/spatial/sort/` | Sorting networks over `R` independent rows of `2^L` PEs holding `K` keys each: `batcher_oddeven_wse3_1D` (Batcher odd-even mergesort), `odd_even_sort_1D_looped` (1D looped odd-even transposition sort), and `shearsort_2D` (2D snake-order shearsort on WSE-3) |
 | `samples/spatial/stencils/` | Stencil examples: `laplacian` (high-level) and `laplacian_routed` (explicit routing) |
 | `samples/spst/` | Stencil IR examples |
 
 ---
 
 ## SDK Version and WSE compatibility
 
-The code has been tested for CSL SDK 1.4 and WSE-2.
+The code has been tested with CSL SDK 1.4 on WSE-2 and WSE-3. The target architecture is selected via `WSE_ARCH` (`wse2`, the default, or `wse3`). CI runs simulator tests for both architectures. Kernels targeting features unique to WSE-3 (such as `shearsort_2D`, which requires four concurrent inbound queues per PE, and `batcher_oddeven_wse3_1D`) automatically skip during WSE-2 test runs.
 
 ## Testing
 
@@ -159,6 +160,7 @@ make -C tests/csl_runtime check-sdk
 
 ```bash
 make -C tests/csl_runtime test
+make -C tests/csl_runtime test WSE_ARCH=wse3
 ```
 
 **Run a single test:**
@@ -189,13 +191,14 @@ tests/csl_runtime/run-in-lima.sh --sdk-url <url>
 ```
 
 This creates the Lima VM on first use (~5–10 min), downloads and extracts the SDK to `tests/csl_runtime/cerebras-sdk/`, installs Python dependencies inside the VM, and runs the full test suite. 
-If the SDK tarball is already downloaded or extracted, use `--sdk /path/to/cs_sdk` instead of `--sdk-url`.
+If the SDK tarball is already downloaded or extracted, use `--sdk /path/to/cs_sdk` instead of `--sdk-url`. Pass `--arch wse3` to compile and simulate for WSE-3.
 
 Other modes:
 
 ```bash
 # Run a single test
 tests/csl_runtime/run-in-lima.sh --sdk <dir> --test test_add.sh
+tests/csl_runtime/run-in-lima.sh --sdk <dir> --arch wse3 --test test_shearsort_2d_looped.sh
 
 # Verify the SDK toolchain only
 tests/csl_runtime/run-in-lima.sh  --sdk <dir> --check

irspec/docs/spatial/routing.md

@@ -192,3 +192,8 @@ to receive.
     If multicasting is used, the correctness conditions must be adapted accordingly, 
     especially when considering multiple phases.
 
+
+## Lowering to Cerebras WSE
+
+How epochs, switch positions, control wavelets, shift bundling, and counter filters are realized on
+the Cerebras Wafer-Scale Engine is described in [Routing Semantics on Cerebras WSE](routing_wse.md).