Commit 8568ef5

mo khan <mo@mokhan.ca>
2026-09-08 02:35:11
add HIP/ROCm backend auto-detection, fix shim link fallback and cp
Adds an auto-detected HIP backend (via rocminfo/hipcc) alongside cuda/vulkan/cpu, ordered before vulkan since ROCm outperforms it on supported GPUs. Also moves the shim link step inside the build/rescue block so a link failure triggers the same CPU fallback a cmake failure does, and replaces the `sh -c "cp -a ..."` step with FileUtils to avoid shelling out and breaking on paths with spaces.
1 parent cb86eaa
Changed files (1)
ext
elelem
ext/elelem/llama/extconf.rb
@@ -8,17 +8,10 @@ VENDOR    = File.join(GEM_ROOT, "vendor", "llama.cpp")
 BUILD_DIR = File.join(EXT_DIR, "build")
 NATIVE    = File.join(GEM_ROOT, "ext", "elelem", "llama", "native")
 STAMP     = File.join(BUILD_DIR, ".elelem_backend") # last backend built here
-VALID_BACKENDS = %w[cpu vulkan cuda].freeze
-
-GPU_FLAGS = {
-  "cpu"    => [],                     # GGML_NATIVE is ON by default -> host CPU SIMD
-  "vulkan" => ["-DGGML_VULKAN=ON"],
-  "cuda"   => ["-DGGML_CUDA=ON"]
-}.freeze
-
+VALID_BACKENDS = %w[cpu vulkan cuda hip].freeze
 
 if RbConfig::CONFIG["host_os"] =~ /darwin/
-  abort "elelem: macOS/Metal is not supported yet -- Linux (CPU/CUDA/Vulkan) only."
+  abort "elelem: macOS/Metal is not supported yet -- Linux (CPU/CUDA/Vulkan/HIP) only."
 end
 
 class BuildError < StandardError; end
@@ -42,6 +35,29 @@ def vulkan_loader?
   end
 end
 
+# Most specific gfxNNNN target reported by rocminfo (e.g. "gfx1151"),
+# skipping the generic "gfxNN" family aliases it also prints.
+def hip_gfx_target
+  out = `rocminfo 2>/dev/null`
+  targets = out.scan(/gfx[0-9a-f]+/).uniq
+  targets.max_by(&:length)
+rescue Errno::ENOENT
+  nil
+end
+
+def hip_available?
+  which("hipcc") && !hip_gfx_target.nil?
+end
+
+def gpu_flags(backend)
+  {
+    "cpu"    => [],                     # GGML_NATIVE is ON by default -> host CPU SIMD
+    "vulkan" => ["-DGGML_VULKAN=ON"],
+    "cuda"   => ["-DGGML_CUDA=ON"],
+    "hip"    => ["-DGGML_HIP=ON", "-DCMAKE_HIP_ARCHITECTURES=#{ENV.fetch("ELELEM_HIP_ARCH") { hip_gfx_target || "gfx1100" }}"]
+  }.fetch(backend) { abort "elelem: unknown backend #{backend.inspect}" }
+end
+
 def detect_backend
   forced = ENV["ELELEM_LLAMA_BACKEND"].to_s.strip.downcase
   unless forced.empty? || forced == "auto" || VALID_BACKENDS.include?(forced)
@@ -50,12 +66,13 @@ def detect_backend
   return forced unless forced.empty? || forced == "auto"
 
   return "cuda" if which("nvcc")
+  return "hip" if hip_available?
   return "vulkan" if which("glslc") && vulkan_loader?
   "cpu"
 end
 
 def configure_and_build(backend)
-  flags = GPU_FLAGS.fetch(backend) { abort "elelem: unknown backend #{backend.inspect}" }
+  flags = gpu_flags(backend)
   run(
     "cmake", "-S", VENDOR, "-B", BUILD_DIR,
     "-DCMAKE_BUILD_TYPE=Release",
@@ -85,33 +102,36 @@ if File.exist?(STAMP) && File.read(STAMP).strip != backend
   FileUtils.rm_rf(BUILD_DIR)
 end
 
+def link_shim(libdir)
+  FileUtils.mkdir_p(NATIVE)
+  FileUtils.rm_f(Dir.glob(File.join(NATIVE, "*.so*")))
+  FileUtils.cp(Dir.glob(File.join(libdir, "*.so*")), NATIVE)
+
+  cxx = ENV["CXX"] || "c++"
+  run(cxx, "-std=c++17", "-O2", "-Wall", "-Wextra", "-shared", "-fPIC",
+      "-I", File.join(VENDOR, "include"),
+      "-I", File.join(VENDOR, "ggml", "include"),
+      "-I", File.join(VENDOR, "common"),   # chat.h (common_chat)
+      "-I", File.join(VENDOR, "vendor"),   # nlohmann/json.hpp
+      File.join(EXT_DIR, "elelem.cpp"),
+      "-o", File.join(NATIVE, "libelelem_llama.so"),
+      "-L", libdir, "-lllama-common", "-lllama",
+      "-Wl,-rpath,$ORIGIN")
+end
+
 begin
   warn "elelem: building llama.cpp backend=#{backend}"
   configure_and_build(backend)
+  link_shim(File.join(BUILD_DIR, "bin"))
 rescue BuildError => e
   raise if backend == "cpu"
   warn "elelem: #{backend} build failed (#{e.message}); falling back to CPU"
   FileUtils.rm_rf(BUILD_DIR)
   backend = "cpu"
   configure_and_build(backend)
+  link_shim(File.join(BUILD_DIR, "bin"))
 end
 
-FileUtils.mkdir_p(NATIVE)
-FileUtils.rm_f(Dir.glob(File.join(NATIVE, "*.so*")))
-libdir = File.join(BUILD_DIR, "bin")
-run("sh", "-c", "cp -a #{libdir}/*.so* #{NATIVE}/")
-
-cxx = ENV["CXX"] || "c++"
-run(cxx, "-std=c++17", "-O2", "-Wall", "-Wextra", "-shared", "-fPIC",
-    "-I", File.join(VENDOR, "include"),
-    "-I", File.join(VENDOR, "ggml", "include"),
-    "-I", File.join(VENDOR, "common"),   # chat.h (common_chat)
-    "-I", File.join(VENDOR, "vendor"),   # nlohmann/json.hpp
-    File.join(EXT_DIR, "elelem.cpp"),
-    "-o", File.join(NATIVE, "libelelem_llama.so"),
-    "-L", libdir, "-lllama-common", "-lllama",
-    "-Wl,-rpath,$ORIGIN")
-
 File.write(File.join(NATIVE, "backend"), backend)
 File.write(File.join(EXT_DIR, "Makefile"), <<~MAKE)
   all: