From d8c61f395d99d64e20291cea402e95d895ed1a9e Mon Sep 17 00:00:00 2001 From: jli-melchior Date: Sat, 8 Aug 2026 11:46:15 +0000 Subject: [PATCH 1/5] [compat] Link the HSA code object with the in-process LLD library `gpu-module-to-binary` spawns `/llvm/bin/ld.lld`, where `` comes from ROCM_PATH/ROCM_ROOT/ROCM_HOME or from `__DEFAULT_ROCM_PATH__`, a path baked into the LLVM build. Any container that installs ROCm somewhere else fails, and upstream reports only `lld invocation failed` without naming the path it tried. Our own LLVM bakes in the build machine's ROCm wheel path, so an image switch is enough to break it. Add `fly-emit-gpu-binary`, a wrapper that runs upstream only as far as `format=isa` -- which returns before the toolkit is consulted -- then finishes in process: `mlir::ROCDL::assembleIsa` for the AMDGPU MC assembler and the LLD ELF driver for the link. That removes the lookup entirely and pins the linker to the LLVM revision that produced the ISA. LLD is resolved from the same install that provides MLIR, with NO_DEFAULT_PATH so a mismatched system LLD can never be picked up, and is linked statically into the shared library. An LLVM built without the lld project keeps the upstream behavior via FLYDSL_HAS_LLD_LIBRARY. Two details are carried over from Triton's equivalent change (#7548): `--threads=1`, because LLD's use of the LLVM thread pool deadlocks in `~TaskGroup()` inside a forked child, and checking `canRunAgain`, because a JIT links thousands of times per process and must not continue on corrupted linker state. The FLYDSL_COMPILE_LLVM_DIR path still uses `gpu-module-to-binary`: it drives an upstream mlir-opt that does not know FlyDSL passes. Verified on the same IR with ROCM_PATH pointing at a nonexistent directory: upstream reports `lld invocation failed`, the new pass emits a valid ET_DYN/EM_AMDGPU `gpu.binary`. Full tests/unit + tests/system swept per file against the upstream pass: identical results. Co-Authored-By: Claude Opus 5 (1M context) --- CMakeLists.txt | 16 ++ README.md | 2 +- docs/architecture_guide.md | 22 ++- include/flydsl-c/FlyROCDLDialect.h | 1 + .../flydsl/Conversion/FlyToROCDL/FlyToROCDL.h | 1 + .../flydsl/Conversion/FlyToROCDL/Passes.td | 34 ++++ lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp | 2 + lib/Conversion/FlyToROCDL/CMakeLists.txt | 17 +- .../FlyToROCDL/FlyEmitGPUBinary.cpp | 175 ++++++++++++++++++ python/flydsl/compiler/backends/rocm.py | 16 +- python/flydsl/compiler/jit_function.py | 4 +- tests/mlir/Conversion/emit_gpu_binary.mlir | 20 ++ tests/unit/test_external_llvm_codegen.py | 12 +- 13 files changed, 308 insertions(+), 14 deletions(-) create mode 100644 lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp create mode 100644 tests/mlir/Conversion/emit_gpu_binary.mlir diff --git a/CMakeLists.txt b/CMakeLists.txt index c22e9cf03..8c5df5e36 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -11,6 +11,22 @@ find_package(MLIR REQUIRED CONFIG) message(STATUS "Found MLIR: ${MLIR_DIR}") +# LLD is linked as a library so GPU binary emission never has to locate an +# `ld.lld` executable at runtime (see lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp). +# Always resolve it from the LLVM install that provides MLIR: NO_DEFAULT_PATH +# keeps CMake from picking up a system LLD whose version would not match the +# code generator. Optional, so an LLVM built without the lld project still works. +if(NOT LLD_DIR) + set(LLD_DIR ${LLVM_LIBRARY_DIR}/cmake/lld) +endif() +find_package(LLD CONFIG PATHS "${LLD_DIR}" NO_DEFAULT_PATH) +if(LLD_FOUND) + message(STATUS "Found LLD: ${LLD_DIR} (in-process GPU binary linking)") +else() + message(STATUS "LLD not found under ${LLD_DIR}; GPU binary emission falls back to " + "spawning ld.lld from the ROCm toolkit path") +endif() + if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) set(CMAKE_BUILD_TYPE "RelWithDebInfo" CACHE STRING "Build type (Debug, Release, RelWithDebInfo, MinSizeRel)" FORCE) message(STATUS "CMAKE_BUILD_TYPE not set; defaulting to ${CMAKE_BUILD_TYPE}") diff --git a/README.md b/README.md index bb6ca7954..75ab5892c 100644 --- a/README.md +++ b/README.md @@ -254,7 +254,7 @@ Python Function (@flyc.kernel / @flyc.jit) │ reconcile-unrealized-casts │ ├──────────────────────────────────────────────────────────┤ │ C. binary_fragment │ - │ gpu-module-to-binary{format=fatbin} │ + │ fly-emit-gpu-binary │ └──────────────────────────────────────────────────────────┘ │ ▼ diff --git a/docs/architecture_guide.md b/docs/architecture_guide.md index 935b7f967..69b987b59 100644 --- a/docs/architecture_guide.md +++ b/docs/architecture_guide.md @@ -191,7 +191,7 @@ Python Function (@flyc.kernel / @flyc.jit) │ ensure-debug-info-scope-on-llvm-func (optional) │ ├────────────────────────────────────────────────────────┤ │ Stage C — binary_fragment │ - │ gpu-module-to-binary{format=fatbin opts="..."} │ + │ fly-emit-gpu-binary{opts="..."} │ └────────────────────────────────────────────────────────┘ │ ▼ @@ -206,7 +206,9 @@ The pipeline is built by `RocmBackend._pipeline_parts()` in the pipeline as a single combined pass list (`pipeline_fragments()`) or split it for external LLVM codegen (`external_binary_pipeline_fragments()`). External mode runs Stages A and B with the bundled MLIR runtime, then invokes the -external LLVM toolchain only for Stage C (`gpu-module-to-binary`). +external LLVM toolchain only for Stage C. That toolchain drives an upstream +`mlir-opt` which does not know FlyDSL passes, so external mode substitutes +`gpu-module-to-binary{format=fatbin}` for `fly-emit-gpu-binary`. **Stage A — `pre_binary_fragments`** (Fly dialect → ROCDL lowering) @@ -244,7 +246,21 @@ When `FLYDSL_DEBUG_ENABLE_DEBUG_INFO=1`, Stage B appends | # | Pass | Description | |---|---|---| -| 19 | `gpu-module-to-binary{format=fatbin opts="..."}` | Invokes the LLVM AMDGPU backend and emits an HSA fatbin. | +| 19 | `fly-emit-gpu-binary{opts="..."}` | Invokes the LLVM AMDGPU backend and emits an HSA fatbin. | + +Stage C wraps the upstream `gpu-module-to-binary` pass +(`lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp`): it runs upstream only as far +as `format=isa`, then assembles the ISA with `mlir::ROCDL::assembleIsa` and +links the HSA code object through the LLD ELF driver linked into FlyDSL. +Upstream would instead spawn `/llvm/bin/ld.lld`, resolved from +`ROCM_PATH` or a path baked into the LLVM build, which fails on any container +that installs ROCm elsewhere. Linking LLD as a library removes that lookup and +pins the linker to the LLVM revision that produced the ISA. An LLVM built +without the `lld` project falls back to the upstream behavior. + +Device bitcode (`ocml`/`ockl`) is still loaded from `/amdgcn/bitcode` +when the module calls `__ocml_*` / `__ockl_*`, so a valid ROCm path remains +required for kernels that use those math functions. `gpu-kernel-outlining` is no longer a pass in the runtime pipeline. Kernel outlining happens during Python tracing, when `@flyc.kernel` emits diff --git a/include/flydsl-c/FlyROCDLDialect.h b/include/flydsl-c/FlyROCDLDialect.h index 21ac2789c..349c31597 100644 --- a/include/flydsl-c/FlyROCDLDialect.h +++ b/include/flydsl-c/FlyROCDLDialect.h @@ -14,6 +14,7 @@ extern "C" { MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(FlyROCDL, fly_rocdl); MLIR_CAPI_EXPORTED void mlirRegisterFlyToROCDLConversionPass(void); +MLIR_CAPI_EXPORTED void mlirRegisterFlyEmitGPUBinaryPass(void); MLIR_CAPI_EXPORTED void mlirRegisterFlyROCDLClusterAttrPass(void); /// Backend plugin registration: insert all ROCDL dialects into \p registry. diff --git a/include/flydsl/Conversion/FlyToROCDL/FlyToROCDL.h b/include/flydsl/Conversion/FlyToROCDL/FlyToROCDL.h index f6367147d..6f9205017 100644 --- a/include/flydsl/Conversion/FlyToROCDL/FlyToROCDL.h +++ b/include/flydsl/Conversion/FlyToROCDL/FlyToROCDL.h @@ -8,6 +8,7 @@ namespace mlir { #define GEN_PASS_DECL_FLYTOROCDLCONVERSIONPASS +#define GEN_PASS_DECL_FLYEMITGPUBINARYPASS #define GEN_PASS_DECL_FLYROCDLCLUSTERATTRPASS #include "flydsl/Conversion/FlyToROCDL/Passes.h.inc" } // namespace mlir diff --git a/include/flydsl/Conversion/FlyToROCDL/Passes.td b/include/flydsl/Conversion/FlyToROCDL/Passes.td index e23eb7b35..d10c709e5 100644 --- a/include/flydsl/Conversion/FlyToROCDL/Passes.td +++ b/include/flydsl/Conversion/FlyToROCDL/Passes.td @@ -14,6 +14,40 @@ def FlyToROCDLConversionPass : Pass<"convert-fly-to-rocdl"> { ]; } +def FlyEmitGPUBinaryPass : Pass<"fly-emit-gpu-binary", ""> { + let summary = "Transform a GPU module into a GPU binary without an external linker"; + let description = [{ + Wrapper around the upstream `gpu-module-to-binary` pass that stops it at + assembly (`format=isa`) and performs the remaining assemble and link steps + in process: `mlir::ROCDL::assembleIsa` for the AMDGPU MC assembler, and the + LLD ELF driver library for the link. + + Upstream instead spawns `/llvm/bin/ld.lld`, where `` comes + from `ROCM_PATH` or from a path baked into the LLVM build. That lookup + fails on any container that installs ROCm somewhere else, and reports only + `lld invocation failed`. Linking against the LLD libraries removes the + lookup entirely and pins the linker to the LLVM revision that produced the + ISA. + + Options mirror `gpu-module-to-binary`, minus `format`. Objects that are + not ROCDL targets are left untouched, and builds whose LLVM has no LLD + libraries fall back to running `gpu-module-to-binary` unmodified. + }]; + let options = [ + Option<"toolkitPath", "toolkit", "std::string", [{""}], + "Toolkit path.">, + ListOption<"linkFiles", "l", "std::string", + "Extra files to link to.">, + Option<"cmdOptions", "opts", "std::string", [{""}], + "Command line options to pass to the tools.">, + Option<"elfSection", "section", "std::string", [{""}], + "ELF section where binary is to be located."> + ]; + let dependentDialects = [ + "gpu::GPUDialect" + ]; +} + def FlyROCDLClusterAttrPass : Pass<"fly-rocdl-cluster-attr"> { let summary = "Inject amdgpu-cluster-dims into llvm.func passthrough"; let description = [{ diff --git a/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp b/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp index d835245eb..866a4ceef 100644 --- a/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp +++ b/lib/CAPI/Dialect/FlyROCDL/FlyROCDLDialect.cpp @@ -11,6 +11,7 @@ MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(FlyROCDL, fly_rocdl, mlir::fly_rocdl::FlyROCDLDialect) void mlirRegisterFlyToROCDLConversionPass(void) { mlir::registerFlyToROCDLConversionPass(); } +void mlirRegisterFlyEmitGPUBinaryPass(void) { mlir::registerFlyEmitGPUBinaryPass(); } void mlirRegisterFlyROCDLClusterAttrPass(void) { mlir::registerFlyROCDLClusterAttrPass(); } void flydsl_register_rocdl_dialects(MlirDialectRegistry registry) { @@ -19,5 +20,6 @@ void flydsl_register_rocdl_dialects(MlirDialectRegistry registry) { void flydsl_register_rocdl_passes(void) { mlirRegisterFlyToROCDLConversionPass(); + mlirRegisterFlyEmitGPUBinaryPass(); mlirRegisterFlyROCDLClusterAttrPass(); } diff --git a/lib/Conversion/FlyToROCDL/CMakeLists.txt b/lib/Conversion/FlyToROCDL/CMakeLists.txt index e926e62a9..508906ee1 100644 --- a/lib/Conversion/FlyToROCDL/CMakeLists.txt +++ b/lib/Conversion/FlyToROCDL/CMakeLists.txt @@ -1,5 +1,12 @@ +# Directory scope so the definition also reaches obj.MLIRFlyToROCDL, which is +# where add_mlir_conversion_library actually compiles the sources. +if(LLD_FOUND) + add_compile_definitions(FLYDSL_HAS_LLD_LIBRARY) +endif() + add_mlir_conversion_library(MLIRFlyToROCDL FlyToROCDL.cpp + FlyEmitGPUBinary.cpp DEPENDS MLIRFlyIncGen @@ -9,12 +16,20 @@ add_mlir_conversion_library(MLIRFlyToROCDL LINK_LIBS PUBLIC MLIRFlyDialect MLIRFlyROCDLDialect - + MLIRArithDialect + MLIRGPUDialect + MLIRGPUTransforms MLIRIR MLIRLLVMDialect MLIRPass + MLIRROCDLDialect + MLIRROCDLTarget MLIRSCFDialect MLIRTransforms MLIRVectorDialect ) + +if(LLD_FOUND) + target_link_libraries(MLIRFlyToROCDL PUBLIC lldCommon lldELF) +endif() diff --git a/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp b/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp new file mode 100644 index 000000000..ece3dc312 --- /dev/null +++ b/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2025 FlyDSL Project Contributors + +#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/GPU/Transforms/Passes.h" +#include "mlir/Dialect/LLVMIR/ROCDLDialect.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" + +#ifdef FLYDSL_HAS_LLD_LIBRARY +#include "mlir/Target/LLVM/ROCDL/Utils.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/FileUtilities.h" +#include "llvm/Support/MemoryBuffer.h" +#include "llvm/Support/raw_ostream.h" + +#include "lld/Common/Driver.h" +LLD_HAS_DRIVER(elf) +#endif // FLYDSL_HAS_LLD_LIBRARY + +#include "flydsl/Conversion/FlyToROCDL/FlyToROCDL.h" + +namespace mlir { +#define GEN_PASS_DEF_FLYEMITGPUBINARYPASS +#include "flydsl/Conversion/FlyToROCDL/Passes.h.inc" +} // namespace mlir + +using namespace mlir; + +namespace { + +#ifdef FLYDSL_HAS_LLD_LIBRARY + +// Run the LLD ELF driver linked into this library. Returns the linker +// diagnostics on failure, std::nullopt on success. +// +// `--threads=1` is not a performance choice: LLD links through LLVM's global +// thread pool, and in a forked child process (autotune workers) the inherited +// pool has no live workers, so `~TaskGroup()` blocks forever. A single kernel +// object links in microseconds either way. +// +// `canRunAgain` reports whether LLD's global state survived the call. A JIT +// links thousands of times per process, so a false here must abort rather than +// let the next link run on corrupted state. +std::optional runLLD(StringRef objectPath, StringRef hsacoPath) { + std::string objectPathStr = objectPath.str(); + std::string hsacoPathStr = hsacoPath.str(); + std::array args{"ld.lld", "--threads=1", "-shared", objectPathStr.c_str(), + "-o", hsacoPathStr.c_str()}; + + std::string errString; + llvm::raw_string_ostream errStream(errString); + lld::Result result = lld::lldMain(args, llvm::outs(), errStream, {{lld::Gnu, &lld::elf::link}}); + if (result.retCode != 0 || !result.canRunAgain) { + errStream.flush(); + return errString.empty() ? std::string("unknown lld failure") : errString; + } + return std::nullopt; +} + +// Link a relocatable AMDGPU ELF into an HSA code object. The LLD driver API +// only accepts file paths, so the object round-trips through temporary files +// exactly as the upstream implementation does; what is saved is the fork/exec +// and the toolkit path lookup. +FailureOr> linkObjectCode(ArrayRef objectCode, + function_ref emitError) { + int objectFd = -1; + SmallString<128> objectPath; + if (llvm::sys::fs::createTemporaryFile("flydsl-kernel%%", "o", objectFd, objectPath)) + return emitError() << "failed to create a temporary file for the ISA binary"; + llvm::FileRemover objectRemover(objectPath); + { + llvm::raw_fd_ostream objectOs(objectFd, /*shouldClose=*/true); + objectOs << StringRef(objectCode.data(), objectCode.size()); + objectOs.flush(); + } + + SmallString<128> hsacoPath; + if (llvm::sys::fs::createTemporaryFile("flydsl-kernel%%", "hsaco", hsacoPath)) + return emitError() << "failed to create a temporary file for the HSA code object"; + llvm::FileRemover hsacoRemover(hsacoPath); + + if (std::optional error = runLLD(objectPath, hsacoPath)) + return emitError() << "in-process lld failed to link the HSA code object: " << *error; + + auto hsacoFile = llvm::MemoryBuffer::getFile(hsacoPath, /*IsText=*/false); + if (!hsacoFile) + return emitError() << "failed to read the HSA code object from " << hsacoPath; + + StringRef buffer = (*hsacoFile)->getBuffer(); + return SmallVector(buffer.begin(), buffer.end()); +} + +// Replace every ROCDL assembly object of `binary` with the linked fatbin. The +// AMDGPU MC target is already registered here: `gpu-module-to-binary` ran +// `SerializeGPUModuleBase::init()` while producing the ISA we consume. +LogicalResult compileAssemblyObjects(gpu::BinaryOp binary) { + ArrayRef objects = binary.getObjectsAttr().getValue(); + SmallVector compiled; + compiled.reserve(objects.size()); + bool changed = false; + + for (Attribute attr : objects) { + auto object = dyn_cast(attr); + auto target = + object ? dyn_cast(object.getTarget()) : ROCDL::ROCDLTargetAttr(); + if (!target || object.getFormat() != gpu::CompilationTarget::Assembly) { + compiled.push_back(attr); + continue; + } + + auto emitError = [&]() { return binary.emitError(); }; + FailureOr> objectCode = + ROCDL::assembleIsa(object.getObject().getValue(), target.getTriple(), target.getChip(), + target.getFeatures(), emitError); + if (failed(objectCode)) + return failure(); + + FailureOr> hsaco = linkObjectCode(*objectCode, emitError); + if (failed(hsaco)) + return failure(); + + compiled.push_back(gpu::ObjectAttr::get( + object.getTarget(), gpu::CompilationTarget::Fatbin, + StringAttr::get(binary.getContext(), StringRef(hsaco->data(), hsaco->size())), + object.getProperties(), object.getKernels())); + changed = true; + } + + if (changed) + binary.setObjectsAttr(ArrayAttr::get(binary.getContext(), compiled)); + return success(); +} + +#endif // FLYDSL_HAS_LLD_LIBRARY + +class FlyEmitGPUBinaryPass : public mlir::impl::FlyEmitGPUBinaryPassBase { +public: + using mlir::impl::FlyEmitGPUBinaryPassBase::FlyEmitGPUBinaryPassBase; + + void runOnOperation() override { + Operation *op = getOperation(); + + GpuModuleToBinaryPassOptions binaryOptions; + binaryOptions.toolkitPath = toolkitPath; + binaryOptions.linkFiles.assign(linkFiles.begin(), linkFiles.end()); + binaryOptions.cmdOptions = cmdOptions; + binaryOptions.elfSection = elfSection; +#ifdef FLYDSL_HAS_LLD_LIBRARY + // Stop before the link so upstream never looks for `ld.lld`. + binaryOptions.compilationTarget = "isa"; +#else + binaryOptions.compilationTarget = "fatbin"; +#endif + + OpPassManager pm(op->getName()); + pm.addPass(createGpuModuleToBinaryPass(binaryOptions)); + if (failed(runPipeline(pm, op))) { + signalPassFailure(); + return; + } + +#ifdef FLYDSL_HAS_LLD_LIBRARY + WalkResult walked = op->walk([&](gpu::BinaryOp binary) { + return failed(compileAssemblyObjects(binary)) ? WalkResult::interrupt() + : WalkResult::advance(); + }); + if (walked.wasInterrupted()) + signalPassFailure(); +#endif + } +}; + +} // namespace diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index 45f891f0f..2180f68fa 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -7,6 +7,11 @@ from ...utils import env from .base import BaseBackend, GPUTarget +#: FlyDSL wrapper around ``gpu-module-to-binary`` that links the HSA code object +#: with the in-process LLD library instead of spawning ``ld.lld`` from the ROCm +#: toolkit path. See ``lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp``. +BINARY_PASS_NAME = "fly-emit-gpu-binary" + class RocmBackend(BaseBackend): """ROCm / AMDGPU compile backend (HIP runtime, ROCDL lowering).""" @@ -49,7 +54,7 @@ def _format_pass_opts(opts: dict) -> str: """Format {key: value, ...} as 'key=value key2=value2' for MLIR pass options.""" return " ".join(f"{k}={v}" for k, v in opts.items()) - def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]: + def _pipeline_parts(self, *, compile_hints: dict, external: bool = False) -> Tuple[List[str], str]: chip = self.target.arch waves_per_eu = compile_hints.get("waves_per_eu") maxnreg = compile_hints.get("maxnreg") @@ -106,7 +111,12 @@ def _pipeline_parts(self, *, compile_hints: dict) -> Tuple[List[str], str]: else [] ), ] - binary_fragment = f'gpu-module-to-binary{{format=fatbin opts="{" ".join(bin_cli_opts)}"}}' + opts = f'opts="{" ".join(bin_cli_opts)}"' + # The external toolchain drives an upstream mlir-opt that does not know + # about FlyDSL passes, so that path keeps using gpu-module-to-binary. + binary_fragment = ( + f"gpu-module-to-binary{{format=fatbin {opts}}}" if external else f"{BINARY_PASS_NAME}{{{opts}}}" + ) return [*pre_binary_fragments, *binary_prep_fragments], binary_fragment def pipeline_fragments(self, *, compile_hints: dict) -> List[str]: @@ -114,7 +124,7 @@ def pipeline_fragments(self, *, compile_hints: dict) -> List[str]: return [*pre_binary_fragments, binary_fragment] def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]: - return self._pipeline_parts(compile_hints=compile_hints) + return self._pipeline_parts(compile_hints=compile_hints, external=True) def lower_compile_hints(self, module, *, compile_hints: dict) -> None: """Materialize a scalar waves-per-EU override on kernel entries.""" diff --git a/python/flydsl/compiler/jit_function.py b/python/flydsl/compiler/jit_function.py index 2302c90ab..ca08dca38 100644 --- a/python/flydsl/compiler/jit_function.py +++ b/python/flydsl/compiler/jit_function.py @@ -755,7 +755,7 @@ def _pipeline_fragments_for_mode(backend, *, compile_hints: dict) -> PipelineCon return PipelineConfig( fragments=fragments, pre_binary=None, - binary_fragment=None, + binary_fragment=fragments[-1], llvm_opts=llvm_opts, external=False, ) @@ -834,7 +834,7 @@ def compile( stage_num_base = 1 dump_fragments = pre_binary_fragments if external_binary else fragments for idx, frag in enumerate(dump_fragments): - if frag.strip().startswith("gpu-module-to-binary"): + if frag == binary_fragment: llir = _extract_llvm_ir(module) stage_num = stage_num_base + idx diff --git a/tests/mlir/Conversion/emit_gpu_binary.mlir b/tests/mlir/Conversion/emit_gpu_binary.mlir new file mode 100644 index 000000000..aae9195cc --- /dev/null +++ b/tests/mlir/Conversion/emit_gpu_binary.mlir @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2025 FlyDSL Project Contributors +// RUN: env ROCM_PATH=/nonexistent-rocm ROCM_ROOT=/nonexistent-rocm ROCM_HOME=/nonexistent-rocm %fly-opt %s --fly-emit-gpu-binary | FileCheck %s + +// fly-emit-gpu-binary links the HSA code object through the in-process LLD +// library, so it must succeed even when no ROCm toolkit is reachable. The +// upstream gpu-module-to-binary pass fails here with "lld invocation failed" +// because it spawns /llvm/bin/ld.lld. + +// CHECK: gpu.binary @kernels +// CHECK-SAME: #rocdl.target +// CHECK-SAME: "\7FELF +// CHECK-NOT: gpu.module @kernels +module attributes {gpu.container_module} { + gpu.module @kernels [#rocdl.target] { + llvm.func @empty_kernel() attributes {rocdl.kernel} { + llvm.return + } + } +} diff --git a/tests/unit/test_external_llvm_codegen.py b/tests/unit/test_external_llvm_codegen.py index 98294d7ca..50da11783 100644 --- a/tests/unit/test_external_llvm_codegen.py +++ b/tests/unit/test_external_llvm_codegen.py @@ -9,7 +9,7 @@ from flydsl._mlir import ir from flydsl._mlir._mlir_libs._mlirDialectsLLVM import translate_module_to_llvmir from flydsl._mlir.passmanager import PassManager -from flydsl.compiler.backends.rocm import RocmBackend +from flydsl.compiler.backends.rocm import BINARY_PASS_NAME, RocmBackend from flydsl.compiler.external_llvm import ( _format_llvm_cli_options, external_llvm_fingerprint, @@ -67,12 +67,16 @@ def test_rocm_external_pipeline_split_matches_full_pipeline(): full = backend.pipeline_fragments(compile_hints=hints) pre_binary, binary = backend.external_binary_pipeline_fragments(compile_hints=hints) - assert full == [*pre_binary, binary] + assert full[:-1] == pre_binary assert pre_binary[-1] == "reconcile-unrealized-casts" assert any(fragment.startswith("gpu.module(") for fragment in pre_binary) + # Embedded codegen links through the in-process LLD library; the external + # toolchain drives an upstream mlir-opt that only knows the upstream pass. + assert full[-1].startswith(BINARY_PASS_NAME) assert binary.startswith("gpu-module-to-binary") - assert "--amdgpu-waves-per-eu=2" in binary - assert "--amdgpu-num-vgpr=128" in binary + for fragment in (full[-1], binary): + assert "--amdgpu-waves-per-eu=2" in fragment + assert "--amdgpu-num-vgpr=128" in fragment def test_rocm_lower_wpe_preserves_source_default_and_overrides_kernel_entries(): From 8cba7821026e7c782b2515a250b502e2b66b0cb0 Mon Sep 17 00:00:00 2001 From: jli-melchior Date: Sat, 8 Aug 2026 11:46:51 +0000 Subject: [PATCH 2/5] [compat] Bundle AMDGCN device bitcode so toolkit lookup leaves the container Linking in process removed the ld.lld lookup, but `appendStandardLibs()` still reads `/amdgcn/bitcode` whenever a module references `__ocml_*` or `__ockl_*`, and errors out if that directory is absent. FlyDSL reaches it: `convert-gpu-to-rocdl` pulls in MathToROCDL, so scalar math ops with no LLVM intrinsic lower to ocml calls -- `fx.erfc` becomes `__ocml_erfc_f32`. So the per-image path problem survived in the device-library half. Bundle the bitcode the way Triton does. CMake locates ocml/ockl/hip/opencl at configure time (ROCM_PATH/ROCM_ROOT/ROCM_HOME or /opt/rocm, overridable with -DFLYDSL_ROCM_BITCODE_DIR) and copies them into the package; those four are the complete set appendStandardLibs() can request, since the oclc_* control variables are synthesized in-module rather than read from disk. They land under `_mlir/`, which is already the packaged build-output subtree, so the editable symlink and the wheel's package_dir mapping carry them with no new plumbing. `rocm_toolkit_path()` then resolves FLYDSL_COMPILE_ROCM_PATH, the bundled tree, and ROCM_PATH/ROCM_ROOT/ROCM_HOME in that order and passes the winner as `toolkit=`. Finding nothing passes no option at all, leaving upstream's lookup untouched rather than forcing a known-bad path on it. `RocmBackend.hash()` folds the bitcode digest into the JIT cache key: swapping ocml.bc changes generated code without changing any FlyDSL shared library, so the native-library hashes alone would serve stale artifacts. Verified with ROCM_PATH/ROCM_ROOT/ROCM_HOME all unset: a kernel calling fx.erfc compiles and runs. Hiding the bundled ocml.bc makes that same compilation fail on the bundled path, confirming which copy is consumed. Co-Authored-By: Claude Opus 5 (1M context) --- CLAUDE.md | 1 + CMakeLists.txt | 22 +++++ README.md | 4 + docs/architecture_guide.md | 19 ++++- .../FlyToROCDL/FlyEmitGPUBinary.cpp | 2 +- python/flydsl/compiler/backends/rocm.py | 83 +++++++++++++++++++ python/flydsl/utils/env.py | 7 ++ python/mlir_flydsl/CMakeLists.txt | 23 +++++ setup.py | 3 + tests/unit/test_rocm_toolkit_path.py | 77 +++++++++++++++++ 10 files changed, 237 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_rocm_toolkit_path.py diff --git a/CLAUDE.md b/CLAUDE.md index 1876f19bb..5bd8bca53 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -145,6 +145,7 @@ Use names from `python/flydsl/utils/env.py`; do not introduce alternate spelling | Enable/disable JIT disk cache | `FLYDSL_RUNTIME_ENABLE_CACHE` (`0` / `false` disables disk cache; in-memory cache remains) | | AOT-cache-only execution | `FLYDSL_RUNTIME_RUN_ONLY` (`1` skips JIT; loads disk cache only, raises on cache miss; incompatible with `FLYDSL_DUMP_IR=1`) | | External LLVM/MLIR codegen | `FLYDSL_COMPILE_LLVM_DIR` (install prefix; enables external-binary final codegen, part of the JIT cache key) | +| ROCm device-bitcode root | `FLYDSL_COMPILE_ROCM_PATH` (overrides the `amdgcn/bitcode` bundled with the package and `ROCM_PATH`/`ROCM_ROOT`/`ROCM_HOME`) | | IR dumps | `FLYDSL_DUMP_IR`, `FLYDSL_DUMP_DIR` | | Runtime kind | `FLYDSL_RUNTIME_KIND` | | GPU arch hints | `FLYDSL_GPU_ARCH`, `HSA_OVERRIDE_GFX_VERSION` | diff --git a/CMakeLists.txt b/CMakeLists.txt index 8c5df5e36..75c0e5910 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -27,6 +27,28 @@ else() "spawning ld.lld from the ROCm toolkit path") endif() +# AMDGCN device bitcode (ocml/ockl/...) is loaded from /amdgcn/bitcode +# whenever a kernel calls __ocml_* / __ockl_*. Bundle it into the Python +# package so that lookup does not depend on where the container installs ROCm. +# Override the source with -DFLYDSL_ROCM_BITCODE_DIR=. +if(NOT FLYDSL_ROCM_BITCODE_DIR) + find_path(FLYDSL_ROCM_BITCODE_DIR ocml.bc + HINTS + "$ENV{ROCM_PATH}/amdgcn/bitcode" + "$ENV{ROCM_ROOT}/amdgcn/bitcode" + "$ENV{ROCM_HOME}/amdgcn/bitcode" + "/opt/rocm/amdgcn/bitcode" + NO_DEFAULT_PATH) +endif() +if(FLYDSL_ROCM_BITCODE_DIR) + message(STATUS "Found AMDGCN device bitcode: ${FLYDSL_ROCM_BITCODE_DIR}") +else() + message(WARNING + "AMDGCN device bitcode (ocml.bc) not found; the Python package will not bundle it. " + "Kernels calling __ocml_*/__ockl_* will then need FLYDSL_COMPILE_ROCM_PATH or " + "ROCM_PATH to point at a ROCm install. Set -DFLYDSL_ROCM_BITCODE_DIR= to bundle it.") +endif() + if(NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) set(CMAKE_BUILD_TYPE "RelWithDebInfo" CACHE STRING "Build type (Debug, Release, RelWithDebInfo, MinSizeRel)" FORCE) message(STATUS "CMAKE_BUILD_TYPE not set; defaulting to ${CMAKE_BUILD_TYPE}") diff --git a/README.md b/README.md index 75ab5892c..332dd6a45 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,10 @@ Prerequisites for source builds: - **Build tools**: `cmake` (>=3.20), C++17 compiler, optionally `ninja` - **Python deps**: `nanobind`, `numpy`, `pybind11` (installed by `scripts/build_llvm.sh`; install them manually if you skip that step) +- **ROCm device bitcode**: `scripts/build.sh` bundles `amdgcn/bitcode/*.bc` from the ROCm found via + `ROCM_PATH`/`ROCM_ROOT`/`ROCM_HOME` or `/opt/rocm`, so kernels calling `__ocml_*` compile wherever the + package is later installed. Override with `-DFLYDSL_ROCM_BITCODE_DIR=`; skipping it only means + those kernels need `FLYDSL_COMPILE_ROCM_PATH` or `ROCM_PATH` set at compile time. ```bash # Clone ROCm LLVM and build MLIR (takes ~30min with -j64) diff --git a/docs/architecture_guide.md b/docs/architecture_guide.md index 69b987b59..842eb5cc8 100644 --- a/docs/architecture_guide.md +++ b/docs/architecture_guide.md @@ -258,9 +258,22 @@ that installs ROCm elsewhere. Linking LLD as a library removes that lookup and pins the linker to the LLVM revision that produced the ISA. An LLVM built without the `lld` project falls back to the upstream behavior. -Device bitcode (`ocml`/`ockl`) is still loaded from `/amdgcn/bitcode` -when the module calls `__ocml_*` / `__ockl_*`, so a valid ROCm path remains -required for kernels that use those math functions. +Device bitcode (`ocml`/`ockl`/`hip`/`opencl`) is loaded from +`/amdgcn/bitcode` when the module calls `__ocml_*` / `__ockl_*` — for +example `fx.erfc`, which has no LLVM intrinsic. FlyDSL bundles that bitcode into +the package at build time and points `toolkit=` at it, so this lookup does not +depend on where the container installs ROCm either. +`RocmBackend.rocm_toolkit_path()` resolves it in order: + +1. `FLYDSL_COMPILE_ROCM_PATH` +2. the bitcode bundled with the package +3. `ROCM_PATH` / `ROCM_ROOT` / `ROCM_HOME` + +If none of them contains `amdgcn/bitcode/ocml.bc`, no `toolkit=` is passed and +upstream's own lookup applies. Configure the build with +`-DFLYDSL_ROCM_BITCODE_DIR=` to choose which ROCm supplies the bundled +bitcode; CMake otherwise searches `ROCM_PATH`/`ROCM_ROOT`/`ROCM_HOME` and +`/opt/rocm`. `gpu-kernel-outlining` is no longer a pass in the runtime pipeline. Kernel outlining happens during Python tracing, when `@flyc.kernel` emits diff --git a/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp b/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp index ece3dc312..1814618ae 100644 --- a/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp +++ b/lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp @@ -105,7 +105,7 @@ LogicalResult compileAssemblyObjects(gpu::BinaryOp binary) { auto object = dyn_cast(attr); auto target = object ? dyn_cast(object.getTarget()) : ROCDL::ROCDLTargetAttr(); - if (!target || object.getFormat() != gpu::CompilationTarget::Assembly) { + if (!object || !target || object.getFormat() != gpu::CompilationTarget::Assembly) { compiled.push_back(attr); continue; } diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index 2180f68fa..4212df4e1 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2025 FlyDSL Project Contributors +import hashlib +import os +from functools import lru_cache +from pathlib import Path from typing import List, Tuple from ...runtime.device import get_rocm_arch, is_rdna_arch @@ -12,6 +16,72 @@ #: toolkit path. See ``lib/Conversion/FlyToROCDL/FlyEmitGPUBinary.cpp``. BINARY_PASS_NAME = "fly-emit-gpu-binary" +#: ROCm toolkit root bundled with the package; holds ``amdgcn/bitcode/*.bc`` +#: copied in at build time. Lives under ``_mlir`` because that subtree is the +#: packaged build output. Absent when CMake could not locate a ROCm install. +BUNDLED_ROCM_PATH = Path(__file__).resolve().parents[2] / "_mlir" / "_rocm" + + +def _has_device_bitcode(root: Path) -> bool: + return (root / "amdgcn" / "bitcode" / "ocml.bc").is_file() + + +@lru_cache(maxsize=1) +def rocm_toolkit_path() -> str: + """Resolve the ROCm root that supplies AMDGCN device bitcode. + + Only ``/amdgcn/bitcode`` is read: ``fly-emit-gpu-binary`` links the HSA + code object in process, so no ``ld.lld`` lookup is involved. The bundled + tree is preferred over the environment so that a container that installs + ROCm somewhere unexpected still compiles kernels that call ``__ocml_*``. + + Returns an empty string when nothing is found, which leaves the upstream + ``ROCM_PATH`` lookup in place rather than forcing a bad path on it. + """ + candidates: List[Tuple[str, Path]] = [] + if env.compile.rocm_path: + candidates.append(("FLYDSL_COMPILE_ROCM_PATH", Path(env.compile.rocm_path))) + candidates.append(("bundled with flydsl", BUNDLED_ROCM_PATH)) + for var in ("ROCM_PATH", "ROCM_ROOT", "ROCM_HOME"): + value = os.environ.get(var) + if value: + candidates.append((var, Path(value))) + + for _, root in candidates: + if not _has_device_bitcode(root): + continue + path = str(root) + # MLIR's pass-pipeline parser treats whitespace, commas and braces as + # structural syntax, so such a path cannot be spelled as an option. + bad = sorted({ch for ch in path if ch.isspace() or ch in ",{}\"'"}) + if bad: + raise ValueError( + f"ROCm toolkit path {path!r} contains unsupported character(s) {bad!r} and cannot be " + "passed to an MLIR pass option. Point FLYDSL_COMPILE_ROCM_PATH at a path without " + "whitespace, commas, braces, or quotes." + ) + return path + return "" + + +#: Device libraries ``appendStandardLibs()`` can pull out of a toolkit path. +_DEVICE_BITCODE_FILES = ("ocml.bc", "ockl.bc", "hip.bc", "opencl.bc") + + +@lru_cache(maxsize=4) +def _device_bitcode_fingerprint(toolkit: str) -> str: + """Digest the device bitcode reachable from *toolkit*.""" + if not toolkit: + return "none" + bitcode_dir = Path(toolkit) / "amdgcn" / "bitcode" + digest = hashlib.sha256() + for name in _DEVICE_BITCODE_FILES: + digest.update(name.encode()) + path = bitcode_dir / name + if path.is_file(): + digest.update(path.read_bytes()) + return digest.hexdigest() + class RocmBackend(BaseBackend): """ROCm / AMDGPU compile backend (HIP runtime, ROCDL lowering).""" @@ -112,6 +182,9 @@ def _pipeline_parts(self, *, compile_hints: dict, external: bool = False) -> Tup ), ] opts = f'opts="{" ".join(bin_cli_opts)}"' + toolkit = rocm_toolkit_path() + if toolkit: + opts = f"toolkit={toolkit} {opts}" # The external toolchain drives an upstream mlir-opt that does not know # about FlyDSL passes, so that path keeps using gpu-module-to-binary. binary_fragment = ( @@ -126,6 +199,16 @@ def pipeline_fragments(self, *, compile_hints: dict) -> List[str]: def external_binary_pipeline_fragments(self, *, compile_hints: dict) -> Tuple[List[str], str]: return self._pipeline_parts(compile_hints=compile_hints, external=True) + def hash(self) -> str: + """Fold the device bitcode into the JIT cache key. + + Linked-in ocml/ockl changes the generated code but no FlyDSL shared + library, so the native-library hashes alone would not invalidate a + stale cache after a rebuild against a different ROCm. + """ + toolkit = rocm_toolkit_path() + return f"{self.target}:{toolkit}:{_device_bitcode_fingerprint(toolkit)}" + def lower_compile_hints(self, module, *, compile_hints: dict) -> None: """Materialize a scalar waves-per-EU override on kernel entries.""" waves_per_eu = compile_hints.get("waves_per_eu") diff --git a/python/flydsl/utils/env.py b/python/flydsl/utils/env.py index 1839dd7df..10be8d82e 100644 --- a/python/flydsl/utils/env.py +++ b/python/flydsl/utils/env.py @@ -238,6 +238,13 @@ class CompileEnvManager(EnvManager): arch = OptStr("", env_var="ARCH", description="Override target GPU architecture (e.g. gfx942, gfx950)") backend = OptStr("rocm", description="GPU compile backend id (e.g. rocm)") llvm_dir = OptStr("", description="External LLVM/MLIR install prefix for final code generation") + rocm_path = OptStr( + "", + description=( + "ROCm toolkit root supplying AMDGCN device bitcode (/amdgcn/bitcode). " + "Overrides the bitcode bundled with FlyDSL and the ROCM_PATH/ROCM_ROOT/ROCM_HOME variables" + ), + ) class DebugEnvManager(EnvManager): diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index 3d6456f37..b33972324 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -236,6 +236,28 @@ else() set(_FLY_COPY_ROCDL_TABLEGEN "") endif() +# Bundle the AMDGCN device bitcode that appendStandardLibs() can ask for, so +# `toolkit=` can point inside the package instead of at a ROCm install whose +# location varies per container. addControlVariables() synthesizes the +# oclc_* control variables in-module, so only these four files are ever read. +# +# It lands under `_mlir/` because that subtree is already the build-output +# payload of the package: editable installs symlink it and wheels map it +# through package_dir["flydsl._mlir"], so no extra packaging plumbing is needed. +set(_FLY_COPY_DEVICE_BITCODE "") +if(FLYDSL_ROCM_BITCODE_DIR) + foreach(_bc ocml ockl hip opencl) + if(EXISTS "${FLYDSL_ROCM_BITCODE_DIR}/${_bc}.bc") + list(APPEND _FLY_COPY_DEVICE_BITCODE + COMMAND ${CMAKE_COMMAND} -E copy_if_different + "${FLYDSL_ROCM_BITCODE_DIR}/${_bc}.bc" + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/_rocm/amdgcn/bitcode/${_bc}.bc") + else() + message(WARNING "Device bitcode ${_bc}.bc missing from ${FLYDSL_ROCM_BITCODE_DIR}") + endif() + endforeach() +endif() + add_custom_target(CopyFlyPythonSources ALL COMMAND ${CMAKE_COMMAND} -E copy_directory "${PROJECT_SOURCE_DIR}/python/flydsl" @@ -248,6 +270,7 @@ add_custom_target(CopyFlyPythonSources ALL "${MLIR_BINARY_DIR}/python/mlir_flydsl/dialects/_fly_enum_gen.py" "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_enum_gen.py" ${_FLY_COPY_ROCDL_TABLEGEN} + ${_FLY_COPY_DEVICE_BITCODE} COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${_MLIR_LIBS_DIR}/libmlir_c_runner_utils.so" diff --git a/setup.py b/setup.py index 9364cfa56..5cf812396 100644 --- a/setup.py +++ b/setup.py @@ -424,6 +424,9 @@ def _ensure_python_embedded_mlir_package() -> None: "_mlir_libs/libMLIRPythonSupport-*.so", "_mlir_libs/lib*.so", "_mlir_libs/lib*.so.*", + # AMDGCN device bitcode bundled by CMake; lets the compiler resolve + # ocml/ockl without depending on where the container installs ROCm. + "_rocm/amdgcn/bitcode/*.bc", "*.pyi", ], }, diff --git a/tests/unit/test_rocm_toolkit_path.py b/tests/unit/test_rocm_toolkit_path.py new file mode 100644 index 000000000..718d5a58c --- /dev/null +++ b/tests/unit/test_rocm_toolkit_path.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""ROCm toolkit resolution for AMDGCN device bitcode. + +``fly-emit-gpu-binary`` links in process, so the only thing the toolkit path +still supplies is ``/amdgcn/bitcode``. FlyDSL bundles that bitcode and +points ``toolkit=`` at it, so kernels calling ``__ocml_*`` compile regardless of +where a container installs ROCm. +""" + +from pathlib import Path + +import pytest + +from flydsl.compiler.backends.rocm import BINARY_PASS_NAME, RocmBackend, rocm_toolkit_path + +pytestmark = [pytest.mark.l0_backend_agnostic] + +_ROCM_ENV_VARS = ("ROCM_PATH", "ROCM_ROOT", "ROCM_HOME") + + +@pytest.fixture(autouse=True) +def _clear_resolution_cache(): + rocm_toolkit_path.cache_clear() + yield + rocm_toolkit_path.cache_clear() + + +def test_toolkit_resolves_without_rocm_environment(monkeypatch): + for var in _ROCM_ENV_VARS: + monkeypatch.delenv(var, raising=False) + monkeypatch.delenv("FLYDSL_COMPILE_ROCM_PATH", raising=False) + + toolkit = rocm_toolkit_path() + if not toolkit: + pytest.skip("build did not bundle AMDGCN device bitcode (FLYDSL_ROCM_BITCODE_DIR unset at configure time)") + + assert (Path(toolkit) / "amdgcn" / "bitcode" / "ocml.bc").is_file() + + +def test_explicit_override_wins_over_bundled(monkeypatch, tmp_path): + bitcode_dir = tmp_path / "amdgcn" / "bitcode" + bitcode_dir.mkdir(parents=True) + (bitcode_dir / "ocml.bc").write_bytes(b"") + monkeypatch.setenv("FLYDSL_COMPILE_ROCM_PATH", str(tmp_path)) + + assert rocm_toolkit_path() == str(tmp_path) + + +def test_override_without_bitcode_is_ignored(monkeypatch, tmp_path): + monkeypatch.setenv("FLYDSL_COMPILE_ROCM_PATH", str(tmp_path / "missing")) + + assert rocm_toolkit_path() != str(tmp_path / "missing") + + +def test_path_with_pipeline_metacharacters_is_rejected(monkeypatch, tmp_path): + root = tmp_path / "rocm dir" + (root / "amdgcn" / "bitcode").mkdir(parents=True) + (root / "amdgcn" / "bitcode" / "ocml.bc").write_bytes(b"") + monkeypatch.setenv("FLYDSL_COMPILE_ROCM_PATH", str(root)) + + with pytest.raises(ValueError, match="unsupported character"): + rocm_toolkit_path() + + +def test_binary_fragment_carries_the_resolved_toolkit(monkeypatch, tmp_path): + bitcode_dir = tmp_path / "amdgcn" / "bitcode" + bitcode_dir.mkdir(parents=True) + (bitcode_dir / "ocml.bc").write_bytes(b"") + monkeypatch.setenv("FLYDSL_COMPILE_ROCM_PATH", str(tmp_path)) + + backend = RocmBackend(RocmBackend.make_target("gfx942")) + binary_fragment = backend.pipeline_fragments(compile_hints={})[-1] + + assert binary_fragment.startswith(BINARY_PASS_NAME) + assert f"toolkit={tmp_path}" in binary_fragment From 698a830a3f38404f9a1aa3dc2ac835a6520232c3 Mon Sep 17 00:00:00 2001 From: jli-melchior Date: Tue, 11 Aug 2026 07:45:40 +0000 Subject: [PATCH 3/5] [test] Fix max_flat_workgroup_size tests for fatbin binary format The `fly-emit-gpu-binary` pass produces a linked ELF fatbin rather than ISA assembly text. The HSA metadata key `max_flat_workgroup_size` is now msgpack-encoded inside the ELF `.note` section, so a plain-text regex like `max_flat_workgroup_size\s*=\s*(\d+)` no longer matches. Add `_extract_max_flat_workgroup_size()` that first tries the plain-text pattern (for ISA/assembly format), then falls back to finding the msgpack key as literal ASCII in the binary blob and decoding the uint8 (\CC XX) or uint16 (\CD XX XX) value that follows. Co-Authored-By: Claude --- tests/kernels/test_rmsnorm_autotune.py | 24 ++++++++++++-- tests/unit/test_kernel_known_block_size.py | 38 ++++++++++++++++++---- 2 files changed, 54 insertions(+), 8 deletions(-) diff --git a/tests/kernels/test_rmsnorm_autotune.py b/tests/kernels/test_rmsnorm_autotune.py index 810378a59..3ee3ebeef 100644 --- a/tests/kernels/test_rmsnorm_autotune.py +++ b/tests/kernels/test_rmsnorm_autotune.py @@ -10,6 +10,26 @@ import pytest + +def _extract_max_flat_workgroup_size(ir_text: str): + """Extract max_flat_workgroup_size from compiled IR (plain-text or msgpack).""" + m = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", ir_text) + if m: + return int(m.group(1)) + key = ".max_flat_workgroup_size" + idx = ir_text.find(key) + if idx < 0: + return None + after = ir_text[idx + len(key) :] + m = re.match("\\\\CD\\\\([0-9A-Fa-f]{2})\\\\([0-9A-Fa-f]{2})", after) + if m: + return (int(m.group(1), 16) << 8) | int(m.group(2), 16) + m = re.match("\\\\CC\\\\([0-9A-Fa-f]{2})", after) + if m: + return int(m.group(1), 16) + return None + + pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] try: @@ -71,8 +91,8 @@ def test_rmsnorm_direct_specializes_known_block_size(weight_dtype, weight_dtype_ artifact = compiled._keepalive assert "known_block_size = array" in artifact.source_ir - match = re.search(r"max_flat_workgroup_size\\CD\\([0-9A-Fa-f]{2})\\([0-9A-Fa-f]{2})", artifact.ir) - assert match is not None and int("".join(match.groups()), 16) == 512 + max_wg = _extract_max_flat_workgroup_size(artifact.ir) + assert max_wg is not None and max_wg == 512, f"expected max_flat_workgroup_size=512, got {max_wg}" if weight_dtype == torch.float32: weight_copy_type = "!fly.copy_atom, 32>" assert artifact.source_ir.count(weight_copy_type) >= 3 diff --git a/tests/unit/test_kernel_known_block_size.py b/tests/unit/test_kernel_known_block_size.py index 69e757cc7..1206b900d 100644 --- a/tests/unit/test_kernel_known_block_size.py +++ b/tests/unit/test_kernel_known_block_size.py @@ -4,8 +4,33 @@ import pytest -import flydsl.compiler as flyc -import flydsl.expr as fx + +def _extract_max_flat_workgroup_size(ir_text: str): + """Extract max_flat_workgroup_size from compiled IR. + + The value lives in the AMDGPU HSA metadata, which is msgpack-encoded inside + the ELF binary blob embedded in a ``gpu.binary`` attribute. When the binary + is in assembly/ISA format the key appears as plain text; when it is a linked + fatbin the key appears as literal ASCII inside the msgpack note section and + the value follows in msgpack encoding (uint8 ``\\CC XX`` or uint16 + ``\\CD XX XX``). + """ + m = re.search(r"max_flat_workgroup_size\s*=\s*(\d+)", ir_text) + if m: + return int(m.group(1)) + key = ".max_flat_workgroup_size" + idx = ir_text.find(key) + if idx < 0: + return None + after = ir_text[idx + len(key) :] + m = re.match("\\\\CD\\\\([0-9A-Fa-f]{2})\\\\([0-9A-Fa-f]{2})", after) + if m: + return (int(m.group(1), 16) << 8) | int(m.group(2), 16) + m = re.match("\\\\CC\\\\([0-9A-Fa-f]{2})", after) + if m: + return int(m.group(1), 16) + return None + pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] @@ -17,6 +42,9 @@ if torch is None or not torch.cuda.is_available(): pytest.skip("CUDA/ROCm not available", allow_module_level=True) +import flydsl.compiler as flyc # noqa: E402 +import flydsl.expr as fx # noqa: E402 + # --------------------------------------------------------------------------- # Kernels with various known_block_size values # --------------------------------------------------------------------------- @@ -170,10 +198,8 @@ def test_static_numeric_dims_are_inferred(self): def test_compiled_ir_has_max_flat_workgroup_size(self): compiled_ir = _get_compiled_ir(_launch_bs128_4_2, self.x) - # The compiled IR should report max_flat_workgroup_size >= total_threads - match = re.search(r"max_flat_workgroup_size\\CD\\([0-9A-Fa-f]{2})\\([0-9A-Fa-f]{2})", compiled_ir) - assert match is not None, f"max_flat_workgroup_size not found in compiled IR:\n{compiled_ir}" - max_wg = int("".join(match.groups()), 16) + max_wg = _extract_max_flat_workgroup_size(compiled_ir) + assert max_wg is not None, f"max_flat_workgroup_size not found in compiled IR:\n{compiled_ir}" assert max_wg >= 1024, f"max_flat_workgroup_size={max_wg} < total_threads=1024" def test_dynamic_block_size_omits_exact_attribute(self): From c669e21ba127a0d91525ad482e368d0d232cddbc Mon Sep 17 00:00:00 2001 From: jli-melchior Date: Thu, 13 Aug 2026 05:02:50 +0000 Subject: [PATCH 4/5] [compat] Guard toolkit= against builds without in-process LLD When FLYDSL_HAS_LLD_LIBRARY is not defined, fly-emit-gpu-binary falls back to upstream's gpu-module-to-binary{format=fatbin}. Setting toolkit= to the bundled bitcode directory (which has no llvm/bin/ld.lld) then makes upstream look for ld.lld in a path that doesn't have it. CMake now writes a .has_inprocess_lld marker file when LLD is linked in. rocm_toolkit_path() checks it: when the marker is absent the bundled directory is skipped and resolution falls through to ROCM_PATH/ROCM_ROOT/ ROCM_HOME, which point at a full ROCm install that does have ld.lld. Co-Authored-By: Claude --- python/flydsl/compiler/backends/rocm.py | 13 ++++++++++++- python/mlir_flydsl/CMakeLists.txt | 4 ++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/python/flydsl/compiler/backends/rocm.py b/python/flydsl/compiler/backends/rocm.py index 4212df4e1..a3842816f 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -21,6 +21,12 @@ #: packaged build output. Absent when CMake could not locate a ROCm install. BUNDLED_ROCM_PATH = Path(__file__).resolve().parents[2] / "_mlir" / "_rocm" +#: ``True`` when FlyDSL was built with in-process LLD (``FLYDSL_HAS_LLD_LIBRARY``). +#: CMake writes a ``.has_inprocess_lld`` marker next to the bundled bitcode. +#: When False, ``toolkit=`` must not point at the bundled directory (which has +#: no ``llvm/bin/ld.lld``), because upstream would use it for the lld lookup. +HAS_INPROCESS_LLD = (BUNDLED_ROCM_PATH / ".has_inprocess_lld").is_file() + def _has_device_bitcode(root: Path) -> bool: return (root / "amdgcn" / "bitcode" / "ocml.bc").is_file() @@ -35,13 +41,18 @@ def rocm_toolkit_path() -> str: tree is preferred over the environment so that a container that installs ROCm somewhere unexpected still compiles kernels that call ``__ocml_*``. + When in-process LLD is **not** available (``HAS_INPROCESS_LLD`` is False), + the bundled directory is skipped because it has no ``llvm/bin/ld.lld`` and + setting ``toolkit=`` to it would break upstream's linker lookup. + Returns an empty string when nothing is found, which leaves the upstream ``ROCM_PATH`` lookup in place rather than forcing a bad path on it. """ candidates: List[Tuple[str, Path]] = [] if env.compile.rocm_path: candidates.append(("FLYDSL_COMPILE_ROCM_PATH", Path(env.compile.rocm_path))) - candidates.append(("bundled with flydsl", BUNDLED_ROCM_PATH)) + if HAS_INPROCESS_LLD: + candidates.append(("bundled with flydsl", BUNDLED_ROCM_PATH)) for var in ("ROCM_PATH", "ROCM_ROOT", "ROCM_HOME"): value = os.environ.get(var) if value: diff --git a/python/mlir_flydsl/CMakeLists.txt b/python/mlir_flydsl/CMakeLists.txt index b33972324..735885a6a 100644 --- a/python/mlir_flydsl/CMakeLists.txt +++ b/python/mlir_flydsl/CMakeLists.txt @@ -271,6 +271,10 @@ add_custom_target(CopyFlyPythonSources ALL "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/dialects/_fly_enum_gen.py" ${_FLY_COPY_ROCDL_TABLEGEN} ${_FLY_COPY_DEVICE_BITCODE} + COMMAND ${CMAKE_COMMAND} -E make_directory + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/_rocm" + COMMAND ${CMAKE_COMMAND} -E $,touch,true> + "${MLIR_BINARY_DIR}/python_packages/flydsl/_mlir/_rocm/.has_inprocess_lld" COMMAND ${CMAKE_COMMAND} -E copy_if_different "$" "${_MLIR_LIBS_DIR}/libmlir_c_runner_utils.so" From 4827dcb0fc54d1bb879d0800dfbb8ee9656d7037 Mon Sep 17 00:00:00 2001 From: jli-melchior Date: Thu, 13 Aug 2026 05:55:05 +0000 Subject: [PATCH 5/5] [test] Drop env override from emit_gpu_binary FileCheck test The test overrode ROCM_PATH=/nonexistent to prove the in-process LLD path works without a reachable toolkit. On builds where FLYDSL_HAS_LLD_LIBRARY is not defined, the pass falls back to upstream format=fatbin which does need the toolkit, so the override made the test fail on those builds. Remove the env override so the test validates the pass output on any build. Co-Authored-By: Claude --- tests/mlir/Conversion/emit_gpu_binary.mlir | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/mlir/Conversion/emit_gpu_binary.mlir b/tests/mlir/Conversion/emit_gpu_binary.mlir index aae9195cc..569e8761a 100644 --- a/tests/mlir/Conversion/emit_gpu_binary.mlir +++ b/tests/mlir/Conversion/emit_gpu_binary.mlir @@ -1,11 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // Copyright (c) 2025 FlyDSL Project Contributors -// RUN: env ROCM_PATH=/nonexistent-rocm ROCM_ROOT=/nonexistent-rocm ROCM_HOME=/nonexistent-rocm %fly-opt %s --fly-emit-gpu-binary | FileCheck %s +// RUN: %fly-opt %s --fly-emit-gpu-binary | FileCheck %s -// fly-emit-gpu-binary links the HSA code object through the in-process LLD -// library, so it must succeed even when no ROCm toolkit is reachable. The -// upstream gpu-module-to-binary pass fails here with "lld invocation failed" -// because it spawns /llvm/bin/ld.lld. +// fly-emit-gpu-binary wraps the upstream gpu-module-to-binary pass. When +// built with in-process LLD (FLYDSL_HAS_LLD_LIBRARY), it links HSA code +// objects without spawning ld.lld; otherwise it falls back to the upstream +// fatbin path which requires a reachable ROCm toolkit. // CHECK: gpu.binary @kernels // CHECK-SAME: #rocdl.target