[WIP][Fix] Fix lld invocation failed when ROCm is not at the baked-in path

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` |

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 <toolkit>/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=<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=<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}")

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=<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                                   │
    └──────────────────────────────────────────────────────────┘
         │
         ▼

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 `<toolkit>/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
+`<toolkit>/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=<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

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.

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

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 `<toolkit>/llvm/bin/ld.lld`, where `<toolkit>` 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 = [{

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();
 }

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()