Skip to content
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

Merged
merged 2 commits into from
Jun 17, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 31 additions & 10 deletions mlir/include/mlir/Target/LLVM/ROCDL/Utils.h
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,19 @@ namespace ROCDL {
/// 5. Returns an empty string.
StringRef getROCMPath();

/// Helper enum for specifying the AMD GCN device libraries required for
/// compilation.
enum class AMDGCNLibraries : uint32_t {
None = 0,
Ockl = 1,
Ocml = 2,
OpenCL = 4,
Hip = 8,
LastLib = Hip,
LLVM_MARK_AS_BITMASK_ENUM(LastLib),
All = (LastLib << 1) - 1
};

/// 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 {
Expand All @@ -49,8 +62,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(AMDGCNLibraries libs);

/// Loads the bitcode files in `fileList`.
virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
Expand All @@ -63,15 +76,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, AMDGCNLibraries 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);
Expand All @@ -84,6 +102,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 none.
AMDGCNLibraries deviceLibs = AMDGCNLibraries::None;
};
} // namespace ROCDL
} // namespace mlir
Expand Down
2 changes: 1 addition & 1 deletion mlir/lib/Dialect/GPU/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ if(MLIR_ENABLE_ROCM_CONVERSIONS)
"Building mlir with ROCm support requires the AMDGPU backend")
endif()

set(DEFAULT_ROCM_PATH "/opt/rocm" CACHE PATH "Fallback path to search for ROCm installs")
set(DEFAULT_ROCM_PATH "" CACHE PATH "Fallback path to search for ROCm installs")
target_compile_definitions(obj.MLIRGPUTransforms
PRIVATE
__DEFAULT_ROCM_PATH__="${DEFAULT_ROCM_PATH}"
Expand Down
7 changes: 1 addition & 6 deletions mlir/lib/Target/LLVM/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Contributor

Choose a reason for hiding this comment

The 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?

Copy link
Contributor Author

Choose a reason for hiding this comment

The 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")
Copy link
Contributor

Choose a reason for hiding this comment

The 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 /opt/rocm for the device libraries if there are no clues about where they are in the user's environment, though?

Copy link
Contributor Author

Choose a reason for hiding this comment

The 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}")

Expand Down
Loading
Loading