-
Notifications
You must be signed in to change notification settings - Fork 11.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[mlir][Target] Improve ROCDL gpu serialization API #95456
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -27,6 +27,62 @@ namespace ROCDL { | |
/// 5. Returns an empty string. | ||
StringRef getROCMPath(); | ||
|
||
/// Helper class for specifying the AMD GCN device libraries required for | ||
/// compilation. | ||
class AMDGCNLibraryList { | ||
public: | ||
typedef enum : uint32_t { | ||
None = 0, | ||
Ockl = 1, | ||
Ocml = 2, | ||
OpenCL = 4, | ||
Hip = 8, | ||
LastLib = Hip, | ||
All = (LastLib << 1) - 1 | ||
} Library; | ||
|
||
explicit AMDGCNLibraryList(uint32_t libs = All) : libList(All & libs) {} | ||
|
||
/// Return a list with no libraries. | ||
static AMDGCNLibraryList getEmpty() { return AMDGCNLibraryList(None); } | ||
|
||
/// Return the libraries needed for compiling code with OpenCL calls. | ||
static AMDGCNLibraryList getOpenCL() { | ||
return AMDGCNLibraryList(Ockl | Ocml | OpenCL); | ||
} | ||
|
||
/// Returns true if the list is empty. | ||
bool isEmpty() const { return libList == None; } | ||
|
||
/// Adds a library to the list. | ||
AMDGCNLibraryList addLibrary(Library lib) { | ||
libList = libList | lib; | ||
return *this; | ||
} | ||
|
||
/// Adds all the libraries in `list` to the library list. | ||
AMDGCNLibraryList addList(AMDGCNLibraryList list) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. For instance, this is |
||
libList = libList | list.libList; | ||
return *this; | ||
} | ||
|
||
/// Removes a library from the list. | ||
AMDGCNLibraryList removeLibrary(Library lib) { | ||
libList = libList & ~lib; | ||
return *this; | ||
} | ||
|
||
/// Returns true if `lib` is in the list of libraries. | ||
bool requiresLibrary(Library lib) const { return (libList & lib) == lib; } | ||
|
||
/// Returns true if `libList` contains any of the libraries in `libs`. | ||
bool requiresAnyOf(uint32_t libs) const { return (libList & libs) != None; } | ||
|
||
private: | ||
/// Library list. | ||
uint32_t libList; | ||
}; | ||
|
||
/// Base class for all ROCDL serializations from GPU modules into binary | ||
/// strings. By default this class serializes into LLVM bitcode. | ||
class SerializeGPUModuleBase : public LLVM::ModuleToObject { | ||
|
@@ -49,8 +105,8 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject { | |
/// Returns the bitcode files to be loaded. | ||
ArrayRef<std::string> getFileList() const; | ||
|
||
/// Appends standard ROCm device libraries like `ocml.bc`, `ockl.bc`, etc. | ||
LogicalResult appendStandardLibs(); | ||
/// Appends standard ROCm device libraries to `fileList`. | ||
LogicalResult appendStandardLibs(AMDGCNLibraryList libs); | ||
|
||
/// Loads the bitcode files in `fileList`. | ||
virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>> | ||
|
@@ -63,15 +119,20 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject { | |
LogicalResult handleBitcodeFile(llvm::Module &module) override; | ||
|
||
protected: | ||
/// Appends the paths of common ROCm device libraries to `libs`. | ||
LogicalResult getCommonBitcodeLibs(llvm::SmallVector<std::string> &libs, | ||
SmallVector<char, 256> &libPath, | ||
StringRef isaVersion); | ||
|
||
/// Adds `oclc` control variables to the LLVM module. | ||
void addControlVariables(llvm::Module &module, bool wave64, bool daz, | ||
bool finiteOnly, bool unsafeMath, bool fastMath, | ||
bool correctSqrt, StringRef abiVer); | ||
void addControlVariables(llvm::Module &module, AMDGCNLibraryList libs, | ||
bool wave64, bool daz, bool finiteOnly, | ||
bool unsafeMath, bool fastMath, bool correctSqrt, | ||
StringRef abiVer); | ||
|
||
/// Compiles assembly to a binary. | ||
virtual std::optional<SmallVector<char, 0>> | ||
compileToBinary(const std::string &serializedISA); | ||
|
||
/// Default implementation of `ModuleToObject::moduleToObject`. | ||
std::optional<SmallVector<char, 0>> | ||
moduleToObjectImpl(const gpu::TargetOptions &targetOptions, | ||
llvm::Module &llvmModule); | ||
|
||
/// Returns the assembled ISA. | ||
std::optional<SmallVector<char, 0>> assembleIsa(StringRef isa); | ||
|
@@ -84,6 +145,9 @@ class SerializeGPUModuleBase : public LLVM::ModuleToObject { | |
|
||
/// List of LLVM bitcode files to link to. | ||
SmallVector<std::string> fileList; | ||
|
||
/// AMD GCN libraries to use when linking, the default is using all. | ||
AMDGCNLibraryList deviceLibs = AMDGCNLibraryList::getEmpty(); | ||
}; | ||
} // namespace ROCDL | ||
} // namespace mlir | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -123,17 +123,12 @@ add_mlir_dialect_library(MLIRROCDLTarget | |
) | ||
|
||
if(MLIR_ENABLE_ROCM_CONVERSIONS) | ||
if (NOT ("AMDGPU" IN_LIST LLVM_TARGETS_TO_BUILD)) | ||
message(SEND_ERROR | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Is this because there's a duplicate check? What happens if the AMDGPU backend isn't available? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. MLIR_ENABLE_ROCM_CONVERSIONS is already an alias for checking whether AMDGPU is being built or not. |
||
"Building mlir with ROCm support requires the AMDGPU backend") | ||
endif() | ||
|
||
if (DEFINED ROCM_PATH) | ||
set(DEFAULT_ROCM_PATH "${ROCM_PATH}" CACHE PATH "Fallback path to search for ROCm installs") | ||
elseif(DEFINED ENV{ROCM_PATH}) | ||
set(DEFAULT_ROCM_PATH "$ENV{ROCM_PATH}" CACHE PATH "Fallback path to search for ROCm installs") | ||
else() | ||
set(DEFAULT_ROCM_PATH "/opt/rocm" CACHE PATH "Fallback path to search for ROCm installs") | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This is meant to be a fallback so that the build goes to look in There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, and the flag is still available, however, having a hard coded value causes issues in windows, as the value is for linux. a better solution would be detecting it with CMake, I'll look into if there's anything like findHip in CMake. |
||
set(DEFAULT_ROCM_PATH "" CACHE PATH "Fallback path to search for ROCm installs") | ||
endif() | ||
message(VERBOSE "MLIR Default ROCM toolkit path: ${DEFAULT_ROCM_PATH}") | ||
|
||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Nit, isn't this just
enum class AMDGCNLibraryList : uint32_t
?... Actually, on top of that - there's already general support for bit enums. This could and should be tablegen'd
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I was trying to avoid going to tablegen for such a small class, but I'll switch it.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's more that there's a lot of common infrastructure around these sorts of flag enums (including, say, printing them) that can be autogenerated. And it'll be useful to have if someone ever wants to put this as an attribute somewhere