[metal] Move to hal/drivers and default build for Apple silicon (#14129)

This commit moves the metal HAL driver from the `experimental/`
directory to the `iree/hal/drivers/`, and turns on building by
default for arm64 Apple devices. x86_64 Apple devices are not
supported right now with missing features and such; that's to
come later.

These existing code is already extensively reviewed and passing
various in-tree tests. Further improvements are for performance
and tracked in #14050.

.github/CODEOWNERS

@@ -75,4 +75,5 @@
 # Runtime
 /runtime/src/iree/ @benvanik
 /runtime/src/iree/hal/cts/ @ScottTodd
+/runtime/src/iree/hal/drivers/metal/ @antiagainst
 /runtime/src/iree/hal/drivers/vulkan/ @antiagainst @ScottTodd

CMakeLists.txt

@@ -231,10 +231,18 @@ if(APPLE)
   set(IREE_HAL_DRIVER_VULKAN_DEFAULT OFF)
 endif()
 
+# Metal support is enabled if it's one of the Apple platforms.
+set(IREE_HAL_DRIVER_METAL_DEFAULT ${IREE_HAL_DRIVER_DEFAULTS})
+# Right now only support Apple silicon devices.
+if(NOT APPLE OR NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm64")
+  set(IREE_HAL_DRIVER_METAL_DEFAULT OFF)
+endif()
+
 option(IREE_HAL_DRIVER_CUDA "Enables the 'cuda' runtime HAL driver" ${IREE_HAL_DRIVER_CUDA_DEFAULT})
 option(IREE_HAL_DRIVER_LOCAL_SYNC "Enables the 'local-sync' runtime HAL driver" ${IREE_HAL_DRIVER_DEFAULTS})
 option(IREE_HAL_DRIVER_LOCAL_TASK "Enables the 'local-task' runtime HAL driver" ${IREE_HAL_DRIVER_DEFAULTS})
 option(IREE_HAL_DRIVER_VULKAN "Enables the 'vulkan' runtime HAL driver" ${IREE_HAL_DRIVER_VULKAN_DEFAULT})
+option(IREE_HAL_DRIVER_METAL "Enables the 'metal' runtime HAL driver" ${IREE_HAL_DRIVER_METAL_DEFAULT})
 
 option(IREE_HAL_EXECUTABLE_LOADER_DEFAULTS "Sets the default value for all runtime HAL executable loaders" ON)
 set(IREE_HAL_EXECUTABLE_LOADER_EMBEDDED_ELF_DEFAULT ${IREE_HAL_EXECUTABLE_LOADER_DEFAULTS})
@@ -295,6 +303,9 @@ endif()
 if(IREE_HAL_DRIVER_VULKAN)
   message(STATUS "  - vulkan")
 endif()
+if(IREE_HAL_DRIVER_METAL)
+  message(STATUS "  - metal")
+endif()
 if(IREE_EXTERNAL_HAL_DRIVERS)
   message(STATUS "  + external: ${IREE_EXTERNAL_HAL_DRIVERS}")
 endif()
@@ -336,11 +347,6 @@ set(IREE_EXTERNAL_ROCM_HAL_DRIVER_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/experi
 set(IREE_EXTERNAL_ROCM_HAL_DRIVER_TARGET "iree::experimental::rocm::registration")
 set(IREE_EXTERNAL_ROCM_HAL_DRIVER_REGISTER "iree_hal_rocm_driver_module_register")
 
-set(IREE_EXTERNAL_METAL_HAL_DRIVER_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/experimental/metal")
-set(IREE_EXTERNAL_METAL_HAL_DRIVER_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/experimental/metal")
-set(IREE_EXTERNAL_METAL_HAL_DRIVER_TARGET "iree::experimental::metal::registration")
-set(IREE_EXTERNAL_METAL_HAL_DRIVER_REGISTER "iree_hal_metal_driver_module_register")
-
 #-------------------------------------------------------------------------------
 # Experimental WebGPU HAL driver
 #-------------------------------------------------------------------------------

docs/website/docs/index.md

@@ -53,7 +53,7 @@ Support for hardware accelerators and APIs is also included:
 
 - [x] Vulkan
 - [x] CUDA
-- [ ] Metal (planned)
+- [x] Metal (for Apple silicon devices)
 - [ ] WebGPU (planned)
 
 ## Project architecture

runtime/src/iree/hal/drivers/CMakeLists.txt

@@ -126,6 +126,10 @@ if(IREE_HAL_DRIVER_VULKAN)
   add_subdirectory(vulkan)
   list(APPEND _INIT_INTERNAL_DEPS iree::hal::drivers::vulkan::registration)
 endif()
+if(IREE_HAL_DRIVER_METAL)
+  add_subdirectory(metal)
+  list(APPEND _INIT_INTERNAL_DEPS iree::hal::drivers::metal::registration)
+endif()
 
 iree_cc_library(
   NAME

runtime/src/iree/hal/drivers/init.c

@@ -22,6 +22,10 @@
 #include "iree/hal/drivers/vulkan/registration/driver_module.h"
 #endif  // IREE_HAVE_HAL_VULKAN_DRIVER_MODULE
 
+#if defined(IREE_HAVE_HAL_METAL_DRIVER_MODULE)
+#include "iree/hal/drivers/metal/registration/driver_module.h"
+#endif  // IREE_HAVE_HAL_METAL_DRIVER_MODULE
+
 #if defined(IREE_HAVE_HAL_EXTERNAL_DRIVERS)
 // Defined in the generated init_external.c file:
 extern iree_status_t iree_hal_register_external_drivers(
@@ -57,6 +61,11 @@ iree_hal_register_all_available_drivers(iree_hal_driver_registry_t* registry) {
       z0, iree_hal_vulkan_driver_module_register(registry));
 #endif  // IREE_HAVE_HAL_VULKAN_DRIVER_MODULE
 
+#if defined(IREE_HAVE_HAL_METAL_DRIVER_MODULE)
+  IREE_RETURN_AND_END_ZONE_IF_ERROR(
+      z0, iree_hal_metal_driver_module_register(registry));
+#endif  // IREE_HAVE_HAL_METAL_DRIVER_MODULE
+
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_hal_register_external_drivers(registry));
 

experimental/metal/CMakeLists.txt → runtime/src/iree/hal/drivers/metal/CMakeLists.txt

@@ -4,10 +4,6 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-# Set the root for package namespacing to the current directory.
-set(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}")
-set(IREE_PACKAGE_ROOT_PREFIX "iree/experimental/metal")
-
 iree_add_all_subdirs()
 
 iree_cc_library(
@@ -43,8 +39,8 @@ iree_cc_library(
     iree::base::internal
     iree::base::internal::arena
     iree::base::internal::flatcc::parsing
-    iree::experimental::metal::builtin
     iree::hal
+    iree::hal::drivers::metal::builtin
     iree::hal::utils::buffer_transfer
     iree::hal::utils::resource_set
     iree::schemas::metal_executable_def_c_fbs

experimental/metal/README.md → runtime/src/iree/hal/drivers/metal/README.md

@@ -1,8 +1,8 @@
 # IREE Metal HAL Driver
 
 This document lists technical details regarding the Metal HAL driver. Note that
-the Metal HAL driver is a work in progress and experimental; this document is
-expected to be updated along the way.
+the Metal HAL driver is a work in progress; this document is expected to be
+updated along the way.
 
 IREE provides a [Hardware Abstraction Layer (HAL)][iree-hal] as a common
 interface to different compute accelerators. IREE HAL's design draws inspiration
@@ -26,10 +26,6 @@ already dominant ([macOS][macos-version-share], [iOS][ios-version-share]).
 
 ### Programming Languages and Libraries
 
-The Metal HAL driver lives under the [`experimental/metal/`][iree-metal]
-directory right now. Once more stable, it will be graduated into the
-`runtime/src/iree/hal/drivers/` directory like other HAL drivers.
-
 The Metal framework only exposes Objective-C or Swift programming language APIs.
 Metal HAL driver needs to inherit from common HAL abstraction definitions, which
 are in C. To minimize dependency and binary size and increase performance, we
@@ -313,7 +309,6 @@ with the current active `MTLComputeCommandEncoder`:
 [macos-version-share]: https://gs.statcounter.com/macos-version-market-share/desktop/worldwide
 [ios-version-share]: https://developer.apple.com/support/app-store/
 [iree-hal]: https://github.com/openxla/iree/tree/main/runtime/src/iree/hal
-[iree-metal]: https://github.com/openxla/iree/tree/main/experimental/metal
 [hal-allocator]: https://github.com/openxla/iree/blob/main/runtime/src/iree/hal/allocator.h
 [hal-buffer]: https://github.com/openxla/iree/blob/main/runtime/src/iree/hal/buffer.h
 [hal-command-buffer]: https://github.com/openxla/iree/blob/main/runtime/src/iree/hal/command_buffer.h
@@ -329,7 +324,7 @@ with the current active `MTLComputeCommandEncoder`:
 [metal-kernel-library]: https://github.com/openxla/iree/tree/main/experimental/metal/kernel_library.h
 [metal-shared-event]: https://github.com/openxla/iree/tree/main/experimental/metal/shared_event.h
 [metal-spirv-target]: https://github.com/openxla/iree/tree/main/compiler/src/iree/compiler/Dialect/HAL/Target/MetalSPIRV
-[metal-builtin-kernels]: https://github.com/openxla/iree/tree/main/experimental/metal/builtin/
+[metal-builtin-kernels]: https://github.com/openxla/iree/tree/main/runtime/src/iree/hal/drivers/metal/builtin/
 [mtl-argument-buffer]: https://developer.apple.com/documentation/metal/buffers/about_argument_buffers?language=objc
 [mtl-argument-encoder]: https://developer.apple.com/documentation/metal/mtlargumentencoder?language=objc
 [mtl-buffer]: https://developer.apple.com/documentation/metal/mtlbuffer?language=objc

experimental/metal/api.h → runtime/src/iree/hal/drivers/metal/api.h

@@ -6,8 +6,8 @@
 
 // See iree/base/api.h for documentation on the API conventions used.
 
-#ifndef IREE_EXPERIMENTAL_METAL_API_H_
-#define IREE_EXPERIMENTAL_METAL_API_H_
+#ifndef IREE_HAL_DRIVERS_METAL_API_H_
+#define IREE_HAL_DRIVERS_METAL_API_H_
 
 #include "iree/base/api.h"
 #include "iree/hal/api.h"
@@ -98,4 +98,4 @@ IREE_API_EXPORT iree_status_t iree_hal_metal_driver_create(
 }  // extern "C"
 #endif  // __cplusplus
 
-#endif  // IREE_EXPERIMENTAL_METAL_API_H_
+#endif  // IREE_HAL_DRIVERS_METAL_API_H_

experimental/metal/builtin_executables.h → runtime/src/iree/hal/drivers/metal/builtin_executables.h

@@ -4,14 +4,14 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#ifndef IREE_EXPERIMENTAL_METAL_BUILTIN_EXECUTABLES_H_
-#define IREE_EXPERIMENTAL_METAL_BUILTIN_EXECUTABLES_H_
+#ifndef IREE_HAL_DRIVERS_METAL_BUILTIN_EXECUTABLES_H_
+#define IREE_HAL_DRIVERS_METAL_BUILTIN_EXECUTABLES_H_
 
 #import <Metal/Metal.h>
 
-#include "experimental/metal/kernel_library.h"
 #include "iree/base/api.h"
 #include "iree/hal/api.h"
+#include "iree/hal/drivers/metal/kernel_library.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -63,4 +63,4 @@ iree_status_t iree_hal_metal_builtin_executable_copy_buffer(
 }  // extern "C"
 #endif  // __cplusplus
 
-#endif  // IREE_EXPERIMENTAL_METAL_BUILTIN_EXECUTABLES_H_
+#endif  // IREE_HAL_DRIVERS_METAL_BUILTIN_EXECUTABLES_H_

experimental/metal/builtin_executables.m → runtime/src/iree/hal/drivers/metal/builtin_executables.m

@@ -4,14 +4,14 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#include "experimental/metal/builtin_executables.h"
+#include "iree/hal/drivers/metal/builtin_executables.h"
 
 #include <string.h>
 
-#include "experimental/metal/builtin/metal_buffer_kernels.h"
 #include "iree/base/api.h"
 #include "iree/base/tracing.h"
 #include "iree/hal/api.h"
+#include "iree/hal/drivers/metal/builtin/metal_buffer_kernels.h"
 
 typedef struct iree_hal_metal_builtin_executable_data_t {
   const char* entry_point;

experimental/metal/cts/CMakeLists.txt → runtime/src/iree/hal/drivers/metal/cts/CMakeLists.txt

@@ -8,15 +8,15 @@ iree_hal_cts_test_suite(
   DRIVER_NAME
     metal
   DRIVER_REGISTRATION_HDR
-    "experimental/metal/registration/driver_module.h"
+    "runtime/src/iree/hal/drivers/metal/registration/driver_module.h"
   DRIVER_REGISTRATION_FN
     "iree_hal_metal_driver_module_register"
   COMPILER_TARGET_BACKEND
     "metal-spirv"
   EXECUTABLE_FORMAT
     "\"MTLE\""
   DEPS
-    iree::experimental::metal::registration
+    iree::hal::drivers::metal::registration
   EXCLUDED_TESTS
     # HAL event is unimplemented for Metal right now.
     "event"

experimental/metal/direct_allocator.h → runtime/src/iree/hal/drivers/metal/direct_allocator.h

@@ -4,14 +4,14 @@
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-#ifndef IREE_EXPERIMENTAL_METAL_DIRECT_ALLOCATOR_H_
-#define IREE_EXPERIMENTAL_METAL_DIRECT_ALLOCATOR_H_
+#ifndef IREE_HAL_DRIVERS_METAL_DIRECT_ALLOCATOR_H_
+#define IREE_HAL_DRIVERS_METAL_DIRECT_ALLOCATOR_H_
 
 #import <Metal/Metal.h>
 
-#include "experimental/metal/api.h"
 #include "iree/base/api.h"
 #include "iree/hal/api.h"
+#include "iree/hal/drivers/metal/api.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -44,4 +44,4 @@ id<MTLCommandQueue> iree_hal_metal_allocator_command_queue(
 }  // extern "C"
 #endif  // __cplusplus
 
-#endif  // IREE_EXPERIMENTAL_METAL_DIRECT_ALLOCATOR_H_
+#endif  // IREE_HAL_DRIVERS_METAL_DIRECT_ALLOCATOR_H_