diff --git a/mlir/tools/mlir-rocm-runner/CMakeLists.txt b/mlir/tools/mlir-rocm-runner/CMakeLists.txt --- a/mlir/tools/mlir-rocm-runner/CMakeLists.txt +++ b/mlir/tools/mlir-rocm-runner/CMakeLists.txt @@ -34,6 +34,9 @@ message(STATUS "ROCm HIP version: ${HIP_VERSION}") endif() + # Set compile-time flags for ROCm path. + add_definitions(-D__ROCM_PATH__="${ROCM_PATH}") + # Locate HIP runtime library. find_library(ROCM_RUNTIME_LIBRARY hip_hcc PATHS "${HIP_PATH}/lib") @@ -108,6 +111,11 @@ rocm-runtime-wrappers ) llvm_update_compile_flags(mlir-rocm-runner) + target_include_directories(mlir-rocm-runner + PRIVATE + "${HIP_PATH}/../include" + "${HIP_PATH}/include" + ) target_link_libraries(mlir-rocm-runner PRIVATE ${LIBS} ${targets_to_link}) endif() diff --git a/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp b/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp --- a/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp +++ b/mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp @@ -33,8 +33,11 @@ #include "mlir/Target/ROCDLIR.h" #include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/Passes.h" +#include "llvm/Support/ErrorOr.h" #include "llvm/Support/FileUtilities.h" #include "llvm/Support/InitLLVM.h" +#include "llvm/Support/LineIterator.h" +#include "llvm/Support/Program.h" #include "llvm/Support/SourceMgr.h" #include "llvm/Support/TargetRegistry.h" #include "llvm/Support/TargetSelect.h" @@ -58,6 +61,9 @@ // lld headers. #include "lld/Common/Driver.h" +// HIP headers. +#include "hip/hip_version.h" + using namespace mlir; using namespace llvm; @@ -67,15 +73,17 @@ cl::value_desc("triple string"), cl::init("amdgcn-amd-amdhsa")); -// TODO(whchung): Add feature to automatically detect available AMD GCN ISA -// version via `rocm-agent-enumerator` utility. static cl::opt targetChip("target", cl::desc("target chip"), cl::value_desc("AMDGPU ISA version"), - cl::init("gfx900")); + cl::init("")); static cl::opt features("feature", cl::desc("target features"), cl::value_desc("AMDGPU target features"), - cl::init("-code-object-v3")); + cl::init("")); + +static constexpr const char kRunnerProgram[] = "mlir-rocm-runner"; +static constexpr const char kRocmAgentEnumerator[] = "rocm_agent_enumerator"; +static constexpr const char kDefaultTargetChip[] = "gfx900"; static LogicalResult assembleIsa(const std::string isa, StringRef name, Blob &result) { @@ -211,10 +219,86 @@ return {}; } +static void configTargetChip() { + // Set targetChip to default value first. + targetChip = kDefaultTargetChip; + + // Locate rocm_agent_enumerator. + llvm::ErrorOr rocmAgentEnumerator = llvm::sys::findProgramByName( + kRocmAgentEnumerator, {__ROCM_PATH__ "/bin"}); + std::error_code ec; + if (ec = rocmAgentEnumerator.getError()) { + WithColor::warning(errs(), kRunnerProgram) + << kRocmAgentEnumerator << " couldn't be located under " + << __ROCM_PATH__ << ", set target as " << kDefaultTargetChip << "\n"; + return; + } + + // Prepare temp file to hold the outputs. + int tempFd = -1; + SmallString<128> tempFilename; + ec = sys::fs::createTemporaryFile("rocm_agent", "txt", tempFd, tempFilename); + if (ec) { + WithColor::warning(errs(), kRunnerProgram) + << "temporary file for " << kRocmAgentEnumerator + << " creation error, set target as " << kDefaultTargetChip << "\n"; + return; + } + FileRemover cleanup(tempFilename); + + // Invoke rocm_agent_enumerator. + std::string errorMessage; + SmallVector args{"-t", "GPU"}; + Optional redirects[3] = {{""}, tempFilename.str(), {""}}; + int result = + llvm::sys::ExecuteAndWait(rocmAgentEnumerator.get(), args, llvm::None, + redirects, 0, 0, &errorMessage); + if (result) { + WithColor::warning(errs(), kRunnerProgram) + << kRocmAgentEnumerator << " invocation error: " << errorMessage + << ", set target as " << kDefaultTargetChip << "\n"; + return; + } + + // Load and parse the result. + auto gfxIsaList = mlir::openInputFile(tempFilename); + if (!gfxIsaList) { + WithColor::error(errs(), kRunnerProgram) + << "read ROCm agent list temp file error, set target as " + << kDefaultTargetChip << "\n"; + return; + } + for (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. + targetChip = lines->str(); + break; + } +} + +static void configTargetFeatures() { + if (features.size() > 0) + features += ","; + // After ROCm 3.5, adopt HSA code object V3. + if (HIP_VERSION_MAJOR >= 3 && HIP_VERSION_MINOR >= 5) + features += "+code-object-v3"; + else + features += "-code-object-v3"; +} + static LogicalResult runMLIRPasses(ModuleOp m) { PassManager pm(m.getContext()); applyPassManagerCLOptions(pm); + // Configure target chip ISA version if it has not been specified. + if (!targetChip.size()) + configTargetChip(); + + // Configure target features per ROCm / HIP version. + configTargetFeatures(); + pm.addPass(createGpuKernelOutliningPass()); auto &kernelPm = pm.nest(); kernelPm.addPass(createStripDebugInfoPass());