[executorch][native-ggml] Run the ggml engine on any linked GPU backend, including HIP

backends/native/runtime/engine/ggml/GgmlEngine.cpp

@@ -6,6 +6,8 @@
 
 #include <executorch/backends/native/runtime/engine/ggml/GgmlEngine.h>
 
+#include <algorithm>
+#include <cctype>
 #include <memory>
 #include <stdexcept>
 #include <string>
@@ -30,10 +32,19 @@ struct GgmlEngineHost::Impl {
     if (options.graph_size == 0) {
       throw std::invalid_argument("ggml: graph size must be positive");
     }
-    name = options.device == GgmlDevice::Cuda ? "ggml-cuda" : "ggml-cpu";
-    device_name = options.device == GgmlDevice::Cuda
-        ? "CUDA device " + std::to_string(options.device_index)
-        : "CPU";
+    if (options.device == GgmlDevice::Gpu) {
+      ggml_backend_dev_t device = select_ggml_gpu_device(options);
+      name = std::string("ggml-") +
+          ggml_backend_reg_name(ggml_backend_dev_backend_reg(device));
+      std::transform(name.begin(), name.end(), name.begin(), [](char c) {
+        return static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
+      });
+      device_name = std::string(ggml_backend_dev_description(device)) + " (" +
+          ggml_backend_dev_name(device) + ")";
+    } else {
+      name = "ggml-cpu";
+      device_name = "CPU";
+    }
   }
 
   GgmlEngineOptions options;

backends/native/runtime/engine/ggml/GgmlEngine.h

@@ -16,7 +16,9 @@ namespace ptn {
 
 enum class GgmlDevice {
   Cpu,
-  Cuda,
+  // The GPU backend is whichever one the linked ggml variant registers (CUDA,
+  // HIP, ...); device_index counts GPU devices only.
+  Gpu,
 };
 
 struct GgmlEngineOptions {

backends/native/runtime/engine/ggml/GgmlExecutor.cpp

@@ -31,9 +31,6 @@
 #include <ggml-cpp.h>
 #include <ggml-cpu.h>
 #include <ggml.h>
-#ifdef PTN_GGML_USE_CUDA
-#include <ggml-cuda.h>
-#endif
 
 #include <executorch/backends/native/runtime/Method.h>
 #include <executorch/backends/native/runtime/deserialize/ByteSpan.h>
@@ -611,6 +608,12 @@ class GgmlExecutable final : public EngineExecutable {
 
     ggml_backend_sched_reset(scheduler_.get());
     GraphInstance instance = build_graph();
+    // The scheduler otherwise keeps graph inputs on the CPU backend and copies
+    // them into GPU splits once per consuming view.
+    for (const Upload& upload : instance.uploads) {
+      ggml_backend_sched_set_tensor_backend(
+          scheduler_.get(), upload.tensor, primary_backend_);
+    }
     if (!ggml_backend_sched_alloc_graph(scheduler_.get(), instance.graph)) {
       ggml_backend_sched_reset(scheduler_.get());
       throw std::runtime_error("ggml: failed to allocate execution graph");
@@ -714,23 +717,17 @@ class GgmlExecutable final : public EngineExecutable {
   }
 
   void initialize_backends() {
-#ifdef PTN_GGML_USE_CUDA
-    if (options_.device == GgmlDevice::Cuda) {
-      if (options_.device_index >= ggml_backend_cuda_get_device_count()) {
-        throw std::runtime_error("ggml: requested CUDA device is unavailable");
-      }
-      ggml_backend_ptr cuda(ggml_backend_cuda_init(options_.device_index));
-      if (!cuda) {
-        throw std::runtime_error("ggml: failed to initialize CUDA backend");
+    if (options_.device == GgmlDevice::Gpu) {
+      ggml_backend_dev_t device = select_ggml_gpu_device(options_);
+      ggml_backend_ptr gpu(ggml_backend_dev_init(device, nullptr));
+      if (!gpu) {
+        throw std::runtime_error(
+            std::string("ggml: failed to initialize ") +
+            ggml_backend_dev_name(device));
       }
-      primary_backend_ = cuda.get();
-      backends_.push_back(std::move(cuda));
+      primary_backend_ = gpu.get();
+      backends_.push_back(std::move(gpu));
     }
-#else
-    if (options_.device == GgmlDevice::Cuda) {
-      throw std::runtime_error("ggml: this target was built without CUDA");
-    }
-#endif
     ggml_backend_ptr cpu(ggml_backend_cpu_init());
     if (!cpu) {
       throw std::runtime_error("ggml: failed to initialize CPU backend");
@@ -2509,6 +2506,25 @@ GraphInstance GgmlExecutable::build_graph() {
 
 } // namespace
 
+ggml_backend_dev_t select_ggml_gpu_device(const GgmlEngineOptions& options) {
+  int remaining = options.device_index;
+  for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+    ggml_backend_dev_t device = ggml_backend_dev_get(i);
+    const enum ggml_backend_dev_type type = ggml_backend_dev_type(device);
+    if (type != GGML_BACKEND_DEVICE_TYPE_GPU &&
+        type != GGML_BACKEND_DEVICE_TYPE_IGPU) {
+      continue;
+    }
+    if (remaining-- == 0) {
+      return device;
+    }
+  }
+  throw std::runtime_error(
+      "ggml: GPU device " + std::to_string(options.device_index) +
+      " is unavailable (" + std::to_string(options.device_index - remaining) +
+      " GPU devices registered by the linked ggml backends)");
+}
+
 // cppcheck-suppress unusedFunction
 std::unique_ptr<EngineExecutable> make_ggml_executable(
     const Method& method,

backends/native/runtime/engine/ggml/GgmlExecutor.h

@@ -11,6 +11,8 @@
 #include <string>
 #include <unordered_map>
 
+#include <ggml-backend.h>
+
 #include <executorch/backends/native/runtime/engine/Engine.h>
 #include <executorch/backends/native/runtime/engine/ggml/GgmlEngine.h>
 
@@ -19,6 +21,10 @@ namespace ptn {
 struct Method;
 class Package;
 
+// Returns the options.device_index-th GPU device in ggml's registry. Throws if
+// the linked ggml variant registers no such device.
+ggml_backend_dev_t select_ggml_gpu_device(const GgmlEngineOptions& options);
+
 // Mutable buffer key -> method whose live executable binds it.
 using GgmlMutableKeyOwners = std::unordered_map<std::string, std::string>;
 

backends/native/runtime/engine/targets.bzl

@@ -21,55 +21,36 @@ def define_common_targets():
         visibility = ["//executorch/backends/native/..."],
     )
 
-    runtime.cxx_library(
-        name = "ggml_engine_cpu",
-        srcs = [
-            "ggml/GgmlEngine.cpp",
-            "ggml/GgmlExecutor.cpp",
-        ],
-        headers = ["ggml/GgmlExecutor.h"],
-        exported_headers = ["ggml/GgmlEngine.h"],
-        exported_deps = [":engine"],
-        deps = [
-            "//executorch/backends/native/runtime:method",
-            "//executorch/backends/native/runtime:runtime",
-            "//executorch/backends/native/runtime/deserialize:byte_span",
-            "//executorch/backends/native/runtime/deserialize:package",
-            "//executorch/backends/native/runtime/graph:argument",
-            "//executorch/backends/native/runtime/graph:graph",
-            "//executorch/backends/native/runtime/graph:ids",
-            "//executorch/backends/native/runtime/graph:node",
-            "//executorch/backends/native/runtime/graph:scalar_type",
-            "//executorch/backends/native/runtime/graph:tensor_meta",
-            "//executorch/backends/native/runtime/graph:value",
-            "fbsource//third-party/llama.cpp:ggml_cpu",
-        ],
-        visibility = ["//executorch/backends/native/..."],
-    )
-
-    runtime.cxx_library(
-        name = "ggml_engine_cuda",
-        srcs = [
-            "ggml/GgmlEngine.cpp",
-            "ggml/GgmlExecutor.cpp",
-        ],
-        headers = ["ggml/GgmlExecutor.h"],
-        exported_headers = ["ggml/GgmlEngine.h"],
-        exported_deps = [":engine"],
-        preprocessor_flags = ["-DPTN_GGML_USE_CUDA=1"],
-        deps = [
-            "//executorch/backends/native/runtime:method",
-            "//executorch/backends/native/runtime:runtime",
-            "//executorch/backends/native/runtime/deserialize:byte_span",
-            "//executorch/backends/native/runtime/deserialize:package",
-            "//executorch/backends/native/runtime/graph:argument",
-            "//executorch/backends/native/runtime/graph:graph",
-            "//executorch/backends/native/runtime/graph:ids",
-            "//executorch/backends/native/runtime/graph:node",
-            "//executorch/backends/native/runtime/graph:scalar_type",
-            "//executorch/backends/native/runtime/graph:tensor_meta",
-            "//executorch/backends/native/runtime/graph:value",
-            "fbsource//third-party/llama.cpp:ggml_cuda",
-        ],
-        visibility = ["//executorch/backends/native/..."],
-    )
+    # The engine source is backend-neutral: it picks devices from ggml's registry,
+    # so each variant differs only in which ggml backend set it links.
+    for variant, ggml_dep in [
+        ("cpu", "fbsource//third-party/llama.cpp:ggml_cpu"),
+        ("cuda", "fbsource//third-party/llama.cpp:ggml_cuda"),
+        ("hip", "fbsource//third-party/llama.cpp:ggml_hip"),
+    ]:
+        runtime.cxx_library(
+            name = "ggml_engine_" + variant,
+            srcs = [
+                "ggml/GgmlEngine.cpp",
+                "ggml/GgmlExecutor.cpp",
+            ],
+            headers = ["ggml/GgmlExecutor.h"],
+            exported_headers = ["ggml/GgmlEngine.h"],
+            exported_deps = [":engine"],
+            deps = [
+                "//executorch/backends/native/runtime:method",
+                "//executorch/backends/native/runtime:runtime",
+                "//executorch/backends/native/runtime/deserialize:byte_span",
+                "//executorch/backends/native/runtime/deserialize:owned_bytes",
+                "//executorch/backends/native/runtime/deserialize:package",
+                "//executorch/backends/native/runtime/graph:argument",
+                "//executorch/backends/native/runtime/graph:graph",
+                "//executorch/backends/native/runtime/graph:ids",
+                "//executorch/backends/native/runtime/graph:node",
+                "//executorch/backends/native/runtime/graph:scalar_type",
+                "//executorch/backends/native/runtime/graph:tensor_meta",
+                "//executorch/backends/native/runtime/graph:value",
+                ggml_dep,
+            ],
+            visibility = ["//executorch/backends/native/..."],
+        )

backends/native/test/runtime/ggml/BUCK

@@ -1,4 +1,5 @@
 load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target")
+load("@fbcode_macros//build_defs/lib:re_test_utils.bzl", "re_test_utils")
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 oncall("executorch")
@@ -18,29 +19,70 @@ fbcode_target(
     ],
 )
 
-fbcode_target(
-    _kind = runtime.cxx_binary,
-    name = "ggml_ptn_runner",
-    srcs = ["ggml_ptn_runner.cpp"],
-    deps = [
-        "//executorch/backends/native/runtime:runtime",
-        "//executorch/backends/native/runtime/deserialize:package",
-        "//executorch/backends/native/runtime/engine:engine",
-        "//executorch/backends/native/runtime/engine:ggml_engine_cpu",
-        "//executorch/backends/native/runtime/graph:scalar_type",
-    ],
-    external_deps = ["gflags"],
-)
+# The op suite runs once per ggml backend set; only the engine variant, the
+# runner's --device, link deps and resources, and the test host differ.
+_VARIANTS = [
+    {"suffix": "", "engine": "ggml_engine_cpu", "device": "cpu"},
+    {
+        "suffix": "_cuda",
+        "engine": "ggml_engine_cuda",
+        "device": "gpu",
+        "remote_execution": re_test_utils.remote_execution(
+            platform = "gpu-remote-execution",
+        ),
+        # The CUDA backend is a shared library that needs the ggml core shared
+        # beside it, and its VMM pool calls the CUDA driver API.
+        "deps": ["fbsource//third-party/llama.cpp:ggml_base[prefer-shared]"],
+        "external_deps": [("cuda", None, "cuda_driver_api-lazy")],
+        # Buck does not put the CUDA driver stub's library on the runner's
+        # loader path, so it ships beside the runner.
+        "resources": {
+            "//tools/lazify:lazy_stub_binder[shared]": "libtools_lazify_lazy_stub_binder.so",
+        },
+    },
+    {
+        "suffix": "_hip",
+        "engine": "ggml_engine_hip",
+        "device": "gpu",
+        "remote_execution": re_test_utils.remote_execution(
+            platform = "amd-gpu",
+            subplatform = "MI300",
+        ),
+    },
+]
 
-fbcode_target(
-    _kind = runtime.python_test,
-    name = "test_ggml_ops",
-    srcs = ["test_ggml_ops.py"],
-    resources = {":ggml_ptn_runner": "ggml_ptn_runner"},
-    deps = [
-        "//caffe2:torch",
-        "//executorch/exir/native:native",
-        "//pytorch/ao:torchao",
-        "fbsource//third-party/pypi/transformers:transformers",
-    ],
-)
+[
+    fbcode_target(
+        _kind = runtime.cxx_binary,
+        name = "ggml_ptn_runner" + v["suffix"],
+        srcs = ["ggml_ptn_runner.cpp"],
+        deps = [
+            "//executorch/backends/native/runtime:runtime",
+            "//executorch/backends/native/runtime/deserialize:package",
+            "//executorch/backends/native/runtime/engine:engine",
+            "//executorch/backends/native/runtime/engine:" + v["engine"],
+            "//executorch/backends/native/runtime/graph:scalar_type",
+        ] + v.get("deps", []),
+        external_deps = ["gflags"] + v.get("external_deps", []),
+    )
+    for v in _VARIANTS
+]
+
+[
+    fbcode_target(
+        _kind = runtime.python_test,
+        name = "test_ggml_ops" + v["suffix"],
+        srcs = ["test_ggml_ops.py"],
+        env = {"GGML_PTN_DEVICE": v["device"]},
+        remote_execution = v.get("remote_execution"),
+        resources = {":ggml_ptn_runner" + v["suffix"]: "ggml_ptn_runner"} |
+                    v.get("resources", {}),
+        deps = [
+            "//caffe2:torch",
+            "//executorch/exir/native:native",
+            "//pytorch/ao:torchao",
+            "fbsource//third-party/pypi/transformers:transformers",
+        ],
+    )
+    for v in _VARIANTS
+]

backends/native/test/runtime/ggml/ggml_ptn_runner.cpp

@@ -4,14 +4,14 @@
 // This source code is licensed under the BSD-style license found in the
 // LICENSE file in the root directory of this source tree.
 
-// Runs one method of a .ptn on the ggml CPU engine. Inputs and outputs are raw
+// Runs one method of a .ptn on the ggml engine. Inputs and outputs are raw
 // little-endian bytes in the method's own dtypes. Semicolons separate steps,
 // which run in order on one executable, so mutable buffers such as KV caches
 // carry over. Prints one line of space-separated sizes per output per step.
 //
 // Usage:
 //   ggml_ptn_runner --package M.ptn --inputs a.bin,b.bin --outputs out0.bin
-//                   [--input_sizes 1x6,6]
+//                   [--input_sizes 1x6,6] [--device gpu]
 //   ggml_ptn_runner --package M.ptn --inputs a0.bin,b0.bin;a1.bin,b1.bin
 //                   --outputs out0.bin;out1.bin --input_sizes 1x6,6;1x1,1
 //
@@ -56,6 +56,11 @@ DEFINE_string(
     "Comma-separated input shapes such as 1x6, in graph order, for methods "
     "with dynamic shapes; semicolons separate steps. Defaults to the "
     "serialized shapes.");
+DEFINE_string(
+    device,
+    "cpu",
+    "ggml device: cpu, or gpu for the GPU backend this binary links.");
+DEFINE_int32(cpu_threads, 0, "ggml CPU threads; 0 keeps ggml's default.");
 
 namespace {
 
@@ -166,7 +171,14 @@ void run() {
   const ptn::ByteSpan program_bytes = package->program_bytes();
   const auto program = std::make_shared<const ptn::Program>(
       ptn::Program::load(program_bytes.data(), program_bytes.size()));
-  const auto host = ptn::GgmlEngineHost::create();
+  ptn::GgmlEngineOptions options;
+  if (FLAGS_device == "gpu") {
+    options.device = ptn::GgmlDevice::Gpu;
+  } else if (FLAGS_device != "cpu") {
+    throw std::invalid_argument("--device must be cpu or gpu");
+  }
+  options.cpu_threads = FLAGS_cpu_threads;
+  const auto host = ptn::GgmlEngineHost::create(options);
   const auto context = host->create_ggml_context(program, package);
   const auto executable = context->compile(FLAGS_method);
 

backends/native/test/runtime/ggml/test_ggml_engine.cpp

@@ -1054,5 +1054,11 @@ TEST(GgmlEngineTest, RejectsEmbeddingIndicesNotFromAnInput) {
   }
 }
 
+TEST(GgmlEngineTest, GpuUnavailableOnCpuVariant) {
+  GgmlEngineOptions options;
+  options.device = GgmlDevice::Gpu;
+  EXPECT_THROW(GgmlEngineHost::create(options), std::runtime_error);
+}
+
 } // namespace
 } // namespace ptn