[HIP] Adds basics to implement HIP HAL driver

This commit starts a HIP backend in HAL. Following the steps of CUDA rewrite (#13245) effort, this backend will provide improvements over the existing ROCm HAL backend. Building this commmit with -DIREE_EXTERNAL_HAL_DRIVERS=hip will initialize this driver and the driver information can be seen using `tools/iree-run-module --dump_devices`.
iree-org · Nov 9, 2023 · f7ebf0e · f7ebf0e
1 parent bc98b9a
commit f7ebf0e
Show file tree

Hide file tree

Showing 14 changed files with 1,115 additions and 0 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -351,6 +351,16 @@ set(IREE_EXTERNAL_CUDA2_HAL_DRIVER_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/exper
 set(IREE_EXTERNAL_CUDA2_HAL_DRIVER_TARGET "iree::experimental::cuda2::registration")
 set(IREE_EXTERNAL_CUDA2_HAL_DRIVER_REGISTER "iree_hal_cuda2_driver_module_register")
 
+#-------------------------------------------------------------------------------
+# Experimental next-generation HIP HAL driver
+# Enable with: -DIREE_EXTERNAL_HAL_DRIVERS=hip
+#-------------------------------------------------------------------------------
+
+set(IREE_EXTERNAL_HIP_HAL_DRIVER_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/experimental/hip")
+set(IREE_EXTERNAL_HIP_HAL_DRIVER_BINARY_DIR "${CMAKE_CURRENT_BINARY_DIR}/experimental/hip")
+set(IREE_EXTERNAL_HIP_HAL_DRIVER_TARGET "iree::experimental::hip::registration")
+set(IREE_EXTERNAL_HIP_HAL_DRIVER_REGISTER "iree_hal_hip_driver_module_register")
+
 #-------------------------------------------------------------------------------
 # Experimental ROCM HAL driver
 # Enable with: -DIREE_EXTERNAL_HAL_DRIVERS=rocm

diff --git a/experimental/hip/CMakeLists.txt b/experimental/hip/CMakeLists.txt
@@ -0,0 +1,78 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# 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/hip")
+
+iree_add_all_subdirs()
+
+if(NOT DEFINED ROCM_HEADERS_API_ROOT)
+  set(ROCM_HEADERS_API_ROOT "${IREE_SOURCE_DIR}/third_party/hip-build-deps/include")
+endif()
+
+if(NOT EXISTS "${ROCM_HEADERS_API_ROOT}/hip/hip_version.h")
+  message(SEND_ERROR "Could not find HIP headers at: ${ROCM_HEADERS_API_ROOT}")
+endif()
+
+iree_cc_library(
+  NAME
+    hip
+  HDRS
+    "api.h"
+  SRCS
+    "api.h"
+    "hip_driver.c"
+  INCLUDES
+    "${ROCM_HEADERS_API_ROOT}"
+  DEPS
+    ::dynamic_symbols
+    iree::base
+    iree::base::core_headers
+    iree::hal
+  COPTS
+    "-D__HIP_PLATFORM_HCC__=1"
+  PUBLIC
+)
+
+iree_cc_library(
+  NAME
+    dynamic_symbols
+  HDRS
+    "dynamic_symbols.h"
+    "status_util.h"
+  TEXTUAL_HDRS
+    "dynamic_symbol_tables.h"
+  SRCS
+    "hip_headers.h"
+    "dynamic_symbols.c"
+    "status_util.c"
+  INCLUDES
+    "${ROCM_HEADERS_API_ROOT}"
+  COPTS
+    "-D__HIP_PLATFORM_HCC__=1"
+  DEPS
+    iree::base
+    iree::base::core_headers
+    iree::base::internal::dynamic_library
+  PUBLIC
+)
+
+iree_cc_test(
+  NAME
+    dynamic_symbols_test
+  SRCS
+    "dynamic_symbols_test.cc"
+  DEPS
+    ::dynamic_symbols
+    iree::base
+    iree::testing::gtest
+    iree::testing::gtest_main
+  LABELS
+    "driver=hip"
+  COPTS
+    "-D__HIP_PLATFORM_HCC__=1"
+)
diff --git a/experimental/hip/api.h b/experimental/hip/api.h
@@ -0,0 +1,47 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// See iree/base/api.h for documentation on the API conventions used.
+
+#ifndef IREE_EXPERIMENTAL_HIP_API_H_
+#define IREE_EXPERIMENTAL_HIP_API_H_
+
+#include "iree/base/api.h"
+#include "iree/hal/api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+//===----------------------------------------------------------------------===//
+// iree_hal_hip_driver_t
+//===----------------------------------------------------------------------===//
+
+// HIP HAL driver creation options.
+typedef struct iree_hal_hip_driver_options_t {
+  // The index of the default HIP device to use within the list of available
+  // devices.
+  int default_device_index;
+} iree_hal_hip_driver_options_t;
+
+// Initializes the given |out_options| with default driver creation options.
+IREE_API_EXPORT void iree_hal_hip_driver_options_initialize(
+    iree_hal_hip_driver_options_t* out_options);
+
+// Creates a HIP HAL driver with the given |options|, from which HIP devices
+// can be enumerated and created with specific parameters.
+//
+// |out_driver| must be released by the caller (see iree_hal_driver_release).
+IREE_API_EXPORT iree_status_t iree_hal_hip_driver_create(
+    iree_string_view_t identifier,
+    const iree_hal_hip_driver_options_t* options,
+    iree_allocator_t host_allocator, iree_hal_driver_t** out_driver);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif  // __cplusplus
+
+#endif  // IREE_EXPERIMENTAL_HIP_API_H_
diff --git a/experimental/hip/dynamic_symbol_tables.h b/experimental/hip/dynamic_symbol_tables.h
@@ -0,0 +1,66 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+//===----------------------------------------------------------------------===//
+// HIP symbols
+//===----------------------------------------------------------------------===//
+IREE_HIP_PFN_DECL(hipCtxCreate, hipCtx_t *, unsigned int, hipDevice_t)
+IREE_HIP_PFN_DECL(hipCtxDestroy, hipCtx_t)
+IREE_HIP_PFN_DECL(hipDeviceGet, hipDevice_t *, int)  // No direct, need to modify
+IREE_HIP_PFN_DECL(hipGetDeviceCount, int *)
+IREE_HIP_PFN_DECL(hipGetDeviceProperties, hipDeviceProp_t *, int)
+IREE_HIP_PFN_DECL(hipDeviceGetName, char *, int,
+            hipDevice_t)  // No direct, need to modify
+IREE_HIP_PFN_STR_DECL(
+    hipGetErrorName,
+    hipError_t)  // Unlike other functions hipGetErrorName(hipError_t) return
+                 // const char* instead of hipError_t so it uses a different
+                 // macro
+IREE_HIP_PFN_STR_DECL(
+    hipGetErrorString,
+    hipError_t)  // Unlike other functions hipGetErrorName(hipError_t) return
+                 // const char* instead of hipError_t so it uses a different
+                 // macro
+IREE_HIP_PFN_DECL(hipInit, unsigned int)
+IREE_HIP_PFN_DECL(hipModuleLaunchKernel, hipFunction_t, unsigned int, unsigned int,
+            unsigned int, unsigned int, unsigned int, unsigned int,
+            unsigned int, hipStream_t, void **, void **)
+IREE_HIP_PFN_DECL(hipMemset, void *, int, size_t)
+IREE_HIP_PFN_DECL(hipMemsetAsync, void *, int, size_t, hipStream_t)
+IREE_HIP_PFN_DECL(hipMemsetD32Async, void *, int, size_t, hipStream_t)
+IREE_HIP_PFN_DECL(hipMemsetD16Async, void *, short, size_t, hipStream_t)
+IREE_HIP_PFN_DECL(hipMemsetD8Async, void *, char, size_t, hipStream_t)
+IREE_HIP_PFN_DECL(hipMemcpy, void *, const void *, size_t, hipMemcpyKind)
+IREE_HIP_PFN_DECL(hipMemcpyAsync, void *, const void *, size_t, hipMemcpyKind,
+            hipStream_t)
+IREE_HIP_PFN_DECL(hipMalloc, void **, size_t)
+IREE_HIP_PFN_DECL(hipMallocManaged, hipDeviceptr_t *, size_t, unsigned int)
+IREE_HIP_PFN_DECL(hipFree, void *)
+IREE_HIP_PFN_DECL(hipHostFree, void *)
+IREE_HIP_PFN_DECL(hipMemAllocHost, void **, size_t, unsigned int)
+IREE_HIP_PFN_DECL(hipHostGetDevicePointer, void **, void *, unsigned int)
+IREE_HIP_PFN_DECL(hipModuleGetFunction, hipFunction_t *, hipModule_t, const char *)
+IREE_HIP_PFN_DECL(hipModuleLoadDataEx, hipModule_t *, const void *, unsigned int,
+            hipJitOption *, void **)
+IREE_HIP_PFN_DECL(hipModuleLoadData, hipModule_t *, const void *)
+IREE_HIP_PFN_DECL(hipModuleUnload, hipModule_t)
+IREE_HIP_PFN_DECL(hipStreamCreateWithFlags, hipStream_t *, unsigned int)
+IREE_HIP_PFN_DECL(hipStreamDestroy, hipStream_t)
+IREE_HIP_PFN_DECL(hipStreamSynchronize, hipStream_t)
+IREE_HIP_PFN_DECL(hipStreamWaitEvent, hipStream_t, hipEvent_t, unsigned int)
+IREE_HIP_PFN_DECL(hipEventCreate, hipEvent_t *)
+IREE_HIP_PFN_DECL(hipEventDestroy, hipEvent_t)
+IREE_HIP_PFN_DECL(hipEventElapsedTime, float *, hipEvent_t, hipEvent_t)
+IREE_HIP_PFN_DECL(hipEventQuery, hipEvent_t)
+IREE_HIP_PFN_DECL(hipEventRecord, hipEvent_t, hipStream_t)
+IREE_HIP_PFN_DECL(hipEventSynchronize, hipEvent_t)
+IREE_HIP_PFN_DECL(hipDeviceGetAttribute, int *, hipDeviceAttribute_t, int)
+IREE_HIP_PFN_DECL(hipFuncSetAttribute, const void *, hipFuncAttribute, int)
+IREE_HIP_PFN_DECL(hipDeviceGetUuid, hipUUID *, hipDevice_t)
+IREE_HIP_PFN_DECL(hipDevicePrimaryCtxRetain, hipCtx_t *, hipDevice_t)
+IREE_HIP_PFN_DECL(hipCtxGetDevice, hipDevice_t *)
+IREE_HIP_PFN_DECL(hipCtxSetCurrent, hipCtx_t)
+IREE_HIP_PFN_DECL(hipDevicePrimaryCtxRelease, hipDevice_t)
diff --git a/experimental/hip/dynamic_symbols.c b/experimental/hip/dynamic_symbols.c
@@ -0,0 +1,83 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "experimental/hip/dynamic_symbols.h"
+
+#include <string.h>
+
+#include "experimental/hip/status_util.h"
+#include "iree/base/assert.h"
+#include "iree/base/internal/dynamic_library.h"
+#include "iree/base/target_platform.h"
+#include "iree/base/tracing.h"
+
+//===----------------------------------------------------------------------===//
+// HIP dynamic symbols
+//===----------------------------------------------------------------------===//
+
+static const char* iree_hal_hip_dylib_names[] = {
+#if defined(IREE_PLATFORM_WINDOWS)
+    "amdhip64.dll",
+#else
+    "libamdhip64.so",
+#endif  // IREE_PLATFORM_WINDOWS
+};
+
+// Resolves all HIP dynamic symbols in `dynamic_symbol_tables.h`
+static iree_status_t iree_hal_hip_dynamic_symbols_resolve_all(
+    iree_hal_hip_dynamic_symbols_t* syms) {
+#define IREE_HIP_PFN_DECL(hip_symbol_name, ...)                  \
+  {                                                              \
+    static const char* name = #hip_symbol_name;                  \
+    IREE_RETURN_IF_ERROR(iree_dynamic_library_lookup_symbol(     \
+        syms->dylib, name, (void**)&syms->hip_symbol_name));     \
+  }
+#define IREE_HIP_PFN_STR_DECL(hip_symbol_name, ...) IREE_HIP_PFN_DECL(hip_symbol_name, ...)
+#include "experimental/hip/dynamic_symbol_tables.h"  // IWYU pragma: keep
+#undef IREE_HIP_PFN_DECL
+#undef IREE_HIP_PFN_STR_DECL
+  return iree_ok_status();
+}
+
+// #undef IREE_CONCAT
+
+iree_status_t iree_hal_hip_dynamic_symbols_initialize(
+    iree_allocator_t host_allocator,
+    iree_hal_hip_dynamic_symbols_t* out_syms) {
+  IREE_ASSERT_ARGUMENT(out_syms);
+  IREE_TRACE_ZONE_BEGIN(z0);
+
+  memset(out_syms, 0, sizeof(*out_syms));
+  iree_status_t status = iree_dynamic_library_load_from_files(
+      IREE_ARRAYSIZE(iree_hal_hip_dylib_names), iree_hal_hip_dylib_names,
+      IREE_DYNAMIC_LIBRARY_FLAG_NONE, host_allocator, &out_syms->dylib);
+  if (iree_status_is_not_found(status)) {
+    iree_status_ignore(status);
+    status = iree_make_status(
+        IREE_STATUS_UNAVAILABLE,
+        "HIP runtime library 'amdhip64.dll'/'libamdhip64.so' not available; please"
+        "ensure installed and in dynamic library search path");
+  }
+  if (iree_status_is_ok(status)) {
+    status = iree_hal_hip_dynamic_symbols_resolve_all(out_syms);
+  }
+  if (!iree_status_is_ok(status)) {
+    iree_hal_hip_dynamic_symbols_deinitialize(out_syms);
+  }
+
+  IREE_TRACE_ZONE_END(z0);
+  return status;
+}
+
+void iree_hal_hip_dynamic_symbols_deinitialize(
+    iree_hal_hip_dynamic_symbols_t* syms) {
+  IREE_TRACE_ZONE_BEGIN(z0);
+
+  iree_dynamic_library_release(syms->dylib);
+  memset(syms, 0, sizeof(*syms));
+
+  IREE_TRACE_ZONE_END(z0);
+}
diff --git a/experimental/hip/dynamic_symbols.h b/experimental/hip/dynamic_symbols.h
@@ -0,0 +1,59 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#ifndef IREE_EXPERIMENTAL_HIP_DYNAMIC_SYMBOLS_H_
+#define IREE_EXPERIMENTAL_HIP_DYNAMIC_SYMBOLS_H_
+
+#include "experimental/hip/hip_headers.h"
+#include "iree/base/api.h"
+#include "iree/base/internal/dynamic_library.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif  // __cplusplus
+
+// iree_dynamic_library_t allows dynamically loading a subset of HIP driver API.
+// We load all the symbols in `dynamic_symbol_tables.h` and
+// fail if any of the symbol is not available. The functions signatures are
+// matching the declarations in `hipruntime.h`.
+
+//===----------------------------------------------------------------------===//
+// HIP dynamic symbols
+//===----------------------------------------------------------------------===//
+
+// HIP driver API dynamic symbols.
+typedef struct iree_hal_hip_dynamic_symbols_t {
+  // The dynamic library handle.
+  iree_dynamic_library_t* dylib;
+
+  // Concrete HIP symbols defined by including the `dynamic_symbol_tables.h`.
+#define IREE_HIP_PFN_DECL(hipSymbolName, ...) \
+  hipError_t (*hipSymbolName)(__VA_ARGS__);
+#define IREE_HIP_PFN_STR_DECL(hipSymbolName, ...) \
+  const char* (*hipSymbolName)(__VA_ARGS__);
+#include "experimental/hip/dynamic_symbol_tables.h"  // IWYU pragma: export
+#undef IREE_HIP_PFN_DECL
+#undef IREE_HIP_PFN_STR_DECL
+} iree_hal_hip_dynamic_symbols_t;
+
+// Initializes |out_syms| in-place with dynamically loaded HIP symbols.
+// iree_hal_hip_dynamic_symbols_deinitialize must be used to release the
+// library resources.
+iree_status_t iree_hal_hip_dynamic_symbols_initialize(
+    iree_allocator_t host_allocator,
+    iree_hal_hip_dynamic_symbols_t* out_syms);
+
+// Deinitializes |syms| by unloading the backing library. All function pointers
+// will be invalidated. They _may_ still work if there are other reasons the
+// library remains loaded so be careful.
+void iree_hal_hip_dynamic_symbols_deinitialize(
+    iree_hal_hip_dynamic_symbols_t* syms);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif  // __cplusplus
+
+#endif  // IREE_EXPERIMENTAL_HIP_DYNAMIC_SYMBOLS_H_