[Mlir-commits] [mlir] [mlir][Target] Support Fatbin target for static nvptxcompiler (PR #118044)
Zichen Lu
llvmlistbot at llvm.org
Sun Dec 1 22:19:23 PST 2024
https://github.com/MikaOvO updated https://github.com/llvm/llvm-project/pull/118044
>From 640d37cf53b607cec368447d60717af966ff1bbc Mon Sep 17 00:00:00 2001
From: Zichen Lu <mikaovo2000 at gmail.com>
Date: Fri, 29 Nov 2024 11:02:05 +0800
Subject: [PATCH] [mlir][Target] Support Fatbin target for static nvptxcompiler
---
mlir/CMakeLists.txt | 3 ++
mlir/lib/Target/LLVM/CMakeLists.txt | 26 ++++++++++++++--
mlir/lib/Target/LLVM/NVVM/Target.cpp | 44 ++++++++++++++++++++++++++++
3 files changed, 71 insertions(+), 2 deletions(-)
diff --git a/mlir/CMakeLists.txt b/mlir/CMakeLists.txt
index 2880dc30bca91f..d885188661e70f 100644
--- a/mlir/CMakeLists.txt
+++ b/mlir/CMakeLists.txt
@@ -135,6 +135,9 @@ set(MLIR_ENABLE_VULKAN_RUNNER 0 CACHE BOOL "Enable building the MLIR Vulkan runn
set(MLIR_ENABLE_NVPTXCOMPILER 0 CACHE BOOL
"Statically link the nvptxlibrary instead of calling ptxas as a subprocess \
for compiling PTX to cubin")
+set(MLIR_ENABLE_NVFATBIN 0 CACHE BOOL
+ "Statically link the nvfatbin library instead of calling fatbinary as a subprocess \
+ for compiling PTX to fatbin")
set(MLIR_ENABLE_PDL_IN_PATTERNMATCH 1 CACHE BOOL "Enable PDL in PatternMatch")
diff --git a/mlir/lib/Target/LLVM/CMakeLists.txt b/mlir/lib/Target/LLVM/CMakeLists.txt
index 422f7e5fa7caec..c483c87c5a9351 100644
--- a/mlir/lib/Target/LLVM/CMakeLists.txt
+++ b/mlir/lib/Target/LLVM/CMakeLists.txt
@@ -66,7 +66,7 @@ if ("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD)
set(MLIR_CUDAToolkit_ROOT ${CUDAToolkit_LIBRARY_ROOT})
endif()
- # Add the `nvptxcompiler` library.
+ # Add the `nvptxcompiler` library
if(MLIR_ENABLE_NVPTXCOMPILER)
# Find the `nvptxcompiler` library.
# TODO: Bump the MLIR CMake version to 3.25 and use `CUDA::nvptxcompiler_static`.
@@ -85,16 +85,38 @@ if ("NVPTX" IN_LIST LLVM_TARGETS_TO_BUILD)
# set_property(TARGET MLIR_NVPTXCOMPILER_LIB PROPERTY IMPORTED_LOCATION ${...})
# where `...` is to be replaced with the path to the library.
set_property(TARGET MLIR_NVPTXCOMPILER_LIB PROPERTY IMPORTED_LOCATION ${MLIR_NVPTXCOMPILER_LIB_PATH})
- # Link against `nvptxcompiler_static`. TODO: use `CUDA::nvptxcompiler_static`.
+ # Link against `nvptxcompiler_static` and `nvfatbin_static`.
+ # TODO: use `CUDA::nvptxcompiler_static` and `CUDA::nvfatbin_static`.
target_link_libraries(MLIRNVVMTarget PRIVATE MLIR_NVPTXCOMPILER_LIB)
target_include_directories(obj.MLIRNVVMTarget PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
endif()
+
+ # Add the `nvfatbin` library
+ if(MLIR_ENABLE_NVFATBIN)
+ find_library(MLIR_NVFATBIN_LIB_PATH nvfatbin_static
+ PATHS ${CUDAToolkit_LIBRARY_DIR} NO_DEFAULT_PATH)
+ # Fail if `nvfatbin_static` couldn't be found.
+ if(MLIR_NVFATBIN_LIB_PATH STREQUAL "MLIR_NVFATBIN_LIB_PATH-NOTFOUND")
+ message(FATAL_ERROR
+ "Requested using the `nvfatbin` library backend but it couldn't be found.")
+ endif()
+
+ add_library(MLIR_NVFATBIN_LIB STATIC IMPORTED GLOBAL)
+ set_property(TARGET MLIR_NVFATBIN_LIB PROPERTY IMPORTED_LOCATION ${MLIR_NVFATBIN_LIB_PATH})
+ target_link_libraries(MLIRNVVMTarget PRIVATE MLIR_NVFATBIN_LIB)
+ target_include_directories(obj.MLIRNVVMTarget PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
+ endif()
else()
# Fail if `MLIR_ENABLE_NVPTXCOMPILER` is enabled and the toolkit couldn't be found.
if(MLIR_ENABLE_NVPTXCOMPILER)
message(FATAL_ERROR
"Requested using the `nvptxcompiler` library backend but it couldn't be found.")
endif()
+ # Fail if `MLIR_ENABLE_NVFATBIN` is enabled and the toolkit couldn't be found.
+ if(MLIR_ENABLE_NVFATBIN)
+ message(FATAL_ERROR
+ "Requested using the `nvfatbin` library backend but it couldn't be found.")
+ endif()
endif()
message(VERBOSE "MLIR default CUDA toolkit path: ${MLIR_CUDAToolkit_ROOT}")
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index bca26e3a0e84a9..0b51074008f5cc 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -473,6 +473,23 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
} \
} while (false)
+
+#if MLIR_ENABLE_NVFATBIN
+
+#include <nvFatbin.h>
+
+#define RETURN_ON_NVFATBIN_ERROR(expr) \
+ do { \
+ auto result = (expr); \
+ if (result != nvFatbinResult::NVFATBIN_SUCCESS) { \
+ emitError(loc) << llvm::Twine(#expr).concat(" failed with error: ") \
+ << nvFatbinGetErrorString(result); \
+ return std::nullopt; \
+ } \
+ } while (false)
+
+#endif // MLIR_ENABLE_NVFATBIN
+
std::optional<SmallVector<char, 0>>
NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
Location loc = getOperation().getLoc();
@@ -486,6 +503,9 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
targetOptions.tokenizeCmdOptions();
cmdOpts.second.append(
{"-arch", getTarget().getChip().data(), "--opt-level", optLevel.c_str()});
+ bool useFatbin32 = llvm::any_of(cmdOpts.second, [](const char *option) {
+ return llvm::StringRef(option) == "-32";
+ });
// Create the compiler handle.
RETURN_ON_NVPTXCOMPILER_ERROR(
@@ -538,6 +558,30 @@ NVPTXSerializer::compileToBinaryNVPTX(const std::string &ptxCode) {
});
#undef DEBUG_TYPE
RETURN_ON_NVPTXCOMPILER_ERROR(nvPTXCompilerDestroy(&compiler));
+
+#if MLIR_ENABLE_NVFATBIN
+ if (targetOptions.getCompilationTarget() == gpu::CompilationTarget::Fatbin) {
+ const char *cubinOpts[1] = {useFatbin32 ? "-32" : "-64"};
+ nvFatbinHandle handle;
+
+ auto chip = getTarget().getChip();
+ chip.consume_front("sm_");
+
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinCreate(&handle, cubinOpts, 1));
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinAddCubin(
+ handle, binary.data(), binary.size(), chip.data(), nullptr));
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinAddPTX(
+ handle, ptxCode.data(), ptxCode.size(), chip.data(), nullptr, nullptr));
+
+ size_t fatbinSize;
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinSize(handle, &fatbinSize));
+ SmallVector<char, 0> fatbin(fatbinSize, 0);
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinGet(handle, (void *)fatbin.data()));
+ RETURN_ON_NVFATBIN_ERROR(nvFatbinDestroy(&handle));
+ return fatbin;
+ }
+#endif // MLIR_ENABLE_NVFATBIN
+
return binary;
}
#endif // MLIR_ENABLE_NVPTXCOMPILER
More information about the Mlir-commits
mailing list