diff --git a/openxla-nvgpu/.gitignore b/openxla-nvgpu/.gitignore index 567609b1..96191325 100644 --- a/openxla-nvgpu/.gitignore +++ b/openxla-nvgpu/.gitignore @@ -1 +1,6 @@ build/ + +# Source indexing files +compile_commands.json +.cache/clangd +.clangd/ diff --git a/openxla-nvgpu/CMakeLists.txt b/openxla-nvgpu/CMakeLists.txt index cd9bc0aa..775ae568 100644 --- a/openxla-nvgpu/CMakeLists.txt +++ b/openxla-nvgpu/CMakeLists.txt @@ -6,6 +6,8 @@ cmake_minimum_required(VERSION 3.21...3.24) +set(CMAKE_EXPORT_COMPILE_COMMANDS ON) + project(OPENXLA_NVGPU) # TODO: Fix this once the project is slotted into place. @@ -22,5 +24,22 @@ option(IREE_TARGET_BACKEND_DEFAULTS "Disables target backend" OFF) option(IREE_TARGET_BACKEND_CUDA "Enables CUDA target backend" ON) option(IREE_COMPILER_BUILD_SHARED_LIBS "Enables shared libraries in the compiler by default" ON) +# TODO: `llvm-cpu` compiler backend and `local-sync` HAL driver are enabled only +# for running tests. Diable them once we'll use proper CUDA target and driver. +option(IREE_HAL_DRIVER_LOCAL_SYNC "Enables the 'local-sync' runtime HAL driver" ON) +option(IREE_TARGET_BACKEND_LLVM_CPU "Enables the 'llvm-cpu' target backend" ON) + set(IREE_COMPILER_PLUGIN_PATHS "${CMAKE_CURRENT_SOURCE_DIR}" CACHE STRING "OpenXLA nvgpu plugins") add_subdirectory("${IREE_ROOT_DIR}" "iree_core") + +#------------------------------------------------------------------------------- +# OpenXLA NVGPU Runtime. +# +# Integration of NVIDIA libraries with IREE runtime via custom VM modules. +#------------------------------------------------------------------------------- + +# TODO: Use same compiler flags for building runtime targets as the IREE core. +set(IREE_CXX_STANDARD 17) + +add_subdirectory(runtime) +add_subdirectory(tools) diff --git a/openxla-nvgpu/runtime/CMakeLists.txt b/openxla-nvgpu/runtime/CMakeLists.txt new file mode 100644 index 00000000..cb51b1e4 --- /dev/null +++ b/openxla-nvgpu/runtime/CMakeLists.txt @@ -0,0 +1,7 @@ +# 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 + +add_subdirectory(src) diff --git a/openxla-nvgpu/runtime/src/CMakeLists.txt b/openxla-nvgpu/runtime/src/CMakeLists.txt new file mode 100644 index 00000000..e8b275fd --- /dev/null +++ b/openxla-nvgpu/runtime/src/CMakeLists.txt @@ -0,0 +1,11 @@ +# 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(IREE_PACKAGE_ROOT_DIR "${CMAKE_CURRENT_LIST_DIR}") +set(IREE_PACKAGE_ROOT_PREFIX "") +set(IREE_COMPILER_TABLEGEN_INCLUDE_DIRS "${CMAKE_CURRENT_SOURCE_DIR}") + +add_subdirectory(openxla/runtime/nvgpu) diff --git a/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/CMakeLists.txt b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/CMakeLists.txt new file mode 100644 index 00000000..6560fc4e --- /dev/null +++ b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/CMakeLists.txt @@ -0,0 +1,29 @@ +# 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 + +add_subdirectory(test) + +iree_cc_library( + NAME + defs + INCLUDES + "$" + "$" + PUBLIC +) + +iree_cc_library( + NAME + cudnn_module + HDRS + "cudnn_module.h" + SRCS + "cudnn_module.cpp" + DEPS + ::defs + iree::runtime + PUBLIC +) diff --git a/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.cpp b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.cpp new file mode 100644 index 00000000..7e6e1bbe --- /dev/null +++ b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.cpp @@ -0,0 +1,67 @@ +// 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 "openxla/runtime/nvgpu/cudnn_module.h" + +#include + +#include + +#include "iree/vm/native_module_cc.h" + +namespace openxla::runtime::nvgpu { + +using namespace iree; + +//===----------------------------------------------------------------------===// +// CuDNN module state encapsulates all the state required for running cuDNN +// operations (launching cuDNN graphs on a stream) at run time. +//===----------------------------------------------------------------------===// + +class CuDNNModuleState { + public: + Status Hello() { + fprintf(stderr, "Hello from OpenXLA CuDNN Module!\n"); + return OkStatus(); + } +}; + +static const vm::NativeFunction kCuDNNModuleFunctions[] = { + vm::MakeNativeFunction("hello", &CuDNNModuleState::Hello), +}; + +//===----------------------------------------------------------------------===// +// CuDNN module instance that will be allocated and reused across contexts. +//===----------------------------------------------------------------------===// + +class CuDNNModule final : public vm::NativeModule { + public: + using vm::NativeModule::NativeModule; + + StatusOr> CreateState( + iree_allocator_t host_allocator) override { + return std::make_unique(); + } +}; + +} // namespace openxla::runtime::nvgpu + +//===----------------------------------------------------------------------===// +// Register cuDNN module with IREE runtime. +//===----------------------------------------------------------------------===// + +using namespace openxla::runtime::nvgpu; + +extern "C" iree_status_t iree_custom_module_cudnn_create( + iree_vm_instance_t* instance, iree_hal_device_t* device, + iree_allocator_t host_allocator, iree_vm_module_t** out_module) { + IREE_ASSERT_ARGUMENT(out_module); + auto module = std::make_unique( + "cudnn", /*version=*/0, instance, host_allocator, + span>(kCuDNNModuleFunctions)); + *out_module = module.release()->interface(); + return iree_ok_status(); +} diff --git a/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.h b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.h new file mode 100644 index 00000000..2d0a75eb --- /dev/null +++ b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/cudnn_module.h @@ -0,0 +1,27 @@ +// 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 OPENXLA_RUNTIME_NVGPU_CUDNN_MODULE_H_ +#define OPENXLA_RUNTIME_NVGPU_CUDNN_MODULE_H_ + +#include "iree/base/api.h" +#include "iree/hal/api.h" +#include "iree/vm/api.h" + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +iree_status_t iree_custom_module_cudnn_create(iree_vm_instance_t* instance, + iree_hal_device_t* device, + iree_allocator_t host_allocator, + iree_vm_module_t** out_module); + +#ifdef __cplusplus +} // extern "C" +#endif // __cplusplus + +#endif // OPENXLA_RUNTIME_NVGPU_CUDNN_MODULE_H_ diff --git a/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/CMakeLists.txt b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/CMakeLists.txt new file mode 100644 index 00000000..b8846e25 --- /dev/null +++ b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/CMakeLists.txt @@ -0,0 +1,18 @@ +# Copyright 2022 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 + +iree_lit_test_suite( + NAME + lit + SRCS + "example.mlir" + TOOLS + FileCheck + iree-compile + openxla-runner + LABELS + "hostonly" +) diff --git a/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/example.mlir b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/example.mlir new file mode 100644 index 00000000..81931cbb --- /dev/null +++ b/openxla-nvgpu/runtime/src/openxla/runtime/nvgpu/test/example.mlir @@ -0,0 +1,13 @@ +// RUN: iree-compile %s --iree-hal-target-backends=llvm-cpu | openxla-runner - example.main | FileCheck %s + +module @example { + + func.func private @cudnn.hello() + + func.func @main() { + // CHECK: Hello from OpenXLA CuDNN Module! + call @cudnn.hello() : () -> () + return + } + +} diff --git a/openxla-nvgpu/tools/CMakeLists.txt b/openxla-nvgpu/tools/CMakeLists.txt new file mode 100644 index 00000000..3810fd16 --- /dev/null +++ b/openxla-nvgpu/tools/CMakeLists.txt @@ -0,0 +1,25 @@ +# 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 + +iree_cc_binary( + NAME + openxla-runner + SRCS + "openxla-runner.c" + DEPS + iree::base + iree::base::internal::flags + iree::base::tracing + iree::hal + iree::modules::hal::types + iree::tooling::comparison + iree::tooling::context_util + iree::tooling::device_util + iree::tooling::instrument_util + iree::tooling::vm_util + iree::vm + openxla::runtime::nvgpu::cudnn_module +) diff --git a/openxla-nvgpu/tools/openxla-runner.c b/openxla-nvgpu/tools/openxla-runner.c new file mode 100644 index 00000000..6c6caa2b --- /dev/null +++ b/openxla-nvgpu/tools/openxla-runner.c @@ -0,0 +1,97 @@ +// 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 + +#include "iree/modules/hal/types.h" +#include "iree/runtime/api.h" +#include "openxla/runtime/nvgpu/cudnn_module.h" + +// TODO: This is a temporary work around missing custom modules integration into +// IREE tools (iree-run-module). We already have flags to enable plugins in +// compiler tools (`iree-compiler` and `iree-opt`), but not yet in "runtime" +// tools. This tool can only run VM function with empty arguments and empty +// results, and intended for testing cuDNN custom module. +int main(int argc, char** argv) { + if (argc != 3) { + fprintf(stderr, + "Usage:\n" + " openxla-runner - # read from stdin\n" + " openxla-runner " + "\n"); + return -1; + } + + // Internally IREE does not (in general) use malloc and instead uses the + // provided allocator to allocate and free memory. Applications can integrate + // their own allocator as-needed. + iree_allocator_t host_allocator = iree_allocator_system(); + + // Create and configure the instance shared across all sessions. + iree_runtime_instance_options_t instance_options; + iree_runtime_instance_options_initialize(&instance_options); + iree_runtime_instance_options_use_all_available_drivers(&instance_options); + iree_runtime_instance_t* instance = NULL; + IREE_CHECK_OK(iree_runtime_instance_create(&instance_options, host_allocator, + &instance)); + + // Try to create the device - it should always succeed as it's a CPU device. + iree_hal_device_t* device = NULL; + IREE_CHECK_OK(iree_runtime_instance_try_create_default_device( + instance, iree_make_cstring_view("local-sync"), &device)); + + // Create one session per loaded module to hold the module state. + iree_runtime_session_options_t session_options; + iree_runtime_session_options_initialize(&session_options); + iree_runtime_session_t* session = NULL; + IREE_CHECK_OK(iree_runtime_session_create_with_device( + instance, &session_options, device, + iree_runtime_instance_host_allocator(instance), &session)); + + // Create the custom module that can be reused across contexts. + iree_vm_module_t* custom_module = NULL; + IREE_CHECK_OK(iree_custom_module_cudnn_create( + iree_runtime_instance_vm_instance(instance), device, host_allocator, + &custom_module)); + IREE_CHECK_OK(iree_runtime_session_append_module(session, custom_module)); + iree_vm_module_release(custom_module); + + // Load the module from stdin or a file on disk. + const char* module_path = argv[1]; + if (strcmp(module_path, "-") == 0) { + IREE_CHECK_OK( + iree_runtime_session_append_bytecode_module_from_stdin(session)); + } else { + IREE_CHECK_OK(iree_runtime_session_append_bytecode_module_from_file( + session, module_path)); + } + + iree_string_view_t entry_point = iree_make_cstring_view(argv[2]); + fprintf(stdout, "INVOKE BEGIN %.*s\n", (int)entry_point.size, + entry_point.data); + fflush(stdout); + + iree_vm_list_t* inputs = NULL; + IREE_CHECK_OK(iree_vm_list_create(NULL, 1, host_allocator, &inputs)); + iree_vm_list_t* outputs = NULL; + IREE_CHECK_OK(iree_vm_list_create(NULL, 1, host_allocator, &outputs)); + + // Synchronously invoke the requested function. + IREE_CHECK_OK( + iree_runtime_session_call_by_name(session, entry_point, inputs, outputs)); + + fprintf(stdout, "INVOKE END %.*s\n", (int)entry_point.size, entry_point.data); + fflush(stdout); + + iree_vm_list_release(inputs); + iree_vm_list_release(outputs); + + iree_runtime_session_release(session); + iree_hal_device_release(device); + iree_runtime_instance_release(instance); + + return 0; +}