[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary #65797

fabianmcg · 2023-09-08T19:48:07Z

This patch adds the option of building an optional symbol table for the top operation in the gpu-module-to-binary pass. The table is not created by default as other targets don't need it; instead, it is lazily built. The table is passed through a callback in TargetOptions.

This patch is required to integrate #65539 .

…oduleToBinary This patch adds the option of building an optional symbol table for the top operation in the `gpu-module-to-binary` pass. This table is required by some target attributes. The table is not created by default, as other targets don't need it. The table is passed through `TargetOptions`.

joker-eph · 2023-09-09T02:06:44Z

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -100,6 +106,10 @@ class TargetOptions {
  /// Compilation process target representation.
  CompilationTarget compilationTarget;

+  /// Parent symbol table of all the GPU modules being serialized. By default
+  /// this member is null as it is not required by most targets.
+  SymbolTable *parentTable;


Can we make this a callback that would lazy initialize it on demand if the target wants it?
Then we don't need a pass option and we can remove this weird coupling where initializing the pass requires to know if the target wants a SymbolTable.

I've switched it.

joker-eph · 2023-09-09T02:06:58Z

(same problem with the title here)

llvmbot · 2023-09-09T16:47:10Z

@llvm/pr-subscribers-mlir

Changes

…oduleToBinary

This patch adds the option of building an optional symbol table for the top operation in the gpu-module-to-binary pass. This table is required by some target attributes. The table is not created by default, as other targets don't need it. The table is passed through TargetOptions.

This patch is required to integrate #65539 .

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

4 Files Affected:

(modified) mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h (+14-3)
(modified) mlir/include/mlir/Dialect/GPU/Transforms/Passes.td (+5-3)
(modified) mlir/lib/Dialect/GPU/IR/GPUDialect.cpp (+10-4)
(modified) mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp (+18-1)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
index e0bf560dbd98b92..cd2287ecf239d36 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
@@ -20,6 +20,7 @@ class IRBuilderBase;
 }
 
 namespace mlir {
+class SymbolTable;
 namespace LLVM {
 class ModuleTranslation;
 }
@@ -56,10 +57,11 @@ class TargetOptions {
 
   /// Constructor initializing the toolkit path, the list of files to link to,
   /// extra command line options & the compilation target. The default
-  /// compilation target is `binary`.
+  /// compilation target is `binOrFatbin`.
   TargetOptions(StringRef toolkitPath = {},
                 ArrayRef linkFiles = {}, StringRef cmdOptions = {},
-                CompilationTarget compilationTarget = binOrFatbin);
+                CompilationTarget compilationTarget = binOrFatbin,
+                function_ref symbolTableCallback = {});
 
   /// Returns the typeID.
   TypeID getTypeID() const;
@@ -80,12 +82,17 @@ class TargetOptions {
   /// Returns the compilation target.
   CompilationTarget getCompilationTarget() const;
 
+  /// Returns the parent symbol table if a callback was provided, else returns
+  /// nullptr.
+  SymbolTable *getParentTable() const;
+
 protected:
   /// Derived classes must use this constructor to initialize `typeID` to the
   /// appropiate value: ie. `TargetOptions(TypeID::get())`.
   TargetOptions(TypeID typeID, StringRef toolkitPath = {},
                 ArrayRef linkFiles = {}, StringRef cmdOptions = {},
-                CompilationTarget compilationTarget = binOrFatbin);
+                CompilationTarget compilationTarget = binOrFatbin,
+                function_ref symbolTableCallback = {});
 
   /// Path to the target toolkit.
   std::string toolkitPath;
@@ -100,6 +107,10 @@ class TargetOptions {
   /// Compilation process target representation.
   CompilationTarget compilationTarget;
 
+  /// Callback for obtaining the parent symbol table of all the GPU modules
+  /// being serialized.
+  function_ref symbolTableCallback;
+
 private:
   TypeID typeID;
 };
diff --git a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
index ba8a6266604e46c..0bfb2750992058f 100644
--- a/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/GPU/Transforms/Passes.td
@@ -64,9 +64,11 @@ def GpuModuleToBinaryPass
     with an object for every target.
 
     The `format` argument can have the following values:
-    1. `offloading`, `llvm`: producing an offloading representation.
-    2. `assembly`, `isa`: producing assembly code.
-    3. `binary`, `bin`: producing binaries.
+    1. `offloading`, `llvm`: produces an offloading representation.
+    2. `assembly`, `isa`: produces assembly code.
+    3. `binary`, `bin`: produces binaries.
+    4. `fatbinary`, `fatbin`: produces fatbinaries.
+    5. `binOrFatbin`: produces bins or fatbins, the target decides which.
   }];
   let options = [
     Option<"offloadingHandler", "handler", "Attribute", "nullptr",
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index f417a083337fcaf..68423e74b7f328f 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -1982,17 +1982,19 @@ gpu::SelectObjectAttr::verify(function_ref emitError,
 TargetOptions::TargetOptions(StringRef toolkitPath,
                              ArrayRef linkFiles,
                              StringRef cmdOptions,
-                             CompilationTarget compilationTarget)
+                             CompilationTarget compilationTarget,
+                             function_ref symbolTableCallback)
     : TargetOptions(TypeID::get(), toolkitPath, linkFiles,
-                    cmdOptions, compilationTarget) {}
+                    cmdOptions, compilationTarget, symbolTableCallback) {}
 
 TargetOptions::TargetOptions(TypeID typeID, StringRef toolkitPath,
                              ArrayRef linkFiles,
                              StringRef cmdOptions,
-                             CompilationTarget compilationTarget)
+                             CompilationTarget compilationTarget,
+                             function_ref symbolTableCallback)
     : toolkitPath(toolkitPath.str()), linkFiles(linkFiles),
       cmdOptions(cmdOptions.str()), compilationTarget(compilationTarget),
-      typeID(typeID) {}
+      symbolTableCallback(symbolTableCallback), typeID(typeID) {}
 
 TypeID TargetOptions::getTypeID() const { return typeID; }
 
@@ -2002,6 +2004,10 @@ ArrayRef TargetOptions::getLinkFiles() const { return linkFiles; }
 
 StringRef TargetOptions::getCmdOptions() const { return cmdOptions; }
 
+SymbolTable *TargetOptions::getParentTable() const {
+  return symbolTableCallback ? symbolTableCallback() : nullptr;
+}
+
 std::pair>
 TargetOptions::tokenizeCmdOptions() const {
   std::pair> options;
diff --git a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
index 06b7dee6941e1f4..e29a1f0c3248d04 100644
--- a/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
+++ b/mlir/lib/Dialect/GPU/Transforms/ModuleToBinary.cpp
@@ -66,9 +66,26 @@ void GpuModuleToBinaryPass::runOnOperation() {
                          .Default(-1);
   if (targetFormat == -1)
     getOperation()->emitError() << "Invalid format specified.";
+
+  // Lazy symbol table builder callback.
+  std::optional parentTable;
+  auto lazyTableBuilder = [&]() -> SymbolTable * {
+    // Build the table if it has not been built.
+    if (!parentTable) {
+      Operation *table = SymbolTable::getNearestSymbolTable(getOperation());
+      // It's up to the target attribute to determine if failing to find a
+      // symbol table is an error.
+      if (!table)
+        return nullptr;
+      parentTable = SymbolTable(table);
+    }
+    return &parentTable.value();
+  };
+
   TargetOptions targetOptions(
       toolkitPath, linkFiles, cmdOptions,
-      static_cast(targetFormat));
+      static_cast(targetFormat),
+      lazyTableBuilder);
   if (failed(transformGpuModulesToBinaries(
           getOperation(),
           offloadingHandler ? dyn_cast(

joker-eph · 2023-09-09T21:31:03Z

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -80,12 +82,17 @@ class TargetOptions {
  /// Returns the compilation target.
  CompilationTarget getCompilationTarget() const;

+  /// Returns the parent symbol table if a callback was provided, else returns
+  /// nullptr.


Seems like the callback itself can return null

joker-eph · 2023-09-09T21:31:27Z

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -56,10 +57,11 @@ class TargetOptions {

  /// Constructor initializing the toolkit path, the list of files to link to,
  /// extra command line options & the compilation target. The default
-  /// compilation target is `binary`.
+  /// compilation target is `binOrFatbin`.


Document the new parameter

joker-eph · 2023-09-09T21:32:16Z

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

  TargetOptions(StringRef toolkitPath = {},
                ArrayRef<std::string> linkFiles = {}, StringRef cmdOptions = {},
-                CompilationTarget compilationTarget = binOrFatbin);
+                CompilationTarget compilationTarget = binOrFatbin,
+                function_ref<SymbolTable *()> symbolTableCallback = {});


I'd call it getSymbolTable (or getSymbolTableCallback if you want to include callback in the name) to keep a verb here.

…oduleToBinary (llvm#65797) This patch adds the option of building an optional symbol table for the top operation in the `gpu-module-to-binary` pass. The table is not created by default as most targets don't need it; instead, it is lazily built. The table is passed through a callback in `TargetOptions`. This patch is required to integrate llvm#65539 .

peledins-zimperium · 2023-12-12T13:29:59Z

mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

@@ -100,6 +111,10 @@ class TargetOptions {
  /// Compilation process target representation.
  CompilationTarget compilationTarget;

+  /// Callback for obtaining the parent symbol table of all the GPU modules
+  /// being serialized.
+  function_ref<SymbolTable *()> getSymbolTableCallback;


Is this safe?
function_ref says:
/// An efficient, type-erasing, non-owning reference to a callable. This is
/// intended for use as the type of a function parameter that is not used
/// after the function in question returns.
///
/// This class does not own the callable, so it is not in general safe to store
/// a function_ref.
So who owns = {} here?

I'm seeing clang-tidy warnings/crashes in similar usage - #71363

It can be changed. However, this class is only meant to be used as a temp to pass arguments with the caller being in charge of making sure the object remains alive and not really to store data.

fabianmcg requested a review from joker-eph September 8, 2023 19:48

fabianmcg requested a review from a team as a code owner September 8, 2023 19:48

fabianmcg removed the request for review from a team September 8, 2023 19:48

joker-eph reviewed Sep 9, 2023

View reviewed changes

Switch to a lazy table builder.

24cd7f5

llvmbot added mlir:core MLIR Core Infrastructure mlir:gpu mlir labels Sep 9, 2023

fabianmcg changed the title ~~[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuM…~~ [mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary Sep 9, 2023

fabianmcg requested a review from joker-eph September 9, 2023 16:49

joker-eph approved these changes Sep 9, 2023

View reviewed changes

Addressing the comments: fixing docs, etc.

7600c49

fabianmcg merged commit 444abb3 into llvm:main Sep 9, 2023

fabianmcg deleted the gpu-mod-to-bin-sym-table branch September 10, 2023 00:11

michaelrj-google mentioned this pull request Sep 12, 2023

[libc] Move long double table option to new config #66151

Merged

peledins-zimperium reviewed Dec 12, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary #65797

[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary #65797

fabianmcg commented Sep 8, 2023 •

edited

Loading

joker-eph Sep 9, 2023

fabianmcg Sep 9, 2023

joker-eph commented Sep 9, 2023

llvmbot commented Sep 9, 2023

This patch is required to integrate #65539 .

joker-eph Sep 9, 2023

joker-eph Sep 9, 2023

joker-eph Sep 9, 2023

peledins-zimperium Dec 12, 2023

peledins-zimperium Dec 12, 2023

fabianmcg Dec 12, 2023

[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary #65797

[mlir][gpu] Add a symbol table field to TargetOptions and adjust GpuModuleToBinary #65797

Conversation

fabianmcg commented Sep 8, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

joker-eph commented Sep 9, 2023

llvmbot commented Sep 9, 2023

This patch is required to integrate #65539 .

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fabianmcg commented Sep 8, 2023 •

edited

Loading