1 //===- Utils.h - MLIR ROCDL target utils ------------------------*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This files declares ROCDL target related utility classes and functions. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #ifndef MLIR_TARGET_LLVM_ROCDL_UTILS_H 14 #define MLIR_TARGET_LLVM_ROCDL_UTILS_H 15 16 #include "mlir/Dialect/GPU/IR/CompilationInterfaces.h" 17 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 18 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" 19 #include "mlir/IR/Attributes.h" 20 #include "mlir/Support/LLVM.h" 21 #include "mlir/Target/LLVM/ModuleToObject.h" 22 23 namespace mlir { 24 namespace ROCDL { 25 /// Searches & returns the path ROCM toolkit path, the search order is: 26 /// 1. The `ROCM_PATH` environment variable. 27 /// 2. The `ROCM_ROOT` environment variable. 28 /// 3. The `ROCM_HOME` environment variable. 29 /// 4. The ROCM path detected by CMake. 30 /// 5. Returns an empty string. 31 StringRef getROCMPath(); 32 33 /// Helper enum for specifying the AMD GCN device libraries required for 34 /// compilation. 35 enum class AMDGCNLibraries : uint32_t { 36 None = 0, 37 Ockl = 1, 38 Ocml = 2, 39 OpenCL = 4, 40 Hip = 8, 41 LastLib = Hip, 42 LLVM_MARK_AS_BITMASK_ENUM(LastLib), 43 All = (LastLib << 1) - 1 44 }; 45 46 /// Base class for all ROCDL serializations from GPU modules into binary 47 /// strings. By default this class serializes into LLVM bitcode. 48 class SerializeGPUModuleBase : public LLVM::ModuleToObject { 49 public: 50 /// Initializes the `toolkitPath` with the path in `targetOptions` or if empty 51 /// with the path in `getROCMPath`. 52 SerializeGPUModuleBase(Operation &module, ROCDLTargetAttr target, 53 const gpu::TargetOptions &targetOptions = {}); 54 55 /// Initializes the LLVM AMDGPU target by safely calling 56 /// `LLVMInitializeAMDGPU*` methods if available. 57 static void init(); 58 59 /// Returns the target attribute. 60 ROCDLTargetAttr getTarget() const; 61 62 /// Returns the ROCM toolkit path. 63 StringRef getToolkitPath() const; 64 65 /// Returns the LLVM bitcode libraries to be linked. 66 ArrayRef<Attribute> getLibrariesToLink() const; 67 68 /// Appends standard ROCm device libraries to `fileList`. 69 LogicalResult appendStandardLibs(AMDGCNLibraries libs); 70 71 /// Loads the bitcode files in `fileList`. 72 virtual std::optional<SmallVector<std::unique_ptr<llvm::Module>>> 73 loadBitcodeFiles(llvm::Module &module) override; 74 75 /// Determines required Device Libraries and adds `oclc` control variables to 76 /// the LLVM Module if needed. Also sets 77 /// `amdhsa_code_object_version` module flag. 78 void handleModulePreLink(llvm::Module &module) override; 79 80 /// Removes unnecessary metadata from the loaded bitcode files. 81 LogicalResult handleBitcodeFile(llvm::Module &module) override; 82 83 protected: 84 /// Adds `oclc` control variables to the LLVM Module if needed. It also sets 85 /// `amdhsa_code_object_version` module flag which is equal to ABI version and 86 /// it uses "llvm::Module::Error" to set that flag. 87 void addControlVariables(llvm::Module &module, AMDGCNLibraries libs, 88 bool wave64, bool daz, bool finiteOnly, 89 bool unsafeMath, bool fastMath, bool correctSqrt, 90 StringRef abiVer); 91 92 /// Compiles assembly to a binary. 93 virtual std::optional<SmallVector<char, 0>> 94 compileToBinary(const std::string &serializedISA); 95 96 /// Default implementation of `ModuleToObject::moduleToObject`. 97 std::optional<SmallVector<char, 0>> 98 moduleToObjectImpl(const gpu::TargetOptions &targetOptions, 99 llvm::Module &llvmModule); 100 101 /// Returns the assembled ISA. 102 std::optional<SmallVector<char, 0>> assembleIsa(StringRef isa); 103 104 /// ROCDL target attribute. 105 ROCDLTargetAttr target; 106 107 /// ROCM toolkit path. 108 std::string toolkitPath; 109 110 /// List of LLVM bitcode files to link to. 111 SmallVector<Attribute> librariesToLink; 112 113 /// AMD GCN libraries to use when linking, the default is using none. 114 AMDGCNLibraries deviceLibs = AMDGCNLibraries::None; 115 }; 116 117 /// Returns a map containing the `amdhsa.kernels` ELF metadata for each of the 118 /// kernels in the binary, or `std::nullopt` if the metadata couldn't be 119 /// retrieved. The map associates the name of the kernel with the list of named 120 /// attributes found in `amdhsa.kernels`. For more information on the ELF 121 /// metadata see: https://llvm.org/docs/AMDGPUUsage.html#amdhsa 122 std::optional<DenseMap<StringAttr, NamedAttrList>> 123 getAMDHSAKernelsELFMetadata(Builder &builder, ArrayRef<char> elfData); 124 125 /// Returns a `#gpu.kernel_table` containing kernel metadata for each of the 126 /// kernels in `gpuModule`. If `elfData` is valid, then the `amdhsa.kernels` ELF 127 /// metadata will be added to the `#gpu.kernel_table`. 128 gpu::KernelTableAttr getKernelMetadata(Operation *gpuModule, 129 ArrayRef<char> elfData = {}); 130 } // namespace ROCDL 131 } // namespace mlir 132 133 #endif // MLIR_TARGET_LLVM_ROCDL_UTILS_H 134