[mlir][Target] ROCDL

Author: fabianmcg

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

@@ -27,6 +27,64 @@ 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) != None; }
+
+  /// Returns true if `libList` contains all the libraries in `libs`.
+  bool containLibraries(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 +107,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 Library to `fileList`.
+  LogicalResult appendStandardLibs(AMDGCNLibraryList libs);
 
   /// Loads the bitcode files in `fileList`.
   virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>>
@@ -63,15 +121,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 +147,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