[mlir][Target] Improve ROCDL gpu serialization API

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

@@ -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) {
+    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

mlir/lib/Dialect/GPU/CMakeLists.txt

@@ -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}"

mlir/lib/Target/LLVM/CMakeLists.txt

@@ -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
-      "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")
+    set(DEFAULT_ROCM_PATH "" CACHE PATH "Fallback path to search for ROCm installs")
   endif()
   message(VERBOSE "MLIR Default ROCM toolkit path: ${DEFAULT_ROCM_PATH}")