Arm Backend: Improve VGF runtime and update MLSDK setup (pytorch#13155)

* Adds portability extensions to instance_info
 * Explicitly sets poolSizes in descriptor pool
* Update Arm MLSDK manifest url to default to newly publish repo during
setup.sh
 * Add new exports to setup_path.sh for mlsdk components

Signed-off-by: Ryan O'Shea <[email protected]>
Co-authored-by: Ryan O'Shea <[email protected]>

Author: 2 people

backends/arm/runtime/VGFBackend.cpp

@@ -264,15 +264,60 @@ VkResult vkml_allocate_basics(
       .engineVersion = 0,
       .apiVersion = VK_API_VERSION_1_3,
   };
+
+  std::vector<const char*> requested_extensions;
+  VkInstanceCreateFlags instance_flags = 0;
+
+#ifdef __APPLE__
+  instance_flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;
+
+  uint32_t extension_count = 0;
+  result = vkEnumerateInstanceExtensionProperties(
+      nullptr, &extension_count, nullptr);
+
+  if (result != VK_SUCCESS) {
+    ET_LOG(Error, "Failed to enumerate instance extensions");
+    return result;
+  }
+
+  std::vector<VkExtensionProperties> extension_properties(extension_count);
+  result = vkEnumerateInstanceExtensionProperties(
+      nullptr, &extension_count, extension_properties.data());
+
+  if (result != VK_SUCCESS) {
+    ET_LOG(Error, "Failed to enumerate instance extensions");
+    return result;
+  }
+
+  if (std::any_of(
+          extension_properties.begin(),
+          extension_properties.end(),
+          [](const auto& extension) {
+            return strcmp(
+                       VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME,
+                       extension.extensionName) == 0;
+          })) {
+    requested_extensions.push_back(
+        VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);
+  }
+
+  if (requested_extensions.empty()) {
+    ET_LOG(Error, "VK_KHR_portability_enumeration not found");
+  }
+
+#endif
+
   VkInstanceCreateInfo instance_info{
       .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
       .pNext = nullptr,
-      .flags = 0,
+      .flags = instance_flags,
       .pApplicationInfo = &app_info,
-      0,
-      nullptr,
-      0,
-      nullptr};
+      .enabledLayerCount = 0,
+      .ppEnabledLayerNames = nullptr,
+      .enabledExtensionCount =
+          static_cast<uint32_t>(requested_extensions.size()),
+      .ppEnabledExtensionNames = requested_extensions.data(),
+  };
   result = vkCreateInstance(&instance_info, nullptr, instance);
   if (result != VK_SUCCESS) {
     ET_LOG(Error, "Failed to create VkInstance");

backends/arm/runtime/VGFSetup.cpp

@@ -517,14 +517,30 @@ bool VgfRepr::process_vgf(const char* vgf_data, ArrayRef<CompileSpec> specs) {
     return false;
   }
 
+  std::vector<VkDescriptorPoolSize> poolSizes;
+  poolSizes.reserve(layout_bindings.size());
+  for (const auto& b : layout_bindings) {
+    bool found = false;
+    for (size_t idx = 0; idx < poolSizes.size(); ++idx) {
+      if (poolSizes[idx].type == b.descriptorType) {
+        poolSizes[idx].descriptorCount += b.descriptorCount;
+        found = true;
+        break;
+      }
+    }
+    if (!found) {
+      poolSizes.push_back({b.descriptorType, b.descriptorCount});
+    }
+  }
+
   // Create descriptor pool and descriptors for pipeline
   const VkDescriptorPoolCreateInfo descriptor_pool_info = {
       .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
       .pNext = nullptr,
       .flags = 0,
       .maxSets = static_cast<uint32_t>(set_count),
-      .poolSizeCount = 0,
-      .pPoolSizes = nullptr,
+      .poolSizeCount = static_cast<uint32_t>(poolSizes.size()),
+      .pPoolSizes = poolSizes.data(),
   };
   result = vkCreateDescriptorPool(
       vk_device, &descriptor_pool_info, nullptr, &vk_descriptor_pool);

backends/arm/scripts/mlsdk_utils.sh

@@ -6,8 +6,7 @@
 
 set -euo pipefail
 
-# TODO
-mlsdk_manifest_url=""
+mlsdk_manifest_url="https://github.com/arm/ai-ml-sdk-manifest.git"
 
 script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd)
 
@@ -55,8 +54,9 @@ function download_ai_mlsdk_manifest() {
 function setup_model_converter() {
     local work_dir="$1"
     local manifest_dir="$2"
-    local enable_vgf_lib="$3"
-    local enable_emulation_layer="$4"
+    local enable_model_converter="$3"
+    local enable_vgf_lib="$4"
+    local enable_emulation_layer="$5"
 
     if [[ -z "$work_dir" ]]; then
         echo "Error: work_dir parameter is required."
@@ -76,29 +76,34 @@ function setup_model_converter() {
     pushd "$manifest_dir"
 
     # model-converter
-    # TODO: Remove macOS patch after mlsdk fully supports macOS
-    if [[ "$(uname)" == "Darwin" ]]; then
+    if [[ "${enable_model_converter}" -eq 1 ]]; then
+        # TODO: Remove this workaround once MLSDK has full Darwin support
+        # Do not indent sed command, the whitespace is significant for the patch to work.
+        if [[ "$(uname)" == "Darwin" ]]; then
     sed -i '' '/^ *print(f"Unsupported host platform/ i\
             if system == "Darwin":\
-                # Use default Apple toolchain (Clang) on macOS\
                 return True\
 \
 ' sw/model-converter/scripts/build.py
+        fi
+        python sw/model-converter/scripts/build.py -j$(nproc)
     fi
-    python sw/model-converter/scripts/build.py -j$(nproc)
 
     # libvgf
     if [[ "${enable_vgf_lib}" -eq 1 ]]; then
-    # TODO: Remove macOS patch after mlsdk fully supports macOS
+        # TODO: Remove this workaround once MLSDK has full Darwin support
+        # Do not indent sed command, the whitespace is significant for the patch to work.
         if [[ "$(uname)" == "Darwin" ]]; then
     sed -i '' '/^ *print(f"ERROR: Unsupported host platform/ i\
             if system == "Darwin":\
-                # Use default Apple toolchain (Clang) on macOS\
                 return True\
 \
 ' sw/vgf-lib/scripts/build.py
         fi
-        python sw/vgf-lib/scripts/build.py -j$(nproc)
+        pushd sw/vgf-lib
+        python scripts/build.py -j$(nproc)
+        cmake --install build --prefix deploy
+        popd
     fi
 
     # emu layer
@@ -110,7 +115,9 @@ function setup_model_converter() {
             -DSPIRV_HEADERS_PATH=../../dependencies/SPIRV-Headers    \
             -DSPIRV_TOOLS_PATH=../../dependencies/SPIRV-Tools        \
             -DVULKAN_HEADERS_PATH=../../dependencies/Vulkan-Headers
+
         cmake --build build
+        cmake --install build --prefix deploy
         popd
     fi
 

examples/arm/setup.sh

@@ -23,9 +23,9 @@ target_toolchain=""
 enable_fvps=1
 enable_vela=1
 enable_model_converter=0   # model-converter tool for VGF output
-enable_vgf_lib=0  # vgf reader - runtime backend dependency 
+enable_vgf_lib=0  # vgf reader - runtime backend dependency
 enable_emulation_layer=0  # Vulkan layer driver - emulates Vulkan ML extensions
-mlsdk_manifest_url=""
+mlsdk_manifest_url="https://github.com/arm/ai-ml-sdk-manifest.git"
 
 
 # Figure out if setup.sh was called or sourced and save it into "is_script_sourced"
@@ -370,14 +370,19 @@ function create_setup_path(){
     # Add Path for vgf-lib and emulation-layer
     if [[ "${enable_vgf_lib}" -eq 1 ]]; then
         cd "${root_dir}"
-        model_vgf_lib_bin_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/build && pwd)"
-        echo "export PATH=\${PATH}:${model_vgf_lib_bin_path}" >> ${setup_path_script}
+        model_vgf_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/deploy && pwd)"
+        echo "export PATH=\${PATH}:${model_vgf_path}/bin" >> ${setup_path_script}
+        echo "export LD_LIBRARY_PATH=\${LD_LIBRARY_PATH}:${model_vgf_path}/lib" >> ${setup_path_script}
+        echo "export DYLD_LIBRARY_PATH=\${DYLD_LIBRARY_PATH}:${model_vgf_path}/lib" >> ${setup_path_script}
     fi
 
     if [[ "${enable_emulation_layer}" -eq 1 ]]; then
         cd "${root_dir}"
-        model_emulation_layer_bin_path="$(cd ${mlsdk_manifest_dir}/sw/vgf-lib/build && pwd)"
-        echo "export PATH=\${PATH}:${model_emulation_layer_bin_path}" >> ${setup_path_script}
+        model_emulation_layer_path="$(cd ${mlsdk_manifest_dir}/sw/emulation-layer/ && pwd)"
+        echo "export LD_LIBRARY_PATH=${model_emulation_layer_path}/deploy/lib:\${LD_LIBRARY_PATH}" >> ${setup_path_script}
+        echo "export DYLD_LIBRARY_PATH=${model_emulation_layer_path}/deploy/lib:\${DYLD_LIBRARY_PATH}" >> ${setup_path_script}
+        echo "export VK_INSTANCE_LAYERS=VK_LAYER_ML_Graph_Emulation:VK_LAYER_ML_Tensor_Emulation:\${VK_INSTANCE_LAYERS}" >> ${setup_path_script}
+        echo "export VK_ADD_LAYER_PATH=${model_emulation_layer_path}/deploy/share/vulkan/explicit_layer.d:\${VK_ADD_LAYER_PATH}" >> ${setup_path_script}
     fi
 }
 
@@ -434,19 +439,11 @@ if [[ $is_script_sourced -eq 0 ]]; then
         setup_fvp
     fi
 
-
-    if [[ -z "$mlsdk_manifest_url" && "${enable_model_converter}" -eq 1 ]]; then
-        echo "Warning: mlsdk-manifest-url is not set, but model converter setup is not skipped."
-        echo "         Please set the --mlsdk-manifest-url option to the correct URL."
-        echo "         Skipping MLSDK model converter setup."
-        enable_model_converter=0  # Q: Can we assume if we enable mlsdk, we will always enable model converter
-        enable_vgf_lib=0
-        enable_emulation_layer=0
-    fi
-
-    if [[ "${enable_model_converter}" -eq 1 ]]; then
+    if [[ "${enable_model_converter}" -eq 1 || \
+          "${enable_vgf_lib}" -eq 1 || \
+          "${enable_emulation_layer}" -eq 1 ]]; then
         source $et_dir/backends/arm/scripts/mlsdk_utils.sh -u "${mlsdk_manifest_url}"
-        setup_model_converter ${root_dir} ${mlsdk_manifest_dir} ${enable_vgf_lib} ${enable_emulation_layer}
+        setup_model_converter ${root_dir} ${mlsdk_manifest_dir} ${enable_model_converter} ${enable_vgf_lib} ${enable_emulation_layer}
     fi
 
     # Create new setup_path script