Arm backend: Fix VKML init and env for MacOS (pytorch#18286)

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell

Signed-off-by: Baris Demir <baris.demir@arm.com>
Co-authored-by: Baris Demir <baris.demir@arm.com>

Author: bdemirb

backends/arm/runtime/VGFBackend.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright 2025 Arm Limited and/or its affiliates.
+ * Copyright 2025-2026 Arm Limited and/or its affiliates.
  *
  * This source code is licensed under the BSD-style license found in the
  * LICENSE file in the root directory of this source tree.
@@ -75,7 +75,9 @@ void vkml_free_basics(
     VkInstance* instance,
     VkDevice* device,
     VkCommandPool* command_pool) {
-  vkDestroyCommandPool(*device, *command_pool, nullptr);
+  if (*device != VK_NULL_HANDLE && *command_pool != VK_NULL_HANDLE) {
+    vkDestroyCommandPool(*device, *command_pool, nullptr);
+  }
   // Note: These primitives are used by the emulation layer for vulkan
   //       object allocation, the vulkan objects are freed in in library
   //       shutdown, so we can't yet destroy these here without causing
@@ -111,21 +113,19 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
           result);
       return;
     }
+
+    is_initialized_ = true;
   }
   ~VGFBackend() {
     vkml_free_basics(&vk_instance, &vk_device, &vk_command_pool);
   }
 
   bool is_available() const override {
-    VkResult result;
-
     ET_LOG(Info, "Checking VGFBackend is available");
-    // Query the device prepared in constructor for needed extensions
-    result = vkml_load_extensions(&vk_device);
-    if (result != VK_SUCCESS)
+    if (!is_initialized_) {
       return false;
-
-    return true;
+    }
+    return vkml_load_extensions(&vk_device) == VK_SUCCESS;
   }
 
   Result<DelegateHandle*> init(
@@ -134,6 +134,13 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
       ArrayRef<CompileSpec> compile_specs) const override {
     ET_LOG(Info, "Entered VGF init");
 
+    if (!is_initialized_) {
+      ET_LOG(
+          Error,
+          "VGF backend is unavailable because Vulkan initialization failed");
+      return Error::NotSupported;
+    }
+
     const char* vgf_data = reinterpret_cast<const char*>(processed->data());
 
     MemoryAllocator* allocator = context.get_runtime_allocator();
@@ -230,11 +237,12 @@ class VGFBackend final : public ::executorch::runtime::BackendInterface {
   }
 
  private:
-  VkInstance vk_instance;
-  VkPhysicalDevice vk_physical_device;
-  VkDevice vk_device;
-  VkQueue vk_queue;
-  VkCommandPool vk_command_pool;
+  VkInstance vk_instance = VK_NULL_HANDLE;
+  VkPhysicalDevice vk_physical_device = VK_NULL_HANDLE;
+  VkDevice vk_device = VK_NULL_HANDLE;
+  VkQueue vk_queue = VK_NULL_HANDLE;
+  VkCommandPool vk_command_pool = VK_NULL_HANDLE;
+  bool is_initialized_ = false;
 };
 
 namespace {
@@ -253,6 +261,7 @@ VkResult vkml_allocate_basics(
 
   if (VK_SUCCESS != volkInitialize()) {
     ET_LOG(Error, "Volk failed to initialize");
+    return VK_ERROR_INITIALIZATION_FAILED;
   }
 
   VkApplicationInfo app_info{

backends/arm/test/runner_utils.py

@@ -10,6 +10,7 @@
 import re
 import shutil
 import subprocess  # nosec B404 - invoked only for trusted toolchain binaries
+import sys
 import tempfile
 from pathlib import Path
 
@@ -364,7 +365,7 @@ def run_vkml_emulation_layer(
         cmd_line += input_string
     cmd_line = cmd_line.split()
 
-    result = _run_cmd(cmd_line)
+    result = _run_cmd(cmd_line, env=_get_vkml_runtime_env())
 
     # TODO: Add regex to check for error or fault messages in stdout from Emulation Layer
     result_stdout = result.stdout.decode()  # noqa: F841
@@ -590,7 +591,80 @@ def save_bytes(
     return file_path
 
 
-def _run_cmd(cmd: List[str], check=True) -> subprocess.CompletedProcess[bytes]:
+def _prepend_env_path(existing: str | None, value: str) -> str:
+    if not existing:
+        return value
+
+    parts = [part for part in existing.split(os.path.pathsep) if part]
+    if value in parts:
+        return existing
+    return os.path.pathsep.join([value, *parts])
+
+
+def _find_local_vulkan_sdk_root() -> Path | None:
+    repo_root = Path(__file__).resolve().parents[3]
+    sdk_base_dir = repo_root / "examples/arm/arm-scratch/vulkan_sdk"
+    if not sdk_base_dir.is_dir():
+        return None
+
+    if sys.platform == "darwin":
+        candidates = sorted(
+            path for path in sdk_base_dir.glob("*/macOS") if path.is_dir()
+        )
+    else:
+        arch = os.uname().machine
+        arch_aliases = [arch]
+        if arch == "arm64":
+            arch_aliases.append("aarch64")
+        candidates = sorted(
+            path
+            for alias in arch_aliases
+            for path in sdk_base_dir.glob(f"*/{alias}")
+            if path.is_dir()
+        )
+
+    if not candidates:
+        return None
+    return candidates[-1]
+
+
+def _get_vkml_runtime_env() -> dict[str, str]:
+    """Return an environment with the Vulkan runtime variables needed for
+    VKML.
+    """
+    env = os.environ.copy()
+    sdk_root = _find_local_vulkan_sdk_root()
+    if sdk_root is None:
+        return env
+
+    env["VULKAN_SDK"] = str(sdk_root)
+    env["PATH"] = _prepend_env_path(env.get("PATH"), str(sdk_root / "bin"))
+
+    if sys.platform == "darwin":
+        env["DYLD_LIBRARY_PATH"] = _prepend_env_path(
+            env.get("DYLD_LIBRARY_PATH"), str(sdk_root / "lib")
+        )
+        moltenvk_icd = sdk_root / "share/vulkan/icd.d/MoltenVK_icd.json"
+        if moltenvk_icd.is_file():
+            env["VK_DRIVER_FILES"] = _prepend_env_path(
+                env.get("VK_DRIVER_FILES"), str(moltenvk_icd)
+            )
+        else:
+            logger.debug(
+                "MoltenVK ICD file not found at %s; leaving VK_DRIVER_FILES unset.",
+                moltenvk_icd,
+            )
+    else:
+        env["LD_LIBRARY_PATH"] = _prepend_env_path(
+            env.get("LD_LIBRARY_PATH"), str(sdk_root / "lib")
+        )
+
+    return env
+
+
+def _run_cmd(
+    cmd: List[str], check=True, env: dict[str, str] | None = None
+) -> subprocess.CompletedProcess[bytes]:
     """Run a command and check for errors.
 
     Args:
@@ -599,7 +673,7 @@ def _run_cmd(cmd: List[str], check=True) -> subprocess.CompletedProcess[bytes]:
     """
     try:
         result = subprocess.run(  # nosec B603 - cmd constructed from trusted inputs
-            cmd, check=check, capture_output=True
+            cmd, check=check, capture_output=True, env=env
         )
         return result
     except subprocess.CalledProcessError as e:

examples/arm/setup.sh

@@ -200,6 +200,12 @@ function create_setup_path(){
     clear_setup_path
     log_step "path" "Generating setup path scripts at ${setup_path_script}"
 
+    if [[ -n "${VIRTUAL_ENV:-}" && -d "${VIRTUAL_ENV}/bin" ]]; then
+        prepend_env_in_setup_path PATH "${VIRTUAL_ENV}/bin"
+    elif [[ -d "${et_dir}/env/bin" ]]; then
+        prepend_env_in_setup_path PATH "${et_dir}/env/bin"
+    fi
+
     local use_mlsdk_pip=0
     if use_mlsdk_pip_package; then
         use_mlsdk_pip=1