[Mlir-commits] [mlir] [MLIR][NVVM][GPU] Do not expect nvidia `fatbinary` to be present (PR #94367)

[Christos Perivolaropoulos]

https://github.com/cperivol created https://github.com/llvm/llvm-project/pull/94367

By default use only ptxas to compile NVVM IR. If the user requests fatbin, only then use `fatbinary` and complain if it's missing.

>From af58a8ad715923421b76a36c8940d57b1701be9f Mon Sep 17 00:00:00 2001
From: Christos Perivolaropoulos <cperivol at google.com>
Date: Tue, 4 Jun 2024 15:32:12 +0000
Subject: [PATCH] [MLIR][NVVM][GPU] Do not expect nvidia `fatbinary` to be
 present

By default use only ptxas to compile NVVM IR. If the user requests
fatbin, only then use `fatbinary` and complain if it's missing.
---
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp |  2 +-
 mlir/lib/Target/LLVM/NVVM/Target.cpp   | 34 +++++++++++++++-----------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index 0c2590d711301..269ffe4b25b75 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2263,7 +2263,7 @@ CompilationTarget TargetOptions::getCompilationTarget() const {
 }
 
 CompilationTarget TargetOptions::getDefaultCompilationTarget() {
-  return CompilationTarget::Fatbin;
+  return CompilationTarget::Binary;
 }
 
 std::pair<llvm::BumpPtrAllocator, SmallVector<const char *>>
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index acb903aa37caf..0a2ccc989df2c 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -22,6 +22,7 @@
 #include "mlir/Target/LLVMIR/Export.h"
 
 #include "llvm/Config/llvm-config.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/FormatVariadic.h"
@@ -191,7 +192,10 @@ class NVPTXSerializer : public SerializeGPUModuleBase {
   /// 1. The toolkit path in `targetOptions`.
   /// 2. In the system PATH.
   /// 3. The path from `getCUDAToolkitPath()`.
-  std::optional<std::string> findTool(StringRef tool);
+  ///
+  /// If warn_on_missing is true we emit a warning to the user when the tool is
+  /// missing.
+  std::optional<std::string> findTool(StringRef tool, bool warn_on_missing);
 
   /// Target options.
   gpu::TargetOptions targetOptions;
@@ -220,7 +224,7 @@ gpu::GPUModuleOp NVPTXSerializer::getOperation() {
   return dyn_cast<gpu::GPUModuleOp>(&SerializeGPUModuleBase::getOperation());
 }
 
-std::optional<std::string> NVPTXSerializer::findTool(StringRef tool) {
+std::optional<std::string> NVPTXSerializer::findTool(StringRef tool, bool warn_on_missing) {
   // Find the `tool` path.
   // 1. Check the toolkit path given in the command line.
   StringRef pathRef = targetOptions.getToolkitPath();
@@ -246,11 +250,15 @@ std::optional<std::string> NVPTXSerializer::findTool(StringRef tool) {
     if (llvm::sys::fs::can_execute(path))
       return StringRef(path.data(), path.size()).str();
   }
-  getOperation().emitError()
-      << "Couldn't find the `" << tool
-      << "` binary. Please specify the toolkit "
-         "path, add the compiler to $PATH, or set one of the environment "
-         "variables in `NVVM::getCUDAToolkitPath()`.";
+
+  // If we reached this point we failed to find the tool.
+  if (warn_on_missing) {
+    getOperation().emitWarning()
+        << "Couldn't find the `" << tool
+        << "` binary. Please specify the toolkit "
+        "path, add the compiler to $PATH, or set one of the environment "
+        "variables in `NVVM::getCUDAToolkitPath()`.";
+  }
   return std::nullopt;
 }
 
@@ -264,15 +272,13 @@ NVPTXSerializer::compileToBinary(const std::string &ptxCode) {
       targetOptions.getCompilationTarget() == gpu::CompilationTarget::Fatbin;
 
   // Find the `ptxas` & `fatbinary` tools.
-  std::optional<std::string> ptxasCompiler = findTool("ptxas");
+  std::optional<std::string> ptxasCompiler = findTool("ptxas", true);
   if (!ptxasCompiler)
     return std::nullopt;
-  std::optional<std::string> fatbinaryTool;
-  if (createFatbin) {
-    fatbinaryTool = findTool("fatbinary");
-    if (!fatbinaryTool)
-      return std::nullopt;
-  }
+  std::optional<std::string> fatbinaryTool =
+      findTool("fatbinary", createFatbin);
+  if (createFatbin && !fatbinaryTool)
+    return std::nullopt;
   Location loc = getOperation().getLoc();
 
   // Base name for all temp files: mlir-<module name>-<target triple>-<chip>.