[TLE] Add device intra node barrier and change the logic of pulling the third-party tileir

.github/workflows/enflame3.6-gcu400-build-and-test.yml

@@ -88,7 +88,8 @@ jobs:
 
           python3 -m pytest -s python/test/tle \
               --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \
-              --ignore=python/test/tle/unit/test_tle_get_local_pe.py
+              --ignore=python/test/tle/unit/test_tle_get_local_pe.py \
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py
 
           ## tle raw test
           python3 -m pytest -s third_party/enflame/python/test/tle/raw

.github/workflows/nv3.6-build-and-test.yml

@@ -72,6 +72,7 @@ jobs:
         run: |
           set -x
           source ~/env-3.6.sh
+          export FLAGCX_RECOMPILE=ON
           export USE_FLAGCX=ON
           MAX_JOBS=32 python3 -m pip install . --no-build-isolation
 
@@ -112,9 +113,11 @@ jobs:
           python3 -m pytest -s python/test/tle/integration
           python3 -m pytest -s python/test/tle/unit \
               --ignore=python/test/tle/unit/test_tle_distributed_d2d.py \
-              --ignore=python/test/tle/unit/test_tle_get_local_pe.py
+              --ignore=python/test/tle/unit/test_tle_get_local_pe.py \
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py
           CUDA_LAUNCH_BLOCKING=1 bash python/test/tle/unit/test_tle_distributed_d2d.sh
           CUDA_LAUNCH_BLOCKING=1 bash python/test/tle/unit/test_tle_get_local_pe.sh
+          CUDA_LAUNCH_BLOCKING=1 bash python/test/tle/unit/test_tle_d2d_barrier.sh
           ## flagtree hints python tutorials
           python3 python/tutorials/hints/01/01-vector-add.py --only_unit_test
           # python3 python/tutorials/hints/02/02-fused-softmax.py --only_unit_test

packaging/rpm/helpers/Dockerfile.rpm

@@ -41,14 +41,12 @@ COPY test /src/test
 COPY unittest /src/unittest
 COPY CMakeLists.txt pyproject.toml setup.py MANIFEST.in LICENSE README.md /src/
 
-# FLAGTREE_DEFAULT_BACKENDS=nvidia,amd matches the historical package
-# contents and keeps the tileir backend (GCC >= 13 + cuda-tile
+# Keep the tileir backend (GCC >= 13 + cuda-tile
 # submodule) out of the distro build.
 # The trailing rm keeps only /wheels in this layer: the CMake build
 # tree and the LLVM tarball cache (~/.triton) would otherwise add
 # several GB and push CI runners into ENOSPC.
 RUN unset FLAGTREE_BACKEND && \
-    FLAGTREE_DEFAULT_BACKENDS=nvidia,amd \
     MAX_JOBS=4 python3 -m pip wheel . --no-build-isolation --no-deps -w /wheels -v && \
     ls -lh /wheels/ && \
     test -n "$(ls /wheels/flagtree-*.whl 2>/dev/null)" && \

python/setup_tools/setup_helper.py

@@ -519,6 +519,8 @@ def uninstall_triton():
 
 download_flagtree_third_party("flagcx", condition=(not flagtree_backend), hook="handle_flagcx", required=True)
 
+download_flagtree_third_party("tileir", condition=(flagtree_backend == "tileir"), required=True)
+
 handle_flagtree_backend()
 
 # iluvatar

python/setup_tools/utils/__init__.py

@@ -60,6 +60,10 @@ def _register_submodule(self, submodule):
     },
     {"name": "flir", "url": "https://github.com/FlagTree/flir.git"},
     {"name": "flagcx", "url": "https://github.com/flagos-ai/FlagCX.git", "relative_path": "tle/third_party/flagcx"},
+    {
+        "name": "tileir", "url": "https://github.com/NVIDIA/cuda-tile", "relative_path": "tileir/third_party/cuda-tile",
+        "commit_id": "2e5ccba66fb3afdba34b26cf358418283027c248"
+    },
 ))
 
 

python/setup_tools/utils/default.py

@@ -27,7 +27,7 @@
 
 
 def printinfo(msgs):
-    print(f" [TLE-DIST-INFO]: {msgs}. \n")
+    print(f" [TLE-DIST-INFO]: {msgs}.")
 
 
 class FlagCXRegistrar:
@@ -37,6 +37,29 @@ def __init__(self, external):
         self.shared_lib_name = "libflagcx.so"
         self._set_path(external)
 
+    def _is_flagcx_recompile_required(self):
+        ENVS = ("FLAGCX_RECOMPILE", "DIST_RECOMPILE")
+        for env in ENVS:
+            env_value = os.environ.get(env, "0")
+            if env_value in ("1", "true", "True", "ON"):
+                printinfo(f"Recompiling FlagCX due to {env}={env_value}\n ")
+                return env_value
+        return False
+
+    def _is_cache_available(self):
+        ENVS = ("FLAGCX_CACHE", "DIST_CACHE")
+
+        for env in ENVS:
+            env_value = os.environ.get(env, "1")
+            if env_value in ("0", "false", "False", "OFF", "clean"):
+                printinfo(f"Skipping using cache at {self.cache_lib_dir} due to {env}={env_value}\n ")
+            if env_value in ("clean"):
+                shutil.rmtree(self.cache_lib_dir, ignore_errors=True)
+                os.makedirs(self.cache_lib_dir, exist_ok=True)
+                printinfo(f"Cache {self.cache_lib_dir} cleaned due to {env}={env_value}\n ")
+            return False
+        return True
+
     def _set_path(self, external):
         submodule = external['backend']
         flagtree_cache = external['cache']
@@ -70,15 +93,18 @@ def get_compile_cmds(self):
 
     def _compile_and_cache(self):
         cmds = self.get_compile_cmds()
+        is_unused_cache = self._is_cache_available()
+        is_recompile = self._is_flagcx_recompile_required()
+        is_unused_cache = is_unused_cache or is_recompile
 
         for lib_name, cmd in cmds.items():
             cache_path = getattr(self, f"{lib_name.split('.')[0]}_cache_path")
             src_path = getattr(self, f"{lib_name.split('.')[0]}_src_path")
             runtime_path = self._get_runtime_path(lib_name)
-            if cache_path.exists():
+            if cache_path.exists() and not is_unused_cache:
                 printinfo(f"{lib_name} already exists in cache, skipping compilation ...")
                 shutil.copy(cache_path, runtime_path)
-            elif src_path.exists():
+            elif src_path.exists() and not is_recompile:
                 printinfo(f"{lib_name} already exists in build directory, copying to cache...")
                 shutil.copy(src_path, cache_path)
                 shutil.copy(src_path, runtime_path)
@@ -93,7 +119,7 @@ def _compile_and_cache(self):
                 shutil.copy(src_path, cache_path)
                 shutil.copy(src_path, runtime_path)
                 printinfo(f"{lib_name} copied from {src_path} to cache at {cache_path}")
-                printinfo(f"[32m{lib_name} copied from {src_path} to cache at {runtime_path}")
+                printinfo(f"{lib_name} copied from {src_path} to cache at {runtime_path}")
 
     def _copy_required_files(self):
         dst = Path(self.flagtree_dir) / "python" / "triton" / "experimental" / "tle" / "language" / "flagcx_wrapper.py"

python/setup_tools/utils/tools.py

@@ -42,12 +42,7 @@ def _get_flagtree_root() -> str:
 
 @dataclass
 class FlagtreeConfigs:
-    # Overridable for distro/packaging builds whose toolchain can't build
-    # every default backend (e.g. tileir requires GCC >= 13 and the
-    # cuda-tile submodule):
-    #   FLAGTREE_DEFAULT_BACKENDS=nvidia,amd pip wheel .
-    default_backends: tuple = field(default_factory=lambda: tuple(
-        b for b in os.environ.get("FLAGTREE_DEFAULT_BACKENDS", "nvidia,amd,tileir").replace(" ", "").split(",") if b))
+    default_backends: tuple = ("nvidia", "amd")
     plugin_backends: tuple = ("cambricon", "ascend", "aipu", "tsingmicro", "enflame", "hcu", "thrive")
     use_cuda_toolkit_backends: tuple = ('aipu', 'tileir')
     language_extra_backends: tuple = ('xpu', 'mthreads', "cambricon")
@@ -69,12 +64,12 @@ class FlagtreeConfigs:
     }))
 
     def __post_init__(self):
+        backends = list(self.default_backends)
+        _backends = [backend for backend in backends if os.environ.get(f"USE_{backend.upper()}", "ON").upper() != "OFF"]
+        self.default_backends = tuple(_backends)
         self.flagtree_submodule_dir = os.path.join(self.flagtree_root_dir, "third_party")
         self.activated_module = self._activate_device_module()
 
-    def non_tileir_default_backends(self):
-        return tuple(backend for backend in self.default_backends if backend != "tileir")
-
     def _activate_device_module(self, suffix=".py"):
         backend = self.flagtree_backend or "default"
         module_path = Path(os.path.dirname(__file__)) / backend

python/test/tle/integration/test_tle_distributed.py

@@ -995,7 +995,7 @@ def test_remote_const_shard_load_high_block_encoding_no_regression(self):
         assert re.search(r"tensor<256x!tt\.ptr<f16,\s*3>,\s*#blocked[0-9]*>", ttgir) is not None
         assert re.search(
             r"\"tle\.remote_pointers\"\(%[^,]+,\s*%[^)]+\)\s*"
-            r"<\{space\s*=\s*\"cluster\"\}>\s*:\s*"
+            r"<\{[^}]*space\s*=\s*\"cluster\"[^}]*\}>\s*:\s*"
             r"\(tensor<256x!tt\.ptr<f16,\s*3>,\s*#blocked[0-9]*>,\s*i32\)\s*->\s*"
             r"tensor<256x!tt\.ptr<f16,\s*7>,\s*#blocked[0-9]*>",
             ttgir,

python/test/tle/unit/test_tle_d2d_barrier.py

@@ -0,0 +1,87 @@
+import triton.experimental.tle.language as tle
+import torch
+import triton
+import triton.language as tl
+import torch.distributed as dist
+
+DEVICE_MESH = tle.device_mesh(tle.MeshConfig(device=2))
+N = 8
+
+
+@triton.jit()
+def _barrier_d2d_kernel(out_ptr, device_dptr: tl.constexpr, mesh: tl.constexpr):
+    pid = tl.program_id(0)
+    local_rank = tle.shard_id(mesh, 'device', device_dptr=device_dptr)
+    n_rank = mesh.shape[0]
+    peer = (local_rank + 1) % n_rank  # noqa: F841
+
+    remote_mem = tle.remote(
+        device_dptr,
+        space="device",
+        dtype=tl.float32,
+        shard_id=peer,
+        offset=pid,
+    )
+    val = tl.load(remote_mem)
+    tl.store(out_ptr + pid, val)
+    tle.distributed_barrier(device_dptr=device_dptr, space="device")
+
+
+def _ir_verify(output, device_dptr, grid):
+    compiled = _barrier_d2d_kernel.warmup(
+        device_dptr=device_dptr,
+        out_ptr=output,
+        mesh=DEVICE_MESH,
+        grid=(grid, ),
+        num_ctas=1,
+        num_warps=4,
+    )
+    assert "distributed_barrier" in compiled.asm["ttgir"]
+    assert "remote_pointer" in compiled.asm["ttgir"]
+    assert "flagcxIntraBarrier" in compiled.asm['ptx']
+    assert "flagcxGetIntraPointerC" in compiled.asm['ptx']
+    assert "flagcxDevCommGetIntraRank" in compiled.asm['ptx']
+
+
+def _runtime_verify(output, device_dptr, grid, rank, world_size):
+    dist.barrier()
+    _barrier_d2d_kernel[grid](device_dptr=device_dptr, out_ptr=output, mesh=DEVICE_MESH)
+
+    torch.cuda.synchronize()
+
+    import sys
+    peer_rank = (rank + 1) % world_size
+    expected = torch.arange(N, dtype=torch.float32, device="cuda") + peer_rank * 1000
+    if torch.allclose(output, expected):
+        print(f"[Rank {rank}] [PASSED] read peer rank {peer_rank}")
+        print(f"[Rank {rank}] sample output[:4] = {output[:4].tolist()}")
+    else:
+        print(f"[Rank {rank}] [FAILED] read peer rank {peer_rank}")
+        print(f"[Rank {rank}] expected[:4] = {expected[:4].tolist()}")
+        print(f"[Rank {rank}] output[:4] = {output[:4].tolist()}")
+        sys.exit(1)
+
+    tle.cleanup_communicator()
+
+
+class TestD2DBarrier:
+
+    def test_tle_d2d_barrier(self):
+        grid = (N, )
+
+        mem_pool = tle.get_mem_pool()
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+        with torch.cuda.use_mem_pool(mem_pool):
+            x = (torch.arange(N, dtype=torch.float32, device="cuda") + rank * 1000).clone()
+
+        device_dptr = tle.create_dist_tensor(x)
+        output = torch.zeros(N, dtype=torch.float32, device="cuda")
+
+        _ir_verify(output, device_dptr, grid)
+
+        _runtime_verify(output, device_dptr, grid, rank, world_size)
+
+
+if __name__ == "__main__":
+    TestD2DBarrier().test_tle_d2d_barrier()

python/test/tle/unit/test_tle_d2d_barrier.sh

@@ -0,0 +1,53 @@
+rm -rf ~/.triton/cache
+#!/bin/bash
+
+# Check if the debug flag is provided as an argument
+if [ "$1" == "debug" ]; then
+    export NCCL_DEBUG=INFO
+    export NCCL_DEBUG_SUBSYS=all
+    echo "NCCL debug information enabled."
+else
+    unset NCCL_DEBUG
+    unset NCCL_DEBUG_SUBSYS
+    echo "NCCL debug information disabled."
+fi
+#export FLAGCX_SOCKET_IFNAME=eth0
+export FLAGCX_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_6,mlx5_7,mlx5_8,mlx5_9
+export FLAGCX_USE_HETERO_COMM=1
+export FLAGCX_MEM_ENABLE=1
+export FLAGCX_VMM_ENABLE=0
+export FLAGCX_P2P_DISABLE=1
+#export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+# Need to preload customized gloo library specified for FlagCX linkage
+#export LD_PRELOAD=/usr/local/lib/libgloo.so
+#export LD_PRELOAD=/usr/local/nccl/build/lib/libnccl.so
+#export TORCH_DISTRIBUTED_DETAIL=DEBUG
+#pytest  FlagTree/python/test/tle/unit/test_tle_distributed_d2d.py
+#export NCCL_DEBUG=INFO
+#export CUDA_LAUNCH_BLOCKING=1
+#export TORCH_USE_CUDA_DSA=1
+#export NCCL_NVLS_ENABLE=0
+#export FLAGCX_DMABUF_ENABLE=1
+#export FLAGCX_DEBUG=TRACE
+#export FLAGCX_DEBUG_SUBSY
+
+run_test() {
+    local script_dir
+    script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+    torchrun \
+        --nproc_per_node=2 \
+        --nnodes=1 \
+        --node_rank=0 \
+        --master_addr=localhost \
+        --master_port=8333 \
+        "${script_dir}/test_tle_d2d_barrier.py"
+}
+
+run_test
+
+
+if [ $? -ne 0 ]; then
+    echo "ERROR: $CMD failed"
+    exit 1
+fi

python/test/tle/unit/test_tle_distributed_d2d.py

@@ -32,8 +32,8 @@
 
 @triton.jit
 def lsa_read_kernel(
-    dev_mem_ptr,
     output_ptr,
+    device_dptr: tl.constexpr,
     N: tl.constexpr,
     MY_RANK: tl.constexpr,
     N_RANKS: tl.constexpr,
@@ -42,7 +42,7 @@ def lsa_read_kernel(
     peer = (MY_RANK + 1) % N_RANKS
 
     remote_mem = tle.remote(
-        dev_mem_ptr,
+        device_dptr,
         space="device",
         dtype=tl.float32,
         shard_id=peer,
@@ -77,18 +77,16 @@ def main():
           f"stride={buf_tensor.stride()}, "
           f"sample={buf_tensor[:4].tolist()}")
 
-    dev_comm_ptr, dev_mem_ptr = tle.create_comm_tensor(buf_tensor)
-
-    print(f"[Rank {rank}] dev_comm_ptr={dev_comm_ptr:#x}, dev_mem_ptr={dev_mem_ptr:#x}")
+    device_dptr = tle.create_dist_tensor(buf_tensor)
 
     output = torch.zeros(N, dtype=torch.float32, device="cuda")
 
     dist.barrier()
 
     grid = (N, )
     lsa_read_kernel[grid](
-        dev_mem_ptr,
         output,
+        device_dptr=device_dptr,
         N=N,
         MY_RANK=rank,
         N_RANKS=world_size,