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 c22e9cf03..75c0e5910 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -11,6 +11,44 @@ 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() + +# 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 bb6ca7954..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) @@ -254,7 +258,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..842eb5cc8 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,34 @@ 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`/`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/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..1814618ae --- /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 (!object || !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..a3842816f 100644 --- a/python/flydsl/compiler/backends/rocm.py +++ b/python/flydsl/compiler/backends/rocm.py @@ -1,12 +1,98 @@ # 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 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" + +#: 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" + +#: ``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() + + +@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_*``. + + 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))) + 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: + 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).""" @@ -49,7 +135,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 +192,15 @@ 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)}"' + 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 = ( + 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 +208,17 @@ 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 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.""" 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/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..735885a6a 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,11 @@ 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 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" 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/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/mlir/Conversion/emit_gpu_binary.mlir b/tests/mlir/Conversion/emit_gpu_binary.mlir new file mode 100644 index 000000000..569e8761a --- /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: %fly-opt %s --fly-emit-gpu-binary | FileCheck %s + +// 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 +// 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(): 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): 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