Merge branch 'main' into add_nv_dist_ci_20260810

Author: zhzhcookie

cmake/FlagTreeOptions.cmake

@@ -454,6 +454,14 @@ function(flagtree_add_tle_generated_header_dependencies)
     list(APPEND _flagtree_tle_codegen_deps TritonTLETransformsIncGen)
   endif()
 
+  set(_flagtree_enflame_tle_header_targets
+      MLIRTritonToGCU_gcu300
+      MLIRTritonToGCU_gcu400
+      MLIRGCUTritonToTritonGPU_gcu400
+      MLIRTritonGCUTransforms_gcu400
+      triton_gcu300_core
+      triton_gcu400_core)
+
   # Native compiler targets include TLE generated headers under __TLE__ guards.
   # The TLE dialect is added after the core libraries, so the dependency must be
   # attached explicitly once the TLE tablegen targets exist; otherwise a clean
@@ -466,16 +474,21 @@ function(flagtree_add_tle_generated_header_dependencies)
       TritonNVIDIAGPUToLLVM
       TritonGPUToLLVM
       NVHopperTransforms
+      ${_flagtree_enflame_tle_header_targets}
       triton
       triton-opt
       triton-reduce
       triton-lsp
       triton-llvm-opt
       triton-tensor-layout)
-    if(TARGET ${_flagtree_tle_header_target})
-      add_dependencies(${_flagtree_tle_header_target}
-        ${_flagtree_tle_codegen_deps})
-    endif()
+    foreach(_flagtree_tle_dependency_target IN ITEMS
+        ${_flagtree_tle_header_target}
+        obj.${_flagtree_tle_header_target})
+      if(TARGET ${_flagtree_tle_dependency_target})
+        add_dependencies(${_flagtree_tle_dependency_target}
+          ${_flagtree_tle_codegen_deps})
+      endif()
+    endforeach()
   endforeach()
 endfunction()
 

include/triton/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -1,3 +1,3 @@
-set(LLVM_TARGET_DEFINITIONS Passes.td)
+flagtree_spec_td_set(LLVM_TARGET_DEFINITIONS Passes.td)
 mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonGPU)
 add_public_tablegen_target(TritonGPUTransformsIncGen)

python/setup_tools/utils/enflame.py

@@ -85,19 +85,6 @@ def install_extension(*args, **kargs):
 
     project_root_dir = cmake_dir.parent.parent
 
-    # Modify nvidia driver's is_active() to return False for enflame backend
-    drvfile = project_root_dir / 'third_party' / 'nvidia' / 'backend' / 'driver.py'
-    if drvfile.exists():
-        with open(drvfile, 'r') as f:
-            lines = f.readlines()
-        for i, line in enumerate(lines):
-            if 'def is_active():' in line:
-                if i + 1 < len(lines) and 'return False' not in lines[i + 1]:
-                    lines.insert(i + 1, '        return False\n')
-                break
-        with open(drvfile, 'w') as f:
-            f.writelines(lines)
-
     dst_dir = project_root_dir / "third_party" / "enflame" / "backend"
     dst_dir.mkdir(parents=True, exist_ok=True)
 

python/test/tle/unit/test_tle_d2d_barrier.sh

@@ -31,6 +31,15 @@ export FLAGCX_P2P_DISABLE=1
 #export FLAGCX_DEBUG=TRACE
 #export FLAGCX_DEBUG_SUBSY
 
+port=8333
+# Check whether port is occupied
+while ss -ltn | grep -q ":${port} "; do
+    echo "Port ${port} is occupied, trying next..."
+    port=$((port + 2))
+done
+
+echo "Using master_port=${port}"
+
 run_test() {
     local script_dir
     script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -40,7 +49,7 @@ run_test() {
         --nnodes=1 \
         --node_rank=0 \
         --master_addr=localhost \
-        --master_port=8333 \
+        --master_port=${port} \
         "${script_dir}/test_tle_d2d_barrier.py"
 }
 

python/test/tle/unit/test_tle_distributed_d2d.sh

@@ -8,6 +8,15 @@ export FLAGCX_MEM_ENABLE=1
 export FLAGCX_VMM_ENABLE=0
 export FLAGCX_P2P_DISABLE=1
 
+port=8333
+# Check whether port is occupied
+while ss -ltn | grep -q ":${port} "; do
+    echo "Port ${port} is occupied, trying next..."
+    port=$((port + 2))
+done
+
+echo "Using master_port=${port}"
+
 run_test() {
     local script_dir
     script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -17,7 +26,7 @@ run_test() {
         --nnodes=1 \
         --node_rank=0 \
         --master_addr=localhost \
-        --master_port=8333 \
+        --master_port=${port} \
         "${script_dir}/test_tle_distributed_d2d.py"
 }
 

python/test/tle/unit/test_tle_get_local_pe.sh

@@ -17,7 +17,6 @@ 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
@@ -30,6 +29,17 @@ export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
 #export FLAGCX_DMABUF_ENABLE=1
 #export FLAGCX_DEBUG=TRACE
 #export FLAGCX_DEBUG_SUBSY
+
+
+port=8333
+# Check whether port is occupied
+while ss -ltn | grep -q ":${port} "; do
+    echo "Port ${port} is occupied, trying next..."
+    port=$((port + 2))
+done
+
+echo "Using master_port=${port}"
+
 run_test() {
     local script_dir
     script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
@@ -39,7 +49,7 @@ run_test() {
         --nnodes=1 \
         --node_rank=0 \
         --master_addr=localhost \
-        --master_port=8333 \
+        --master_port=${port} \
         "${script_dir}/test_tle_get_local_pe.py"
 }
 

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

@@ -157,7 +157,7 @@ def cleanup_communicator():
 
 def init_communicator():
     global comm, rank, _init_communicator_
-    if _init_communicator_:
+    if enabled and _init_communicator_:
         return
     dist.init_process_group(backend="nccl")
     rank = dist.get_rank()