build(flydsl): upgrade FlyDSL to 0.4.0

.github/workflows/ci_cco.yml

@@ -166,7 +166,9 @@ jobs:
             mpirun --allow-run-as-root -np 2 python 02_lsa_put/main.py && \
             mpirun --allow-run-as-root -np 2 python 04_flydsl_lsa_put/main.py && \
             mpirun --allow-run-as-root -np 2 python 05_flydsl_lsa_allreduce/main.py && \
-            mpirun --allow-run-as-root -np 2 python 08_torch_symm_import/main.py"
+            MORI_ENABLE_SDMA=1 mpirun --allow-run-as-root -np 2 -x MORI_ENABLE_SDMA python 07_flydsl_sdma/main.py && \
+            mpirun --allow-run-as-root -np 2 python 08_torch_symm_import/main.py && \
+            MORI_ENABLE_SDMA=1 torchrun --standalone --nproc_per_node=2 10_flydsl_sdma_allgather_async/main.py --iters 5 --timing-iters 5"
             # GDA (device-initiated RDMA) examples fail on the runner (cross-rail
             # RDMA); re-enable once that is fixed.
             # mpirun --allow-run-as-root -np 2 python 03_flydsl_put/main.py
@@ -191,6 +193,13 @@ jobs:
             cd $GITHUB_WORKSPACE && \
             timeout 300 pytest tests/python/cco/test_sdma_api.py -v"
 
+      - name: Run CCO FlyDSL buffer cache-policy check
+        run: |
+          # A wrong cache policy still runs, so check the bits in the emitted ISA.
+          $CT exec $CONTAINER bash -c "\
+            cd /tmp && \
+            timeout 600 python -m pytest -v $GITHUB_WORKSPACE/tests/python/cco/test_flydsl_buffer.py"
+
       - name: Run CCO GEMM+AR numerics (up to 8 ranks)
         run: |
           # The op ships in mori.ops.gemm_ar, so run pytest from /tmp: from a

.github/workflows/nightly.yml

@@ -985,3 +985,103 @@ jobs:
         env:
           TWINE_USERNAME: __token__
           TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+
+  flydsl-compat:
+    # CI proper runs the FlyDSL that aiter pins (docker/Dockerfile.dev). mori's
+    # FlyDSL shims (mori.cco.device.flydsl.buffer and
+    # ops/dispatch_combine_v2/flydsl_compat.py) still carry 0.2.x and 0.3.x paths
+    # for environments pinned to an older one; this runs the FlyDSL-backed tests
+    # on those so the paths stay exercised: 0.2.4 is flydsl.expr.buffer_ops, and
+    # 0.3.0 still takes `aux` on the raw buffer ops positionally (0.3.1 made it the
+    # keyword attribute that CI's 0.4.0 covers).
+    name: FlyDSL ${{ matrix.flydsl }} compat (MI355X_AINIC)
+    runs-on: [self-hosted, MI355X-AINIC]
+    strategy:
+      fail-fast: false
+      max-parallel: 1  # each leg uses all 8 GPUs; never share a host
+      matrix:
+        flydsl: ["0.2.4", "0.3.0"]
+    env:
+      MORI_RDMA_DEVICES: rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7  # as ci_cco.yml
+      MORI_RDMA_SL: 3
+      MORI_RDMA_TC: 104
+      IMAGE: rocm/mori:ci-flydsl-${{ matrix.flydsl }}
+      BASE_IMAGE: rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.12.0
+      CONTAINER: mori_flydsl_compat_${{ github.run_id }}_${{ strategy.job-index }}
+      CT: podman  # MI355X-AINIC runner is podman rootless (no docker daemon)
+    timeout-minutes: 90
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          submodules: true
+
+      - name: Build image
+        run: |
+          $CT build --network=host --build-arg BASE_IMAGE=$BASE_IMAGE \
+            --build-arg GPU_ARCHS="gfx950" --build-arg WITH_FLYDSL=1 \
+            --build-arg FLYDSL_VERSION=${{ matrix.flydsl }} \
+            -t $IMAGE -f docker/Dockerfile.dev .
+
+      - name: Start container
+        run: |
+          $CT rm -f $CONTAINER 2>/dev/null || true
+          CONTAINER_RUNTIME=$CT ./docker/ci_run.sh --name $CONTAINER \
+            -e MORI_RDMA_DEVICES=$MORI_RDMA_DEVICES \
+            -e MORI_RDMA_SL=$MORI_RDMA_SL \
+            -e MORI_RDMA_TC=$MORI_RDMA_TC \
+            -v $GITHUB_WORKSPACE:$GITHUB_WORKSPACE \
+            -w $GITHUB_WORKSPACE \
+            $IMAGE sleep infinity
+          $CT exec $CONTAINER \
+            git config --global --add safe.directory $GITHUB_WORKSPACE
+
+      - name: Build mori
+        run: |
+          $CT exec $CONTAINER bash -c \
+            "cd $GITHUB_WORKSPACE && BUILD_UMBP=OFF BUILD_CCO_SDMA=ON pip install . && pip show flydsl"
+
+      - name: CCO SDMA Python API
+        run: |
+          $CT exec $CONTAINER bash -c "\
+            cd $GITHUB_WORKSPACE && \
+            timeout 300 pytest tests/python/cco/test_sdma_api.py -v"
+
+      - name: FlyDSL buffer cache-policy check
+        run: |
+          $CT exec $CONTAINER bash -c "\
+            cd /tmp && \
+            timeout 600 python -m pytest -v $GITHUB_WORKSPACE/tests/python/cco/test_flydsl_buffer.py"
+
+      - name: GEMM+AR numerics
+        run: |
+          $CT exec -e MORI_ENABLE_SDMA=1 -e MORI_SOCKET_IFNAME=lo $CONTAINER bash -c "\
+            cd /tmp && \
+            timeout 1800 python -m pytest -v \
+              $GITHUB_WORKSPACE/tests/python/cco/test_gemm_ar.py \
+              $GITHUB_WORKSPACE/tests/python/cco/test_gemm_ar_op.py \
+              $GITHUB_WORKSPACE/tests/python/cco/test_flydsl_ar.py"
+
+      - name: ops-v2 dispatch/combine, FlyDSL backend (EP8)
+        run: |
+          $CT exec $CONTAINER bash -c "\
+            cd $GITHUB_WORKSPACE && \
+            python -m pytest tests/python/ops/dispatch_combine_v2/test_dispatch_combine_v2_intranode.py -v"
+
+      - name: CCO FlyDSL examples (2 ranks)
+        run: |
+          $CT exec -e MORI_ENABLE_SDMA=1 $CONTAINER bash -c "pip install mpi4py && \
+            cd $GITHUB_WORKSPACE/examples/cco/python && \
+            mpirun --allow-run-as-root -np 2 python 04_flydsl_lsa_put/main.py && \
+            mpirun --allow-run-as-root -np 2 python 05_flydsl_lsa_allreduce/main.py && \
+            mpirun --allow-run-as-root -np 2 -x MORI_ENABLE_SDMA python 07_flydsl_sdma/main.py && \
+            mpirun --allow-run-as-root -np 2 python 08_torch_symm_import/main.py && \
+            torchrun --standalone --nproc_per_node=2 10_flydsl_sdma_allgather_async/main.py --iters 5 --timing-iters 5"
+
+      - name: Cleanup
+        if: always()
+        run: |
+          $CT exec $CONTAINER chown -R $(id -u):$(id -g) $GITHUB_WORKSPACE 2>/dev/null || true
+          CONTAINER_RUNTIME=$CT bash docker/ci_stop.sh $CONTAINER || true
+          # Only this leg's FlyDSL layer is unique; the rest is shared with ci-cco.
+          $CT rmi -f $IMAGE 2>/dev/null || true

docker/Dockerfile.dev

@@ -16,8 +16,6 @@ ARG GPU_ARCHS="gfx942 gfx950"
 ARG ROCM_APT_VERSION=7.14
 ARG ROCM_APT_GPG=https://repo.amd.com/rocm/packages-multi-arch/gpg/rocm.gpg
 ARG ROCM_APT_LIST="https://repo.amd.com/rocm/packages-multi-arch/ubuntu2404 stable main"
-# 1 installs flydsl, which only the cco Python examples need.
-ARG WITH_FLYDSL=0
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
@@ -74,6 +72,12 @@ RUN set -ex; \
 # distro kernel's amdgpu, so bumping the Ubuntu kernel alone does nothing.
 ENV MALLOC_MMAP_THRESHOLD_=131072
 
-# Pinned <0.3: 0.3.x dropped `flydsl.expr.buffer_ops`, still imported directly
-# by the cco examples (04/05/07/08). Unpin once they move to the 0.3.x API.
-RUN if [ "${WITH_FLYDSL}" = "1" ]; then pip install --no-cache-dir "flydsl<0.3"; fi
+# Declared here, next to their only use: every RUN after an ARG has it in its
+# cache key, so up top a new FlyDSL version would rebuild the ROCm layers too.
+# 1 installs flydsl, which only the cco Python examples need.
+ARG WITH_FLYDSL=0
+# The version aiter pins, so CI runs mori against the FlyDSL it ships with.
+# mori itself also runs on 0.2.x and 0.3.x; nightly.yml's flydsl-compat job
+# builds this image with an older one to keep those paths covered.
+ARG FLYDSL_VERSION=0.4.0
+RUN if [ "${WITH_FLYDSL}" = "1" ]; then pip install --no-cache-dir "flydsl==${FLYDSL_VERSION}"; fi

examples/cco/README.md

@@ -50,13 +50,13 @@ Then install the two extra runtime deps the examples need (not pulled in by
 `pip install .`):
 
 ```bash
-pip install mpi4py "flydsl==0.2.2"
+pip install mpi4py "flydsl==0.4.0"
 ```
 
 - `mpi4py` — every example bootstraps the cco `UniqueId` over MPI.
-- `flydsl==0.2.2` — required by the Python **FlyDSL** examples (03–06). Pinned to
-  the FlyDSL ABI the device bitcode targets. (Also available as the optional
-  extra `pip install amd_mori[flydsl]`.) Not needed for `01`, `02`, or the C++
+- `flydsl` — required by the Python **FlyDSL** examples (03–08). 0.4.0 is what CI
+  and aiter use; 0.2.x and 0.3.x work too. (Also available as the optional extra
+  `pip install amd_mori[flydsl]`.) Not needed for `01`, `02`, `09`, or the C++
   examples.
 
 After `pip install .` you do **not** need `PYTHONPATH` / `LD_LIBRARY_PATH` /

examples/cco/python/04_flydsl_lsa_put/main.py

@@ -44,12 +44,12 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl.expr import buffer_ops
 from flydsl.expr.typing import Int64, T
 from flydsl._mlir.dialects import rocdl
 
 from mori.cco import Communicator, CCODevCommRequirements, GDA_CONNECTION_NONE
 import mori.cco.device.flydsl as cco
+from mori.cco.device.flydsl import buffer
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from cco_example_common import set_device, sync, fill, zero, read  # noqa: E402
@@ -60,7 +60,6 @@
 NUM_PACKS = NBYTES // 16  # 16 B (vector<4xi32>) per pack
 THREADS = 256
 DST_RANK = 1  # rank 1's LSA rank
-_CM_UNCACHED = 3  # SC0|SC1: store bypasses L1+L2 -> reaches peer HBM
 
 
 @flyc.kernel(known_block_size=[THREADS, 1, 1])
@@ -72,12 +71,13 @@ def lsa_put_kernel(dev_comm: Int64, win_handle: Int64):
     win = cco.Window(win_handle)
     src = win.lsa_ptr(dc.lsa_rank, 0)  # my own slot (local VA)
     dst = win.lsa_ptr(DST_RANK, 0)  # rank 1's slot (peer VA) — load/store directly
-    src_rsrc = buffer_ops.create_buffer_resource_from_addr(src)
-    dst_rsrc = buffer_ops.create_buffer_resource_from_addr(dst)
+    src_rsrc = buffer.create_buffer_resource_from_addr(src)
+    dst_rsrc = buffer.create_buffer_resource_from_addr(dst)
     for pk in range(tid, NUM_PACKS, THREADS):
         off = pk * 4  # i32-element offset of this 16 B pack
-        v = buffer_ops.buffer_load(src_rsrc, off, vec_width=4, dtype=T.i32)
-        buffer_ops.buffer_store(v, dst_rsrc, off, cache_modifier=_CM_UNCACHED)
+        v = buffer.buffer_load(src_rsrc, off, vec_width=4, dtype=T.i32)
+        # Past L1+L2 so the store reaches the peer's HBM.
+        buffer.buffer_store(v, dst_rsrc, off, cache_modifier=buffer.CM_SC0_SC1)
     rocdl.s_waitcnt(0)
 
 

examples/cco/python/05_flydsl_lsa_allreduce/main.py

@@ -53,21 +53,17 @@
 import flydsl.expr as fx
 from flydsl._mlir import ir
 from flydsl._mlir.dialects import scf
-from flydsl.expr import buffer_ops, range_constexpr
+from flydsl.expr import range_constexpr
 from flydsl.expr import gpu as fgpu
 from flydsl.expr.typing import Int32, Int64, T
 
 from mori.cco import Communicator, CCODevCommRequirements, GDA_CONNECTION_NONE
 import mori.cco.device.flydsl as cco
+from mori.cco.device.flydsl import buffer
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from cco_example_common import set_device, sync, fill, zero, read, F32  # noqa: E402
 
-# Cache-modifier aux bits (GFX942): SC0=bypass L1, SC1=bypass L2.
-_CM_CACHED = 0
-_CM_SC1 = 2  # uncached read  (see peers' fresh signal writes)
-_CM_SC0_SC1 = 3  # uncached write (signal stores)
-
 WS = 2  # world size (ranks, one GPU each)
 THREADS = 256
 NUM_ELEMS = 256 * 1024  # f32 elements to all-reduce (1 MiB)
@@ -83,16 +79,16 @@
 
 
 def _rsrc(addr_i64):
-    return buffer_ops.create_buffer_resource_from_addr(addr_i64)
+    return buffer.create_buffer_resource_from_addr(addr_i64)
 
 
 def _store_u32_uncached(rsrc, val):
-    buffer_ops.buffer_store(val, rsrc, 0, cache_modifier=_CM_SC0_SC1)
+    buffer.buffer_store(val, rsrc, 0, cache_modifier=buffer.CM_SC0_SC1)
 
 
 def _load_u32_uncached(rsrc):
-    return buffer_ops.buffer_load(
-        rsrc, 0, vec_width=1, dtype=T.i32, cache_modifier=_CM_SC1
+    return buffer.buffer_load(
+        rsrc, 0, vec_width=1, dtype=T.i32, cache_modifier=buffer.CM_SC1
     )
 
 
@@ -142,12 +138,12 @@ def custom_ar_kernel(dev_comm: Int64, win: Int64, flag: Int32):
         acc = None
         for p in range_constexpr(WS):
             raw = fx.Vector(
-                buffer_ops.buffer_load(in_rsrc[p], elem_off, vec_width=4, dtype=T.i32)
+                buffer.buffer_load(in_rsrc[p], elem_off, vec_width=4, dtype=T.i32)
             )
             vf = raw.bitcast(fx.Float32)
             acc = vf if acc is None else acc + vf
-        buffer_ops.buffer_store(
-            acc.bitcast(fx.Int32), out_rsrc, elem_off, cache_modifier=_CM_CACHED
+        buffer.buffer_store(
+            acc.bitcast(fx.Int32), out_rsrc, elem_off, cache_modifier=buffer.CM_CACHED
         )
 
 

examples/cco/python/07_flydsl_sdma/main.py

@@ -56,11 +56,12 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl.expr import buffer_ops, range_constexpr
+from flydsl.expr import range_constexpr
 from flydsl.expr.typing import Int64, T
 
 from mori.cco import Communicator, CCODevCommRequirements, GDA_CONNECTION_NONE
 import mori.cco.device.flydsl as cco
+from mori.cco.device.flydsl import buffer
 
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from cco_example_common import set_device, sync, fill, zero, read, F32  # noqa: E402
@@ -127,25 +128,21 @@ def reduce_kernel(win: Int64):
         w = cco.Window(win)
         # All reads are local memory — my own symmetric-window VA.
         src_rsrc = [
-            buffer_ops.create_buffer_resource_from_addr(
-                fx.Int64(w.lsa_ptr(my_rank, off))
-            )
+            buffer.create_buffer_resource_from_addr(fx.Int64(w.lsa_ptr(my_rank, off)))
             for off in src_offs
         ]
-        out_rsrc = buffer_ops.create_buffer_resource_from_addr(
+        out_rsrc = buffer.create_buffer_resource_from_addr(
             fx.Int64(w.lsa_ptr(my_rank, OUT_OFF))
         )
         for pk in range(tid, NUM_PACKS, THREADS):
             elem_off = pk * ELEMS_PER_PACK
             acc = None
             for s in range_constexpr(len(src_rsrc)):
                 v = fx.Vector(
-                    buffer_ops.buffer_load(
-                        src_rsrc[s], elem_off, vec_width=4, dtype=T.f32
-                    )
+                    buffer.buffer_load(src_rsrc[s], elem_off, vec_width=4, dtype=T.f32)
                 )
                 acc = v if acc is None else acc + v
-            buffer_ops.buffer_store(acc, out_rsrc, elem_off)
+            buffer.buffer_store(acc, out_rsrc, elem_off)
 
     @flyc.jit
     def run(win: Int64, stream=fx.Stream(None)):

examples/cco/python/08_torch_symm_import/main.py

@@ -54,13 +54,13 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl.expr import buffer_ops
 from flydsl.expr.typing import Int32, Int64, T
 from flydsl._mlir.dialects import rocdl
 
 from mori.cco import Communicator, CCODevCommRequirements, GDA_CONNECTION_NONE
 from mori.jit.hip_driver import _check, _get_hip_lib
 import mori.cco.device.flydsl as cco
+from mori.cco.device.flydsl import buffer
 
 PER_RANK_VMM = 1 * 1024 * 1024 * 1024
 N = 1024  # int32 elements
@@ -84,12 +84,12 @@ def lsa_gather_kernel(dev_comm: Int64, win_handle: Int64, src_rank: Int32):
     win = cco.Window(win_handle)
     src = win.lsa_ptr(src_rank, 0)  # peer's torch tensor (peer VA) — load directly
     dst = win.lsa_ptr(dc.lsa_rank, 0)  # my own torch tensor (local VA)
-    src_rsrc = buffer_ops.create_buffer_resource_from_addr(src)
-    dst_rsrc = buffer_ops.create_buffer_resource_from_addr(dst)
+    src_rsrc = buffer.create_buffer_resource_from_addr(src)
+    dst_rsrc = buffer.create_buffer_resource_from_addr(dst)
     for pk in range(tid, NUM_PACKS, THREADS):
         off = pk * 4  # i32-element offset of this 16 B pack
-        v = buffer_ops.buffer_load(src_rsrc, off, vec_width=4, dtype=T.i32)
-        buffer_ops.buffer_store(v, dst_rsrc, off)
+        v = buffer.buffer_load(src_rsrc, off, vec_width=4, dtype=T.i32)
+        buffer.buffer_store(v, dst_rsrc, off)
     rocdl.s_waitcnt(0)
 
 

examples/cco/python/10_flydsl_sdma_allgather_async/main.py

@@ -85,9 +85,8 @@
 from flydsl.expr.typing import Int64
 from mori.cco import GDA_CONNECTION_NONE, CCODevCommRequirements, Communicator, UniqueId
 
-# flydsl moved its buffer/atomic helpers between 0.2 and 0.3; this module is the
-# one place that papers over it, so the example runs on either.
-from mori.ops.gemm_ar._compat import (
+# Signal-slot accessors that behave the same on every FlyDSL mori supports.
+from mori.cco.device.flydsl.signals import (
     i32_type,
     local_load_u32,
     local_store_u32,

python/mori/cco/device/README.md

@@ -26,6 +26,8 @@ CCO_DEVICE_FUNCTIONS (one scalar ABI table)
 | Triton extern wrappers | `python/mori/ir/triton/cco/ops.py` |
 | OO API (`DevComm/Window/Gda`) + enums (`CoopScope/SignalOp/ThreadMode`) | `python/mori/cco/device/flydsl/handles.py` |
 | `_ffi` factory + `cco_struct` (handle keeps methods across `if`/`for`) | `flydsl/_internal.py` |
+| Buffer load/store with a cache policy that holds on every FlyDSL version | `flydsl/buffer.py` |
+| u32 signal-slot accessors (peer flags, rank-private epochs) | `flydsl/signals.py` |
 | How the `.bc` is located at runtime | `python/mori/cco/device/bitcode.py` (`find_cco_bitcode`) |
 | How the `.bc` is built (JIT, per arch+NIC+cov) | `bitcode.py::_jit_compile` (reuses `mori.jit`) |
 | How to prebuild the `.bc` manually | `tools/build_cco_bitcode.sh` |
@@ -80,7 +82,7 @@ then install MORI and the example runtime deps:
 
 ```bash
 pip install .                          # builds + co-locates all libmori_*.so
-pip install mpi4py "flydsl==0.2.2"     # mpi4py: bootstrap; flydsl: the device bindings
+pip install mpi4py "flydsl==0.4.0"     # mpi4py: bootstrap; flydsl: the device bindings
 ```
 
 After `pip install .` no `PYTHONPATH` / `LD_LIBRARY_PATH` / `MORI_CCO_BC` is