Skip to content

Commit 70fb1e3

Browse files
authored
Reland [mlir][Target] Improve ROCDL gpu serialization API (#96198)
Reland: #95456 This patch improves the ROCDL gpu serialization API by: - Introducing the enum `AMDGCNLibraries` for specifying the AMD GCN device code libraries to use during linking. - Removing `getCommonBitcodeLibs` in favor of `AMDGCNLibraries`. Previously `getCommonBitcodeLibs` would try to load all AMD GCN bitcode librariesm now it will only load the requested libraries. - Exposing the `compileToBinary` method and making it virtual, allowing downstream users to re-use this method. - Exposing `moduleToObjectImpl`, this method provides a prototype flow for compiling to binary, allowing downstream users to re-use this method. - It also avoids constructing the control variables if no device libraries are being used. - Changes the style of the error messages to be composable, ie no full stops. - Adds an error message for when the ROCm toolkit can't be found but it was required.
1 parent 5c9513a commit 70fb1e3

File tree

3 files changed

+205
-151
lines changed

3 files changed

+205
-151
lines changed

mlir/include/mlir/Target/LLVM/ROCDL/Utils.h

Lines changed: 31 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,19 @@ namespace ROCDL {
2727
/// 5. Returns an empty string.
2828
StringRef getROCMPath();
2929

30+
/// Helper enum for specifying the AMD GCN device libraries required for
31+
/// compilation.
32+
enum class AMDGCNLibraries : uint32_t {
33+
None = 0,
34+
Ockl = 1,
35+
Ocml = 2,
36+
OpenCL = 4,
37+
Hip = 8,
38+
LastLib = Hip,
39+
LLVM_MARK_AS_BITMASK_ENUM(LastLib),
40+
All = (LastLib << 1) - 1
41+
};
42+
3043
/// Base class for all ROCDL serializations from GPU modules into binary
3144
/// strings. By default this class serializes into LLVM bitcode.
3245
class SerializeGPUModuleBase : public LLVM::ModuleToObject {
@@ -49,8 +62,8 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
4962
/// Returns the bitcode files to be loaded.
5063
ArrayRef<std::string> getFileList() const;
5164

52-
/// Appends standard ROCm device libraries like `ocml.bc`, `ockl.bc`, etc.
53-
LogicalResult appendStandardLibs();
65+
/// Appends standard ROCm device libraries to `fileList`.
66+
LogicalResult appendStandardLibs(AMDGCNLibraries libs);
5467

5568
/// Loads the bitcode files in `fileList`.
5669
virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
@@ -63,15 +76,20 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
6376
LogicalResult handleBitcodeFile(llvm::Module &module) override;
6477

6578
protected:
66-
/// Appends the paths of common ROCm device libraries to `libs`.
67-
LogicalResult getCommonBitcodeLibs(llvm::SmallVector<std::string> &libs,
68-
SmallVector<char, 256> &libPath,
69-
StringRef isaVersion);
70-
7179
/// Adds `oclc` control variables to the LLVM module.
72-
void addControlVariables(llvm::Module &module, bool wave64, bool daz,
73-
bool finiteOnly, bool unsafeMath, bool fastMath,
74-
bool correctSqrt, StringRef abiVer);
80+
void addControlVariables(llvm::Module &module, AMDGCNLibraries libs,
81+
bool wave64, bool daz, bool finiteOnly,
82+
bool unsafeMath, bool fastMath, bool correctSqrt,
83+
StringRef abiVer);
84+
85+
/// Compiles assembly to a binary.
86+
virtual std::optional<SmallVector<char, 0>>
87+
compileToBinary(const std::string &serializedISA);
88+
89+
/// Default implementation of `ModuleToObject::moduleToObject`.
90+
std::optional<SmallVector<char, 0>>
91+
moduleToObjectImpl(const gpu::TargetOptions &targetOptions,
92+
llvm::Module &llvmModule);
7593

7694
/// Returns the assembled ISA.
7795
std::optional<SmallVector<char, 0>> assembleIsa(StringRef isa);
@@ -84,6 +102,9 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject {
84102

85103
/// List of LLVM bitcode files to link to.
86104
SmallVector<std::string> fileList;
105+
106+
/// AMD GCN libraries to use when linking, the default is using none.
107+
AMDGCNLibraries deviceLibs = AMDGCNLibraries::None;
87108
};
88109
} // namespace ROCDL
89110
} // namespace mlir

mlir/lib/Target/LLVM/CMakeLists.txt

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -127,17 +127,18 @@ add_mlir_dialect_library(MLIRROCDLTarget
127127
)
128128

129129
if(MLIR_ENABLE_ROCM_CONVERSIONS)
130-
if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD))
131-
message(SEND_ERROR
132-
"Building mlir with ROCm support requires the AMDGPU backend")
133-
endif()
134-
135130
if (DEFINED ROCM_PATH)
136131
set(DEFAULT_ROCM_PATH "${ROCM_PATH}" CACHE PATH "Fallback path to search for ROCm installs")
137132
elseif(DEFINED ENV{ROCM_PATH})
138133
set(DEFAULT_ROCM_PATH "$ENV{ROCM_PATH}" CACHE PATH "Fallback path to search for ROCm installs")
139134
else()
140-
set(DEFAULT_ROCM_PATH "/opt/rocm" CACHE PATH "Fallback path to search for ROCm installs")
135+
if (WIN32)
136+
# Avoid setting an UNIX path for Windows.
137+
# TODO: Eventually migrate to FindHIP once it becomes a part of CMake.
138+
set(DEFAULT_ROCM_PATH "" CACHE PATH "Fallback path to search for ROCm installs")
139+
else()
140+
set(DEFAULT_ROCM_PATH "/opt/rocm" CACHE PATH "Fallback path to search for ROCm installs")
141+
endif()
141142
endif()
142143
message(VERBOSE "MLIR Default ROCM toolkit path: ${DEFAULT_ROCM_PATH}")
143144

0 commit comments

Comments
 (0)