[MLIR][GPU] Make the path to ROCm a runtime option

[krzysz00]

Summary

Our current build assumes that the path to ROCm we find at build time
will be the path at which ROCm is located when the built code is
executed. This commit adds a --rocm-path option to SerializeToHsaco,
and removes the HIP dependency that the SerializeToHsaco previously had.

Depends on D114113

(though the dependency is to ensure the diffs apply cleanly and to capture the dependency on D114107)

[mehdi_amini]

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

71

I think you only need this for the optLevel because it isn't a pass option right now.

89

I think this should be in the parent revision?

93–94

Can you make this a pass option like the other options?

94–100

I'd really rather avoid using environment variable when we have pass options available.

[krzysz00]

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

93–94

Sure, I can do that. I'll fold that stuff into the parent revision

94–100

This is the clang behavior, which I'm mostly replicating

[mehdi_amini]

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

94–100

If clang or any tool want to load from the environment, that's fine: but from a layering point of view I see it as their responsibility to do so and set it as a pass option when they create the passes.

[krzysz00]

• Address review comments

[krzysz00]

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

89

Yep, and the only reason any of this worked is that this whole file was #ifdef'd off with how the CMakeLists.txt was written.

[krzysz00]

Re-push to account for earlier merged patches.

mlir/lib/Dialect/GPU/CMakeLists.txt

message(SEND_ERROR
       "Building mlir with ROCm support requires the AMDGPU backend")
   endif()
 
   # Ensure lld is enabled.
   if (NOT "lld" IN_LIST LLVM_ENABLE_PROJECTS)
     message(SEND_ERROR "lld is not enabled. Please revise LLVM_ENABLE_PROJECTS")
   endif()
 
-  # Configure ROCm support.
-  if (NOT DEFINED ROCM_PATH)
-    if (NOT DEFINED ENV{ROCM_PATH})
-      set(ROCM_PATH "/opt/rocm" CACHE PATH "Path to which ROCm has been installed")
-    else()
-      set(ROCM_PATH $ENV{ROCM_PATH} CACHE PATH "Path to which ROCm has been installed")
-    endif()
-    set(HIP_PATH "${ROCM_PATH}/hip" CACHE PATH " Path to which HIP has been installed")
-  endif()
-  set(CMAKE_MODULE_PATH "${HIP_PATH}/cmake" ${CMAKE_MODULE_PATH})
-  find_package(HIP)
-  if (NOT HIP_FOUND)
-    message(SEND_ERROR "Building mlir with ROCm support requires a working ROCm and HIP install")
-  else()
-    message(STATUS "ROCm HIP version: ${HIP_VERSION}")
-  endif()
-
+  set(DEFAULT_ROCM_PATH "/opt/rocm" CACHE PATH "Fallback path to search for ROCm installs")
   target_compile_definitions(obj.MLIRGPUTransforms
     PRIVATE
-    __HIP_PLATFORM_HCC__
-    __ROCM_PATH__="${ROCM_PATH}"
+    __DEFAULT_ROCM_PATH__="${DEFAULT_ROCM_PATH}"
     MLIR_GPU_TO_HSACO_PASS_ENABLE=1
   )
 
   target_include_directories(obj.MLIRGPUTransforms
     PRIVATE
     ${MLIR_SOURCE_DIR}/../lld/include
-    ${HIP_PATH}/include
-    ${ROCM_PATH}/include
   )
 
   target_link_libraries(MLIRGPUOps
     PRIVATE
     lldELF
     MLIRExecutionEngine
     MLIRROCDLToLLVMIRTranslation
   )
 
   # Link lldELF also to libmlir.so. Create an alias that starts with LLVM
   # because LINK_COMPONENTS elements are implicitly prefixed with LLVM.
   add_library(LLVMAliasTolldELF ALIAS lldELF)
   set_property(GLOBAL APPEND PROPERTY MLIR_LLVM_LINK_COMPONENTS AliasTolldELF)
 
 endif()

mlir/lib/Dialect/GPU/Transforms/SerializeToHsaco.cpp

 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/MCTargetAsmParser.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSubtargetInfo.h"
-
 #include "llvm/MC/TargetRegistry.h"
+
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
-#include "llvm/Support/LineIterator.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TargetSelect.h"
 #include "llvm/Support/WithColor.h"
 
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 
 #include "lld/Common/Driver.h"
 
-#include "hip/hip_version.h"
-
 #include <mutex>
 
 using namespace mlir;
 
 namespace {
 class SerializeToHsacoPass
     : public PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass> {
 public:
   SerializeToHsacoPass(StringRef triple, StringRef arch, StringRef features,
                        int optLevel);
   SerializeToHsacoPass(const SerializeToHsacoPass &other);
   StringRef getArgument() const override { return "gpu-to-hsaco"; }
   StringRef getDescription() const override {
     return "Lower GPU kernel function to HSACO binary annotations";
   }
 
 protected:
   Option<int> optLevel{
       *this, "opt-level",
       llvm::cl::desc("Optimization level for HSACO compilation"),
       llvm::cl::init(2)};
 
+  Option<std::string> rocmPath{*this, "rocm-path",
+                               llvm::cl::desc("Path to ROCm install")};
+
   /// Adds LLVM optimization passes
   LogicalResult optimizeLlvm(llvm::Module &llvmModule,
                              llvm::TargetMachine &targetMachine) override;
 
 private:
   void getDependentDialects(DialectRegistry &registry) const override;
 
   // Serializes ROCDL to HSACO.
   std::unique_ptr<std::vector<char>>
   serializeISA(const std::string &isa) override;
 
   std::unique_ptr<SmallVectorImpl<char>> assembleIsa(const std::string &isa);
   std::unique_ptr<std::vector<char>>
   createHsaco(const SmallVectorImpl<char> &isaBinary);
+
+  std::string getRocmPath();
 };
-} // namespace
+} // end namespace
 
 SerializeToHsacoPass::SerializeToHsacoPass(const SerializeToHsacoPass &other)
     : PassWrapper<SerializeToHsacoPass, gpu::SerializeToBlobPass>(other) {}
-static std::string getDefaultChip() {
-  const char kDefaultChip[] = "gfx900";
 
-  // Locate rocm_agent_enumerator.
-  const char kRocmAgentEnumerator[] = "rocm_agent_enumerator";
-  llvm::ErrorOr<std::string> rocmAgentEnumerator = llvm::sys::findProgramByName(
-      kRocmAgentEnumerator, {__ROCM_PATH__ "/bin"});
-  if (!rocmAgentEnumerator) {
-    llvm::WithColor::warning(llvm::errs())
-        << kRocmAgentEnumerator << "couldn't be located under " << __ROCM_PATH__
-        << "/bin\n";
-    return kDefaultChip;
-  }
-
-  // Prepare temp file to hold the outputs.
-  int tempFd = -1;
-  SmallString<128> tempFilename;
-  if (llvm::sys::fs::createTemporaryFile("rocm_agent", "txt", tempFd,
-                                         tempFilename)) {
-    llvm::WithColor::warning(llvm::errs())
-        << "temporary file for " << kRocmAgentEnumerator << " creation error\n";
-    return kDefaultChip;
-  }
-  llvm::FileRemover cleanup(tempFilename);
-
-  // Invoke rocm_agent_enumerator.
-  std::string errorMessage;
-  SmallVector<StringRef, 2> args{"-t", "GPU"};
-  Optional<StringRef> redirects[3] = {{""}, tempFilename.str(), {""}};
-  int result =
-      llvm::sys::ExecuteAndWait(rocmAgentEnumerator.get(), args, llvm::None,
-                                redirects, 0, 0, &errorMessage);
-  if (result) {
-    llvm::WithColor::warning(llvm::errs())
-        << kRocmAgentEnumerator << " invocation error: " << errorMessage
-        << "\n";
-    return kDefaultChip;
-  }
-
-  // Load and parse the result.
-  auto gfxIsaList = openInputFile(tempFilename);
-  if (!gfxIsaList) {
-    llvm::WithColor::error(llvm::errs())
-        << "read ROCm agent list temp file error\n";
-    return kDefaultChip;
-  }
-  for (llvm::line_iterator lines(*gfxIsaList); !lines.is_at_end(); ++lines) {
-    // Skip the line with content "gfx000".
-    if (*lines == "gfx000")
-      continue;
-    // Use the first ISA version found.
-    return lines->str();
-  }
+/// Get a user-specified path to ROCm
+// Tries, in order, the --rocm-path option, the ROCM_PATH environment variable
+// and a compile-time default
+std::string SerializeToHsacoPass::getRocmPath() {
+  if (rocmPath.getNumOccurrences() > 0)
+    return rocmPath.getValue();
 
-  return kDefaultChip;
+  return __DEFAULT_ROCM_PATH__;
 }
 
 // Sets the 'option' to 'value' unless it already has a value.
 static void maybeSetOption(Pass::Option<std::string> &option,
                            function_ref<std::string()> getValue) {
   if (!option.hasValue())
     option = getValue();
 }