[KMCompiler] [TLERaw] Add WS thread budget repro cases

include/triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h

@@ -102,6 +102,11 @@ void populatePrintOpToLLVMPattern(LLVMTypeConverter &typeConverter,
                                   const TargetInfoBase &targetInfo,
                                   PatternBenefit benefit);
 
+void populateExternCallOpToLLVMPattern(LLVMTypeConverter &typeConverter,
+                                       RewritePatternSet &patterns,
+                                       const TargetInfoBase &targetInfo,
+                                       PatternBenefit benefit);
+
 void populateInstrumentationToLLVMPatterns(LLVMTypeConverter &typeConverter,
                                            const TargetInfoBase &targetInfo,
                                            RewritePatternSet &patterns,

include/triton/Dialect/Triton/IR/TritonOps.td

@@ -845,6 +845,32 @@ def TT_MapElementwiseReturnOp: TT_Op<"map_elementwise.return",
     let assemblyFormat = "attr-dict ($result^ `:` type($result))?";
 }
 
+//
+// External Call op
+//
+def TT_ExternCallOp : TT_Op<"extern_call", [
+  DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+  ConditionallySpeculatable,
+]> {
+
+    let description = [{
+        call an external function $symbol implemented in $libpath/$libname with $args
+        return $libpath/$libname:$symbol($args...)
+    }];
+
+    let arguments = (ins Variadic<TT_Type>:$srcs, StrAttr:$libname, StrAttr:$libpath, StrAttr:$symbol, BoolAttr:$pure);
+
+    let results = (outs Variadic<TT_Type>:$result);
+
+    let assemblyFormat = "operands attr-dict `:` functional-type(operands, $result)";
+
+    let extraClassDeclaration = [{
+      // Interface method for ConditionallySpeculatable.
+      Speculation::Speculatability getSpeculatability();
+    }];
+
+}
+
 //
 // External Elementwise op
 //

lib/Conversion/TritonGPUToLLVM/CMakeLists.txt

@@ -7,6 +7,7 @@ add_triton_library(TritonGPUToLLVM
     AssertOpToLLVM.cpp
     ControlFlowOpToLLVM.cpp
     ConvertLayoutOpToLLVM.cpp
+    ExternCallOpToLLVM.cpp
     ElementwiseOpToLLVM.cpp
     FuncOpToLLVM.cpp
     GatherOpToLLVM.cpp

lib/Conversion/TritonGPUToLLVM/ExternCallOpToLLVM.cpp

@@ -0,0 +1,61 @@
+#include "mlir/Conversion/LLVMCommon/Pattern.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Support/LLVM.h"
+#include "triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h"
+#include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
+#include "triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h"
+#include "triton/Conversion/TritonGPUToLLVM/Utility.h"
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+
+namespace {
+
+class ExternCallOpConversion
+    : public ConvertOpToLLVMPattern<triton::ExternCallOp> {
+public:
+  ExternCallOpConversion(const LLVMTypeConverter &converter,
+                         const PatternBenefit &benefit)
+      : ConvertOpToLLVMPattern<triton::ExternCallOp>(converter, benefit) {}
+
+  LogicalResult
+  matchAndRewrite(triton::ExternCallOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = op->getLoc();
+
+    if (op->getNumResults() > 1) {
+      llvm::errs() << "ExternCallConversion does not support multi outs.";
+      return failure();
+    }
+
+    LLVM::LLVMVoidType voidTy = void_ty(op->getContext());
+    auto newOperands = adaptor.getOperands();
+    Type retType =
+        op->getNumResults() == 0
+            ? voidTy
+            : this->getTypeConverter()->convertType(op->getResult(0).getType());
+    std::string funcName = op.getSymbol().str();
+    StringRef libname = op.getLibname();
+    StringRef libpath = op.getLibpath();
+
+    Operation *externCallOp;
+    Type funcType = mlir::triton::gpu::getFunctionType(retType, newOperands);
+    LLVM::LLVMFuncOp funcOp = mlir::triton::gpu::appendOrGetExternFuncOp(
+        rewriter, op, funcName, funcType, libname, libpath);
+    externCallOp = LLVM::createLLVMCallOp(rewriter, loc, funcOp, newOperands);
+
+    if (op->getNumResults() == 0) {
+      rewriter.eraseOp(op);
+    } else {
+      rewriter.replaceOp(op, externCallOp->getResult(0));
+    }
+
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::triton::populateExternCallOpToLLVMPattern(
+    LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+    const TargetInfoBase &targetInfo, PatternBenefit benefit) {
+  patterns.add<ExternCallOpConversion>(typeConverter, benefit);
+}

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -618,6 +618,7 @@ void populateTritonPatterns(TritonGPUTypeConverter &typeConverter,
       GenericOpPattern<triton::HistogramOp>,
       GenericOpPattern<triton::GatherOp>,
       GenericOpPattern<triton::ExternElementwiseOp>,
+      GenericOpPattern<triton::ExternCallOp>,
       GenericOpPattern<triton::PrintOp>,
       GenericOpPattern<triton::AssertOp>,
       GenericOpPattern<triton::AtomicCASOp>,

lib/Dialect/Triton/IR/Ops.cpp

@@ -1313,6 +1313,24 @@ Speculation::Speculatability ExternElementwiseOp::getSpeculatability() {
   return Speculation::NotSpeculatable;
 }
 
+// -- ExternCallOp --
+void ExternCallOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  if (getPure())
+    return;
+  effects.emplace_back(MemoryEffects::Write::get(),
+                       SideEffects::DefaultResource::get());
+  effects.emplace_back(MemoryEffects::Read::get(),
+                       SideEffects::DefaultResource::get());
+}
+
+Speculation::Speculatability ExternCallOp::getSpeculatability() {
+  if (getPure())
+    return Speculation::Speculatable;
+  return Speculation::NotSpeculatable;
+}
+
 // -- GatherOp --
 LogicalResult GatherOp::verify() {
   RankedTensorType indicesTy = getIndices().getType();

python/src/ir.cc

@@ -1673,6 +1673,14 @@ void init_triton_ir(py::module &&m) {
              return self.create<ExternElementwiseOp>(retType, argList, libName,
                                                      libPath, symbol, isPure);
            })
+      .def("create_extern_call",
+           [](TritonOpBuilder &self, const std::string &libName,
+              const std::string &libPath, const std::string &symbol,
+              std::vector<Value> &argList, const std::vector<Type> &retTypes,
+              bool isPure) -> OpState {
+             return self.create<ExternCallOp>(retTypes, argList, libName,
+                                              libPath, symbol, isPure);
+           })
       // Built-in instruction
       .def("create_get_program_id",
            [](TritonOpBuilder &self, int axis) -> Value {

python/test/tle/integration/isolated_repros/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.py[cod]
+*.o
+*.so
+*.lock
+*.tmp

python/test/tle/integration/isolated_repros/no_nvshmem_ws_444/README_CN.md

@@ -0,0 +1,47 @@
+# no-NVSHMEM TLE WS 4+4+4 复现
+
+这个目录用于检查 `Required: 384, Hardware limit: 256` 是否可以在完全不使用
+NVSHMEM/raw extern 的情况下复现。
+
+本目录只使用：
+
+- 单进程
+- 单 GPU
+- 普通 CUDA tensor
+- `tle.gpu.warp_specialize`
+- 普通 `tl.store` / `tl.load`
+
+不使用：
+
+- `mpirun`
+- `NVSHMEM_HOME`
+- raw dialect extern
+- `.cu` device/host helper
+
+## 运行命令
+
+```bash
+cd /workspace/megakernel/triton_flagos_support_nvshmem/python/test/tle/integration/isolated_repros/no_nvshmem_ws_444
+
+PYTHONNOUSERSITE=1 \
+PYTHONDONTWRITEBYTECODE=1 \
+MEGAMOE_TORCH_SITE_PACKAGES=/path/to/site-packages \
+TRITON_CACHE_DIR=/tmp/tle_ws_444_no_nvshmem \
+${PYTHON_BIN:-python} repro_ws_444_no_nvshmem.py
+```
+
+注意：如果使用 editable install 的 FlagTree/Triton 环境，不要额外把
+`/path/to/FlagTree/python` 放进 `PYTHONPATH`，否则可能绕过 editable backend
+finder，导致 `triton.backends.*` 或 `triton.experimental.tle` 解析异常。
+
+## 验证结果
+
+已验证这个 case 会触发：
+
+```text
+triton.runtime.errors.OutOfResources:
+out of resource: threads, Required: 384, Hardware limit: 256.
+```
+
+结论：这个 thread-budget OOR 不依赖 NVSHMEM/raw extern。单进程、单 GPU、
+普通 tensor、无 raw dialect extern 的 TLE WS 4+4+4 role 结构已经足够复现。

python/test/tle/integration/isolated_repros/no_nvshmem_ws_444/repro_ws_444_no_nvshmem.py

@@ -0,0 +1,86 @@
+"""Minimal no-NVSHMEM TLE WS 4+4+4 thread-budget probe.
+
+This file intentionally avoids raw dialect externs and NVSHMEM.  It checks
+whether a TLE warp-specialized kernel with:
+
+  default partition: 4 warps
+  worker0 partition: 4 warps
+  worker1 partition: 4 warps
+
+can reproduce the same thread-budget OOR seen in the MegaMoE-derived repros.
+"""
+
+from __future__ import annotations
+
+import os
+import site
+from pathlib import Path
+
+
+cuda_home = os.environ.get("CUDA_HOME", "/usr/local/cuda-12.8")
+os.environ.setdefault("CUDA_HOME", cuda_home)
+os.environ["CPATH"] = (
+    f"{cuda_home}/targets/x86_64-linux/include:" + os.environ.get("CPATH", "")
+)
+os.environ["LD_LIBRARY_PATH"] = (
+    f"{cuda_home}/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")
+)
+
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+torch_site_packages = os.environ.get("MEGAMOE_TORCH_SITE_PACKAGES")
+if torch_site_packages and Path(torch_site_packages).exists():
+    site.addsitedir(torch_site_packages)
+
+import torch
+
+
+@triton.jit
+def _default_partition(out, VALUE: tl.constexpr):
+    tl.store(out + 0, VALUE)
+
+
+@triton.jit
+def _worker0_partition(out, BLOCK: tl.constexpr):
+    offs = tl.arange(0, BLOCK)
+    vals = offs + 10
+    tl.store(out + 1 + offs, vals)
+
+
+@triton.jit
+def _worker1_partition(out, BLOCK: tl.constexpr):
+    offs = tl.arange(0, BLOCK)
+    vals = tl.load(out + 1 + offs, mask=offs < BLOCK, other=0)
+    tl.store(out + 1 + BLOCK + offs, vals + 20)
+
+
+@triton.jit
+def _ws_444_no_nvshmem_kernel(out, BLOCK: tl.constexpr):
+    tle.gpu.warp_specialize(
+        [
+            (_default_partition, (out, 0x444)),
+            (_worker0_partition, (out, BLOCK)),
+            (_worker1_partition, (out, BLOCK)),
+        ],
+        [4, 4],
+        [80, 180],
+    )
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+
+    torch.cuda.set_device(0)
+    out = torch.empty((257,), device="cuda", dtype=torch.int32)
+    out.zero_()
+    _ws_444_no_nvshmem_kernel[(1,)](out, 128, num_warps=4, maxnreg=240)
+    torch.cuda.synchronize()
+    got = out[:5].detach().cpu().tolist()
+    print(f"unexpected PASS: no-NVSHMEM WS 4+4+4 launched successfully, out[:5]={got}")
+
+
+if __name__ == "__main__":
+    main()