Revert "[MLIR] Add callback functions for ModuleToObject"

[joker-eph]

Reverts #116007

Bot is broken.

[llvmbot]

@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Mehdi Amini (joker-eph)

Changes

Reverts llvm/llvm-project#116007

Bot is broken.

---

Full diff: https://github.com/llvm/llvm-project/pull/116811.diff

6 Files Affected:

• (modified) mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h (+2-42)
• (modified) mlir/include/mlir/Target/LLVM/ModuleToObject.h (+2-22)
• (modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+4-37)
• (modified) mlir/lib/Target/LLVM/ModuleToObject.cpp (+3-19)
• (modified) mlir/lib/Target/LLVM/NVVM/Target.cpp (-3)
• (modified) mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp (-62)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
index d4b16a1de8eddc..6d7cb5ca7a7f81 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
@@ -14,7 +14,6 @@
 #define MLIR_DIALECT_GPU_IR_COMPILATIONINTERFACES_H
 
 #include "mlir/IR/Attributes.h"
-#include "llvm/IR/Module.h"
 
 namespace llvm {
 class IRBuilderBase;
@@ -53,11 +52,7 @@ class TargetOptions {
       StringRef toolkitPath = {}, ArrayRef<std::string> linkFiles = {},
       StringRef cmdOptions = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
-      function_ref<SymbolTable *()> getSymbolTableCallback = {},
-      function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
-      function_ref<void(StringRef)> isaCallback = {});
+      function_ref<SymbolTable *()> getSymbolTableCallback = {});
 
   /// Returns the typeID.
   TypeID getTypeID() const;
@@ -85,22 +80,6 @@ class TargetOptions {
   /// table.
   SymbolTable *getSymbolTable() const;
 
-  /// Returns the callback invoked with the initial LLVM IR for the device
-  /// module.
-  function_ref<void(llvm::Module &)> getInitialLlvmIRCallback() const;
-
-  /// Returns the callback invoked with LLVM IR for the device module
-  /// after linking the device libraries.
-  function_ref<void(llvm::Module &)> getLinkedLlvmIRCallback() const;
-
-  /// Returns the callback invoked with LLVM IR for the device module after
-  /// LLVM optimizations but before codegen.
-  function_ref<void(llvm::Module &)> getOptimizedLlvmIRCallback() const;
-
-  /// Returns the callback invoked with the target ISA for the device,
-  /// for example PTX assembly.
-  function_ref<void(StringRef)> getISACallback() const;
-
   /// Returns the default compilation target: `CompilationTarget::Fatbin`.
   static CompilationTarget getDefaultCompilationTarget();
 
@@ -111,11 +90,7 @@ class TargetOptions {
       TypeID typeID, StringRef toolkitPath = {},
       ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
       CompilationTarget compilationTarget = getDefaultCompilationTarget(),
-      function_ref<SymbolTable *()> getSymbolTableCallback = {},
-      function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
-      function_ref<void(StringRef)> isaCallback = {});
+      function_ref<SymbolTable *()> getSymbolTableCallback = {});
 
   /// Path to the target toolkit.
   std::string toolkitPath;
@@ -134,21 +109,6 @@ class TargetOptions {
   /// being serialized.
   function_ref<SymbolTable *()> getSymbolTableCallback;
 
-  /// Callback invoked with the initial LLVM IR for the device module.
-  function_ref<void(llvm::Module &)> initialLlvmIRCallback;
-
-  /// Callback invoked with LLVM IR for the device module after
-  /// linking the device libraries.
-  function_ref<void(llvm::Module &)> linkedLlvmIRCallback;
-
-  /// Callback invoked with LLVM IR for the device module after
-  /// LLVM optimizations but before codegen.
-  function_ref<void(llvm::Module &)> optimizedLlvmIRCallback;
-
-  /// Callback invoked with the target ISA for the device,
-  /// for example PTX assembly.
-  function_ref<void(StringRef)> isaCallback;
-
 private:
   TypeID typeID;
 };
diff --git a/mlir/include/mlir/Target/LLVM/ModuleToObject.h b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
index 07fc55b41ae9c5..e40d7e9a43dd6b 100644
--- a/mlir/include/mlir/Target/LLVM/ModuleToObject.h
+++ b/mlir/include/mlir/Target/LLVM/ModuleToObject.h
@@ -29,13 +29,8 @@ class ModuleTranslation;
 /// operations being transformed must be translatable into LLVM IR.
 class ModuleToObject {
 public:
-  ModuleToObject(
-      Operation &module, StringRef triple, StringRef chip,
-      StringRef features = {}, int optLevel = 3,
-      function_ref<void(llvm::Module &)> initialLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> linkedLlvmIRCallback = {},
-      function_ref<void(llvm::Module &)> optimizedLlvmIRCallback = {},
-      function_ref<void(StringRef)> isaCallback = {});
+  ModuleToObject(Operation &module, StringRef triple, StringRef chip,
+                 StringRef features = {}, int optLevel = 3);
   virtual ~ModuleToObject();
 
   /// Returns the operation being serialized.
@@ -119,21 +114,6 @@ class ModuleToObject {
   /// Optimization level.
   int optLevel;
 
-  /// Callback invoked with the initial LLVM IR for the device module.
-  function_ref<void(llvm::Module &)> initialLlvmIRCallback;
-
-  /// Callback invoked with LLVM IR for the device module after
-  /// linking the device libraries.
-  function_ref<void(llvm::Module &)> linkedLlvmIRCallback;
-
-  /// Callback invoked with LLVM IR for the device module after
-  /// LLVM optimizations but before codegen.
-  function_ref<void(llvm::Module &)> optimizedLlvmIRCallback;
-
-  /// Callback invoked with the target ISA for the device,
-  /// for example PTX assembly.
-  function_ref<void(StringRef)> isaCallback;
-
 private:
   /// The TargetMachine created for the given Triple, if available.
   /// Accessible through `getOrCreateTargetMachine()`.
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index d62ea72dcea2f6..956877497d9338 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -2302,31 +2302,17 @@ KernelMetadataAttr KernelTableAttr::lookup(StringAttr key) const {
 TargetOptions::TargetOptions(
     StringRef toolkitPath, ArrayRef<std::string> linkFiles,
     StringRef cmdOptions, CompilationTarget compilationTarget,
-    function_ref<SymbolTable *()> getSymbolTableCallback,
-    function_ref<void(llvm::Module &)> initialLlvmIRCallback,
-    function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
-    function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
-    function_ref<void(StringRef)> isaCallback)
+    function_ref<SymbolTable *()> getSymbolTableCallback)
     : TargetOptions(TypeID::get<TargetOptions>(), toolkitPath, linkFiles,
-                    cmdOptions, compilationTarget, getSymbolTableCallback,
-                    initialLlvmIRCallback, linkedLlvmIRCallback,
-                    optimizedLlvmIRCallback, isaCallback) {}
+                    cmdOptions, compilationTarget, getSymbolTableCallback) {}
 
 TargetOptions::TargetOptions(
     TypeID typeID, StringRef toolkitPath, ArrayRef<std::string> linkFiles,
     StringRef cmdOptions, CompilationTarget compilationTarget,
-    function_ref<SymbolTable *()> getSymbolTableCallback,
-    function_ref<void(llvm::Module &)> initialLlvmIRCallback,
-    function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
-    function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
-    function_ref<void(StringRef)> isaCallback)
+    function_ref<SymbolTable *()> getSymbolTableCallback)
     : toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
       cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
-      getSymbolTableCallback(getSymbolTableCallback),
-      initialLlvmIRCallback(initialLlvmIRCallback),
-      linkedLlvmIRCallback(linkedLlvmIRCallback),
-      optimizedLlvmIRCallback(optimizedLlvmIRCallback),
-      isaCallback(isaCallback), typeID(typeID) {}
+      getSymbolTableCallback(getSymbolTableCallback), typeID(typeID) {}
 
 TypeID TargetOptions::getTypeID() const { return typeID; }
 
@@ -2340,25 +2326,6 @@ SymbolTable *TargetOptions::getSymbolTable() const {
   return getSymbolTableCallback ? getSymbolTableCallback() : nullptr;
 }
 
-function_ref<void(llvm::Module &)>
-TargetOptions::getInitialLlvmIRCallback() const {
-  return initialLlvmIRCallback;
-}
-
-function_ref<void(llvm::Module &)>
-TargetOptions::getLinkedLlvmIRCallback() const {
-  return linkedLlvmIRCallback;
-}
-
-function_ref<void(llvm::Module &)>
-TargetOptions::getOptimizedLlvmIRCallback() const {
-  return optimizedLlvmIRCallback;
-}
-
-function_ref<void(StringRef)> TargetOptions::getISACallback() const {
-  return isaCallback;
-}
-
 CompilationTarget TargetOptions::getCompilationTarget() const {
   return compilationTarget;
 }
diff --git a/mlir/lib/Target/LLVM/ModuleToObject.cpp b/mlir/lib/Target/LLVM/ModuleToObject.cpp
index 3f5b3d5e31864b..77391341adaad2 100644
--- a/mlir/lib/Target/LLVM/ModuleToObject.cpp
+++ b/mlir/lib/Target/LLVM/ModuleToObject.cpp
@@ -34,17 +34,10 @@
 using namespace mlir;
 using namespace mlir::LLVM;
 
-ModuleToObject::ModuleToObject(
-    Operation &module, StringRef triple, StringRef chip, StringRef features,
-    int optLevel, function_ref<void(llvm::Module &)> initialLlvmIRCallback,
-    function_ref<void(llvm::Module &)> linkedLlvmIRCallback,
-    function_ref<void(llvm::Module &)> optimizedLlvmIRCallback,
-    function_ref<void(StringRef)> isaCallback)
+ModuleToObject::ModuleToObject(Operation &module, StringRef triple,
+                               StringRef chip, StringRef features, int optLevel)
     : module(module), triple(triple), chip(chip), features(features),
-      optLevel(optLevel), initialLlvmIRCallback(initialLlvmIRCallback),
-      linkedLlvmIRCallback(linkedLlvmIRCallback),
-      optimizedLlvmIRCallback(optimizedLlvmIRCallback),
-      isaCallback(isaCallback) {}
+      optLevel(optLevel) {}
 
 ModuleToObject::~ModuleToObject() = default;
 
@@ -222,9 +215,6 @@ std::optional<SmallVector<char, 0>> ModuleToObject::run() {
   }
   setDataLayoutAndTriple(*llvmModule);
 
-  if (initialLlvmIRCallback)
-    initialLlvmIRCallback(*llvmModule);
-
   // Link bitcode files.
   handleModulePreLink(*llvmModule);
   {
@@ -237,16 +227,10 @@ std::optional<SmallVector<char, 0>> ModuleToObject::run() {
     handleModulePostLink(*llvmModule);
   }
 
-  if (linkedLlvmIRCallback)
-    linkedLlvmIRCallback(*llvmModule);
-
   // Optimize the module.
   if (failed(optimizeModule(*llvmModule, optLevel)))
     return std::nullopt;
 
-  if (optimizedLlvmIRCallback)
-    optimizedLlvmIRCallback(*llvmModule);
-
   // Return the serialized object.
   return moduleToObject(*llvmModule);
 }
diff --git a/mlir/lib/Target/LLVM/NVVM/Target.cpp b/mlir/lib/Target/LLVM/NVVM/Target.cpp
index 2a95f343bb2f84..69602af8563aa0 100644
--- a/mlir/lib/Target/LLVM/NVVM/Target.cpp
+++ b/mlir/lib/Target/LLVM/NVVM/Target.cpp
@@ -572,9 +572,6 @@ NVPTXSerializer::moduleToObject(llvm::Module &llvmModule) {
     getOperation().emitError() << "Failed translating the module to ISA.";
     return std::nullopt;
   }
-  if (isaCallback)
-    isaCallback(serializedISA.value());
-
 #define DEBUG_TYPE "serialize-to-isa"
   LLVM_DEBUG({
     llvm::dbgs() << "PTX for module: " << getOperation().getNameAttr() << "\n";
diff --git a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
index eb0c5358ab3530..642aa045178095 100644
--- a/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
+++ b/mlir/unittests/Target/LLVM/SerializeNVVMTarget.cpp
@@ -156,65 +156,3 @@ TEST_F(MLIRTargetLLVMNVVM, SKIP_WITHOUT_NVPTX(SerializeNVVMToBinary)) {
     ASSERT_TRUE(!object->empty());
   }
 }
-
-// Test callback functions invoked with LLVM IR and ISA.
-TEST_F(MLIRTargetLLVMNVVM,
-       SKIP_WITHOUT_NVPTX(CallbackInvokedWithLLVMIRAndISA)) {
-  if (!hasPtxas())
-    GTEST_SKIP() << "PTXAS compiler not found, skipping test.";
-
-  MLIRContext context(registry);
-
-  OwningOpRef<ModuleOp> module =
-      parseSourceString<ModuleOp>(moduleStr, &context);
-  ASSERT_TRUE(!!module);
-
-  NVVM::NVVMTargetAttr target = NVVM::NVVMTargetAttr::get(&context);
-
-  auto serializer = dyn_cast<gpu::TargetAttrInterface>(target);
-  ASSERT_TRUE(!!serializer);
-
-  std::string initialLLVMIR;
-  auto initialCallback = [&initialLLVMIR](llvm::Module &module) {
-    llvm::raw_string_ostream ros(initialLLVMIR);
-    module.print(ros, nullptr);
-  };
-
-  std::string linkedLLVMIR;
-  auto linkedCallback = [&linkedLLVMIR](llvm::Module &module) {
-    llvm::raw_string_ostream ros(linkedLLVMIR);
-    module.print(ros, nullptr);
-  };
-
-  std::string optimizedLLVMIR;
-  auto optimizedCallback = [&optimizedLLVMIR](llvm::Module &module) {
-    llvm::raw_string_ostream ros(optimizedLLVMIR);
-    module.print(ros, nullptr);
-  };
-
-  std::string isaResult;
-  auto isaCallback = [&isaResult](llvm::StringRef isa) {
-    isaResult = isa.str();
-  };
-
-  gpu::TargetOptions options({}, {}, {}, gpu::CompilationTarget::Binary, {},
-                             initialCallback, linkedCallback, optimizedCallback,
-                             isaCallback);
-
-  for (auto gpuModule : (*module).getBody()->getOps<gpu::GPUModuleOp>()) {
-    std::optional<SmallVector<char, 0>> object =
-        serializer.serializeToObject(gpuModule, options);
-
-    ASSERT_TRUE(object != std::nullopt);
-    ASSERT_TRUE(!object->empty());
-    ASSERT_TRUE(!initialLLVMIR.empty());
-    ASSERT_TRUE(!linkedLLVMIR.empty());
-    ASSERT_TRUE(!optimizedLLVMIR.empty());
-    ASSERT_TRUE(!isaResult.empty());
-
-    initialLLVMIR.clear();
-    linkedLLVMIR.clear();
-    optimizedLLVMIR.clear();
-    isaResult.clear();
-  }
-}