Merge branch 'main' into feature/gpu_set_layout&gpu_alloc

Author: ph0375

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

@@ -106,7 +106,8 @@ jobs:
           python3 -m pytest 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_d2d_barrier.py
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py \
+              --ignore=python/test/tle/unit/test_tle_get_node_rank.py
 
           ## tle raw test
           python3 -m pytest third_party/enflame/python/test/tle/raw

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

@@ -130,7 +130,8 @@ jobs:
           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_d2d_barrier.py
+              --ignore=python/test/tle/unit/test_tle_d2d_barrier.py \
+              --ignore=python/test/tle/unit/test_tle_get_node_rank.py
           ## 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

python/test/tle/unit/test_tle_distributed.py

@@ -226,6 +226,16 @@ def __init__(self):
         self.builder = _LegacyDistributedBarrierBuilder()
 
 
+class TestShardId:
+
+    @pytest.mark.parametrize("axis", ("device", "node"))
+    def test_rank_axis_requires_device_dptr(self, axis):
+        mesh = tle.device_mesh({"node": 2, "device": 4})
+        semantic = _FakeSemantic()
+        with pytest.raises(ValueError, match=rf"device_dptr is required for axis '{axis}'"):
+            tle.shard_id(mesh, axis, _semantic=semantic)
+
+
 class TestDistributedBarrierScope:
 
     def test_distributed_barrier_full_cluster_mesh(self):

python/test/tle/unit/test_tle_get_node_rank.py

@@ -0,0 +1,77 @@
+import os
+
+import torch
+import torch.distributed as dist
+import triton
+import triton.experimental.tle.language as tle
+import triton.language as tl
+
+LOCAL_WORLD_SIZE = int(os.environ["LOCAL_WORLD_SIZE"])
+WORLD_SIZE = int(os.environ["WORLD_SIZE"])
+if WORLD_SIZE % LOCAL_WORLD_SIZE != 0:
+    raise ValueError("WORLD_SIZE must be divisible by LOCAL_WORLD_SIZE")
+
+DEVICE_MESH = tle.device_mesh(tle.MeshConfig(node=WORLD_SIZE // LOCAL_WORLD_SIZE, device=LOCAL_WORLD_SIZE))
+
+
+@triton.jit
+def _tle_node_rank_kernel(out_ptr, device_dptr: tl.constexpr, mesh: tl.constexpr):
+    pid = tl.program_id(0)
+    node_rank = tle.shard_id(mesh, "node", device_dptr=device_dptr)
+    tl.store(out_ptr + pid, node_rank)
+
+
+def test_tle_get_node_rank():
+    grid = 2
+    with torch.cuda.use_mem_pool(tle.get_mem_pool()):
+        source = torch.empty((1, ), dtype=torch.float32, device="cuda")
+    device_dptr = tle.create_dist_tensor(source)
+    node_rank_out = torch.empty((grid, ), dtype=torch.int32, device="cuda")
+
+    compiled = _tle_node_rank_kernel.warmup(
+        out_ptr=node_rank_out,
+        device_dptr=device_dptr,
+        mesh=DEVICE_MESH,
+        grid=(grid, ),
+        num_ctas=1,
+        num_warps=4,
+    )
+    assert "get_world_rank" in compiled.asm["ttgir"]
+    assert "get_num_pes" in compiled.asm["ttgir"]
+    assert "flagcxDevCommGetRank" in compiled.asm["ptx"]
+    assert "flagcxDevCommGetIntraSize" in compiled.asm["ptx"]
+
+    _tle_node_rank_kernel[(grid, )](
+        out_ptr=node_rank_out,
+        device_dptr=device_dptr,
+        mesh=DEVICE_MESH,
+    )
+    torch.cuda.synchronize()
+
+    rank = dist.get_rank()
+    expected_node_rank = rank // LOCAL_WORLD_SIZE
+    actual_node_ranks = node_rank_out.cpu().tolist()
+    try:
+        torch.testing.assert_close(
+            node_rank_out,
+            torch.full_like(node_rank_out, expected_node_rank),
+        )
+    except AssertionError:
+        print(
+            f"[Rank {rank}] FAILED: node ranks={actual_node_ranks}, "
+            f"expected={expected_node_rank}",
+            flush=True,
+        )
+        raise
+    else:
+        print(
+            f"[Rank {rank}] PASSED: node ranks={actual_node_ranks}, "
+            f"expected={expected_node_rank}",
+            flush=True,
+        )
+    finally:
+        tle.cleanup_communicator()
+
+
+if __name__ == "__main__":
+    test_tle_get_node_rank()

python/test/tle/unit/test_tle_get_node_rank.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+if [ "$1" = "debug" ]; then
+    export NCCL_DEBUG=INFO
+    export NCCL_DEBUG_SUBSYS=all
+else
+    unset NCCL_DEBUG
+    unset NCCL_DEBUG_SUBSYS
+fi
+
+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
+
+nproc_per_node=${NPROC_PER_NODE:-2}
+nnodes=${NNODES:-2}
+node_rank=${NODE_RANK:-0}
+master_addr=${MASTER_ADDR:-10.0.9.3}
+port=${MASTER_PORT:-8335}
+
+if [ "${nnodes}" -eq 1 ]; then
+    while ss -ltn | grep -q ":${port} "; do
+        echo "Port ${port} is occupied, trying next..."
+        port=$((port + 2))
+    done
+fi
+
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+echo "Using master ${master_addr}:${port}, node ${node_rank}/${nnodes}"
+
+torchrun \
+    --nproc_per_node="${nproc_per_node}" \
+    --nnodes="${nnodes}" \
+    --node_rank="${node_rank}" \
+    --master_addr="${master_addr}" \
+    --master_port="${port}" \
+    "${script_dir}/test_tle_get_node_rank.py"

python/triton/experimental/tle/language/distributed.py

@@ -71,12 +71,23 @@ def _get_local_rank(device_dptr, _semantic=None, ret_dtype=tl.int32):
     return tl.tensor(result, ret_dtype)
 
 
-# The number of devices in the world
+# Get the current world rank
+@tl.builtin
+def _get_world_rank(device_dptr, _semantic=None, ret_dtype=tl.int32):
+    builder = _semantic.builder
+    ret_ir_ty = ret_dtype.to_ir(builder)
+    ptr = _parse_src_arg(builder, device_dptr, 1)
+    result = builder.get_world_rank(ret_ir_ty, ptr)
+    return tl.tensor(result, ret_dtype)
+
+
+# The number of devices on the current node
 @tl.builtin
 def n_pes(dev_mem_ptr, _semantic=None, ret_dtype=tl.int32):
     builder = _semantic.builder
     ret_ir_ty = ret_dtype.to_ir(builder)
-    result = builder.get_n_pes(ret_ir_ty, dev_mem_ptr.handle)
+    ptr = _parse_src_arg(builder, dev_mem_ptr, 1)
+    result = builder.get_n_pes(ret_ir_ty, ptr)
     return tl.tensor(result, ret_dtype)
 
 
@@ -639,13 +650,21 @@ def shard_id(
     Return current shard coordinate on the given launch mesh axis.
 
     `axis` can be axis name (`str`) or axis index (`int`, supports negative).
+    `device` returns the intra-node rank; `node` returns the inter-node rank.
     The returned value is a scalar int32 tensor.
     """
     mesh = tl._unwrap_if_constexpr(mesh)
     axis = tl._unwrap_if_constexpr(axis)
 
-    if axis in ("device", "node"):
+    if axis in ("device", "node") and device_dptr is None:
+        raise ValueError(f"device_dptr is required for axis {axis!r}")
+
+    if axis == "device":
         return _get_local_rank(device_dptr, _semantic=_semantic, ret_dtype=tl.int32)
+    if axis == "node":
+        world_rank = _get_world_rank(device_dptr, _semantic=_semantic, ret_dtype=tl.int32)
+        local_world_size = n_pes(device_dptr, _semantic=_semantic, ret_dtype=tl.int32)
+        return _semantic.floordiv(world_rank, local_world_size)
 
     if not isinstance(mesh, device_mesh):
         raise TypeError(f"mesh must be device_mesh, got {type(mesh).__name__}")

third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h

@@ -23,6 +23,7 @@
 
 #include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h"
 #include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h"
+#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetWorldRankOpToLLVM.h"
 
 namespace mlir::triton::tle {
 void populateFlagCxOpToLLVMPatterns(LLVMTypeConverter &typeConverter,

third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetWorldRankOpToLLVM.h

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2025-     FlagOS Contributors
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining
+ * a copy of this software and associated documentation files
+ * (the "Software"), to deal in the Software without restriction,
+ * including without limitation the rights to use, copy, modify, merge,
+ * publish, distribute, sublicense, and/or sell copies of the Software,
+ * and to permit persons to whom the Software is furnished to do so,
+ * subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be
+ * included in all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+ * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+ * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+ * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
+ * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+ * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
+ * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef TLE_CONVERSION_TLETOLLVM_FLAGCXOPTOLLVM_GETWORLDRANKOPTOLLVM_H
+#define TLE_CONVERSION_TLETOLLVM_FLAGCXOPTOLLVM_GETWORLDRANKOPTOLLVM_H
+
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
+
+namespace mlir::triton::tle {
+
+void populateGetWorldRankOpToLLVMPatterns(LLVMTypeConverter &typeConverter,
+                                          RewritePatternSet &patterns,
+                                          PatternBenefit benefit);
+
+} // namespace mlir::triton::tle
+
+#endif

third_party/tle/dialect/include/IR/TleOps.td

@@ -384,7 +384,7 @@ def Tle_RemotePointersOp : Tle_Op<"remote_pointers", [Pure, AttrSizedOperandSegm
   let hasVerifier = 1;
 }
 def Tle_GetNumPesOp : Tle_Op<"get_num_pes"> {
-  let summary = "Get nume pes";
+  let summary = "Get the number of intra-node PEs";
   let arguments = (ins
     Tle_LocalPointerResultType:$src
   );
@@ -401,6 +401,15 @@ def Tle_GetDeviceIdOp : Tle_Op<"get_device_id"> {
   let hasVerifier = 1;
 }
 
+def Tle_GetWorldRankOp : Tle_Op<"get_world_rank"> {
+  let summary = "Get global PE rank";
+  let arguments = (ins
+     Tle_LocalPointerResultType:$input
+  );
+  let results = (outs I32:$result);
+  let hasVerifier = 1;
+}
+
 
 def Tle_DSLRegionOp : Tle_Op<"dsl_region", [IsolatedFromAbove, MemDescViewTrait,
                                             DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {

third_party/tle/dialect/include/Tools/FlagcxUtils.h

@@ -35,6 +35,10 @@ LLVM::CallOp getLocalPeFuncCall(mlir::Location loc,
                                 ConversionPatternRewriter &rewriter,
                                 Value memPtrInt);
 
+LLVM::CallOp getWorldRankFuncCall(mlir::Location loc,
+                                  ConversionPatternRewriter &rewriter,
+                                  Value memPtrInt);
+
 LLVM::CallOp getNumPesFunCall(mlir::Location loc,
                               ConversionPatternRewriter &rewriter,
                               Value memPtrInt);