From eabb6aeda207e7be8cfe2cdc7dae0484efaee82f Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Thu, 2 Jul 2026 17:49:03 +0800 Subject: [PATCH 01/32] init device intra node barrier --- third_party/tle/dialect/include/IR/TleOps.td | 7 ++- .../tle/dialect/include/IR/VerfiyUtils.h | 4 ++ third_party/tle/dialect/lib/IR/Ops.cpp | 4 ++ .../tle/dialect/lib/IR/VerfiyUtils.cpp | 47 +++++++++++++++++++ .../tle/dialect/lib/Tools/FlagcxUtils.cpp | 41 ++++++++++++---- third_party/tle/triton_tle.cc | 29 +++++++++++- 6 files changed, 120 insertions(+), 12 deletions(-) diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index c927258f7..1b38636db 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/third_party/tle/dialect/include/IR/TleOps.td @@ -227,13 +227,17 @@ def Tle_TMAStoreCommitGroupOp def Tle_DistributedBarrierOp : Tle_Op<"distributed_barrier", [MemoryEffects<[MemRead, MemWrite]>]> { let arguments = (ins + Optional:$src, + OptionalAttr:$space, + OptionalAttr:$barrier_type, + OptionalAttr:$order, OptionalAttr:$group_kind, + OptionalAttr:$barrier_index, OptionalAttr:$group_rank, OptionalAttr:$group_shape, OptionalAttr:$group_axes, OptionalAttr:$group_mask ); - let assemblyFormat = "attr-dict"; let hasVerifier = 1; } @@ -265,6 +269,7 @@ def Tle_GetDeviceIdOp : Tle_Op<"get_device_id"> { let hasVerifier = 1; } + def Tle_DSLRegionOp : Tle_Op<"dsl_region", [IsolatedFromAbove, MemDescViewTrait, DeclareOpInterfaceMethods]> { let arguments = (ins diff --git a/third_party/tle/dialect/include/IR/VerfiyUtils.h b/third_party/tle/dialect/include/IR/VerfiyUtils.h index 24d2e45e8..dc1dbe0fb 100644 --- a/third_party/tle/dialect/include/IR/VerfiyUtils.h +++ b/third_party/tle/dialect/include/IR/VerfiyUtils.h @@ -17,4 +17,8 @@ namespace RemotePointers { llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result); } +namespace DistributedBarrier { +llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op); +} + } // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index 334c56a06..745eb9c13 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -680,6 +680,10 @@ LogicalResult ExclusiveCumsumOp::verify() { LogicalResult DistributedBarrierOp::verify() { auto *op = getOperation(); + auto spaceAttr = op->getAttrOfType("space"); + if (spaceAttr && spaceAttr.getValue() != "device") + return DistributedBarrier::verifyDeviceSpace(op); + auto kindAttr = op->getAttrOfType("group_kind"); auto rankAttr = op->getAttrOfType("group_rank"); auto shapeAttr = op->getAttrOfType("group_shape"); diff --git a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp index 753faccae..9e6bbe395 100644 --- a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp +++ b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp @@ -32,4 +32,51 @@ llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result) { } } // namespace RemotePointers +namespace DistributedBarrier { +llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op) { + auto emitInvalidAttr = [&](StringRef attrName, StringRef value, + StringRef expected) -> LogicalResult { + return op->emitOpError() << "invalid " << attrName << " '" << value + << "', expected one of: " << expected; + }; + auto kindAttr = op->getAttrOfType("group_kind"); + auto barrierTypeAttr = op->getAttrOfType("barrier_type"); + auto orderAttr = op->getAttrOfType("order"); + + bool kindValid = llvm::StringSwitch(kindAttr.getValue()) + .Case("thread", true) + .Case("warp", true) + .Case("block", true) + .Case("tile_span", true) + .Case("lanes", true) + .Default(false); + + if (!kindValid) { + return emitInvalidAttr("group_kind", kindAttr.getValue(), + "thread, warp, block, tile_span, lanes"); + } + bool barrierTypeValid = llvm::StringSwitch(barrierTypeAttr.getValue()) + .Case("arrive", true) + .Case("wait", true) + .Case("sync", true) + .Default(false); + if (!barrierTypeValid) { + return emitInvalidAttr("barrier_type", barrierTypeAttr.getValue(), + "arrive, wait, sync"); + } + bool orderValid = llvm::StringSwitch(orderAttr.getValue()) + .Case("acqrel", true) + .Case("acquire", true) + .Case("release", true) + .Case("relaxed", true) + .Default(false); + if (!orderValid) { + return emitInvalidAttr("order", orderAttr.getValue(), + "relaxed, acqrel, acquire, release"); + } + return success(); +} + +} // namespace DistributedBarrier + } // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp index bc7bdd3ea..f25433557 100644 --- a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp +++ b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp @@ -13,7 +13,9 @@ using namespace mlir; static const llvm::StringMap runtimeNames = { {"getLocalPeFunction", "flagcxDevCommGetIntraRank"}, {"getNumPesFunction", "flagcxDevCommGetIntraSize"}, - {"getReadSignalFunction", "flagcxDevNetReadSignal"}}; + {"getIntraBarrierArriveSignalFunction", "flagcxIntraBarrierArriveS"}, + {"getIntraBarrierWaitSignalFunction", "flagcxIntraBarrierWaitS"}, + {"getIntraBarrierSyncSignalFunction", "flagcxIntraBarrierSyncS"}}; static inline LLVM::LLVMFuncOp createFuncInstance(const char *funcName, ModuleOp module, @@ -57,22 +59,43 @@ LLVM::CallOp getNumPesFunCall(mlir::Location loc, FlatSymbolRefAttr::get(func), ValueRange{comm_dev_ptr}); } -LLVM::CallOp getReadSignalFunCall(mlir::Location loc, - ConversionPatternRewriter &rewriter, - Value memPtrInt) { +LLVM::CallOp getBarrierFuncCall(mlir::Location loc, + ConversionPatternRewriter &rewriter, Value comm, + IntegerAttr barrier_index, IntegerAttr coopKind, + IntegerAttr order, StringAttr barrierType) { auto ctx = rewriter.getContext(); ModuleOp module = rewriter.getInsertionPoint()->getParentOp()->getParentOfType(); auto PtrTy = LLVM::LLVMPointerType::get(ctx, 1); auto i32Ty = IntegerType::get(ctx, 32); - auto func = createFuncInstance( - runtimeNames.lookup("getNumPesFunction").data(), module, {PtrTy}, i32Ty); - - auto comm_dev_ptr = getFlagcxMemPtr(loc, rewriter, memPtrInt); + auto i1Ty = IntegerType::get(ctx, 1); + auto funcName = ""; + if (barrierType.getValue() == "arrive") { + funcName = "getIntraBarrierArriveSignalFunction"; + } else if (barrierType.getValue() == "wait") { + funcName = "getIntraBarrierWaitSignalFunction"; + } else if (barrierType.getValue() == "sync") { + funcName = "getIntraBarrierSyncSignalFunction"; + } else { + llvm_unreachable("Unknown barrier type"); + } + + auto func = createFuncInstance(runtimeNames.lookup(funcName).data(), module, + {PtrTy, i32Ty, i32Ty, i1Ty, i32Ty}, i32Ty); + + auto comm_dev_ptr = getFlagcxMemPtr(loc, rewriter, comm); + auto falseVal = + rewriter.create(loc, i1Ty, rewriter.getBoolAttr(false)); + auto barrierIndexVal = + rewriter.create(loc, i32Ty, barrier_index); + auto coopKindVal = rewriter.create(loc, i32Ty, coopKind); + auto orderVal = rewriter.create(loc, i32Ty, order); return rewriter.create( loc, TypeRange{func.getFunctionType().getReturnType()}, - FlatSymbolRefAttr::get(func), ValueRange{comm_dev_ptr}); + FlatSymbolRefAttr::get(func), + ValueRange{comm_dev_ptr, coopKindVal, barrierIndexVal, falseVal, + orderVal}); } LLVM::CallOp getLocalPeFuncCall(mlir::Location loc, diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index f26cb1158..fc7f9068a 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -361,7 +361,31 @@ void init_triton_tle_ir(py::module &&m) { .def("create_distributed_barrier", [](TritonOpBuilder &self) -> void { self.create( - StringAttr(), IntegerAttr(), DenseI32ArrayAttr(), + Value(), StringAttr(), StringAttr(), StringAttr(), + StringAttr(), IntegerAttr(), IntegerAttr(), + DenseI32ArrayAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr()); + }) + .def("create_distributed_barrier", + [](TritonOpBuilder &self, Value src, size_t barrier_index = 0, + const std::string &space = "device", + const std::string &groupKind = "block", + const std::string &order = "acqrel", + const std::string &barrierType = "sync") -> void { + auto &builder = self.getBuilder(); + auto *ctx = builder.getContext(); + auto getOptStrAttr = [&](const std::string &s) -> StringAttr { + return s.empty() ? StringAttr() : builder.getStringAttr(s); + }; + auto spaceAttr = getOptStrAttr(space); + auto kindAttr = getOptStrAttr(groupKind); + auto orderAttr = getOptStrAttr(order); + auto barrierTypeAttr = getOptStrAttr(barrierType); + auto barrierIndexAttr = + builder.getI32IntegerAttr(static_cast(barrier_index)); + + self.create( + src, spaceAttr, barrierTypeAttr, orderAttr, kindAttr, + barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr()); }) .def( @@ -399,7 +423,8 @@ void init_triton_tle_ir(py::module &&m) { } self.create( - kindAttr, rankAttr, shapeAttr, axesAttr, maskAttr); + Value(), StringAttr(), StringAttr(), StringAttr(), kindAttr, + IntegerAttr(), rankAttr, shapeAttr, axesAttr, maskAttr); }, py::arg("group_kind"), py::arg("group_shape"), py::arg("group_axes"), py::arg("group_mask")) From d4563fd02df861c6d6d2810f1a6e2180b6a49b6e Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 3 Jul 2026 17:56:47 +0800 Subject: [PATCH 02/32] add conversions --- .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 4 +- .../DeviceIntraBarrierOpToLLVM.h | 15 +++++ .../tle/dialect/include/IR/FlagCxOps.td | 33 ++++++++++ .../tle/dialect/include/IR/VerfiyUtils.h | 2 +- .../tle/dialect/include/Tools/FlagcxUtils.h | 5 ++ .../lib/Conversion/TleToLLVM/CMakeLists.txt | 2 + .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 51 ++++++++++++++++ .../TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt | 15 +++++ .../DeviceIntraBarrierOpToLLVM.cpp | 61 +++++++++++++++++++ third_party/tle/dialect/lib/IR/FlagCxOps.cpp | 48 +++++++++++++++ third_party/tle/dialect/lib/IR/Ops.cpp | 5 +- .../tle/dialect/lib/IR/VerfiyUtils.cpp | 46 +++----------- .../tle/dialect/lib/Tools/FlagcxUtils.cpp | 10 +-- 13 files changed, 249 insertions(+), 48 deletions(-) create mode 100644 third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h create mode 100644 third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt create mode 100644 third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index b3c9b193c..0b6e57653 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -180,8 +180,6 @@ struct ConvertTritonGPUToLLVM benefit); mlir::triton::tle::populateDistributedBarrierOpToLLVMPatterns( typeConverter, patterns, benefit); - // mlir::triton::tle::populateGetNumPesOpToLLVMPatterns( - // typeConverter, patterns, benefit + 1); mlir::triton::tle::populateLocalPointersOpToLLVMPatterns( typeConverter, targetInfo, patterns, benefit); mlir::triton::tle::populateExtractTileOpToLLVMPatterns( @@ -208,6 +206,8 @@ struct ConvertTritonGPUToLLVM typeConverter, patterns, benefit); mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, patterns, benefit); + mlir::triton::tle::populateDistributedBarrierOpToLLVMPatterns( + typeConverter, patterns, benefit); } #endif diff --git a/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h b/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h new file mode 100644 index 000000000..fe002abed --- /dev/null +++ b/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h @@ -0,0 +1,15 @@ +#ifndef TLE_CONVERSION_TLE_TO_LLVM_DEVICE_INTRA_BARRIER_OP_TO_LLVM_H +#define TLE_CONVERSION_TLE_TO_LLVM_DEVICE_INTRA_BARRIER_OP_TO_LLVM_H + +#include "mlir/Conversion/LLVMCommon/TypeConverter.h" +#include "mlir/IR/PatternMatch.h" + +namespace mlir::triton::tle { + +void populateDeviceIntraBarrierOpToLLVMPatterns( + LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, + PatternBenefit benefit); + +} // namespace mlir::triton::tle + +#endif // TLE_CONVERSION_TLE_TO_LLVM_DEVICE_INTRA_BARRIER_OP_TO_LLVM_H diff --git a/third_party/tle/dialect/include/IR/FlagCxOps.td b/third_party/tle/dialect/include/IR/FlagCxOps.td index 7c283d78b..e882c8d7f 100644 --- a/third_party/tle/dialect/include/IR/FlagCxOps.td +++ b/third_party/tle/dialect/include/IR/FlagCxOps.td @@ -23,3 +23,36 @@ def Tle_GetLocalRankOp : Tle_Op<"get local rank"> { let results = (outs I32:$result); let hasVerifier = 1; } + + +def Tle_DeviceIntraBarrierOp : Tle_Op<"device intra barrier"> { + let summary = "Device-local synchronization barrier"; + let description = [{ + Performs a device-local synchronization barrier by lowering to one of the + following runtime calls: + + * flagcxIntraBarrierArriveS(comm, coopKind, index, multimem, order) + * flagcxIntraBarrierWaitS(comm, coopKind, index, multimem, order) + * flagcxIntraBarrierSyncS(comm, coopKind, index, multimem, order) + + The specific runtime routine is selected according to the barrier operation + (arrive, wait, or sync). + + The barrier is identified by `index`, scoped by `coopKind`, and obeys the + specified memory ordering semantics (`order`). `multimem` indicates whether + the barrier operates on multimem storage. + + Returns the runtime status code. + }]; + + let arguments = (ins + Tle_LocalPointerResultType:$comm, + OptionalAttr:$barrier_type, + OptionalAttr:$coop_kind, + OptionalAttr:$index, + OptionalAttr:$multimem, + OptionalAttr:$order + ); + + let hasVerifier = 1; +} diff --git a/third_party/tle/dialect/include/IR/VerfiyUtils.h b/third_party/tle/dialect/include/IR/VerfiyUtils.h index dc1dbe0fb..bdbf2fdf2 100644 --- a/third_party/tle/dialect/include/IR/VerfiyUtils.h +++ b/third_party/tle/dialect/include/IR/VerfiyUtils.h @@ -18,7 +18,7 @@ llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result); } namespace DistributedBarrier { -llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op); +llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op, mlir::Value src); } } // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/include/Tools/FlagcxUtils.h b/third_party/tle/dialect/include/Tools/FlagcxUtils.h index 84b5ecbfa..144f0add8 100644 --- a/third_party/tle/dialect/include/Tools/FlagcxUtils.h +++ b/third_party/tle/dialect/include/Tools/FlagcxUtils.h @@ -16,4 +16,9 @@ LLVM::CallOp getNumPesFunCall(mlir::Location loc, ConversionPatternRewriter &rewriter, Value memPtrInt); +LLVM::CallOp getBarrierFuncCall(mlir::Location loc, + ConversionPatternRewriter &rewriter, Value comm, + size_t barrier_index, size_t coopKind, + size_t order, llvm::StringRef barrierType); + } // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt index 0a634494a..6fced8154 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt @@ -11,6 +11,8 @@ set(TLE_TO_LLVM_SRC WGMMAFenceOpToLLVM.cpp ) +add_subdirectory(FlagCxOpToLLVM) + if(FLAGCX_ENABLED) list(APPEND TLE_TO_LLVM_SRC GetDeviceIdToFlagCX.cpp diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index 14a9950e1..bf61013cf 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -1,4 +1,5 @@ #include "tle/dialect/include/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.h" +#include "tle/dialect/include/Tools/FlagcxUtils.h" #include "mlir/Conversion/LLVMCommon/Pattern.h" #include "mlir/Dialect/GPU/IR/GPUDialect.h" @@ -18,6 +19,9 @@ namespace { using namespace mlir; namespace tle = mlir::triton::tle; +constexpr llvm::StringLiteral kSpaceAttr = "space"; +constexpr llvm::StringLiteral kOrderAttr = "order"; +constexpr llvm::StringLiteral kIndexAttr = "index"; constexpr llvm::StringLiteral kGroupKindAttr = "group_kind"; constexpr llvm::StringLiteral kGroupShapeAttr = "group_shape"; constexpr llvm::StringLiteral kGroupMaskAttr = "group_mask"; @@ -422,9 +426,56 @@ struct DistributedBarrierOpConversion return success(); } + LogicalResult + lowerDeviceSpaceBarrier(tle::DistributedBarrierOp op, + ConversionPatternRewriter &rewriter) const { + auto kindAttr = op->getAttrOfType(kGroupKindAttr); + auto orderAttr = op->getAttrOfType(kOrderAttr); + auto indexAttr = op->getAttrOfType(kIndexAttr); + auto getCoopKindValue = [](StringRef kind) -> int32_t { + return llvm::StringSwitch(kind) + .Case("thread", 0) + .Case("warp", 1) + .Case("block", 2) + .Case("grid", 3) + .Default(-1); + }; + auto getOrderValue = [](StringRef order) -> int32_t { + return llvm::StringSwitch(order) + .Case("relaxed", 0) + .Case("acquire", 1) + .Case("release", 2) + .Case("acqrel", 3) + .Default(-1); + }; + + int32_t coopKind = getCoopKindValue(kindAttr.getValue()); + int32_t order = getOrderValue(orderAttr.getValue()); + if (coopKind < 0) + return rewriter.notifyMatchFailure(op, "invalid coop_kind"); + + if (order < 0) + return rewriter.notifyMatchFailure(op, "invalid order"); + + auto comm = op.getSrc(); + auto coopKindAttr = rewriter.getI32IntegerAttr(coopKind); + auto newOrderAttr = rewriter.getI32IntegerAttr(order); + auto barrierTypeAttr = op.getBarrierTypeAttr(); + auto multimemAttr = rewriter.getBoolAttr(false); + + rewriter.replaceOpWithNewOp( + op, comm, barrierTypeAttr, coopKindAttr, indexAttr, multimemAttr, + newOrderAttr); + + return success(); + } LogicalResult matchAndRewrite(tle::DistributedBarrierOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + if (auto spaceAttr = op->getAttrOfType(kSpaceAttr)) + if (spaceAttr.getValue() == "device") + return lowerDeviceSpaceBarrier(op, rewriter); + if (auto kindAttr = op->getAttrOfType(kGroupKindAttr)) { if (kindAttr.getValue() == "grid") return lowerGridBarrier(op, rewriter); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt new file mode 100644 index 000000000..bdc429145 --- /dev/null +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt @@ -0,0 +1,15 @@ +set(FLAGCX_TO_LLVM_SRC + DeviceIntraBarrierOpToLLVM.cpp +) + +add_triton_library(FlagCxToLLVM + ${FLAGCX_TO_LLVM_SRC} + + DEPENDS + TleTableGen + TritonGPUTypeInterfacesIncGen + TleTools + + LINK_LIBS PUBLIC + TleTools +) diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp new file mode 100644 index 000000000..0f08cea47 --- /dev/null +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp @@ -0,0 +1,61 @@ +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h" +#include "tle/dialect/include/Tools/FlagcxUtils.h" + +#include "mlir/Conversion/LLVMCommon/Pattern.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMTypes.h" +#include "mlir/Dialect/LLVMIR/NVVMDialect.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/Transforms/DialectConversion.h" +#include "tle/dialect/include/IR/Dialect.h" +#include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Tools/LayoutUtils.h" +#include "llvm/Support/raw_ostream.h" + +namespace { +using namespace mlir; +namespace ttg = mlir::triton::gpu; +namespace tle = mlir::triton::tle; + +struct DeviceIntraBarrierOpConversion + : public ConvertOpToLLVMPattern { + DeviceIntraBarrierOpConversion(LLVMTypeConverter &typeConverter, + PatternBenefit benefit) + : ConvertOpToLLVMPattern(typeConverter, benefit) {} + + LogicalResult + matchAndRewrite(tle::DeviceIntraBarrierOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op.getLoc(); + + auto indexValue = op.getIndexAttr().getInt(); + ; + auto coopValue = op.getCoopKindAttr().getInt(); + ; + auto orderValue = op.getOrderAttr().getInt(); + ; + auto barrierType = op.getBarrierTypeAttr().getValue(); + + if (!llvm::is_contained(std::array{0, 1, 2, 3, 4}, coopValue)) + return rewriter.notifyMatchFailure(op, "invalid coop_kind"); + + if (!llvm::is_contained(std::array{0, 1, 2, 3}, orderValue)) + return rewriter.notifyMatchFailure(op, "invalid coop_kind"); + + tle::getBarrierFuncCall(loc, rewriter, adaptor.getComm(), indexValue, + coopValue, orderValue, barrierType); + rewriter.eraseOp(op); + return success(); + } +}; + +} // namespace + +void tle::populateDeviceIntraBarrierOpToLLVMPatterns( + LLVMTypeConverter &typeConverter, RewritePatternSet &patterns, + PatternBenefit benefit) { + patterns.add(typeConverter, benefit); +} diff --git a/third_party/tle/dialect/lib/IR/FlagCxOps.cpp b/third_party/tle/dialect/lib/IR/FlagCxOps.cpp index 1af15b524..67a8d5ff2 100644 --- a/third_party/tle/dialect/lib/IR/FlagCxOps.cpp +++ b/third_party/tle/dialect/lib/IR/FlagCxOps.cpp @@ -25,4 +25,52 @@ LogicalResult GetLocalRankOp::verify() { return success(); } +LogicalResult DeviceIntraBarrierOp::verify() { + auto *op = getOperation(); + auto spaceAttr = op->getAttrOfType("space"); + if (spaceAttr && spaceAttr.getValue() != "device") + return success(); + auto emitInvalidAttr = [&](StringRef attrName, StringRef value, + StringRef expected) -> LogicalResult { + return op->emitOpError() << "invalid " << attrName << " '" << value + << "', expected one of: " << expected; + }; + auto kindAttr = op->getAttrOfType("group_kind"); + auto barrierTypeAttr = op->getAttrOfType("barrier_type"); + auto orderAttr = op->getAttrOfType("order"); + + bool kindValid = llvm::StringSwitch(kindAttr.getValue()) + .Case("thread", true) + .Case("warp", true) + .Case("block", true) + .Case("tile_span", true) + .Case("lanes", true) + .Default(false); + + if (!kindValid) { + return emitInvalidAttr("group_kind", kindAttr.getValue(), + "thread, warp, block, tile_span, lanes"); + } + bool barrierTypeValid = llvm::StringSwitch(barrierTypeAttr.getValue()) + .Case("arrive", true) + .Case("wait", true) + .Case("sync", true) + .Default(false); + if (!barrierTypeValid) { + return emitInvalidAttr("barrier_type", barrierTypeAttr.getValue(), + "arrive, wait, sync"); + } + bool orderValid = llvm::StringSwitch(orderAttr.getValue()) + .Case("acqrel", true) + .Case("acquire", true) + .Case("release", true) + .Case("relaxed", true) + .Default(false); + if (!orderValid) { + return emitInvalidAttr("order", orderAttr.getValue(), + "relaxed, acqrel, acquire, release"); + } + return success(); +} + } // namespace mlir::triton::tle diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index 745eb9c13..532b4c7ef 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -681,8 +681,9 @@ LogicalResult ExclusiveCumsumOp::verify() { LogicalResult DistributedBarrierOp::verify() { auto *op = getOperation(); auto spaceAttr = op->getAttrOfType("space"); - if (spaceAttr && spaceAttr.getValue() != "device") - return DistributedBarrier::verifyDeviceSpace(op); + + if (spaceAttr && spaceAttr.getValue() == "device") + return DistributedBarrier::verifyDeviceSpace(op, getSrc()); auto kindAttr = op->getAttrOfType("group_kind"); auto rankAttr = op->getAttrOfType("group_rank"); diff --git a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp index 9e6bbe395..c1d9177fa 100644 --- a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp +++ b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp @@ -33,48 +33,18 @@ llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result) { } // namespace RemotePointers namespace DistributedBarrier { -llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op) { - auto emitInvalidAttr = [&](StringRef attrName, StringRef value, - StringRef expected) -> LogicalResult { - return op->emitOpError() << "invalid " << attrName << " '" << value - << "', expected one of: " << expected; - }; +llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op, mlir::Value src) { + auto kindAttr = op->getAttrOfType("group_kind"); auto barrierTypeAttr = op->getAttrOfType("barrier_type"); auto orderAttr = op->getAttrOfType("order"); - bool kindValid = llvm::StringSwitch(kindAttr.getValue()) - .Case("thread", true) - .Case("warp", true) - .Case("block", true) - .Case("tile_span", true) - .Case("lanes", true) - .Default(false); - - if (!kindValid) { - return emitInvalidAttr("group_kind", kindAttr.getValue(), - "thread, warp, block, tile_span, lanes"); - } - bool barrierTypeValid = llvm::StringSwitch(barrierTypeAttr.getValue()) - .Case("arrive", true) - .Case("wait", true) - .Case("sync", true) - .Default(false); - if (!barrierTypeValid) { - return emitInvalidAttr("barrier_type", barrierTypeAttr.getValue(), - "arrive, wait, sync"); - } - bool orderValid = llvm::StringSwitch(orderAttr.getValue()) - .Case("acqrel", true) - .Case("acquire", true) - .Case("release", true) - .Case("relaxed", true) - .Default(false); - if (!orderValid) { - return emitInvalidAttr("order", orderAttr.getValue(), - "relaxed, acqrel, acquire, release"); - } - return success(); + if (src && kindAttr && barrierTypeAttr && orderAttr) + return success(); + else + return op->emitOpError() + << "expects src, group_kind, barrier_type and order attributes to " + "be present for device space distributed barrier"; } } // namespace DistributedBarrier diff --git a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp index f25433557..19c7d4076 100644 --- a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp +++ b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp @@ -61,8 +61,8 @@ LLVM::CallOp getNumPesFunCall(mlir::Location loc, LLVM::CallOp getBarrierFuncCall(mlir::Location loc, ConversionPatternRewriter &rewriter, Value comm, - IntegerAttr barrier_index, IntegerAttr coopKind, - IntegerAttr order, StringAttr barrierType) { + size_t barrier_index, size_t coopKind, + size_t order, llvm::StringRef barrierType) { auto ctx = rewriter.getContext(); ModuleOp module = rewriter.getInsertionPoint()->getParentOp()->getParentOfType(); @@ -71,11 +71,11 @@ LLVM::CallOp getBarrierFuncCall(mlir::Location loc, auto i32Ty = IntegerType::get(ctx, 32); auto i1Ty = IntegerType::get(ctx, 1); auto funcName = ""; - if (barrierType.getValue() == "arrive") { + if (barrierType == "arrive") { funcName = "getIntraBarrierArriveSignalFunction"; - } else if (barrierType.getValue() == "wait") { + } else if (barrierType == "wait") { funcName = "getIntraBarrierWaitSignalFunction"; - } else if (barrierType.getValue() == "sync") { + } else if (barrierType == "sync") { funcName = "getIntraBarrierSyncSignalFunction"; } else { llvm_unreachable("Unknown barrier type"); From 95029aaa0c96c25bbd3423692237bc03018bca78 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 3 Jul 2026 18:07:46 +0800 Subject: [PATCH 03/32] fix conversion bug --- third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/CMakeLists.txt | 3 +++ .../nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 3 ++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/CMakeLists.txt b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/CMakeLists.txt index 23137749c..3bc3be2bf 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/CMakeLists.txt +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/CMakeLists.txt @@ -4,6 +4,9 @@ if(FLAGTREE_TLE) TleToLLVM # flagtree tle raw TritonTLETransforms ) + if(FLAGCX_ENABLED) + list(APPEND _TLE_LIBS FlagCxToLLVM) + endif() else() set(_TLE_LIBS "") endif() diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 0b6e57653..81964e145 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -25,6 +25,7 @@ #include "tle/dialect/include/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/ExtractOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/GetDeviceIdToFlagCX.h" #include "tle/dialect/include/Conversion/TleToLLVM/GetLocalRankOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/LocalPointersOpToLLVM.h" @@ -206,7 +207,7 @@ struct ConvertTritonGPUToLLVM typeConverter, patterns, benefit); mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, patterns, benefit); - mlir::triton::tle::populateDistributedBarrierOpToLLVMPatterns( + mlir::triton::tle::populateDeviceIntraBarrierOpToLLVMPatterns( typeConverter, patterns, benefit); } #endif From 1a311902ec392569e5c3221ba6dc78fab31a461e Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Mon, 6 Jul 2026 10:16:29 +0000 Subject: [PATCH 04/32] add python api args --- python/test/tle/unit/test_tle_d2d_barrier.py | 47 ++++++++++++++++ python/test/tle/unit/test_tle_d2d_barrier.sh | 53 ++++++++++++++++++ .../experimental/tle/language/distributed.py | 56 ++++++++++++++++++- .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 17 ++---- .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 5 +- .../DeviceIntraBarrierOpToLLVM.cpp | 9 +-- third_party/tle/triton_tle.cc | 12 ++-- 7 files changed, 174 insertions(+), 25 deletions(-) create mode 100644 python/test/tle/unit/test_tle_d2d_barrier.py create mode 100644 python/test/tle/unit/test_tle_d2d_barrier.sh diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py new file mode 100644 index 000000000..9caeee7d6 --- /dev/null +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -0,0 +1,47 @@ +import triton.experimental.tle.language as tle +import torch +import triton +import triton.language as tl + +DEVICE_MESH = tle.device_mesh(tle.MeshConfig(device=2)) + + +@triton.jit +def _remote_peer_d2d_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): + tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", barrier_kind="sync") + + +class TestDeviceToDevice: + + def test_tle_d2d_remote(self): + block = 64 + grid = 2 + + N = 64 + + with torch.cuda.use_mem_pool(tle.get_mem_pool()): + x = torch.randn((N, N), dtype=torch.float32, device="cuda") + y = torch.empty_like(x) + + dev_comm_dptr, dev_mem_dptr = tle.create_comm_tensor(x) + + compiled = _remote_peer_d2d_kernel.warmup( + dev_comm_dptr=dev_comm_dptr, + dev_mem_dptr=dev_mem_dptr, + out_ptr=y, + mesh=DEVICE_MESH, + BLOCK=block, + grid=(grid, ), + num_ctas=1, + num_warps=4, + ) + # assert "remote_pointers" in compiled.asm["ttgir"] + # assert "flagcxGetIntraPointerC" in compiled.asm['ptx'] + + _remote_peer_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, + mesh=DEVICE_MESH, BLOCK=block) + + tle.cleanup_communicator() + + +TestDeviceToDevice().test_tle_d2d_remote() diff --git a/python/test/tle/unit/test_tle_d2d_barrier.sh b/python/test/tle/unit/test_tle_d2d_barrier.sh new file mode 100644 index 000000000..bb0a89a95 --- /dev/null +++ b/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 diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 1454d275a..06771604a 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, asdict from itertools import product from typing import Any, Iterable, Mapping, Sequence, List, Tuple, Union, Optional, Dict - +from enum import Enum import triton.language.core as tl Axis = Tuple[str, int] @@ -45,6 +45,24 @@ def n_pes(dev_mem_ptr, _semantic=None, ret_dtype=tl.int32): return tl.tensor(result, ret_dtype) +class BarrierKind(str, Enum): + ARRIVE = "arrive" + WAIT = "wait" + SYNC = "sync" + +class MemoryOrder(str, Enum): + RELAXED = "relaxed" + ACQUIRE = "acquire" + RELEASE = "release" + ACQ_REL = "acq_rel" + +class GroupKind(str, Enum): + THREAD = "thread" + WARP = "warp" + BLOCK = "block" + TILE_SPAN = "tile_span" + LANES = "lanes" + @dataclass class MeshConfig: """ @@ -612,9 +630,25 @@ def shard_id( coord = _semantic.mod(coord, dim) return coord +def _parse_device_barrier_args(argType) -> str: + argType = tl._unwrap_if_constexpr(argType) + argTypes = (BarrierKind, GroupKind, MemoryOrder) + if isinstance(argType, argTypes): + return argType.value + else: + return str(argType).lower() @tl.builtin -def distributed_barrier(mesh: device_mesh | None = None, _semantic=None): +def distributed_barrier( + mesh: device_mesh | None = None, + space: str = None, + comm_ptr=None, + barrier_kind: BarrierKind | str = BarrierKind.SYNC, + group_kind : str | GroupKind = GroupKind.BLOCK, + index: int | None = 0, + order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, + _semantic=None +): """ M3 entrypoint: distributed synchronization primitive. @@ -626,7 +660,23 @@ def distributed_barrier(mesh: device_mesh | None = None, _semantic=None): raise TypeError(f"mesh must be device_mesh or None, got {type(mesh).__name__}") subgroup = None use_grid = mesh is not None and _mesh_uses_grid_barrier(mesh) - + + if space or comm_ptr: + if space and space != "device" and comm_ptr: + raise ValueError(f"{space} space and comm_ptr cannot be used together") + if space and space == "device" and not comm_ptr: + raise NotImplementedError(f"distributed_barrier with space={space} must have comm_ptr arg") + builder = _semantic.builder + builder.create_distributed_barrier( + src=comm_ptr.handle, + barrier_index=index or 0, + space=_parse_device_barrier_args(space), + group_kind=_parse_device_barrier_args(group_kind), + order=_parse_device_barrier_args(order), + barrier_kind=_parse_device_barrier_args(barrier_kind), + ) + return None + if use_grid: if mesh is not None: _apply_mesh_grid_launch(mesh, _semantic) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 81964e145..dfb594b64 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -114,9 +114,7 @@ class TleLLVMConversionTarget : public ConversionTarget { }); addLegalOp(); addLegalOp(); - // addIllegalOp(); - // // addLegalOp(); - // addLegalOp(); + addLegalOp(); // Allow non-TLE ops to remain during this partial conversion. markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; }); } @@ -195,12 +193,7 @@ struct ConvertTritonGPUToLLVM typeConverter, patterns, benefit); mlir::triton::tle::populateTMAStoreCommitGroupOpToLLVMPatterns( typeConverter, patterns, benefit); - if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { - return signalPassFailure(); - } - } #ifdef FLAGCX_ENABLED - { mlir::triton::tle::populateGetDeviceIdOpToFlagCxPatterns( typeConverter, patterns, benefit); mlir::triton::tle::populateGetLocalRankOpToLLVMPatterns( @@ -208,10 +201,12 @@ struct ConvertTritonGPUToLLVM mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, patterns, benefit); mlir::triton::tle::populateDeviceIntraBarrierOpToLLVMPatterns( - typeConverter, patterns, benefit); - } + typeConverter, patterns, benefit + 1); #endif - + if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { + return signalPassFailure(); + } + } #endif mlir::triton::NVIDIA::populateConvertLayoutOpToLLVMPatterns( diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index bf61013cf..6e0df54a9 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -462,11 +462,12 @@ struct DistributedBarrierOpConversion auto newOrderAttr = rewriter.getI32IntegerAttr(order); auto barrierTypeAttr = op.getBarrierTypeAttr(); auto multimemAttr = rewriter.getBoolAttr(false); - + # ifdef FLAGCX_ENABLED + llvm::errs() << "replaceOpWithNewOp Lowering DeviceIntraBarrierOp to FlagCx function call\n"; rewriter.replaceOpWithNewOp( op, comm, barrierTypeAttr, coopKindAttr, indexAttr, multimemAttr, newOrderAttr); - + # endif return success(); } LogicalResult diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp index 0f08cea47..d87bb47e5 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp @@ -32,11 +32,11 @@ struct DeviceIntraBarrierOpConversion auto loc = op.getLoc(); auto indexValue = op.getIndexAttr().getInt(); - ; + auto coopValue = op.getCoopKindAttr().getInt(); - ; + auto orderValue = op.getOrderAttr().getInt(); - ; + auto barrierType = op.getBarrierTypeAttr().getValue(); if (!llvm::is_contained(std::array{0, 1, 2, 3, 4}, coopValue)) @@ -44,7 +44,8 @@ struct DeviceIntraBarrierOpConversion if (!llvm::is_contained(std::array{0, 1, 2, 3}, orderValue)) return rewriter.notifyMatchFailure(op, "invalid coop_kind"); - + + llvm::errs() << "Lowering DeviceIntraBarrierOp to FlagCx function call\n"; tle::getBarrierFuncCall(loc, rewriter, adaptor.getComm(), indexValue, coopValue, orderValue, barrierType); rewriter.eraseOp(op); diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index fc7f9068a..6664fe6db 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -368,18 +368,18 @@ void init_triton_tle_ir(py::module &&m) { .def("create_distributed_barrier", [](TritonOpBuilder &self, Value src, size_t barrier_index = 0, const std::string &space = "device", - const std::string &groupKind = "block", + const std::string &group_kind = "block", const std::string &order = "acqrel", - const std::string &barrierType = "sync") -> void { + const std::string &barrier_kind = "sync") -> void { auto &builder = self.getBuilder(); auto *ctx = builder.getContext(); auto getOptStrAttr = [&](const std::string &s) -> StringAttr { return s.empty() ? StringAttr() : builder.getStringAttr(s); }; auto spaceAttr = getOptStrAttr(space); - auto kindAttr = getOptStrAttr(groupKind); + auto kindAttr = getOptStrAttr(group_kind); auto orderAttr = getOptStrAttr(order); - auto barrierTypeAttr = getOptStrAttr(barrierType); + auto barrierTypeAttr = getOptStrAttr(barrier_kind); auto barrierIndexAttr = builder.getI32IntegerAttr(static_cast(barrier_index)); @@ -387,7 +387,9 @@ void init_triton_tle_ir(py::module &&m) { src, spaceAttr, barrierTypeAttr, orderAttr, kindAttr, barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr()); - }) + }, + py::arg("src"), py::arg("barrier_index"), py::arg("space"), + py::arg("group_kind"), py::arg("order"), py::arg("barrier_kind")) .def( "create_distributed_barrier", [](TritonOpBuilder &self, const std::string &groupKind, From eee13c4f17b2c8e963ca269752dc81f19cd8a6f9 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Mon, 6 Jul 2026 10:18:07 +0000 Subject: [PATCH 05/32] add python api args --- python/test/tle/unit/test_tle_d2d_barrier.py | 5 ++- .../experimental/tle/language/distributed.py | 23 +++++----- .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 7 +-- .../DeviceIntraBarrierOpToLLVM.cpp | 8 ++-- third_party/tle/triton_tle.cc | 45 ++++++++++--------- 5 files changed, 45 insertions(+), 43 deletions(-) diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index 9caeee7d6..08750c181 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -8,7 +8,8 @@ @triton.jit def _remote_peer_d2d_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): - tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", barrier_kind="sync") + tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", + barrier_kind="sync") class TestDeviceToDevice: @@ -35,7 +36,7 @@ def test_tle_d2d_remote(self): num_ctas=1, num_warps=4, ) - # assert "remote_pointers" in compiled.asm["ttgir"] + assert "remote_pointers" in compiled.asm["ttgir"] # assert "flagcxGetIntraPointerC" in compiled.asm['ptx'] _remote_peer_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 06771604a..872b5d5a9 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -50,12 +50,14 @@ class BarrierKind(str, Enum): WAIT = "wait" SYNC = "sync" + class MemoryOrder(str, Enum): RELAXED = "relaxed" ACQUIRE = "acquire" RELEASE = "release" ACQ_REL = "acq_rel" + class GroupKind(str, Enum): THREAD = "thread" WARP = "warp" @@ -63,6 +65,7 @@ class GroupKind(str, Enum): TILE_SPAN = "tile_span" LANES = "lanes" + @dataclass class MeshConfig: """ @@ -630,6 +633,7 @@ def shard_id( coord = _semantic.mod(coord, dim) return coord + def _parse_device_barrier_args(argType) -> str: argType = tl._unwrap_if_constexpr(argType) argTypes = (BarrierKind, GroupKind, MemoryOrder) @@ -638,17 +642,12 @@ def _parse_device_barrier_args(argType) -> str: else: return str(argType).lower() + @tl.builtin -def distributed_barrier( - mesh: device_mesh | None = None, - space: str = None, - comm_ptr=None, - barrier_kind: BarrierKind | str = BarrierKind.SYNC, - group_kind : str | GroupKind = GroupKind.BLOCK, - index: int | None = 0, - order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, - _semantic=None -): +def distributed_barrier(mesh: device_mesh | None = None, space: str = None, comm_ptr=None, + barrier_kind: BarrierKind | str = BarrierKind.SYNC, + group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, + order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): """ M3 entrypoint: distributed synchronization primitive. @@ -660,7 +659,7 @@ def distributed_barrier( raise TypeError(f"mesh must be device_mesh or None, got {type(mesh).__name__}") subgroup = None use_grid = mesh is not None and _mesh_uses_grid_barrier(mesh) - + if space or comm_ptr: if space and space != "device" and comm_ptr: raise ValueError(f"{space} space and comm_ptr cannot be used together") @@ -676,7 +675,7 @@ def distributed_barrier( barrier_kind=_parse_device_barrier_args(barrier_kind), ) return None - + if use_grid: if mesh is not None: _apply_mesh_grid_launch(mesh, _semantic) diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index 6e0df54a9..d674f4b79 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -462,12 +462,13 @@ struct DistributedBarrierOpConversion auto newOrderAttr = rewriter.getI32IntegerAttr(order); auto barrierTypeAttr = op.getBarrierTypeAttr(); auto multimemAttr = rewriter.getBoolAttr(false); - # ifdef FLAGCX_ENABLED - llvm::errs() << "replaceOpWithNewOp Lowering DeviceIntraBarrierOp to FlagCx function call\n"; +#ifdef FLAGCX_ENABLED + llvm::errs() << "replaceOpWithNewOp Lowering DeviceIntraBarrierOp to " + "FlagCx function call\n"; rewriter.replaceOpWithNewOp( op, comm, barrierTypeAttr, coopKindAttr, indexAttr, multimemAttr, newOrderAttr); - # endif +#endif return success(); } LogicalResult diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp index d87bb47e5..55a6e6589 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp @@ -32,11 +32,11 @@ struct DeviceIntraBarrierOpConversion auto loc = op.getLoc(); auto indexValue = op.getIndexAttr().getInt(); - + auto coopValue = op.getCoopKindAttr().getInt(); - + auto orderValue = op.getOrderAttr().getInt(); - + auto barrierType = op.getBarrierTypeAttr().getValue(); if (!llvm::is_contained(std::array{0, 1, 2, 3, 4}, coopValue)) @@ -44,7 +44,7 @@ struct DeviceIntraBarrierOpConversion if (!llvm::is_contained(std::array{0, 1, 2, 3}, orderValue)) return rewriter.notifyMatchFailure(op, "invalid coop_kind"); - + llvm::errs() << "Lowering DeviceIntraBarrierOp to FlagCx function call\n"; tle::getBarrierFuncCall(loc, rewriter, adaptor.getComm(), indexValue, coopValue, orderValue, barrierType); diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 6664fe6db..aa00c305d 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -365,29 +365,30 @@ void init_triton_tle_ir(py::module &&m) { StringAttr(), IntegerAttr(), IntegerAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr()); }) - .def("create_distributed_barrier", - [](TritonOpBuilder &self, Value src, size_t barrier_index = 0, - const std::string &space = "device", - const std::string &group_kind = "block", - const std::string &order = "acqrel", - const std::string &barrier_kind = "sync") -> void { - auto &builder = self.getBuilder(); - auto *ctx = builder.getContext(); - auto getOptStrAttr = [&](const std::string &s) -> StringAttr { - return s.empty() ? StringAttr() : builder.getStringAttr(s); - }; - auto spaceAttr = getOptStrAttr(space); - auto kindAttr = getOptStrAttr(group_kind); - auto orderAttr = getOptStrAttr(order); - auto barrierTypeAttr = getOptStrAttr(barrier_kind); - auto barrierIndexAttr = - builder.getI32IntegerAttr(static_cast(barrier_index)); + .def( + "create_distributed_barrier", + [](TritonOpBuilder &self, Value src, size_t barrier_index = 0, + const std::string &space = "device", + const std::string &group_kind = "block", + const std::string &order = "acqrel", + const std::string &barrier_kind = "sync") -> void { + auto &builder = self.getBuilder(); + auto *ctx = builder.getContext(); + auto getOptStrAttr = [&](const std::string &s) -> StringAttr { + return s.empty() ? StringAttr() : builder.getStringAttr(s); + }; + auto spaceAttr = getOptStrAttr(space); + auto kindAttr = getOptStrAttr(group_kind); + auto orderAttr = getOptStrAttr(order); + auto barrierTypeAttr = getOptStrAttr(barrier_kind); + auto barrierIndexAttr = + builder.getI32IntegerAttr(static_cast(barrier_index)); - self.create( - src, spaceAttr, barrierTypeAttr, orderAttr, kindAttr, - barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), - DenseI32ArrayAttr(), DenseI32ArrayAttr()); - }, + self.create( + src, spaceAttr, barrierTypeAttr, orderAttr, kindAttr, + barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), + DenseI32ArrayAttr(), DenseI32ArrayAttr()); + }, py::arg("src"), py::arg("barrier_index"), py::arg("space"), py::arg("group_kind"), py::arg("order"), py::arg("barrier_kind")) .def( From f95e1a956f5320478e7a59a8cdf6efc8a6a0dc81 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 7 Jul 2026 08:39:14 +0000 Subject: [PATCH 06/32] fix conversion bug --- .../enflame3.6-gcu400-build-and-test.yml | 3 +- .github/workflows/hopper-build-and-test.yml | 4 +- python/test/tle/unit/test_tle_d2d_barrier.py | 18 +-- .../experimental/tle/language/distributed.py | 38 +++--- .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 6 +- .../lib/Conversion/TleToLLVM/CMakeLists.txt | 4 +- .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 12 +- .../DeviceIntraBarrierOpToLLVM.cpp | 4 - third_party/tle/dialect/lib/IR/FlagCxOps.cpp | 114 ++++++++++++------ 9 files changed, 126 insertions(+), 77 deletions(-) diff --git a/.github/workflows/enflame3.6-gcu400-build-and-test.yml b/.github/workflows/enflame3.6-gcu400-build-and-test.yml index 831f0b7a7..22495d920 100644 --- a/.github/workflows/enflame3.6-gcu400-build-and-test.yml +++ b/.github/workflows/enflame3.6-gcu400-build-and-test.yml @@ -68,7 +68,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 diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index d8b2eb243..408a7873a 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -188,9 +188,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 diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index 08750c181..469616c90 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -7,14 +7,14 @@ @triton.jit -def _remote_peer_d2d_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): +def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", barrier_kind="sync") -class TestDeviceToDevice: +class TestD2DBarrier: - def test_tle_d2d_remote(self): + def test_tle_d2d_barrier(self): block = 64 grid = 2 @@ -26,7 +26,7 @@ def test_tle_d2d_remote(self): dev_comm_dptr, dev_mem_dptr = tle.create_comm_tensor(x) - compiled = _remote_peer_d2d_kernel.warmup( + compiled = _barrier_d2d_kernel.warmup( dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, @@ -36,13 +36,13 @@ def test_tle_d2d_remote(self): num_ctas=1, num_warps=4, ) - assert "remote_pointers" in compiled.asm["ttgir"] - # assert "flagcxGetIntraPointerC" in compiled.asm['ptx'] + assert "distributed_barrier" in compiled.asm["ttgir"] + assert "flagcxIntraBarrier" in compiled.asm['ptx'] - _remote_peer_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, - mesh=DEVICE_MESH, BLOCK=block) + _barrier_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, + mesh=DEVICE_MESH, BLOCK=block) tle.cleanup_communicator() -TestDeviceToDevice().test_tle_d2d_remote() +TestD2DBarrier().test_tle_d2d_barrier() diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 872b5d5a9..b97e1aa6e 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -643,6 +643,28 @@ def _parse_device_barrier_args(argType) -> str: return str(argType).lower() +def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, + barrier_kind: BarrierKind | str = BarrierKind.SYNC, + group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, + order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): + if space or comm_ptr: + if space and space != "device" and comm_ptr: + raise ValueError(f"{space} space and comm_ptr cannot be used together") + if space and space == "device" and not comm_ptr: + raise NotImplementedError(f"distributed_barrier with space={space} must have comm_ptr arg") + builder = _semantic.builder + builder.create_distributed_barrier( + src=comm_ptr.handle, + barrier_index=index or 0, + space=_parse_device_barrier_args("device"), + group_kind=_parse_device_barrier_args(group_kind), + order=_parse_device_barrier_args(order), + barrier_kind=_parse_device_barrier_args(barrier_kind), + ) + return True + return False + + @tl.builtin def distributed_barrier(mesh: device_mesh | None = None, space: str = None, comm_ptr=None, barrier_kind: BarrierKind | str = BarrierKind.SYNC, @@ -660,20 +682,8 @@ def distributed_barrier(mesh: device_mesh | None = None, space: str = None, comm subgroup = None use_grid = mesh is not None and _mesh_uses_grid_barrier(mesh) - if space or comm_ptr: - if space and space != "device" and comm_ptr: - raise ValueError(f"{space} space and comm_ptr cannot be used together") - if space and space == "device" and not comm_ptr: - raise NotImplementedError(f"distributed_barrier with space={space} must have comm_ptr arg") - builder = _semantic.builder - builder.create_distributed_barrier( - src=comm_ptr.handle, - barrier_index=index or 0, - space=_parse_device_barrier_args(space), - group_kind=_parse_device_barrier_args(group_kind), - order=_parse_device_barrier_args(order), - barrier_kind=_parse_device_barrier_args(barrier_kind), - ) + if check_and_handle_device_intra_barrier(space=space, comm_ptr=comm_ptr, barrier_kind=barrier_kind, + group_kind=group_kind, index=index, order=order, _semantic=_semantic): return None if use_grid: diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index dfb594b64..acc0645f9 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -114,7 +114,9 @@ class TleLLVMConversionTarget : public ConversionTarget { }); addLegalOp(); addLegalOp(); - addLegalOp(); + // addIllegalOp(); + // addIllegalOp(); + // addLegalOp(); // Allow non-TLE ops to remain during this partial conversion. markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; }); } @@ -201,7 +203,7 @@ struct ConvertTritonGPUToLLVM mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, patterns, benefit); mlir::triton::tle::populateDeviceIntraBarrierOpToLLVMPatterns( - typeConverter, patterns, benefit + 1); + typeConverter, patterns, benefit); #endif if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { return signalPassFailure(); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt index 6fced8154..d1772d125 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt @@ -11,10 +11,10 @@ set(TLE_TO_LLVM_SRC WGMMAFenceOpToLLVM.cpp ) -add_subdirectory(FlagCxOpToLLVM) if(FLAGCX_ENABLED) - list(APPEND TLE_TO_LLVM_SRC +add_subdirectory(FlagCxOpToLLVM) +list(APPEND TLE_TO_LLVM_SRC GetDeviceIdToFlagCX.cpp GetLocalRankOpToLLVM.cpp ) diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index d674f4b79..d477246d1 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -21,7 +21,7 @@ namespace tle = mlir::triton::tle; constexpr llvm::StringLiteral kSpaceAttr = "space"; constexpr llvm::StringLiteral kOrderAttr = "order"; -constexpr llvm::StringLiteral kIndexAttr = "index"; +constexpr llvm::StringLiteral kIndexAttr = "barrier_index"; constexpr llvm::StringLiteral kGroupKindAttr = "group_kind"; constexpr llvm::StringLiteral kGroupShapeAttr = "group_shape"; constexpr llvm::StringLiteral kGroupMaskAttr = "group_mask"; @@ -427,11 +427,12 @@ struct DistributedBarrierOpConversion } LogicalResult - lowerDeviceSpaceBarrier(tle::DistributedBarrierOp op, + lowerDeviceSpaceBarrier(tle::DistributedBarrierOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const { auto kindAttr = op->getAttrOfType(kGroupKindAttr); auto orderAttr = op->getAttrOfType(kOrderAttr); auto indexAttr = op->getAttrOfType(kIndexAttr); + auto loc = op.getLoc(); auto getCoopKindValue = [](StringRef kind) -> int32_t { return llvm::StringSwitch(kind) .Case("thread", 0) @@ -457,14 +458,13 @@ struct DistributedBarrierOpConversion if (order < 0) return rewriter.notifyMatchFailure(op, "invalid order"); - auto comm = op.getSrc(); + auto srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); + auto comm = srcElems[0]; auto coopKindAttr = rewriter.getI32IntegerAttr(coopKind); auto newOrderAttr = rewriter.getI32IntegerAttr(order); auto barrierTypeAttr = op.getBarrierTypeAttr(); auto multimemAttr = rewriter.getBoolAttr(false); #ifdef FLAGCX_ENABLED - llvm::errs() << "replaceOpWithNewOp Lowering DeviceIntraBarrierOp to " - "FlagCx function call\n"; rewriter.replaceOpWithNewOp( op, comm, barrierTypeAttr, coopKindAttr, indexAttr, multimemAttr, newOrderAttr); @@ -476,7 +476,7 @@ struct DistributedBarrierOpConversion ConversionPatternRewriter &rewriter) const override { if (auto spaceAttr = op->getAttrOfType(kSpaceAttr)) if (spaceAttr.getValue() == "device") - return lowerDeviceSpaceBarrier(op, rewriter); + return lowerDeviceSpaceBarrier(op, adaptor, rewriter); if (auto kindAttr = op->getAttrOfType(kGroupKindAttr)) { if (kindAttr.getValue() == "grid") diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp index 55a6e6589..879566541 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp @@ -32,11 +32,8 @@ struct DeviceIntraBarrierOpConversion auto loc = op.getLoc(); auto indexValue = op.getIndexAttr().getInt(); - auto coopValue = op.getCoopKindAttr().getInt(); - auto orderValue = op.getOrderAttr().getInt(); - auto barrierType = op.getBarrierTypeAttr().getValue(); if (!llvm::is_contained(std::array{0, 1, 2, 3, 4}, coopValue)) @@ -45,7 +42,6 @@ struct DeviceIntraBarrierOpConversion if (!llvm::is_contained(std::array{0, 1, 2, 3}, orderValue)) return rewriter.notifyMatchFailure(op, "invalid coop_kind"); - llvm::errs() << "Lowering DeviceIntraBarrierOp to FlagCx function call\n"; tle::getBarrierFuncCall(loc, rewriter, adaptor.getComm(), indexValue, coopValue, orderValue, barrierType); rewriter.eraseOp(op); diff --git a/third_party/tle/dialect/lib/IR/FlagCxOps.cpp b/third_party/tle/dialect/lib/IR/FlagCxOps.cpp index 67a8d5ff2..1cca185a5 100644 --- a/third_party/tle/dialect/lib/IR/FlagCxOps.cpp +++ b/third_party/tle/dialect/lib/IR/FlagCxOps.cpp @@ -14,6 +14,21 @@ #include "triton/Dialect/TritonGPU/IR/Dialect.h" #include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +enum class CoopKind : int32_t { + Thread = 0, + Warp = 1, + Block = 2, + TileSpan = 3, + Lanes = 4, +}; + +enum class MemoryOrder : int32_t { + Relaxed = 0, + Acquire = 1, + Release = 2, + AcqRel = 3, +}; + namespace mlir::triton::tle { LogicalResult GetLocalRankOp::verify() { @@ -27,50 +42,73 @@ LogicalResult GetLocalRankOp::verify() { LogicalResult DeviceIntraBarrierOp::verify() { auto *op = getOperation(); - auto spaceAttr = op->getAttrOfType("space"); - if (spaceAttr && spaceAttr.getValue() != "device") - return success(); - auto emitInvalidAttr = [&](StringRef attrName, StringRef value, - StringRef expected) -> LogicalResult { + + auto barrierTypeAttr = getBarrierTypeAttr(); + auto coopKindAttr = getCoopKindAttr(); + auto orderAttr = getOrderAttr(); + + auto emitInvalidIntAttr = [&](StringRef attrName, int64_t value, + StringRef expected) -> LogicalResult { + return op->emitOpError() << "invalid " << attrName << " (" << value + << "), expected one of: " << expected; + }; + + auto emitInvalidStrAttr = [&](StringRef attrName, StringRef value, + StringRef expected) -> LogicalResult { return op->emitOpError() << "invalid " << attrName << " '" << value << "', expected one of: " << expected; }; - auto kindAttr = op->getAttrOfType("group_kind"); - auto barrierTypeAttr = op->getAttrOfType("barrier_type"); - auto orderAttr = op->getAttrOfType("order"); - - bool kindValid = llvm::StringSwitch(kindAttr.getValue()) - .Case("thread", true) - .Case("warp", true) - .Case("block", true) - .Case("tile_span", true) - .Case("lanes", true) - .Default(false); - - if (!kindValid) { - return emitInvalidAttr("group_kind", kindAttr.getValue(), - "thread, warp, block, tile_span, lanes"); + + // barrier_type + if (barrierTypeAttr) { + StringRef barrierType = barrierTypeAttr.getValue(); + + bool valid = llvm::StringSwitch(barrierType) + .Case("arrive", true) + .Case("wait", true) + .Case("sync", true) + .Default(false); + + if (!valid) + return emitInvalidStrAttr("barrier_type", barrierType, + "arrive, wait, sync"); } - bool barrierTypeValid = llvm::StringSwitch(barrierTypeAttr.getValue()) - .Case("arrive", true) - .Case("wait", true) - .Case("sync", true) - .Default(false); - if (!barrierTypeValid) { - return emitInvalidAttr("barrier_type", barrierTypeAttr.getValue(), - "arrive, wait, sync"); + + // coop_kind + if (coopKindAttr) { + auto coopKind = static_cast(coopKindAttr.getInt()); + + switch (coopKind) { + case CoopKind::Thread: + case CoopKind::Warp: + case CoopKind::Block: + case CoopKind::TileSpan: + case CoopKind::Lanes: + break; + default: + return emitInvalidIntAttr( + "coop_kind", coopKindAttr.getInt(), + "Thread(0), Warp(1), Block(2), TileSpan(3), Lanes(4)"); + } } - bool orderValid = llvm::StringSwitch(orderAttr.getValue()) - .Case("acqrel", true) - .Case("acquire", true) - .Case("release", true) - .Case("relaxed", true) - .Default(false); - if (!orderValid) { - return emitInvalidAttr("order", orderAttr.getValue(), - "relaxed, acqrel, acquire, release"); + + // order + if (orderAttr) { + auto order = static_cast(orderAttr.getInt()); + + switch (order) { + case MemoryOrder::Relaxed: + case MemoryOrder::Acquire: + case MemoryOrder::Release: + case MemoryOrder::AcqRel: + break; + default: + return emitInvalidIntAttr( + "order", orderAttr.getInt(), + "Relaxed(0), Acquire(1), Release(2), AcqRel(3)"); + } } + return success(); } - } // namespace mlir::triton::tle From a30a571768da04848e682de1e94d0a09b0e96ea5 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 7 Jul 2026 09:46:17 +0000 Subject: [PATCH 07/32] add test case --- python/test/tle/unit/test_tle_d2d_barrier.py | 18 +++++++++++++++--- .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 4 ---- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index 469616c90..b44491e56 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -7,7 +7,19 @@ @triton.jit -def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): +def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): + pid = tl.program_id(0) + local_rank = tle.shard_id(mesh, 'device', comm_ptr=dev_comm_dptr) + n_rank = mesh.shape[0] + peer = (local_rank + 1) % n_rank # noqa: F841 + remote_mem = tle.remote( + dev_mem_ptr, + space="device", + dtype=tl.float32, + shard_id=peer, + offset=pid, + ) + tl.load(remote_mem) tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", barrier_kind="sync") @@ -28,7 +40,7 @@ def test_tle_d2d_barrier(self): compiled = _barrier_d2d_kernel.warmup( dev_comm_dptr=dev_comm_dptr, - dev_mem_dptr=dev_mem_dptr, + dev_mem_ptr=dev_mem_dptr, out_ptr=y, mesh=DEVICE_MESH, BLOCK=block, @@ -39,7 +51,7 @@ def test_tle_d2d_barrier(self): assert "distributed_barrier" in compiled.asm["ttgir"] assert "flagcxIntraBarrier" in compiled.asm['ptx'] - _barrier_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, + _barrier_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_ptr=dev_mem_dptr, out_ptr=y, mesh=DEVICE_MESH, BLOCK=block) tle.cleanup_communicator() diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index acc0645f9..221caa3eb 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -113,10 +113,6 @@ class TleLLVMConversionTarget : public ConversionTarget { return hasLegalRegions && typeConverter.isLegal(op); }); addLegalOp(); - addLegalOp(); - // addIllegalOp(); - // addIllegalOp(); - // addLegalOp(); // Allow non-TLE ops to remain during this partial conversion. markUnknownOpDynamicallyLegal([](Operation *) -> bool { return true; }); } From cbde6f26ae477bc2d691fd28b80f405b7c526f9a Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 7 Jul 2026 10:06:04 +0000 Subject: [PATCH 08/32] adjust dir --- .../nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 2 +- .../TleToLLVM/{ => FlagCxOpToLLVM}/GetLocalRankOpToLLVM.h | 0 third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt | 1 - .../lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt | 1 + .../TleToLLVM/{ => FlagCxOpToLLVM}/GetLocalRankOpToLLVM.cpp | 2 +- 5 files changed, 3 insertions(+), 3 deletions(-) rename third_party/tle/dialect/include/Conversion/TleToLLVM/{ => FlagCxOpToLLVM}/GetLocalRankOpToLLVM.h (100%) rename third_party/tle/dialect/lib/Conversion/TleToLLVM/{ => FlagCxOpToLLVM}/GetLocalRankOpToLLVM.cpp (97%) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 221caa3eb..43c1c7e88 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -26,8 +26,8 @@ #include "tle/dialect/include/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/ExtractOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/GetDeviceIdToFlagCX.h" -#include "tle/dialect/include/Conversion/TleToLLVM/GetLocalRankOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/LocalPointersOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/PackOpToLLVM.h" #include "tle/dialect/include/IR/Dialect.h" diff --git a/third_party/tle/dialect/include/Conversion/TleToLLVM/GetLocalRankOpToLLVM.h b/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h similarity index 100% rename from third_party/tle/dialect/include/Conversion/TleToLLVM/GetLocalRankOpToLLVM.h rename to third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt index d1772d125..d59aef3a5 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/CMakeLists.txt @@ -16,7 +16,6 @@ if(FLAGCX_ENABLED) add_subdirectory(FlagCxOpToLLVM) list(APPEND TLE_TO_LLVM_SRC GetDeviceIdToFlagCX.cpp - GetLocalRankOpToLLVM.cpp ) endif() diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt index bdc429145..b7f2e3d3b 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt @@ -1,5 +1,6 @@ set(FLAGCX_TO_LLVM_SRC DeviceIntraBarrierOpToLLVM.cpp + GetLocalRankOpToLLVM.cpp ) add_triton_library(FlagCxToLLVM diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetLocalRankOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp similarity index 97% rename from third_party/tle/dialect/lib/Conversion/TleToLLVM/GetLocalRankOpToLLVM.cpp rename to third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp index 9a63b05f9..26bb7cca4 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetLocalRankOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp @@ -1,4 +1,4 @@ -#include "tle/dialect/include/Conversion/TleToLLVM/GetLocalRankOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h" #include "tle/dialect/include/Tools/FlagcxUtils.h" #include "mlir/Conversion/LLVMCommon/Pattern.h" From be3ba634dd90a8a83e26e82af29bf66ec058a620 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Wed, 8 Jul 2026 06:34:36 +0000 Subject: [PATCH 09/32] fix test case bug --- .github/workflows/hopper-build-and-test.yml | 4 +- python/setup_tools/utils/default.py | 34 +++++++- python/test/tle/unit/test_tle_d2d_barrier.py | 81 +++++++++++++------ .../tle/dialect/lib/Tools/FlagcxUtils.cpp | 15 ++-- 4 files changed, 97 insertions(+), 37 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 408a7873a..91772c8fc 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -75,7 +75,9 @@ jobs: source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true rm -rf ~/.triton/cache - USE_FLAGCX=ON MAX_JOBS=32 python3 -m pip install . --no-build-isolation + export FLAGCX_RECOMPILE=ON + export USE_FLAGCX=ON + MAX_JOBS=32 python3 -m pip install . --no-build-isolation - name: FlagTree Test on NVidia (triton_v3.4.x branch) if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.4.x' }} diff --git a/python/setup_tools/utils/default.py b/python/setup_tools/utils/default.py index 7e1d3494e..c0d6f638d 100644 --- a/python/setup_tools/utils/default.py +++ b/python/setup_tools/utils/default.py @@ -7,7 +7,7 @@ def printinfo(msgs): - print(f" [TLE-DIST-INFO]: {msgs}. \n") + print(f" [TLE-DIST-INFO]: {msgs}.") class FlagCXRegistrar: @@ -17,6 +17,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'] @@ -50,15 +73,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) @@ -73,7 +99,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" diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index b44491e56..db2422fd8 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -2,16 +2,19 @@ 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(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): +def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr): pid = tl.program_id(0) local_rank = tle.shard_id(mesh, 'device', comm_ptr=dev_comm_dptr) n_rank = mesh.shape[0] peer = (local_rank + 1) % n_rank # noqa: F841 + remote_mem = tle.remote( dev_mem_ptr, space="device", @@ -19,42 +22,68 @@ def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr, shard_id=peer, offset=pid, ) - tl.load(remote_mem) + val = tl.load(remote_mem) + tl.store(out_ptr + pid, val) tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", barrier_kind="sync") +def _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid): + compiled = _barrier_d2d_kernel.warmup( + dev_comm_dptr=dev_comm_dptr, + dev_mem_ptr=dev_mem_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(dev_comm_dptr, dev_mem_dptr, output, grid, rank, world_size): + dist.barrier() + _barrier_d2d_kernel[grid](dev_comm_dptr=dev_comm_dptr, dev_mem_ptr=dev_mem_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): - block = 64 - grid = 2 + grid = (N, ) - N = 64 - - with torch.cuda.use_mem_pool(tle.get_mem_pool()): - x = torch.randn((N, N), dtype=torch.float32, device="cuda") - y = torch.empty_like(x) + 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() dev_comm_dptr, dev_mem_dptr = tle.create_comm_tensor(x) + output = torch.zeros(N, dtype=torch.float32, device="cuda") - compiled = _barrier_d2d_kernel.warmup( - dev_comm_dptr=dev_comm_dptr, - dev_mem_ptr=dev_mem_dptr, - out_ptr=y, - mesh=DEVICE_MESH, - BLOCK=block, - grid=(grid, ), - num_ctas=1, - num_warps=4, - ) - assert "distributed_barrier" in compiled.asm["ttgir"] - assert "flagcxIntraBarrier" in compiled.asm['ptx'] - - _barrier_d2d_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_ptr=dev_mem_dptr, out_ptr=y, - mesh=DEVICE_MESH, BLOCK=block) + _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid) - tle.cleanup_communicator() + _runtime_verify(dev_comm_dptr, dev_mem_dptr, output, grid, rank, world_size) -TestD2DBarrier().test_tle_d2d_barrier() +if __name__ == "__main__": + TestD2DBarrier().test_tle_d2d_barrier() diff --git a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp index 19c7d4076..e43034b63 100644 --- a/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp +++ b/third_party/tle/dialect/lib/Tools/FlagcxUtils.cpp @@ -33,9 +33,12 @@ static inline LLVM::LLVMFuncOp createFuncInstance(const char *funcName, return func; } -static inline Value getFlagcxMemPtr(mlir::Location loc, - ConversionPatternRewriter &rewriter, - Value memPtrInt) { +// The frontend passes the FlagCX global memory/communication pointer as an +// integer. Convert it back to an LLVM pointer in global address space (AS1) +// before passing it to device/runtime functions. +static inline Value getFlagcxMemOrCommPtr(mlir::Location loc, + ConversionPatternRewriter &rewriter, + Value memPtrInt) { auto ctx = rewriter.getContext(); auto ptrTy = LLVM::LLVMPointerType::get(ctx, 1); return rewriter.create(loc, ptrTy, memPtrInt); @@ -53,7 +56,7 @@ LLVM::CallOp getNumPesFunCall(mlir::Location loc, auto func = createFuncInstance( runtimeNames.lookup("getNumPesFunction").data(), module, {PtrTy}, i32Ty); - auto comm_dev_ptr = getFlagcxMemPtr(loc, rewriter, memPtrInt); + auto comm_dev_ptr = getFlagcxMemOrCommPtr(loc, rewriter, memPtrInt); return rewriter.create( loc, TypeRange{func.getFunctionType().getReturnType()}, FlatSymbolRefAttr::get(func), ValueRange{comm_dev_ptr}); @@ -84,7 +87,7 @@ LLVM::CallOp getBarrierFuncCall(mlir::Location loc, auto func = createFuncInstance(runtimeNames.lookup(funcName).data(), module, {PtrTy, i32Ty, i32Ty, i1Ty, i32Ty}, i32Ty); - auto comm_dev_ptr = getFlagcxMemPtr(loc, rewriter, comm); + auto comm_dev_ptr = getFlagcxMemOrCommPtr(loc, rewriter, comm); auto falseVal = rewriter.create(loc, i1Ty, rewriter.getBoolAttr(false)); auto barrierIndexVal = @@ -110,7 +113,7 @@ LLVM::CallOp getLocalPeFuncCall(mlir::Location loc, auto func = createFuncInstance( runtimeNames.lookup("getLocalPeFunction").data(), module, {PtrTy}, i32Ty); - auto comm_dev_ptr = getFlagcxMemPtr(loc, rewriter, memPtrInt); + auto comm_dev_ptr = getFlagcxMemOrCommPtr(loc, rewriter, memPtrInt); return rewriter.create( loc, TypeRange{func.getFunctionType().getReturnType()}, FlatSymbolRefAttr::get(func), ValueRange{comm_dev_ptr}); From 57482d8123fc0db81398c2ff49db81063a881576 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Wed, 8 Jul 2026 06:59:58 +0000 Subject: [PATCH 10/32] solve conflict --- .github/workflows/hopper-build-and-test.yml | 117 ++------------------ 1 file changed, 11 insertions(+), 106 deletions(-) diff --git a/.github/workflows/hopper-build-and-test.yml b/.github/workflows/hopper-build-and-test.yml index 91772c8fc..9c4645cd2 100644 --- a/.github/workflows/hopper-build-and-test.yml +++ b/.github/workflows/hopper-build-and-test.yml @@ -4,9 +4,9 @@ on: schedule: - cron: '0 21 * * *' push: - branches: [ "triton_v3.4.x", "triton_v3.5.x", "triton_v3.6.x" ] + branches: [ "triton_v3.6.x" ] pull_request: - branches: [ "triton_v3.4.x", "triton_v3.5.x", "triton_v3.6.x" ] + branches: [ "triton_v3.6.x" ] concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -46,26 +46,6 @@ jobs: echo "TARGET_BRANCH=$TARGET_BRANCH" >> $GITHUB_ENV echo "TARGET_BRANCH=$TARGET_BRANCH" - - name: FlagTree Build on NVidia (triton_v3.4.x branch) - if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.4.x' }} - shell: bash - run: | - set -x - pip uninstall -y triton - source ~/env-3.4.sh - env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true - MAX_JOBS=32 python3 -m pip install . --no-build-isolation - - - name: FlagTree Build on NVidia (triton_v3.5.x branch) - if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.5.x' }} - shell: bash - run: | - set -x - pip uninstall -y triton - source ~/env-3.5.sh - env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true - MAX_JOBS=32 python3 -m pip install . --no-build-isolation - - name: FlagTree Build on NVidia (triton_v3.6.x branch) if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.6.x' }} shell: bash @@ -74,91 +54,17 @@ jobs: pip uninstall -y triton source ~/env-3.6.sh env | grep -E '^(LLVM_SYSPATH)=' >> $GITHUB_ENV || true - rm -rf ~/.triton/cache export FLAGCX_RECOMPILE=ON export USE_FLAGCX=ON MAX_JOBS=32 python3 -m pip install . --no-build-isolation - - name: FlagTree Test on NVidia (triton_v3.4.x branch) - if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.4.x' }} - shell: bash - run: | - set -x - # python unit test - python3 -m pytest -s python/test/unit/cuda - python3 -m pytest -s python/test/unit/instrumentation - python3 -m pytest -s python/test/unit/language \ - --ignore=python/test/unit/language/test_line_info.py - python3 -m pytest -s python/test/unit/runtime - if [ -d "python/test/operators" ]; then - python3 -m pytest -s python/test/operators - fi - - - name: FlagTree Test on NVidia (triton_v3.5.x branch) - if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.5.x' }} - shell: bash - run: | - set -x - source ~/env.sh - # python tutorials - python3 python/tutorials/01-vector-add.py --only_unit_test - python3 python/tutorials/02-fused-softmax.py --only_unit_test - python3 python/tutorials/03-matrix-multiplication.py --only_unit_test - python3 python/tutorials/04-low-memory-dropout.py --only_unit_test - python3 python/tutorials/05-layer-norm.py --only_unit_test - python3 python/tutorials/06-fused-attention.py --only_unit_test - python3 python/tutorials/07-extern-functions.py --only_unit_test - python3 python/tutorials/08-grouped-gemm.py --only_unit_test - python3 python/tutorials/09-persistent-matmul.py --only_unit_test - python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test - # python unit test - python3 -m pytest -s python/test/unit/cuda - python3 -m pytest -s python/test/unit/instrumentation - python3 -m pytest -s python/test/unit/language \ - --ignore=python/test/unit/language/test_line_info.py - python3 -m pytest -s python/test/unit/runtime - if [ -d "python/test/operators" ]; then - python3 -m pytest -s python/test/operators - fi - # flagtree tle - # python tutorials - python3 python/tutorials/tle/01-fft.py - python3 python/tutorials/tle/02-moe_align_block_size.py - CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/03-topk.py - CUDA_LAUNCH_BLOCKING=1 python3 python/tutorials/tle/04-cluster-gemm.py - python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py - python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py --mode test - python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py - # python unit test - python3 -m pytest -s python/test/tle/integration - python3 -m pytest -s python/test/tle/unit - # 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 - python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test - python3 python/tutorials/hints/04/04-low-memory-dropout.py --only_unit_test - python3 python/tutorials/hints/05/05-layer-norm.py --only_unit_test - python3 python/tutorials/hints/06/06-fused-attention.py --only_unit_test - python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test - python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test - python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test - # flagtree tle raw - python3 python/tutorials/tle/raw/mlir/01-vector-add.py - python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py - python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py - python3 python/tutorials/tle/raw/mlir/04-hello-world.py - python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton - python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle - python3 python/tutorials/tle/raw/mlir/06-test-vassert.py - - name: FlagTree Test on NVidia (triton_v3.6.x branch) if: ${{ steps.check_backend.outputs.should_skip != 'true' && env.TARGET_BRANCH == 'triton_v3.6.x' }} shell: bash run: | set -x source ~/env.sh - # python tutorials + ## python tutorials python3 python/tutorials/01-vector-add.py --only_unit_test python3 python/tutorials/02-fused-softmax.py --only_unit_test python3 python/tutorials/03-matrix-multiplication.py --only_unit_test @@ -169,7 +75,7 @@ jobs: python3 python/tutorials/08-grouped-gemm.py --only_unit_test python3 python/tutorials/09-persistent-matmul.py --only_unit_test python3 python/tutorials/11-programmatic-dependent-launch.py --only_unit_test - # python unit test + ## python unit test python3 -m pytest -s python/test/unit/cuda python3 -m pytest -s python/test/unit/instrumentation python3 -m pytest -s python/test/unit/language \ @@ -178,15 +84,14 @@ jobs: if [ -d "python/test/operators" ]; then python3 -m pytest -s python/test/operators fi - # flagtree tle - # python tutorials + ## flagtree tle python tutorials python3 python/tutorials/tle/01-fft.py python3 python/tutorials/tle/02-moe_align_block_size.py python3 python/tutorials/tle/03-topk.py - python3 python/tutorials/tle/04-cluster-gemm.py + # python3 python/tutorials/tle/04-cluster-gemm.py # TODO: sometimes failed python3 python/tutorials/tle/deepseek_v32/01-topk_selector.py python3 python/tutorials/tle/deepseek_v32/02-sparse-mla.py - # python unit test + ## flagtree tle python unit test 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 \ @@ -195,8 +100,8 @@ jobs: 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 + ## 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 python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test @@ -206,7 +111,7 @@ jobs: python3 python/tutorials/hints/07/07-extern-functions.py --only_unit_test python3 python/tutorials/hints/08/08-grouped-gemm.py --only_unit_test python3 python/tutorials/hints/11/11-programmatic-dependent-launch.py --only_unit_test - # flagtree tle raw + ## flagtree tle raw python3 python/tutorials/tle/raw/mlir/01-vector-add.py python3 python/tutorials/tle/raw/mlir/02-fused-softmax.py python3 python/tutorials/tle/raw/mlir/03-matrix-multiplication.py @@ -214,7 +119,7 @@ jobs: python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel triton python3 python/tutorials/tle/raw/mlir/05-topk.py --kernel tle python3 python/tutorials/tle/raw/mlir/06-test-vassert.py - # flagtree tle cuda + ## flagtree tle cuda # TODO: These tests are currently skipped because the CLANG environment variable cannot be set. # python3 python/tutorials/tle/raw/cuda/01-vector-add.py # python3 python/tutorials/tle/raw/cuda/02-fused-softmax.py From cbf2bd71b69867820a151006a27cb0231c5b7119 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Wed, 8 Jul 2026 09:11:20 +0000 Subject: [PATCH 11/32] polish code --- .../TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 17 +++++---------- .../TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h | 8 +++++++ .../TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt | 1 + .../FlagCxOpToLLVM/FlagCxOpToLLVM.cpp | 21 +++++++++++++++++++ 4 files changed, 35 insertions(+), 12 deletions(-) create mode 100644 third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h create mode 100644 third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.cpp diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 43c1c7e88..54d1cedfd 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -25,8 +25,7 @@ #include "tle/dialect/include/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/ExclusiveCumsumOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/ExtractOpToLLVM.h" -#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/FlagCxOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/GetDeviceIdToFlagCX.h" #include "tle/dialect/include/Conversion/TleToLLVM/LocalPointersOpToLLVM.h" #include "tle/dialect/include/Conversion/TleToLLVM/PackOpToLLVM.h" @@ -191,16 +190,10 @@ struct ConvertTritonGPUToLLVM typeConverter, patterns, benefit); mlir::triton::tle::populateTMAStoreCommitGroupOpToLLVMPatterns( typeConverter, patterns, benefit); -#ifdef FLAGCX_ENABLED - mlir::triton::tle::populateGetDeviceIdOpToFlagCxPatterns( - typeConverter, patterns, benefit); - mlir::triton::tle::populateGetLocalRankOpToLLVMPatterns( - typeConverter, patterns, benefit); - mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, - patterns, benefit); - mlir::triton::tle::populateDeviceIntraBarrierOpToLLVMPatterns( - typeConverter, patterns, benefit); -#endif + // FlagCX ops are lowered to LLVM. + mlir::triton::tle::populateFlagCxOpToLLVMPatterns(typeConverter, patterns, + benefit); + if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { return signalPassFailure(); } diff --git a/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h b/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h new file mode 100644 index 000000000..e2704e8e5 --- /dev/null +++ b/third_party/tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.h @@ -0,0 +1,8 @@ +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h" + +namespace mlir::triton::tle { +void populateFlagCxOpToLLVMPatterns(LLVMTypeConverter &typeConverter, + RewritePatternSet &patterns, + PatternBenefit benefit); +} diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt index b7f2e3d3b..124815cb3 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/CMakeLists.txt @@ -1,6 +1,7 @@ set(FLAGCX_TO_LLVM_SRC DeviceIntraBarrierOpToLLVM.cpp GetLocalRankOpToLLVM.cpp + FlagCxOpToLLVM.cpp ) add_triton_library(FlagCxToLLVM diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.cpp new file mode 100644 index 000000000..214bcc590 --- /dev/null +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/FlagCxOpToLLVM.cpp @@ -0,0 +1,21 @@ +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.h" +#include "tle/dialect/include/Conversion/TleToLLVM/GetDeviceIdToFlagCX.h" + +namespace mlir::triton::tle { +void populateFlagCxOpToLLVMPatterns(LLVMTypeConverter &typeConverter, + RewritePatternSet &patterns, + PatternBenefit benefit) { +#ifdef FLAGCX_ENABLED + mlir::triton::tle::populateGetDeviceIdOpToFlagCxPatterns(typeConverter, + patterns, benefit); + mlir::triton::tle::populateGetLocalRankOpToLLVMPatterns(typeConverter, + patterns, benefit); + mlir::triton::tle::populateGetNumPesOpToLLVMPatterns(typeConverter, patterns, + benefit); + mlir::triton::tle::populateDeviceIntraBarrierOpToLLVMPatterns( + typeConverter, patterns, benefit); +#endif +} + +} // namespace mlir::triton::tle From 10a20dfd2154dc58d6538a3acb6e8174d0d451fc Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Wed, 8 Jul 2026 10:03:37 +0000 Subject: [PATCH 12/32] fix third_party/nvidia/backend/distributed.py --- third_party/nvidia/backend/distributed.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/nvidia/backend/distributed.py b/third_party/nvidia/backend/distributed.py index 798500573..159a4d902 100644 --- a/third_party/nvidia/backend/distributed.py +++ b/third_party/nvidia/backend/distributed.py @@ -69,8 +69,8 @@ def __init__(self, path_order=0): self.include_path = self._get_include_paths()[path_order] def _check_path_available(self, paths): - available_paths = [Path(p) for p in paths if p and p.exists()] - if len(paths) == 0: + available_paths = [Path(p) for p in paths if p and Path(p).exists()] + if len(available_paths) == 0: raise RuntimeError(f"There are no available {self.bt_name} path in this {available_paths}") return available_paths From 89e5383e60af1dd9fbc062d883e1834aa784dca0 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Thu, 9 Jul 2026 02:18:14 +0000 Subject: [PATCH 13/32] fix bugs --- .../nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp index 54d1cedfd..552dc79ed 100644 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp @@ -191,9 +191,10 @@ struct ConvertTritonGPUToLLVM mlir::triton::tle::populateTMAStoreCommitGroupOpToLLVMPatterns( typeConverter, patterns, benefit); // FlagCX ops are lowered to LLVM. +#ifdef FLAGCX_ENABLED mlir::triton::tle::populateFlagCxOpToLLVMPatterns(typeConverter, patterns, benefit); - +#endif if (failed(applyPartialConversion(mod, target, std::move(patterns)))) { return signalPassFailure(); } From 64e54bcdc60fe8fc4742257bcb088246282b84af Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Thu, 9 Jul 2026 02:21:42 +0000 Subject: [PATCH 14/32] add arg for GroupKind --- python/triton/experimental/tle/language/distributed.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index f72fafebd..6ae0b0526 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -64,6 +64,7 @@ class GroupKind(str, Enum): BLOCK = "block" TILE_SPAN = "tile_span" LANES = "lanes" + GRID = "grid" @dataclass From 57fe69eac9472bbd9b0bff0c6a495cfc23241ff5 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 14 Jul 2026 04:35:33 +0000 Subject: [PATCH 15/32] add flag to control tileir --- .gitmodules | 6 +-- CMakeLists.txt | 2 +- python/setup_tools/setup_helper.py | 2 + python/setup_tools/utils/__init__.py | 4 ++ python/setup_tools/utils/tools.py | 7 ++++ python/triton/compiler/code_generator.py | 1 + setup.py | 1 - third_party/tileir/README.md | 6 ++- .../tileir/lib/TritonToTileIR/Utils.cpp | 2 +- third_party/tileir/third_party/cuda-tile | 1 - .../tle/dialect/include/Transforms/Passes.td | 26 ++++++++++++ .../tle/dialect/lib/Transforms/CMakeLists.txt | 1 + .../Transforms/TleAddDistributedParams.cpp | 40 +++++++++++++++++++ third_party/tle/triton_tle.cc | 2 + 14 files changed, 92 insertions(+), 9 deletions(-) delete mode 160000 third_party/tileir/third_party/cuda-tile create mode 100644 third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp diff --git a/.gitmodules b/.gitmodules index fd32e978f..c63a55bf6 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ -[submodule "third_party/tileir/third_party/cuda-tile"] - path = third_party/tileir/third_party/cuda-tile - url = https://github.com/NVIDIA/cuda-tile +# [submodule "third_party/tileir/third_party/cuda-tile"] +# path = third_party/tileir/third_party/cuda-tile +# url = https://github.com/NVIDIA/cuda-tile diff --git a/CMakeLists.txt b/CMakeLists.txt index 52b0bee20..024281b0b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -459,7 +459,7 @@ if(TRITON_BUILD_PYTHON_MODULE) # ${TRITON_BINARY_DIR}/third_party/${PLUGIN_NAME} cmake_path(APPEND TRITON_BINARY_DIR "third_party" ${PLUGIN_NAME} OUTPUT_VARIABLE PLUGIN_DIR_BUILD_OUTPUT) message(STATUS "Building plugin '${PLUGIN_NAME}' from ${PLUGIN_DIR} with output ${PLUGIN_DIR_BUILD_OUTPUT}") - add_subdirectory(${PLUGIN_DIR} ${PLUGIN_DIR_BUILD_OUTPUT}) + # add_subdirectory(${PLUGIN_DIR} ${PLUGIN_DIR_BUILD_OUTPUT}) endforeach() endif() diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 4f101733c..f73d3e985 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -503,6 +503,8 @@ def uninstall_triton(): download_flagtree_third_party("flagcx", condition=(not flagtree_backend), hook="handle_flagcx", required=True) +download_flagtree_third_party("tileir", condition=(not flagtree_backend), required=True) + handle_flagtree_backend() # iluvatar diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py index 4c618155b..48d7d6888 100644 --- a/python/setup_tools/utils/__init__.py +++ b/python/setup_tools/utils/__init__.py @@ -40,6 +40,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" + }, )) diff --git a/python/setup_tools/utils/tools.py b/python/setup_tools/utils/tools.py index 918ee56c1..752a168d2 100644 --- a/python/setup_tools/utils/tools.py +++ b/python/setup_tools/utils/tools.py @@ -44,6 +44,13 @@ 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() diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index e79845f51..ea8fd56b8 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -656,6 +656,7 @@ def visit_FunctionDef(self, node): for arg_name, arg_value in zip(arg_names, arg_values): self._maybe_set_loc_to_name(arg_value, arg_name) self.set_value(arg_name, arg_value) + print(f"{arg_name} : {arg_value}") insert_pt = self.builder.get_insertion_block() self.builder.set_insertion_point_to_start(entry) # visit function body diff --git a/setup.py b/setup.py index af96928d7..0d7fe420c 100644 --- a/setup.py +++ b/setup.py @@ -648,7 +648,6 @@ def download_and_copy_dependencies(): f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz", ) - if helper.flagtree_backend: if helper.flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir"): default_backends = helper.configs.non_tileir_default_backends() diff --git a/third_party/tileir/README.md b/third_party/tileir/README.md index 2e1c1fe48..e6ed4d907 100644 --- a/third_party/tileir/README.md +++ b/third_party/tileir/README.md @@ -123,7 +123,7 @@ shown above. **FlagTree compatibility** -- Target FlagTree/Triton base: `triton_v3.6.x` / Triton 3.6.0 +- Target FlagTree/Triton base: `main` / Triton 3.6.0 - FlagTree's pinned LLVM is older than the LLVM expected by cuda-tile v13.3.0. - `scripts/patch_bytecode_utils.sh` is a FlagTree-local compatibility patch for that LLVM skew. It rewrites cuda-tile source before build, including: @@ -133,7 +133,7 @@ shown above. **FlagTree-local changes** -Compared with the imported upstream commit, FlagTree modifies 8 files and adds +Compared with the imported upstream commit, FlagTree modifies 9 files and adds 4 entries: - **Build and provenance:** Modified `README.md`, `CMakeLists.txt`, and @@ -151,6 +151,8 @@ Compared with the imported upstream commit, FlagTree modifies 8 files and adds - **Lowering:** Modified `lib/TritonToTileIR/TritonToTileIRPass.cpp` to add `math.erf` lowering and align flattened TensorDescriptor arguments with FlagTree's ABI. +- **Release builds:** Modified `lib/TritonToTileIR/Utils.cpp` so its unsupported + reduction-type path remains unreachable when `NDEBUG` removes assertions. - **Runtime:** Upstream activates `TileIRDriver` globally with `ENABLE_TILE=1` and creates it during import. Modified `backend/driver.py` so it is not a global active driver and is created only for a kernel routed to TileIR. The diff --git a/third_party/tileir/lib/TritonToTileIR/Utils.cpp b/third_party/tileir/lib/TritonToTileIR/Utils.cpp index f1a34a9d1..6132e6614 100644 --- a/third_party/tileir/lib/TritonToTileIR/Utils.cpp +++ b/third_party/tileir/lib/TritonToTileIR/Utils.cpp @@ -97,7 +97,7 @@ Attribute getIdentitiesAttr(MLIRContext *context, } } - assert(false && "unexpected data type for reduceOp."); + llvm_unreachable("unexpected data type for reduceOp."); } bool isI8OrI1ElementTensor(Type type) { diff --git a/third_party/tileir/third_party/cuda-tile b/third_party/tileir/third_party/cuda-tile deleted file mode 160000 index 2e5ccba66..000000000 --- a/third_party/tileir/third_party/cuda-tile +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 2e5ccba66fb3afdba34b26cf358418283027c248 diff --git a/third_party/tle/dialect/include/Transforms/Passes.td b/third_party/tle/dialect/include/Transforms/Passes.td index 2ddd3b1bf..a9e295d73 100644 --- a/third_party/tle/dialect/include/Transforms/Passes.td +++ b/third_party/tle/dialect/include/Transforms/Passes.td @@ -397,4 +397,30 @@ def TleRemoveRedundantCopy def TleDSLRegionInline : Pass<"tle-dslregion-inline", "mlir::ModuleOp"> {} +def TritonTleAddDistributedParams + : Pass<"triton-tle-add-distributed-params", "mlir::ModuleOp"> { + + let summary = "Add distributed runtime parameters to Triton kernel functions"; + + let description = [{ + This pass updates Triton kernel function signatures by inserting two + runtime parameters at the beginning of each kernel function: + + * device_comm_ptr + * device_mem_ptr + + These parameters provide device-side communication and runtime state + required by distributed TLE operations. + + All call sites and function argument uses are updated accordingly. + }]; + + let dependentDialects = [ + "mlir::func::FuncDialect", + "mlir::arith::ArithDialect", + "mlir::triton::TritonDialect" + ]; +} + + #endif // TRITON_TLE_PASSES diff --git a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt index f77a3b937..f6ee4f946 100644 --- a/third_party/tle/dialect/lib/Transforms/CMakeLists.txt +++ b/third_party/tle/dialect/lib/Transforms/CMakeLists.txt @@ -28,6 +28,7 @@ add_triton_library(TritonTLETransforms RemoveRedundantCopy.cpp DSLRegionInline.cpp TleUtility.cpp + TleAddDistributedParams.cpp DEPENDS TritonTLETransformsIncGen diff --git a/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp new file mode 100644 index 000000000..c24473848 --- /dev/null +++ b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp @@ -0,0 +1,40 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" +#include "tle/dialect/include/IR/Dialect.h" +#include "tle/dialect/include/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" + +namespace mlir::triton::tle { + +namespace ttg = mlir::triton::gpu; +namespace ttng = mlir::triton::nvidia_gpu; + +#define GEN_PASS_DEF_TRITONTLEADDDISTRIBUTEDPARAMS +#include "tle/dialect/include/Transforms/Passes.h.inc" + + +namespace{ + +struct TritonTleAddDistributedParams + : public impl::TritonTleAddDistributedParamsBase< + TritonTleAddDistributedParams> { + void runOnOperation() override { + ModuleOp m = getOperation(); + OpBuilder builder(m.getContext()); + + m->walk([&, this](triton::FuncOp func) { + auto type = func.getFunctionType(); + llvm::errs() << type << "\n"; + }); + + } + +}; + + + +} +} // mlir::triton::tle \ No newline at end of file diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index caa435eaf..44ae34fbd 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -660,6 +660,8 @@ void init_triton_tle_ir(py::module &&m) { } void init_triton_tle_passes(py::module &&m) { + ADD_PASS_WRAPPER_0("add_params_for_distribution", + tle::createTritonTleAddDistributedParams); ADD_PASS_WRAPPER_0("add_early_assign_memory_space", tle::createTritonTleEarlyAssignMemorySpace); ADD_PASS_WRAPPER_0("add_select_encodings", From f613b218e0c2e9b776f2531e0974a6d0d76aae40 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 14 Jul 2026 04:36:28 +0000 Subject: [PATCH 16/32] add flag to control tileir --- python/setup_tools/utils/__init__.py | 4 ++-- python/setup_tools/utils/tools.py | 6 +----- setup.py | 1 + .../lib/Transforms/TleAddDistributedParams.cpp | 17 ++++++----------- 4 files changed, 10 insertions(+), 18 deletions(-) diff --git a/python/setup_tools/utils/__init__.py b/python/setup_tools/utils/__init__.py index 48d7d6888..a69c8bd7d 100644 --- a/python/setup_tools/utils/__init__.py +++ b/python/setup_tools/utils/__init__.py @@ -41,8 +41,8 @@ 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" + "name": "tileir", "url": "https://github.com/NVIDIA/cuda-tile", "relative_path": "tileir/third_party/cuda-tile", + "commit_id": "2e5ccba66fb3afdba34b26cf358418283027c248" }, )) diff --git a/python/setup_tools/utils/tools.py b/python/setup_tools/utils/tools.py index 752a168d2..4134c9238 100644 --- a/python/setup_tools/utils/tools.py +++ b/python/setup_tools/utils/tools.py @@ -45,11 +45,7 @@ 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" - ] + _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() diff --git a/setup.py b/setup.py index 0d7fe420c..af96928d7 100644 --- a/setup.py +++ b/setup.py @@ -648,6 +648,7 @@ def download_and_copy_dependencies(): f"https://developer.download.nvidia.com/compute/cuda/redist/cuda_cupti/{system}-{arch}/cuda_cupti-{system}-{arch}-{version}-archive.tar.xz", ) + if helper.flagtree_backend: if helper.flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir"): default_backends = helper.configs.non_tileir_default_backends() diff --git a/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp index c24473848..5e3223b40 100644 --- a/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp @@ -15,26 +15,21 @@ namespace ttng = mlir::triton::nvidia_gpu; #define GEN_PASS_DEF_TRITONTLEADDDISTRIBUTEDPARAMS #include "tle/dialect/include/Transforms/Passes.h.inc" - -namespace{ +namespace { struct TritonTleAddDistributedParams : public impl::TritonTleAddDistributedParamsBase< - TritonTleAddDistributedParams> { + TritonTleAddDistributedParams> { void runOnOperation() override { ModuleOp m = getOperation(); OpBuilder builder(m.getContext()); m->walk([&, this](triton::FuncOp func) { - auto type = func.getFunctionType(); - llvm::errs() << type << "\n"; + auto type = func.getFunctionType(); + llvm::errs() << type << "\n"; }); - } - }; - - -} -} // mlir::triton::tle \ No newline at end of file +} // namespace +} // namespace mlir::triton::tle From bc95f16ccd26f8cc372410e62241db06b6530acc Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Tue, 14 Jul 2026 06:20:03 +0000 Subject: [PATCH 17/32] polish code --- .gitmodules | 3 --- CMakeLists.txt | 2 +- 2 files changed, 1 insertion(+), 4 deletions(-) delete mode 100644 .gitmodules diff --git a/.gitmodules b/.gitmodules deleted file mode 100644 index c63a55bf6..000000000 --- a/.gitmodules +++ /dev/null @@ -1,3 +0,0 @@ -# [submodule "third_party/tileir/third_party/cuda-tile"] -# path = third_party/tileir/third_party/cuda-tile -# url = https://github.com/NVIDIA/cuda-tile diff --git a/CMakeLists.txt b/CMakeLists.txt index 024281b0b..52b0bee20 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -459,7 +459,7 @@ if(TRITON_BUILD_PYTHON_MODULE) # ${TRITON_BINARY_DIR}/third_party/${PLUGIN_NAME} cmake_path(APPEND TRITON_BINARY_DIR "third_party" ${PLUGIN_NAME} OUTPUT_VARIABLE PLUGIN_DIR_BUILD_OUTPUT) message(STATUS "Building plugin '${PLUGIN_NAME}' from ${PLUGIN_DIR} with output ${PLUGIN_DIR_BUILD_OUTPUT}") - # add_subdirectory(${PLUGIN_DIR} ${PLUGIN_DIR_BUILD_OUTPUT}) + add_subdirectory(${PLUGIN_DIR} ${PLUGIN_DIR_BUILD_OUTPUT}) endforeach() endif() From ee5e5d48a5a5ab5fa9c44a2bedf07481cf19d243 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Thu, 16 Jul 2026 02:59:34 +0000 Subject: [PATCH 18/32] enable distribued params pass --- python/triton/compiler/code_generator.py | 1 - .../tle/language/communication.py | 3 + .../experimental/tle/language/distributed.py | 39 ++++++----- python/triton/runtime/__init__.py | 2 + python/triton/runtime/_distributed.py | 70 +++++++++++++++++++ python/triton/runtime/jit.py | 15 ++-- third_party/nvidia/backend/compiler.py | 3 + third_party/nvidia/backend/driver.py | 5 ++ third_party/tle/dialect/include/IR/TleOps.td | 6 +- .../tle/dialect/include/Transforms/Passes.td | 2 - .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 10 ++- .../DeviceIntraBarrierOpToLLVM.cpp | 1 - .../FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp | 6 +- .../TleToLLVM/GetDeviceIdToFlagCX.cpp | 17 ++++- .../TleToLLVM/LocalPointersOpToLLVM.cpp | 42 +++++++---- third_party/tle/dialect/lib/IR/Ops.cpp | 4 +- .../tle/dialect/lib/IR/VerfiyUtils.cpp | 4 +- .../Transforms/TleAddDistributedParams.cpp | 49 +++++++++++-- third_party/tle/triton_tle.cc | 31 ++++---- 19 files changed, 240 insertions(+), 70 deletions(-) create mode 100644 python/triton/runtime/_distributed.py diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py index ea8fd56b8..e79845f51 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py @@ -656,7 +656,6 @@ def visit_FunctionDef(self, node): for arg_name, arg_value in zip(arg_names, arg_values): self._maybe_set_loc_to_name(arg_value, arg_name) self.set_value(arg_name, arg_value) - print(f"{arg_name} : {arg_value}") insert_pt = self.builder.get_insertion_block() self.builder.set_insertion_point_to_start(entry) # visit function body diff --git a/python/triton/experimental/tle/language/communication.py b/python/triton/experimental/tle/language/communication.py index fc8a44e28..d57a34343 100644 --- a/python/triton/experimental/tle/language/communication.py +++ b/python/triton/experimental/tle/language/communication.py @@ -205,4 +205,7 @@ def create_comm_tensor(buf_tensor): # Synchronize all ranks before kernel launch dist.barrier() + + from triton.runtime import DistributedRtContext + DistributedRtContext(dev_comm_dptr.value, dev_mem_dptr.value) return dev_comm_dptr.value, dev_mem_dptr.value diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 6ae0b0526..2af4ed53a 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -27,12 +27,20 @@ def _as_positive_int(value: Any, label: str) -> int: return value +def _parse_src_arg(builder, src): + if src is None: + return src + else: + return src.handle + + # Get the current device id @tl.builtin -def _get_local_rank(dev_mem_ptr, _semantic=None, ret_dtype=tl.int32): +def _get_local_rank(dev_comm_ptr, _semantic=None, ret_dtype=tl.int32): builder = _semantic.builder ret_ir_ty = ret_dtype.to_ir(builder) - result = builder.get_device_id(ret_ir_ty, dev_mem_ptr.handle) + ptr = _parse_src_arg(builder, dev_comm_ptr) + result = builder.get_device_id(ret_ir_ty, ptr) return tl.tensor(result, ret_dtype) @@ -609,7 +617,7 @@ def shard_id( mesh = tl._unwrap_if_constexpr(mesh) axis = tl._unwrap_if_constexpr(axis) - if comm_ptr is not None: + if axis in ("device", "node"): return _get_local_rank(comm_ptr, _semantic=_semantic, ret_dtype=tl.int32) if not isinstance(mesh, device_mesh): @@ -648,14 +656,12 @@ def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, barrier_kind: BarrierKind | str = BarrierKind.SYNC, group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): - if space or comm_ptr: - if space and space != "device" and comm_ptr: - raise ValueError(f"{space} space and comm_ptr cannot be used together") - if space and space == "device" and not comm_ptr: - raise NotImplementedError(f"distributed_barrier with space={space} must have comm_ptr arg") + + if space and space in ("device", "node"): builder = _semantic.builder + ptr = _parse_src_arg(builder, comm_ptr) builder.create_distributed_barrier( - src=comm_ptr.handle, + src=ptr, barrier_index=index or 0, space=_parse_device_barrier_args("device"), group_kind=_parse_device_barrier_args(group_kind), @@ -813,7 +819,7 @@ def _create_remote_pointers_tensor( "cluster": (dtype, 7), "device": (dtype, 1), }.get(space)) - if tensor.type.is_block(): + if tensor and tensor.type.is_block(): remote_type = tl.block_type(remote_ptr_dtype, list(tensor.shape)).to_ir(builder) else: remote_type = remote_ptr_dtype.to_ir(builder) @@ -838,11 +844,12 @@ def _create_remote_pointers_tensor( # automatic injection (e.g. inside this helper). if offset_tensor.dtype != tl.int64: offset_tensor = tl.cast(offset_tensor, tl.int64, _semantic=_semantic) - remote_op = builder.create_remote_pointers(remote_type, tensor.handle, shard_id_tensor.handle, space, + ptr = _parse_src_arg(builder, tensor) + remote_op = builder.create_remote_pointers(remote_type, ptr, shard_id_tensor.handle, space, offset_tensor.handle) else: remote_op = builder.create_remote_pointers(remote_type, tensor.handle, shard_id_tensor.handle, space) - if tensor.type.is_block(): + if tensor and tensor.type.is_block(): return tl.tensor(remote_op.get_result(0), tl.block_type(remote_ptr_dtype, list(tensor.shape))) return tl.tensor(remote_op.get_result(0), remote_ptr_dtype) @@ -888,7 +895,7 @@ def _remote_pointer( _semantic=None, ) -> tl.tensor: - if not isinstance(tensor, tl.tensor): + if not isinstance(tensor, tl.tensor) and space not in ("device", "node"): raise TypeError(f"tensor must be tl.tensor, got {type(tensor).__name__}") space = tl._unwrap_if_constexpr(space) @@ -917,8 +924,8 @@ def _remote_pointer( @tl.builtin def remote( - tensor, - shard_id, + tensor=None, + shard_id=None, scope: device_mesh | None = None, space: str = "cluster", dtype: tl.dtype = None, @@ -953,7 +960,7 @@ def remote( # Direct pointer path: support local_ptr scalar/tensor values and return # remote pointer with preserved shape. - if isinstance(tensor, tl.tensor): + if isinstance(tensor, tl.tensor) or (tensor is None and space in ("device", "node")): return _remote_pointer(tensor, shard_id, scope=scope, space=space, _semantic=_semantic, dtype=dtype, offset=offset) diff --git a/python/triton/runtime/__init__.py b/python/triton/runtime/__init__.py index 0b3979d28..85b7b00bc 100644 --- a/python/triton/runtime/__init__.py +++ b/python/triton/runtime/__init__.py @@ -3,6 +3,7 @@ from .driver import driver from .jit import JITFunction, KernelInterface, MockTensor, TensorWrapper, reinterpret from .errors import OutOfResources, InterpreterError +from ._distributed import DistributedRtContext __all__ = [ "autotune", @@ -20,4 +21,5 @@ "reinterpret", "RemoteCacheBackend", "TensorWrapper", + "DistributedRtContext", ] diff --git a/python/triton/runtime/_distributed.py b/python/triton/runtime/_distributed.py new file mode 100644 index 000000000..67ccd87a3 --- /dev/null +++ b/python/triton/runtime/_distributed.py @@ -0,0 +1,70 @@ +class DistributedRtContext: + _instance = None + _initialized = False + _init_count = 0 + + def __new__(cls, *args, **kwargs): + if cls._instance is None: + cls._instance = super().__new__(cls) + cls._instance._comm_ptr = None + cls._instance._mem_ptr = None + return cls._instance + + def __init__(self, _comm_ptr=None, _mem_ptr=None): + if _comm_ptr and _mem_ptr: + type(self)._init_count += 1 + + if self._initialized: + return + self._comm_ptr = _comm_ptr + self._mem_ptr = _mem_ptr + self._initialized = True + + @property + def comm_ptr(self) -> int | None: + """Communication runtime pointer.""" + return self._comm_ptr + + def _get_needed_params(self): + return {"device_comm_ptr": self._comm_ptr, "device_mem_ptr": self._mem_ptr} + + @property + def mem_ptr(self) -> int | None: + """Distributed memory runtime pointer.""" + return self._mem_ptr + + @property + def is_lite_mode(self) -> bool: + import os + user_action = os.getenv("FLAGTREE_LITE_DIST", "").strip().upper() in { + "1", + "ON", + "TRUE", + } + inner_action = self._init_count == 1 + return inner_action and user_action + + def add_args_to_jitfunction(self, **kwargs): + params = kwargs['params'] # list + _kwargs = kwargs['kwargs'] + if (isinstance(params, list) and len(params) > 0): + needed_params = self._get_needed_params() + needed_params_size = len(needed_params) + template_ele = params[0] + KernelParam = type(template_ele) # KernelParam type + Parameter = type(template_ele._param) # inspect.Parameter type + + dist_params = [] + for i, (name, val) in enumerate(needed_params.items()): + _kwargs[name] = val + param = Parameter(name, kind=template_ele._param.kind) + dist_params.append(KernelParam(i, param, False, False)) + + new_params = [] + for param in params: + new_loc = param.num + needed_params_size + #num: int, param: inspect.Parameter, do_not_specialize: bool, do_not_specialize_on_alignment: bool + new_params.append( + KernelParam(new_loc, param._param, param.do_not_specialize, param.do_not_specialize_on_alignment)) + new_params = dist_params + new_params + params[:] = new_params diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index ecd8dc8d7..b62bbc3b5 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -20,6 +20,7 @@ from .._utils import find_paths_if, get_iterable_path, type_canonicalisation_dict, is_namedtuple from .cache import get_cache_key from triton._C.libtriton import get_cache_invalidating_env_vars, native_specialize_impl +from ._distributed import DistributedRtContext TRITON_MODULE = "triton.language" GLUON_MODULE = "triton.experimental.gluon.language" @@ -706,6 +707,7 @@ def _pack_args(self, backend, kwargs, bound_args, specialization, options): def run(self, *args, grid, warmup, **kwargs): kwargs["debug"] = kwargs.get("debug", self.debug) or knobs.runtime.debug kwargs["instrumentation_mode"] = knobs.compilation.instrumentation_mode + # DistributedRtContext().add_args_to_jitfunction(kwargs=kwargs, params=self.params) # parse options device = driver.active.get_current_device() @@ -727,7 +729,6 @@ def run(self, *args, grid, warmup, **kwargs): if kernel is None: options, signature, constexprs, attrs = self._pack_args(backend, kwargs, bound_args, specialization, options) - kernel = self._do_compile(key, signature, device, constexprs, options, attrs, warmup) if kernel is None: return None @@ -752,8 +753,15 @@ def run(self, *args, grid, warmup, **kwargs): kernel = kernel.result() # launch kernel launch_metadata = kernel.launch_metadata(grid, stream, *bound_args.values()) - kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, launch_metadata, - knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *bound_args.values()) + + dist_param = [] + ctx = DistributedRtContext() + if ctx.is_lite_mode: + dist_param += [ctx.comm_ptr, ctx.mem_ptr] + kernel.run(grid_0, grid_1, grid_2, stream, kernel.function, kernel.packed_metadata, + launch_metadata, knobs.runtime.launch_enter_hook, knobs.runtime.launch_exit_hook, *dist_param, + *bound_args.values()) + return kernel def repr(self, _): @@ -771,7 +779,6 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o self.do_not_specialize_on_alignment = do_not_specialize_on_alignment self._repr = repr self.launch_metadata = launch_metadata - self.params = [] for i, param in enumerate(self.signature.parameters.values()): dns = i in do_not_specialize or param.name in do_not_specialize diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index fe4ecde0c..43a47181e 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -3,6 +3,7 @@ from triton._C.libtriton import tle from triton import knobs from triton.runtime.errors import PTXASError +from triton.runtime._distributed import DistributedRtContext from dataclasses import dataclass import functools @@ -279,6 +280,8 @@ def make_ttgir(mod, metadata, opt, capability): pm = ir.pass_manager(mod.context) dump_enabled = pm.enable_debug() emuTF32 = (capability // 10 >= 8) + if DistributedRtContext().is_lite_mode: + tle.passes.add_params_for_distribution(pm) passes.ttir.add_convert_to_ttgpuir(pm, f"cuda:{capability}", opt.num_warps, 32, opt.num_ctas) # flagtree tle raw tle.raw_passes.add_tle_convert_arg_to_memdesc(pm) diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py index 8e9c97acf..7745b4439 100644 --- a/third_party/nvidia/backend/driver.py +++ b/third_party/nvidia/backend/driver.py @@ -215,6 +215,9 @@ def format_of(ty): flat_signature = [] for sig in signature.values(): _flatten_signature(sig, flat_signature) + #adxx + # flat_signature.insert(0, "*i64") # dev_comm_ptr + # flat_signature.insert(1, "*i64") # dev_mem_ptr signature = {i: s for i, s in enumerate(flat_signature)} args_list = ', ' + ', '.join(f"&_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else '' # Record the end of regular arguments; @@ -257,7 +260,9 @@ def format_of(ty): for i, ty in signature.items() if ty in FLOAT_STORAGE_TYPE ] + params = [f"&arg{i}" for i, ty in signature.items() if ty != "constexpr"] + params.append("&global_scratch") params.append("&profile_scratch") tle_cpp_define = "#define __TLE__ 1" diff --git a/third_party/tle/dialect/include/IR/TleOps.td b/third_party/tle/dialect/include/IR/TleOps.td index f6fe248c2..c72884494 100644 --- a/third_party/tle/dialect/include/IR/TleOps.td +++ b/third_party/tle/dialect/include/IR/TleOps.td @@ -333,9 +333,9 @@ def Tle_DistributedBarrierOp : Tle_Op<"distributed_barrier", let hasVerifier = 1; } -def Tle_RemotePointersOp : Tle_Op<"remote_pointers", [Pure]> { +def Tle_RemotePointersOp : Tle_Op<"remote_pointers", [Pure, AttrSizedOperandSegments]> { let arguments = (ins - Tle_LocalPointerResultType:$src, + Optional:$src, TT_Int:$shard_id, StrAttr:$space, Optional:$offset @@ -356,7 +356,7 @@ def Tle_GetNumPesOp : Tle_Op<"get_num_pes"> { def Tle_GetDeviceIdOp : Tle_Op<"get_device_id"> { let summary = "Get local PE id"; let arguments = (ins - AnyType:$input + Optional:$input ); let results = (outs I32:$result); let hasVerifier = 1; diff --git a/third_party/tle/dialect/include/Transforms/Passes.td b/third_party/tle/dialect/include/Transforms/Passes.td index a9e295d73..50523383b 100644 --- a/third_party/tle/dialect/include/Transforms/Passes.td +++ b/third_party/tle/dialect/include/Transforms/Passes.td @@ -416,8 +416,6 @@ def TritonTleAddDistributedParams }]; let dependentDialects = [ - "mlir::func::FuncDialect", - "mlir::arith::ArithDialect", "mlir::triton::TritonDialect" ]; } diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index d477246d1..9ee9670be 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -458,8 +458,14 @@ struct DistributedBarrierOpConversion if (order < 0) return rewriter.notifyMatchFailure(op, "invalid order"); - auto srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); - auto comm = srcElems[0]; + auto comm = Value(); + if (auto src = adaptor.getSrc()) { + auto srcElems = unpackLLElements(loc, src, rewriter); + comm = srcElems[0]; + } else { + auto func = op->getParentOfType(); + comm = func.getArgument(0); + } auto coopKindAttr = rewriter.getI32IntegerAttr(coopKind); auto newOrderAttr = rewriter.getI32IntegerAttr(order); auto barrierTypeAttr = op.getBarrierTypeAttr(); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp index 879566541..972dd07c6 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/DeviceIntraBarrierOpToLLVM.cpp @@ -35,7 +35,6 @@ struct DeviceIntraBarrierOpConversion auto coopValue = op.getCoopKindAttr().getInt(); auto orderValue = op.getOrderAttr().getInt(); auto barrierType = op.getBarrierTypeAttr().getValue(); - if (!llvm::is_contained(std::array{0, 1, 2, 3, 4}, coopValue)) return rewriter.notifyMatchFailure(op, "invalid coop_kind"); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp index 26bb7cca4..ae793adb7 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/FlagCxOpToLLVM/GetLocalRankOpToLLVM.cpp @@ -53,20 +53,18 @@ struct GetLocalRankOpConversion LogicalResult matchAndRewrite(tle::GetLocalRankOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto reportFailure = [&](StringRef msg) -> LogicalResult { llvm::errs() << "[GetLocalRankOpConversion] " << msg << "\n"; return failure(); }; auto loc = op.getLoc(); - auto srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); - auto getLocalPeCall = tle::getLocalPeFuncCall(loc, rewriter, srcElems[0]); + auto comm = op.getSrc(); + auto getLocalPeCall = tle::getLocalPeFuncCall(loc, rewriter, comm); Value localPe = getLocalPeCall.getResult(); if (!localPe.getType().isInteger(32)) return reportFailure("expected i32 result"); rewriter.replaceOp(op, localPe); - return success(); } }; diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp index 79061205e..8eaee1c6d 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp @@ -30,10 +30,21 @@ struct GetDeviceIdOpConversion matchAndRewrite(tle::GetDeviceIdOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Value src = adaptor.getInput(); - + auto comm = Value(); + auto func = op->getParentOfType(); + if (!func) { + return rewriter.notifyMatchFailure( + op, "expected parent LLVM::LLVMFuncOp, but none was found. "); + } + + if (auto src = adaptor.getInput()) + comm = src; + else + comm = func.getArgument(0); + + rewriter.modifyOpInPlace(op, [&]() { op->insertOperands(0, comm); }); auto localRank = rewriter.create( - op.getLoc(), rewriter.getI32Type(), src); - + op.getLoc(), rewriter.getI32Type(), comm); rewriter.replaceOp(op, localRank.getResult()); return success(); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp index 4dd033c45..331c8ab1c 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp @@ -374,7 +374,7 @@ LogicalResult lowerClusterSpace(Location loc, ValueRange srcElems, return success(); } -LogicalResult lowerDeviceSpace(Location loc, ValueRange srcElems, +LogicalResult lowerDeviceSpace(Location loc, Value mem_ptr, ValueRange shardElems, Value offsetVal, int elemBytes, ConversionPatternRewriter &rewriter, @@ -387,7 +387,7 @@ LogicalResult lowerDeviceSpace(Location loc, ValueRange srcElems, } auto func = getOrInsertGetPeerPointer(module, rewriter.getContext()); - Value memPtr = srcElems[0]; + Value memPtr = mem_ptr; Value peer = shardElems[0]; auto i64Ty = rewriter.getI64Type(); @@ -452,6 +452,15 @@ LogicalResult lowerNodeSpace(Location loc, ValueRange srcElems, return failure(); // Not implemented yet } +Value getDistDevicePtr(tle::RemotePointersOp op, SmallVector &srcElems) { + if (!srcElems.empty()) + return srcElems[0]; + else { + auto func = op->getParentOfType(); + return func.getArgument(1); + } +} + struct RemotePointersOpConversion : public ConvertOpToLLVMPattern { RemotePointersOpConversion(LLVMTypeConverter &typeConverter, @@ -467,19 +476,24 @@ struct RemotePointersOpConversion llvm::errs() << "[RemotePointersOpConversion] " << msg << "\n"; return rewriter.notifyMatchFailure(op, msg); }; - auto srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); - if (srcElems.empty()) + + SmallVector srcElems; + auto space = adaptor.getSpace(); + + if (adaptor.getSrc()) + srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); + + if (space != "device" && srcElems.empty()) return reportFailure("expected non-empty source pointer elements"); auto shardElems = unpackLLElements(loc, adaptor.getShardId(), rewriter); if (shardElems.empty()) return reportFailure("expected non-empty shard_id elements"); - if (shardElems.size() != 1 && shardElems.size() != srcElems.size()) + if (space != "device" && shardElems.size() != 1 && + shardElems.size() != srcElems.size()) return reportFailure( "shard_id must be scalar or match source pointer element count"); - auto space = adaptor.getSpace(); - Value offsetVal; if (adaptor.getOffset()) { auto offsetElems = unpackLLElements(loc, adaptor.getOffset(), rewriter); @@ -489,9 +503,11 @@ struct RemotePointersOpConversion } SmallVector mappedPtrs; + auto mem = getDistDevicePtr(op, srcElems); + if (space == "cluster") { - if (failed(lowerClusterSpace(loc, srcElems, shardElems, rewriter, - mappedPtrs))) { + if (failed( + lowerClusterSpace(loc, mem, shardElems, rewriter, mappedPtrs))) { return rewriter.notifyMatchFailure(op, "cluster lowering failed"); } } else if (space == "device") { @@ -506,19 +522,17 @@ struct RemotePointersOpConversion "result pointee type must be scalar int or float"); elemBytes = elemBits.value() / 8; } - if (failed(lowerDeviceSpace(loc, srcElems, shardElems, offsetVal, - elemBytes, rewriter, mappedPtrs))) { + if (failed(lowerDeviceSpace(loc, mem, shardElems, offsetVal, elemBytes, + rewriter, mappedPtrs))) { return rewriter.notifyMatchFailure(op, "device lowering failed"); } } else if (space == "node") { - if (failed(lowerNodeSpace(loc, srcElems, shardElems, rewriter, - mappedPtrs))) { + if (failed(lowerNodeSpace(loc, mem, shardElems, rewriter, mappedPtrs))) { return rewriter.notifyMatchFailure(op, "node lowering failed"); } } else { return reportFailure("unsupported remote space: " + space.str()); } - Value packed = packLLElements(loc, typeConverter, mappedPtrs, rewriter, op.getType()); rewriter.replaceOp(op, packed); diff --git a/third_party/tle/dialect/lib/IR/Ops.cpp b/third_party/tle/dialect/lib/IR/Ops.cpp index cf563b50a..2856d82bc 100644 --- a/third_party/tle/dialect/lib/IR/Ops.cpp +++ b/third_party/tle/dialect/lib/IR/Ops.cpp @@ -897,13 +897,13 @@ LogicalResult DistributedBarrierOp::verify() { } LogicalResult RemotePointersOp::verify() { - Type srcTy = getSrc().getType(); - Type resultTy = getResult().getType(); auto spaceAttr = getSpace(); if (spaceAttr == "device") { if (failed(RemotePointers::verifyDeviceSpace(getSrc(), getResult()))) return failure(); } else { + Type srcTy = getSrc().getType(); + Type resultTy = getResult().getType(); auto getPtrInfo = [&](Type ty, triton::PointerType &ptr, bool &isTensor, ArrayRef &shape, Attribute &encoding) -> LogicalResult { diff --git a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp index c1d9177fa..74dd36853 100644 --- a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp +++ b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp @@ -19,6 +19,8 @@ namespace RemotePointers { llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result) { // flagcxGetIntraPointerC accept raw device pointers represented as signless // i64 values. + if (!src) + return success(); if (!src.getType().isSignlessInteger(64)) return failure(); @@ -39,7 +41,7 @@ llvm::LogicalResult verifyDeviceSpace(mlir::Operation *op, mlir::Value src) { auto barrierTypeAttr = op->getAttrOfType("barrier_type"); auto orderAttr = op->getAttrOfType("order"); - if (src && kindAttr && barrierTypeAttr && orderAttr) + if (kindAttr && barrierTypeAttr && orderAttr) return success(); else return op->emitOpError() diff --git a/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp index 5e3223b40..49fd83f50 100644 --- a/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp +++ b/third_party/tle/dialect/lib/Transforms/TleAddDistributedParams.cpp @@ -17,17 +17,56 @@ namespace ttng = mlir::triton::nvidia_gpu; namespace { +void addDistributedParams(MLIRContext *ctx, mlir::triton::FuncOp func) { + auto i64Ty = IntegerType::get(ctx, 64); + + // + // 1. Update function type. + // + auto oldType = func.getFunctionType(); + + SmallVector inputTypes; + inputTypes.push_back(i64Ty); // device_comm_ptr + inputTypes.push_back(i64Ty); // device_mem_ptr + llvm::append_range(inputTypes, oldType.getInputs()); + + func.setFunctionType( + FunctionType::get(ctx, inputTypes, oldType.getResults())); + + // + // 2. Update argument attributes. + // + auto makeNameAttr = [&](StringRef name) { + return DictionaryAttr::get(ctx, + {NamedAttribute(StringAttr::get(ctx, "tt.name"), + StringAttr::get(ctx, name))}); + }; + + SmallVector argAttrs; + argAttrs.push_back(makeNameAttr("device_comm_ptr")); + argAttrs.push_back(makeNameAttr("device_mem_ptr")); + + if (auto oldAttrs = func.getArgAttrsAttr()) + llvm::append_range(argAttrs, oldAttrs); + + func.setArgAttrsAttr(ArrayAttr::get(ctx, argAttrs)); + + // + // 3. Insert block arguments. + // + Block &entry = func.front(); + entry.insertArgument(static_cast(0), i64Ty, func.getLoc()); + entry.insertArgument(static_cast(1), i64Ty, func.getLoc()); +} + struct TritonTleAddDistributedParams : public impl::TritonTleAddDistributedParamsBase< TritonTleAddDistributedParams> { void runOnOperation() override { ModuleOp m = getOperation(); - OpBuilder builder(m.getContext()); + MLIRContext *ctx = m.getContext(); - m->walk([&, this](triton::FuncOp func) { - auto type = func.getFunctionType(); - llvm::errs() << type << "\n"; - }); + m.walk([&](mlir::triton::FuncOp func) { addDistributedParams(ctx, func); }); } }; diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 44ae34fbd..1f02385a7 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -531,8 +531,8 @@ void init_triton_tle_ir(py::module &&m) { }) .def( "create_distributed_barrier", - [](TritonOpBuilder &self, Value src, size_t barrier_index = 0, - const std::string &space = "device", + [](TritonOpBuilder &self, std::optional src, + size_t barrier_index = 0, const std::string &space = "device", const std::string &group_kind = "block", const std::string &order = "acqrel", const std::string &barrier_kind = "sync") -> void { @@ -549,12 +549,13 @@ void init_triton_tle_ir(py::module &&m) { builder.getI32IntegerAttr(static_cast(barrier_index)); self.create( - src, spaceAttr, barrierTypeAttr, orderAttr, kindAttr, - barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), + src.value_or(Value()), spaceAttr, barrierTypeAttr, orderAttr, + kindAttr, barrierIndexAttr, IntegerAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr(), DenseI32ArrayAttr()); }, - py::arg("src"), py::arg("barrier_index"), py::arg("space"), - py::arg("group_kind"), py::arg("order"), py::arg("barrier_kind")) + py::arg("src") = py::none(), py::arg("barrier_index"), + py::arg("space"), py::arg("group_kind"), py::arg("order"), + py::arg("barrier_kind")) .def( "create_distributed_barrier", [](TritonOpBuilder &self, const std::string &groupKind, @@ -597,8 +598,8 @@ void init_triton_tle_ir(py::module &&m) { py::arg("group_mask")) .def( "create_remote_pointers", - [](TritonOpBuilder &self, Type resultTy, Value src, Value shardId, - const std::string &space, + [](TritonOpBuilder &self, Type resultTy, std::optional &src, + Value shardId, const std::string &space, std::optional &offset) -> OpState { auto &builder = self.getBuilder(); static const std::unordered_set valid = { @@ -609,15 +610,21 @@ void init_triton_tle_ir(py::module &&m) { ". Expected one of: cluster, device, node."); } auto space_attr = builder.getStringAttr(space); + return self.create( - resultTy, src, shardId, space_attr, offset.value_or(Value())); + resultTy, src.value_or(Value()), shardId, space_attr, + offset.value_or(Value())); }, - py::arg("resultTy"), py::arg("src"), py::arg("shardId"), + py::arg("resultTy"), py::arg("src") = py::none(), py::arg("shardId"), py::arg("space"), py::arg("offset") = py::none()) .def("get_device_id", - [](TritonOpBuilder &self, Type resultTy, Value src) -> Value { + [](TritonOpBuilder &self, Type resultTy, + std::optional src) -> Value { auto &builder = self.getBuilder(); - return self.create(resultTy, src); + auto ptr = Value(); + if (src.has_value()) + ptr = src.value(); + return self.create(resultTy, ptr); }) .def("get_n_pes", [](TritonOpBuilder &self, Type resultTy, Value src) -> Value { From f10b6e6fc96a617e3f96b400083b85ee71445ffb Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Thu, 16 Jul 2026 06:34:46 +0000 Subject: [PATCH 19/32] fix bug --- python/test/tle/integration/test_tle_distributed.py | 2 +- python/test/tle/unit/test_tle_d2d_barrier.py | 3 +-- python/triton/experimental/tle/language/__init__.py | 6 ++++++ python/triton/experimental/tle/language/distributed.py | 10 +++++----- python/triton/runtime/jit.py | 1 - third_party/nvidia/backend/driver.py | 7 ++++--- .../lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp | 4 ++-- 7 files changed, 19 insertions(+), 14 deletions(-) diff --git a/python/test/tle/integration/test_tle_distributed.py b/python/test/tle/integration/test_tle_distributed.py index 067371ec1..e6df4bfb1 100644 --- a/python/test/tle/integration/test_tle_distributed.py +++ b/python/test/tle/integration/test_tle_distributed.py @@ -975,7 +975,7 @@ def test_remote_const_shard_load_high_block_encoding_no_regression(self): assert re.search(r"tensor<256x!tt\.ptr,\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,\s*#blocked[0-9]*>,\s*i32\)\s*->\s*" r"tensor<256x!tt\.ptr,\s*#blocked[0-9]*>", ttgir, diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index db2422fd8..837e17b21 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -24,8 +24,7 @@ def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr) ) val = tl.load(remote_mem) tl.store(out_ptr + pid, val) - tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device", group_kind="block", order="acqrel", - barrier_kind="sync") + tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device") def _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid): diff --git a/python/triton/experimental/tle/language/__init__.py b/python/triton/experimental/tle/language/__init__.py index 932b931b4..d441e239e 100644 --- a/python/triton/experimental/tle/language/__init__.py +++ b/python/triton/experimental/tle/language/__init__.py @@ -23,6 +23,9 @@ ShardingSpec, device_mesh, MeshConfig, + BarrierKind, + MemoryOrder, + GroupKind, distributed_barrier, distributed_dot, _infer_submesh_barrier_group, @@ -92,6 +95,9 @@ def __getattr__(name): "get_mem_pool", "create_comm_tensor", "cleanup_communicator", + "BarrierKind", + "MemoryOrder", + "GroupKind", ] from . import distributed, gpu, raw diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 2af4ed53a..fe452ff4a 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -63,7 +63,7 @@ class MemoryOrder(str, Enum): RELAXED = "relaxed" ACQUIRE = "acquire" RELEASE = "release" - ACQ_REL = "acq_rel" + ACQ_REL = "acqrel" class GroupKind(str, Enum): @@ -656,7 +656,6 @@ def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, barrier_kind: BarrierKind | str = BarrierKind.SYNC, group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): - if space and space in ("device", "node"): builder = _semantic.builder ptr = _parse_src_arg(builder, comm_ptr) @@ -673,10 +672,11 @@ def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, @tl.builtin -def distributed_barrier(mesh: device_mesh | None = None, space: str = None, comm_ptr=None, +def distributed_barrier(mesh: device_mesh | None = None, comm_ptr=None, space: str = None, + group_kind: str | GroupKind = GroupKind.BLOCK, barrier_kind: BarrierKind | str = BarrierKind.SYNC, - group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, - order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): + order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None, + index: int | None = 0): """ M3 entrypoint: distributed synchronization primitive. diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 68749cd7f..91370171f 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -716,7 +716,6 @@ def _pack_args(self, backend, kwargs, bound_args, specialization, options): def run(self, *args, grid, warmup, **kwargs): kwargs["debug"] = kwargs.get("debug", self.debug) or knobs.runtime.debug kwargs["instrumentation_mode"] = knobs.compilation.instrumentation_mode - # DistributedRtContext().add_args_to_jitfunction(kwargs=kwargs, params=self.params) # parse options device = driver.active.get_current_device() diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py index 7745b4439..9c1f87185 100644 --- a/third_party/nvidia/backend/driver.py +++ b/third_party/nvidia/backend/driver.py @@ -9,6 +9,7 @@ from triton.runtime import _allocation from triton.backends.compiler import GPUTarget from triton.backends.driver import GPUDriver +from triton.runtime._distributed import DistributedRtContext dirname = os.path.dirname(os.path.realpath(__file__)) include_dirs = [os.path.join(dirname, "include")] @@ -215,9 +216,9 @@ def format_of(ty): flat_signature = [] for sig in signature.values(): _flatten_signature(sig, flat_signature) - #adxx - # flat_signature.insert(0, "*i64") # dev_comm_ptr - # flat_signature.insert(1, "*i64") # dev_mem_ptr + if DistributedRtContext().is_lite_mode: + flat_signature.insert(0, "*i64") # flagcx_dev_comm_ptr + flat_signature.insert(1, "*i64") # flagcx_dev_comm_ptr signature = {i: s for i, s in enumerate(flat_signature)} args_list = ', ' + ', '.join(f"&_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else '' # Record the end of regular arguments; diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp index 331c8ab1c..a3b5f2575 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp @@ -506,8 +506,8 @@ struct RemotePointersOpConversion auto mem = getDistDevicePtr(op, srcElems); if (space == "cluster") { - if (failed( - lowerClusterSpace(loc, mem, shardElems, rewriter, mappedPtrs))) { + if (failed(lowerClusterSpace(loc, srcElems, shardElems, rewriter, + mappedPtrs))) { return rewriter.notifyMatchFailure(op, "cluster lowering failed"); } } else if (space == "device") { From 5a039c83e402b498a56663a67d1865b482c7a0ee Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 03:24:18 +0000 Subject: [PATCH 20/32] update --- python/setup_tools/setup_helper.py | 2 +- python/test/tle/unit/test_tle_d2d_barrier.py | 25 ++++++------ .../test/tle/unit/test_tle_distributed_d2d.py | 2 +- python/test/tle/unit/test_tle_get_local_pe.py | 2 +- .../experimental/tle/language/__init__.py | 4 +- .../tle/language/communication.py | 6 +-- .../experimental/tle/language/distributed.py | 38 ++++++++++--------- python/triton/runtime/_distributed.py | 8 ++++ .../TleToLLVM/DistributedBarrierOpToLLVM.cpp | 23 +++++++---- .../TleToLLVM/GetDeviceIdToFlagCX.cpp | 22 +++++++---- .../TleToLLVM/LocalPointersOpToLLVM.cpp | 2 +- .../tle/dialect/lib/IR/VerfiyUtils.cpp | 4 -- third_party/tle/triton_tle.cc | 6 +-- 13 files changed, 81 insertions(+), 63 deletions(-) diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 571853574..0cc62714b 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -499,7 +499,7 @@ def uninstall_triton(): download_flagtree_third_party("flagcx", condition=(not flagtree_backend), hook="handle_flagcx", required=True) -download_flagtree_third_party("tileir", condition=(not flagtree_backend), required=True) +download_flagtree_third_party("tileir", condition=(flagtree_backend == "tileir"), required=True) handle_flagtree_backend() diff --git a/python/test/tle/unit/test_tle_d2d_barrier.py b/python/test/tle/unit/test_tle_d2d_barrier.py index 837e17b21..cb2a78d72 100644 --- a/python/test/tle/unit/test_tle_d2d_barrier.py +++ b/python/test/tle/unit/test_tle_d2d_barrier.py @@ -8,15 +8,15 @@ N = 8 -@triton.jit -def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr): +@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', comm_ptr=dev_comm_dptr) + 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( - dev_mem_ptr, + device_dptr, space="device", dtype=tl.float32, shard_id=peer, @@ -24,13 +24,12 @@ def _barrier_d2d_kernel(dev_comm_dptr, dev_mem_ptr, out_ptr, mesh: tl.constexpr) ) val = tl.load(remote_mem) tl.store(out_ptr + pid, val) - tle.distributed_barrier(comm_ptr=dev_comm_dptr, space="device") + tle.distributed_barrier(device_dptr=device_dptr, space="device") -def _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid): +def _ir_verify(output, device_dptr, grid): compiled = _barrier_d2d_kernel.warmup( - dev_comm_dptr=dev_comm_dptr, - dev_mem_ptr=dev_mem_dptr, + device_dptr=device_dptr, out_ptr=output, mesh=DEVICE_MESH, grid=(grid, ), @@ -44,9 +43,9 @@ def _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid): assert "flagcxDevCommGetIntraRank" in compiled.asm['ptx'] -def _runtime_verify(dev_comm_dptr, dev_mem_dptr, output, grid, rank, world_size): +def _runtime_verify(output, device_dptr, grid, rank, world_size): dist.barrier() - _barrier_d2d_kernel[grid](dev_comm_dptr=dev_comm_dptr, dev_mem_ptr=dev_mem_dptr, out_ptr=output, mesh=DEVICE_MESH) + _barrier_d2d_kernel[grid](device_dptr=device_dptr, out_ptr=output, mesh=DEVICE_MESH) torch.cuda.synchronize() @@ -76,12 +75,12 @@ def test_tle_d2d_barrier(self): with torch.cuda.use_mem_pool(mem_pool): x = (torch.arange(N, dtype=torch.float32, device="cuda") + rank * 1000).clone() - dev_comm_dptr, dev_mem_dptr = tle.create_comm_tensor(x) + device_dptr = tle.create_dist_tensor(x) output = torch.zeros(N, dtype=torch.float32, device="cuda") - _ir_verify(dev_comm_dptr, dev_mem_dptr, output, grid) + _ir_verify(output, device_dptr, grid) - _runtime_verify(dev_comm_dptr, dev_mem_dptr, output, grid, rank, world_size) + _runtime_verify(output, device_dptr, grid, rank, world_size) if __name__ == "__main__": diff --git a/python/test/tle/unit/test_tle_distributed_d2d.py b/python/test/tle/unit/test_tle_distributed_d2d.py index e738124fd..545d2ecfa 100644 --- a/python/test/tle/unit/test_tle_distributed_d2d.py +++ b/python/test/tle/unit/test_tle_distributed_d2d.py @@ -57,7 +57,7 @@ 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) + dev_comm_ptr, dev_mem_ptr = tle.create_dist_tensor(buf_tensor) print(f"[Rank {rank}] dev_comm_ptr={dev_comm_ptr:#x}, dev_mem_ptr={dev_mem_ptr:#x}") diff --git a/python/test/tle/unit/test_tle_get_local_pe.py b/python/test/tle/unit/test_tle_get_local_pe.py index fbccde5ac..1e20622a2 100644 --- a/python/test/tle/unit/test_tle_get_local_pe.py +++ b/python/test/tle/unit/test_tle_get_local_pe.py @@ -23,7 +23,7 @@ def test_tle_local_pe_kernel(self): with torch.cuda.use_mem_pool(tle.get_mem_pool()): x = torch.randn((N, N), dtype=torch.float32, device="cuda") y = torch.empty_like(x) - dev_comm_dptr, dev_mem_dptr = tle.create_comm_tensor(x) + dev_comm_dptr, dev_mem_dptr = tle.create_dist_tensor(x) compiled = _tle_local_pe_kernel.warmup( dev_comm_dptr=dev_comm_dptr, diff --git a/python/triton/experimental/tle/language/__init__.py b/python/triton/experimental/tle/language/__init__.py index d441e239e..6362381be 100644 --- a/python/triton/experimental/tle/language/__init__.py +++ b/python/triton/experimental/tle/language/__init__.py @@ -39,7 +39,7 @@ sharding, ) from . import communication -from .communication import get_mem_pool, create_comm_tensor, cleanup_communicator +from .communication import get_mem_pool, create_dist_tensor, cleanup_communicator _EXTENSION_APIS = frozenset({ "make_tensor_view", @@ -93,7 +93,7 @@ def __getattr__(name): "raw", "mem_pool", "get_mem_pool", - "create_comm_tensor", + "create_dist_tensor", "cleanup_communicator", "BarrierKind", "MemoryOrder", diff --git a/python/triton/experimental/tle/language/communication.py b/python/triton/experimental/tle/language/communication.py index d57a34343..382fa1316 100644 --- a/python/triton/experimental/tle/language/communication.py +++ b/python/triton/experimental/tle/language/communication.py @@ -168,7 +168,7 @@ def init_communicator(): _init_communicator_ = True -def create_comm_tensor(buf_tensor): +def create_dist_tensor(buf_tensor): global comm, rank, dev_mem, dev_comm, win buf_ptr = buf_tensor.data_ptr() buf_size = buf_tensor.numel() * buf_tensor.element_size() @@ -207,5 +207,5 @@ def create_comm_tensor(buf_tensor): dist.barrier() from triton.runtime import DistributedRtContext - DistributedRtContext(dev_comm_dptr.value, dev_mem_dptr.value) - return dev_comm_dptr.value, dev_mem_dptr.value + ctx = DistributedRtContext(dev_mem_dptr.value, dev_comm_dptr.value) + return ctx diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index fe452ff4a..12144f4b9 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -7,6 +7,7 @@ from typing import Any, Iterable, Mapping, Sequence, List, Tuple, Union, Optional, Dict from enum import Enum import triton.language.core as tl +from triton.runtime import DistributedRtContext Axis = Tuple[str, int] AxesLike = Union[int, List[Axis]] @@ -27,19 +28,22 @@ def _as_positive_int(value: Any, label: str) -> int: return value -def _parse_src_arg(builder, src): - if src is None: - return src - else: +def _parse_src_arg(builder, src, index=0): + src = tl._unwrap_if_constexpr(src) + if isinstance(src, DistributedRtContext): + return builder.get_int64(src[index]) + elif src: return src.handle + else: + return None # Get the current device id @tl.builtin -def _get_local_rank(dev_comm_ptr, _semantic=None, ret_dtype=tl.int32): +def _get_local_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, dev_comm_ptr) + ptr = _parse_src_arg(builder, device_dptr, 1) result = builder.get_device_id(ret_ir_ty, ptr) return tl.tensor(result, ret_dtype) @@ -605,7 +609,7 @@ def _resolve_launch_axis(mesh: device_mesh, axis: str | int) -> int: def shard_id( mesh: device_mesh, axis: str | int, - comm_ptr=None, + device_dptr=None, _semantic=None, ): """ @@ -618,7 +622,7 @@ def shard_id( axis = tl._unwrap_if_constexpr(axis) if axis in ("device", "node"): - return _get_local_rank(comm_ptr, _semantic=_semantic, ret_dtype=tl.int32) + return _get_local_rank(device_dptr, _semantic=_semantic, ret_dtype=tl.int32) if not isinstance(mesh, device_mesh): raise TypeError(f"mesh must be device_mesh, got {type(mesh).__name__}") @@ -652,13 +656,13 @@ def _parse_device_barrier_args(argType) -> str: return str(argType).lower() -def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, +def check_and_handle_device_intra_barrier(space: str = None, device_dptr=None, barrier_kind: BarrierKind | str = BarrierKind.SYNC, group_kind: str | GroupKind = GroupKind.BLOCK, index: int | None = 0, order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None): if space and space in ("device", "node"): builder = _semantic.builder - ptr = _parse_src_arg(builder, comm_ptr) + ptr = _parse_src_arg(builder, device_dptr, 1) builder.create_distributed_barrier( src=ptr, barrier_index=index or 0, @@ -672,7 +676,7 @@ def check_and_handle_device_intra_barrier(space: str = None, comm_ptr=None, @tl.builtin -def distributed_barrier(mesh: device_mesh | None = None, comm_ptr=None, space: str = None, +def distributed_barrier(mesh: device_mesh | None = None, device_dptr=None, space: str = None, group_kind: str | GroupKind = GroupKind.BLOCK, barrier_kind: BarrierKind | str = BarrierKind.SYNC, order: MemoryOrder | str | int | None = MemoryOrder.ACQ_REL, _semantic=None, @@ -689,7 +693,7 @@ def distributed_barrier(mesh: device_mesh | None = None, comm_ptr=None, space: s subgroup = None use_grid = mesh is not None and _mesh_uses_grid_barrier(mesh) - if check_and_handle_device_intra_barrier(space=space, comm_ptr=comm_ptr, barrier_kind=barrier_kind, + if check_and_handle_device_intra_barrier(space=space, device_dptr=device_dptr, barrier_kind=barrier_kind, group_kind=group_kind, index=index, order=order, _semantic=_semantic): return None @@ -819,7 +823,7 @@ def _create_remote_pointers_tensor( "cluster": (dtype, 7), "device": (dtype, 1), }.get(space)) - if tensor and tensor.type.is_block(): + if space == 'cluster' and tensor and tensor.type.is_block(): remote_type = tl.block_type(remote_ptr_dtype, list(tensor.shape)).to_ir(builder) else: remote_type = remote_ptr_dtype.to_ir(builder) @@ -844,12 +848,12 @@ def _create_remote_pointers_tensor( # automatic injection (e.g. inside this helper). if offset_tensor.dtype != tl.int64: offset_tensor = tl.cast(offset_tensor, tl.int64, _semantic=_semantic) - ptr = _parse_src_arg(builder, tensor) + ptr = _parse_src_arg(builder, tensor, 0) remote_op = builder.create_remote_pointers(remote_type, ptr, shard_id_tensor.handle, space, offset_tensor.handle) else: remote_op = builder.create_remote_pointers(remote_type, tensor.handle, shard_id_tensor.handle, space) - if tensor and tensor.type.is_block(): + if space == "cluster" and tensor and tensor.type.is_block(): return tl.tensor(remote_op.get_result(0), tl.block_type(remote_ptr_dtype, list(tensor.shape))) return tl.tensor(remote_op.get_result(0), remote_ptr_dtype) @@ -924,7 +928,7 @@ def _remote_pointer( @tl.builtin def remote( - tensor=None, + tensor=tl.tensor | None, shard_id=None, scope: device_mesh | None = None, space: str = "cluster", @@ -960,7 +964,7 @@ def remote( # Direct pointer path: support local_ptr scalar/tensor values and return # remote pointer with preserved shape. - if isinstance(tensor, tl.tensor) or (tensor is None and space in ("device", "node")): + if isinstance(tensor, tl.tensor) or (space in ("device", "node")): return _remote_pointer(tensor, shard_id, scope=scope, space=space, _semantic=_semantic, dtype=dtype, offset=offset) diff --git a/python/triton/runtime/_distributed.py b/python/triton/runtime/_distributed.py index 67ccd87a3..a1665927e 100644 --- a/python/triton/runtime/_distributed.py +++ b/python/triton/runtime/_distributed.py @@ -20,6 +20,11 @@ def __init__(self, _comm_ptr=None, _mem_ptr=None): self._mem_ptr = _mem_ptr self._initialized = True + def get_packed_data(self): + return int(self._mem_ptr), int(self._comm_ptr) + import torch + return torch.tensor([int(self._mem_ptr), int(self._comm_ptr)], device='cuda', dtype=torch.int64) + @property def comm_ptr(self) -> int | None: """Communication runtime pointer.""" @@ -44,6 +49,9 @@ def is_lite_mode(self) -> bool: inner_action = self._init_count == 1 return inner_action and user_action + def __getitem__(self, index=0): + return list(self._get_needed_params().values())[index] + def add_args_to_jitfunction(self, **kwargs): params = kwargs['params'] # list _kwargs = kwargs['kwargs'] diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp index 9ee9670be..5bd5ec7c9 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/DistributedBarrierOpToLLVM.cpp @@ -42,6 +42,16 @@ constexpr int32_t kGridScratchAlignment = 4; constexpr int32_t kGridScratchBytes = 4; constexpr int32_t kGridArrivedOffsetBytes = 0; +Value getDistDevicePtr(tle::DistributedBarrierOp op, + SmallVector &srcElems) { + if (!srcElems.empty()) + return srcElems[0]; + else { + auto func = op->getParentOfType(); + return func.getArgument(1); + } +} + FailureOr getOrCreateSubmeshScratchOffset(ModuleOp mod) { if (auto existing = mod->getAttrOfType(kSubmeshScratchOffsetAttr)) { @@ -433,6 +443,7 @@ struct DistributedBarrierOpConversion auto orderAttr = op->getAttrOfType(kOrderAttr); auto indexAttr = op->getAttrOfType(kIndexAttr); auto loc = op.getLoc(); + SmallVector srcElems; auto getCoopKindValue = [](StringRef kind) -> int32_t { return llvm::StringSwitch(kind) .Case("thread", 0) @@ -458,14 +469,10 @@ struct DistributedBarrierOpConversion if (order < 0) return rewriter.notifyMatchFailure(op, "invalid order"); - auto comm = Value(); - if (auto src = adaptor.getSrc()) { - auto srcElems = unpackLLElements(loc, src, rewriter); - comm = srcElems[0]; - } else { - auto func = op->getParentOfType(); - comm = func.getArgument(0); - } + if (auto src = adaptor.getSrc()) + srcElems = unpackLLElements(loc, src, rewriter); + + auto comm = getDistDevicePtr(op, srcElems); auto coopKindAttr = rewriter.getI32IntegerAttr(coopKind); auto newOrderAttr = rewriter.getI32IntegerAttr(order); auto barrierTypeAttr = op.getBarrierTypeAttr(); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp index 8eaee1c6d..dff002aea 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/GetDeviceIdToFlagCX.cpp @@ -20,6 +20,15 @@ using namespace mlir; namespace ttg = mlir::triton::gpu; namespace tle = mlir::triton::tle; +Value getDistDevicePtr(tle::GetDeviceIdOp op, SmallVector &srcElems) { + if (!srcElems.empty()) + return srcElems[0]; + else { + auto func = op->getParentOfType(); + return func.getArgument(1); + } +} + struct GetDeviceIdOpConversion : public ConvertOpToLLVMPattern { GetDeviceIdOpConversion(LLVMTypeConverter &typeConverter, @@ -29,19 +38,16 @@ struct GetDeviceIdOpConversion LogicalResult matchAndRewrite(tle::GetDeviceIdOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - Value src = adaptor.getInput(); - auto comm = Value(); + auto loc = op.getLoc(); + SmallVector srcElems; + if (auto src = adaptor.getInput()) + srcElems = unpackLLElements(loc, src, rewriter); auto func = op->getParentOfType(); if (!func) { return rewriter.notifyMatchFailure( op, "expected parent LLVM::LLVMFuncOp, but none was found. "); } - - if (auto src = adaptor.getInput()) - comm = src; - else - comm = func.getArgument(0); - + auto comm = getDistDevicePtr(op, srcElems); rewriter.modifyOpInPlace(op, [&]() { op->insertOperands(0, comm); }); auto localRank = rewriter.create( op.getLoc(), rewriter.getI32Type(), comm); diff --git a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp index a3b5f2575..28ef347da 100644 --- a/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp +++ b/third_party/tle/dialect/lib/Conversion/TleToLLVM/LocalPointersOpToLLVM.cpp @@ -480,7 +480,7 @@ struct RemotePointersOpConversion SmallVector srcElems; auto space = adaptor.getSpace(); - if (adaptor.getSrc()) + if (auto src = adaptor.getSrc()) srcElems = unpackLLElements(loc, adaptor.getSrc(), rewriter); if (space != "device" && srcElems.empty()) diff --git a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp index 74dd36853..5452c7b19 100644 --- a/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp +++ b/third_party/tle/dialect/lib/IR/VerfiyUtils.cpp @@ -17,12 +17,8 @@ namespace mlir::triton::tle { namespace RemotePointers { llvm::LogicalResult verifyDeviceSpace(mlir::Value src, mlir::Value result) { - // flagcxGetIntraPointerC accept raw device pointers represented as signless - // i64 values. if (!src) return success(); - if (!src.getType().isSignlessInteger(64)) - return failure(); if (auto tensorTy = dyn_cast(result.getType())) { auto ptr = dyn_cast(tensorTy.getElementType()); diff --git a/third_party/tle/triton_tle.cc b/third_party/tle/triton_tle.cc index 1f02385a7..8d595947a 100644 --- a/third_party/tle/triton_tle.cc +++ b/third_party/tle/triton_tle.cc @@ -621,10 +621,8 @@ void init_triton_tle_ir(py::module &&m) { [](TritonOpBuilder &self, Type resultTy, std::optional src) -> Value { auto &builder = self.getBuilder(); - auto ptr = Value(); - if (src.has_value()) - ptr = src.value(); - return self.create(resultTy, ptr); + return self.create(resultTy, + src.value_or(Value())); }) .def("get_n_pes", [](TritonOpBuilder &self, Type resultTy, Value src) -> Value { From b15993e112867a1b7381bbb0f8a30e2777e4c78b Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 03:25:37 +0000 Subject: [PATCH 21/32] update --- .github/workflows/nv3.6-build-and-test.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/nv3.6-build-and-test.yml b/.github/workflows/nv3.6-build-and-test.yml index 3d5e7b9b9..30460802e 100644 --- a/.github/workflows/nv3.6-build-and-test.yml +++ b/.github/workflows/nv3.6-build-and-test.yml @@ -99,7 +99,6 @@ jobs: 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 python3 python/tutorials/hints/03/03-matrix-multiplication.py --only_unit_test From 2801a8f06f0d6b5d28cabf50b7f64fc29b926421 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 06:19:54 +0000 Subject: [PATCH 22/32] fix comment --- python/test/tle/unit/test_tle_distributed_d2d.py | 10 ++++------ python/test/tle/unit/test_tle_get_local_pe.py | 12 +++++------- python/triton/runtime/__init__.py | 1 + third_party/nvidia/backend/driver.py | 3 +-- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/python/test/tle/unit/test_tle_distributed_d2d.py b/python/test/tle/unit/test_tle_distributed_d2d.py index 545d2ecfa..dbe26f209 100644 --- a/python/test/tle/unit/test_tle_distributed_d2d.py +++ b/python/test/tle/unit/test_tle_distributed_d2d.py @@ -12,8 +12,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, @@ -22,7 +22,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, @@ -57,9 +57,7 @@ def main(): f"stride={buf_tensor.stride()}, " f"sample={buf_tensor[:4].tolist()}") - dev_comm_ptr, dev_mem_ptr = tle.create_dist_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") @@ -67,8 +65,8 @@ def main(): grid = (N, ) lsa_read_kernel[grid]( - dev_mem_ptr, output, + device_dptr=device_dptr, N=N, MY_RANK=rank, N_RANKS=world_size, diff --git a/python/test/tle/unit/test_tle_get_local_pe.py b/python/test/tle/unit/test_tle_get_local_pe.py index 1e20622a2..cb82e567f 100644 --- a/python/test/tle/unit/test_tle_get_local_pe.py +++ b/python/test/tle/unit/test_tle_get_local_pe.py @@ -7,9 +7,9 @@ @triton.jit -def _tle_local_pe_kernel(dev_comm_dptr, dev_mem_dptr, out_ptr, mesh: tl.constexpr, BLOCK: tl.constexpr): +def _tle_local_pe_kernel(out_ptr, device_dptr: tl.constexpr, mesh: tl.constexpr, BLOCK: tl.constexpr): pid = tl.program_id(0) # noqa: F841 - local_rank = tle.shard_id(mesh, 'device', comm_ptr=dev_comm_dptr) + local_rank = tle.shard_id(mesh, 'device', device_dptr=device_dptr) n_rank = mesh.shape[0] peer = (local_rank + 1) % n_rank # noqa: F841 @@ -23,12 +23,11 @@ def test_tle_local_pe_kernel(self): with torch.cuda.use_mem_pool(tle.get_mem_pool()): x = torch.randn((N, N), dtype=torch.float32, device="cuda") y = torch.empty_like(x) - dev_comm_dptr, dev_mem_dptr = tle.create_dist_tensor(x) + device_dptr = tle.create_dist_tensor(x) compiled = _tle_local_pe_kernel.warmup( - dev_comm_dptr=dev_comm_dptr, - dev_mem_dptr=dev_mem_dptr, out_ptr=y, + device_dptr=device_dptr, mesh=DEVICE_MESH, BLOCK=block, grid=(grid, ), @@ -37,8 +36,7 @@ def test_tle_local_pe_kernel(self): ) assert "get_device_id" in compiled.asm["ttgir"] - _tle_local_pe_kernel[(grid, )](dev_comm_dptr=dev_comm_dptr, dev_mem_dptr=dev_mem_dptr, out_ptr=y, - mesh=DEVICE_MESH, BLOCK=block) + _tle_local_pe_kernel[(grid, )](out_ptr=y, device_dptr=device_dptr, mesh=DEVICE_MESH, BLOCK=block) tle.cleanup_communicator() diff --git a/python/triton/runtime/__init__.py b/python/triton/runtime/__init__.py index 85b7b00bc..f9e12f61f 100644 --- a/python/triton/runtime/__init__.py +++ b/python/triton/runtime/__init__.py @@ -21,5 +21,6 @@ "reinterpret", "RemoteCacheBackend", "TensorWrapper", + # flagtree tle distributed "DistributedRtContext", ] diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py index 9c1f87185..8c6456f3e 100644 --- a/third_party/nvidia/backend/driver.py +++ b/third_party/nvidia/backend/driver.py @@ -209,13 +209,12 @@ def format_of(ty): expand_signature = _expand_signature(signature.values()) signature = {i: s for i, s in enumerate(expand_signature)} - args_format = ''.join([format_of(ty) for ty in signature.values()]) format = _BASE_ARGS_FORMAT + args_format - flat_signature = [] for sig in signature.values(): _flatten_signature(sig, flat_signature) + # flagtree tle distributed if DistributedRtContext().is_lite_mode: flat_signature.insert(0, "*i64") # flagcx_dev_comm_ptr flat_signature.insert(1, "*i64") # flagcx_dev_comm_ptr From 462062c8ec2a5580b50cbca7744a3634e322d2bf Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 07:08:35 +0000 Subject: [PATCH 23/32] polist code --- python/triton/runtime/_distributed.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/triton/runtime/_distributed.py b/python/triton/runtime/_distributed.py index a1665927e..70fb3fa75 100644 --- a/python/triton/runtime/_distributed.py +++ b/python/triton/runtime/_distributed.py @@ -22,8 +22,6 @@ def __init__(self, _comm_ptr=None, _mem_ptr=None): def get_packed_data(self): return int(self._mem_ptr), int(self._comm_ptr) - import torch - return torch.tensor([int(self._mem_ptr), int(self._comm_ptr)], device='cuda', dtype=torch.int64) @property def comm_ptr(self) -> int | None: From 8980a7b8933e7439264fb4f192a3b9a1fc4cd970 Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 07:40:20 +0000 Subject: [PATCH 24/32] fix backend bugs --- .../experimental/tle/language/distributed.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/python/triton/experimental/tle/language/distributed.py b/python/triton/experimental/tle/language/distributed.py index 12144f4b9..61e278455 100644 --- a/python/triton/experimental/tle/language/distributed.py +++ b/python/triton/experimental/tle/language/distributed.py @@ -7,7 +7,6 @@ from typing import Any, Iterable, Mapping, Sequence, List, Tuple, Union, Optional, Dict from enum import Enum import triton.language.core as tl -from triton.runtime import DistributedRtContext Axis = Tuple[str, int] AxesLike = Union[int, List[Axis]] @@ -29,13 +28,17 @@ def _as_positive_int(value: Any, label: str) -> int: def _parse_src_arg(builder, src, index=0): - src = tl._unwrap_if_constexpr(src) - if isinstance(src, DistributedRtContext): - return builder.get_int64(src[index]) - elif src: + try: + from triton.runtime import DistributedRtContext + src = tl._unwrap_if_constexpr(src) + if isinstance(src, DistributedRtContext): + return builder.get_int64(src[index]) + elif src: + return src.handle + else: + return None + except Exception: return src.handle - else: - return None # Get the current device id From d1d899c3db204a059d35a202ae00496e61d7006e Mon Sep 17 00:00:00 2001 From: yrl <2535184404@qq.com> Date: Fri, 17 Jul 2026 07:45:31 +0000 Subject: [PATCH 25/32] fix comment --- third_party/nvidia/backend/driver.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/third_party/nvidia/backend/driver.py b/third_party/nvidia/backend/driver.py index 8c6456f3e..0c5e98f56 100644 --- a/third_party/nvidia/backend/driver.py +++ b/third_party/nvidia/backend/driver.py @@ -209,8 +209,10 @@ def format_of(ty): expand_signature = _expand_signature(signature.values()) signature = {i: s for i, s in enumerate(expand_signature)} + args_format = ''.join([format_of(ty) for ty in signature.values()]) format = _BASE_ARGS_FORMAT + args_format + flat_signature = [] for sig in signature.values(): _flatten_signature(sig, flat_signature) @@ -260,9 +262,7 @@ def format_of(ty): for i, ty in signature.items() if ty in FLOAT_STORAGE_TYPE ] - params = [f"&arg{i}" for i, ty in signature.items() if ty != "constexpr"] - params.append("&global_scratch") params.append("&profile_scratch") tle_cpp_define = "#define __TLE__ 1" From 6e116270d03add46d99663a0b79c31019e6c3670 Mon Sep 17 00:00:00 2001 From: Galaxy1458 <55453380+Galaxy1458@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:06:08 +0800 Subject: [PATCH 26/32] Remove commented line in jit.py Removed commented line related to flagtree in distributed context. --- python/triton/runtime/jit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 91370171f..79667196c 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -761,7 +761,7 @@ def run(self, *args, grid, warmup, **kwargs): kernel = kernel.result() # launch kernel launch_metadata = kernel.launch_metadata(grid, stream, *bound_args.values()) - + # flagtree tle distributed dist_param = [] ctx = DistributedRtContext() if ctx.is_lite_mode: From dbbaa48b9a2d80e37ee01b07d5b8c72bd09dee15 Mon Sep 17 00:00:00 2001 From: Galaxy1458 <55453380+Galaxy1458@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:08:38 +0800 Subject: [PATCH 27/32] Update jit.py --- python/triton/runtime/jit.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 79667196c..77de76655 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -737,6 +737,7 @@ def run(self, *args, grid, warmup, **kwargs): if kernel is None: options, signature, constexprs, attrs = self._pack_args(backend, kwargs, bound_args, specialization, options) + kernel = self._do_compile(key, signature, device, constexprs, options, attrs, warmup) if kernel is None: return None @@ -761,7 +762,7 @@ def run(self, *args, grid, warmup, **kwargs): kernel = kernel.result() # launch kernel launch_metadata = kernel.launch_metadata(grid, stream, *bound_args.values()) - # flagtree tle distributed + # flagtree tle distributed: Add dist_param to kernel.run dist_param = [] ctx = DistributedRtContext() if ctx.is_lite_mode: From 61c8e81a5483f465d47660c98ed89c1be8e8a0c3 Mon Sep 17 00:00:00 2001 From: Galaxy1458 <55453380+Galaxy1458@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:10:07 +0800 Subject: [PATCH 28/32] Update compiler.py --- third_party/nvidia/backend/compiler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/third_party/nvidia/backend/compiler.py b/third_party/nvidia/backend/compiler.py index 20dbadce8..64ae02436 100644 --- a/third_party/nvidia/backend/compiler.py +++ b/third_party/nvidia/backend/compiler.py @@ -280,6 +280,7 @@ def make_ttgir(mod, metadata, opt, capability): pm = ir.pass_manager(mod.context) dump_enabled = pm.enable_debug() emuTF32 = (capability // 10 >= 8) + # flagtree tle distributed if DistributedRtContext().is_lite_mode: tle.passes.add_params_for_distribution(pm) passes.ttir.add_convert_to_ttgpuir(pm, f"cuda:{capability}", opt.num_warps, 32, opt.num_ctas) From 6edac5188fe7dcf0998543ae63957cf05cfac9e6 Mon Sep 17 00:00:00 2001 From: flagtree-bot Date: Fri, 17 Jul 2026 10:10:39 +0000 Subject: [PATCH 29/32] Apply code-format changes --- python/triton/runtime/jit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 77de76655..7c35791d6 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -737,7 +737,7 @@ def run(self, *args, grid, warmup, **kwargs): if kernel is None: options, signature, constexprs, attrs = self._pack_args(backend, kwargs, bound_args, specialization, options) - + kernel = self._do_compile(key, signature, device, constexprs, options, attrs, warmup) if kernel is None: return None From 2a884aeb5606d68903d058a8b88dcf9449d51931 Mon Sep 17 00:00:00 2001 From: Galaxy1458 <55453380+Galaxy1458@users.noreply.github.com> Date: Fri, 17 Jul 2026 18:11:18 +0800 Subject: [PATCH 30/32] Initialize params list in JIT class --- python/triton/runtime/jit.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 7c35791d6..48e5b49d7 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -788,6 +788,7 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o self.do_not_specialize_on_alignment = do_not_specialize_on_alignment self._repr = repr self.launch_metadata = launch_metadata + self.params = [] for i, param in enumerate(self.signature.parameters.values()): dns = i in do_not_specialize or param.name in do_not_specialize From 342f00088ed06e0e4626f449c20d05720ee11328 Mon Sep 17 00:00:00 2001 From: flagtree-bot Date: Fri, 17 Jul 2026 10:14:15 +0000 Subject: [PATCH 31/32] Apply code-format changes --- python/triton/runtime/jit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/triton/runtime/jit.py b/python/triton/runtime/jit.py index 48e5b49d7..60702174b 100644 --- a/python/triton/runtime/jit.py +++ b/python/triton/runtime/jit.py @@ -788,7 +788,7 @@ def __init__(self, fn, version=None, do_not_specialize=None, do_not_specialize_o self.do_not_specialize_on_alignment = do_not_specialize_on_alignment self._repr = repr self.launch_metadata = launch_metadata - + self.params = [] for i, param in enumerate(self.signature.parameters.values()): dns = i in do_not_specialize or param.name in do_not_specialize From 3ec900f71fc50b179c7b5f283eb2e15260326a35 Mon Sep 17 00:00:00 2001 From: zhzhcookie Date: Sat, 18 Jul 2026 00:19:29 +0800 Subject: [PATCH 32/32] [Build] Fix default_backends without tileir --- packaging/rpm/helpers/Dockerfile.rpm | 4 +--- python/setup_tools/utils/tools.py | 10 +--------- setup.py | 3 +-- third_party/tileir/CMakeLists.txt | 3 +-- third_party/tileir/README.md | 8 +------- 5 files changed, 5 insertions(+), 23 deletions(-) diff --git a/packaging/rpm/helpers/Dockerfile.rpm b/packaging/rpm/helpers/Dockerfile.rpm index dbaea8087..df778a2d1 100644 --- a/packaging/rpm/helpers/Dockerfile.rpm +++ b/packaging/rpm/helpers/Dockerfile.rpm @@ -41,11 +41,9 @@ 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. 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)" diff --git a/python/setup_tools/utils/tools.py b/python/setup_tools/utils/tools.py index 5e34e3436..138e8b5cd 100644 --- a/python/setup_tools/utils/tools.py +++ b/python/setup_tools/utils/tools.py @@ -22,12 +22,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") @@ -55,9 +50,6 @@ def __post_init__(self): 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 diff --git a/setup.py b/setup.py index 8a1d84585..03da7b912 100644 --- a/setup.py +++ b/setup.py @@ -675,9 +675,8 @@ def download_and_copy_dependencies(): if helper.flagtree_backend: if helper.flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir"): - default_backends = helper.configs.non_tileir_default_backends() backends = [ - *BackendInstaller.copy(default_backends + tuple(helper.configs.extend_backends)), + *BackendInstaller.copy(helper.configs.default_backends + tuple(helper.configs.extend_backends)), *BackendInstaller.copy_externals(), ] else: diff --git a/third_party/tileir/CMakeLists.txt b/third_party/tileir/CMakeLists.txt index 96ce81423..de8277ecf 100644 --- a/third_party/tileir/CMakeLists.txt +++ b/third_party/tileir/CMakeLists.txt @@ -40,8 +40,7 @@ else() set(CUDA_TILE_SUBMODULE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/cuda-tile) if(NOT EXISTS ${CUDA_TILE_SUBMODULE_DIR}/CMakeLists.txt) message(FATAL_ERROR - "cuda-tile submodule not initialized at ${CUDA_TILE_SUBMODULE_DIR}. " - "Run: git submodule update --init --recursive") + "cuda-tile submodule not initialized at ${CUDA_TILE_SUBMODULE_DIR}. ") endif() execute_process( diff --git a/third_party/tileir/README.md b/third_party/tileir/README.md index e6ed4d907..72fdf89a9 100644 --- a/third_party/tileir/README.md +++ b/third_party/tileir/README.md @@ -35,13 +35,7 @@ FlagTree's normal NVIDIA assemblers, managed by `setup.py` at: - `third_party/nvidia/backend/bin/ptxas-blackwell` `pip install` and `pip install -e` initialize the pinned cuda-tile submodule -automatically through `setup.py`; no manual submodule command is needed. If you -invoke CMake directly without `setup.py`, initialize the submodule first: - -```bash -git submodule update --init --recursive \ - third_party/tileir/third_party/cuda-tile -``` +automatically through `setup.py`; no manual submodule command is needed. At runtime, TileIR looks for `tileiras` in this order: