[Mlir-commits] [mlir] 4760ea0 - [mlir][gpu] Add the GPU offloading handler attribute interface.

[Fabian Mora]

Author: Fabian Mora
Date: 2023-08-11T19:46:17Z
New Revision: 4760ea029a942ef8da8760cdb1148ec11984a1bd

URL: https://github.com/llvm/llvm-project/commit/4760ea029a942ef8da8760cdb1148ec11984a1bd
DIFF: https://github.com/llvm/llvm-project/commit/4760ea029a942ef8da8760cdb1148ec11984a1bd.diff

LOG: [mlir][gpu] Add the GPU offloading handler attribute interface.

**For an explanation of these patches see D154153.**

Commit message:
This patch adds the `OffloadingTranslationAttrTrait` trait and the
`GPUOffloadingLLVMTranslationAttrInterface` attribute interface.

The interface indicates that an attribute is a GPU offloading translation
attribute. These kinds of attributes must implement an interface for handling
the translation of GPU offloading operations like `gpu.binary` & `gpu.launch_func`.

The interface is meant to be used as an opaque interface for embedding GPU
binaries into LLVM modules and launching GPU kernels.

Depends on D154129

Reviewed By: krzysz00

Differential Revision: https://reviews.llvm.org/D154108

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
    mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
index a3c7e5a5284f06..f30f9d172b08ba 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
@@ -6,7 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines interfaces for GPU target attributes.
+// This file defines interfaces for GPU compilation attributes.
 //
 //===----------------------------------------------------------------------===//
 
@@ -48,4 +48,74 @@ def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttrInterface,
 def GPUNonEmptyTargetArrayAttr :
   ConfinedAttr<GPUTargetArrayAttr, [ArrayMinCount<1>]>;
 
+//===----------------------------------------------------------------------===//
+// GPU offloading translation attribute trait.
+//===----------------------------------------------------------------------===//
+def OffloadingTranslationAttrTrait :
+   NativeTrait<"OffloadingTranslationAttrTrait", ""> {
+  let cppNamespace = "::mlir::gpu";
+}
+
+def HasOffloadingTranslationAttrTrait : AttrConstraint<
+  CPred<"$_self.hasTrait<::mlir::gpu::OffloadingTranslationAttrTrait>()">,
+  "with the `OffloadingTranslationAttrTrait` trait."
+>;
+
+def OffloadingTranslationAttr :
+    ConfinedAttr<AnyAttr, [HasOffloadingTranslationAttrTrait]> {
+  let description = [{
+    Generic GPU offloading translation attribute. These attributes must
+    implement an interface for handling the translation of PU offloading
+    operations like `gpu.binary` & `gpu.launch_func`. An example of such
+    interface is the `OffloadingLLVMTranslationAttrInterface` interface.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// GPU offloading LLVM translation handler attribute interface.
+//===----------------------------------------------------------------------===//
+
+def OffloadingLLVMTranslationAttrInterface :
+    AttrInterface<"OffloadingLLVMTranslationAttrInterface"> {
+  let description = [{
+    Interface for GPU offloading LLVM translation attributes. Attributes
+    implementing this interface manage the interaction between GPU offloading
+    operations and host IR.
+  }];
+  let cppNamespace = "::mlir::gpu";
+  let methods = [
+    InterfaceMethod<[{
+        Translates a `gpu.binary` Op into a sequence of LLVM IR target-specific
+        instructions, embedding the binary into a host LLVM module.
+
+        The LLVM translation mechanism invokes this function when translating a
+        `gpu.binary`.
+
+        The first argument has to be a GPU binary operation.
+        If the function fails at any point, it must return `failure`.
+      }],
+      "LogicalResult", "embedBinary",
+      (ins "Operation*":$binaryOp, "llvm::IRBuilderBase&":$hostBuilder,
+           "LLVM::ModuleTranslation&":$hostModuleTranslation)
+    >,
+    InterfaceMethod<[{
+        Translates a `gpu.launch_func` op into a sequence of LLVM IR
+        target-specific instructions, resulting in a kernel launch on host IR.
+
+        The LLVM translation mechanism invokes this function when translating a
+        `gpu.launch_func` operation; it searches the appropriate binary and uses
+        its offloading handler.
+
+        The first two arguments must be GPU launch and binary operations,
+        respectively. If the function fails at any point, it must return
+        `failure`.
+      }],
+      "LogicalResult", "launchKernel",
+      (ins "Operation*":$launchFunc, "Operation*":$binaryOp,
+           "llvm::IRBuilderBase&":$hostBuilder,
+           "LLVM::ModuleTranslation&":$hostModuleTranslation)
+    >
+  ];
+}
+
 #endif // GPU_COMPILATIONATTRINTERFACES

diff  --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
index 090e272cb17e7c..adf950d7b119d6 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationInterfaces.h
@@ -15,8 +15,26 @@
 
 #include "mlir/IR/Attributes.h"
 
+namespace llvm {
+class IRBuilderBase;
+}
+
 namespace mlir {
+namespace LLVM {
+class ModuleTranslation;
+}
 namespace gpu {
+/// This class indicates that the attribute associated with this trait is a GPU
+/// offloading translation attribute. These kinds of attributes must implement
+/// an interface for handling the translation of GPU offloading operations like
+/// `gpu.binary` & `gpu.launch_func`.
+template <typename ConcreteType>
+class OffloadingTranslationAttrTrait
+    : public AttributeTrait::TraitBase<ConcreteType,
+                                       OffloadingTranslationAttrTrait> {
+  // TODO: Verify the attribute promises or implements the interface.
+};
+
 /// This class serves as an opaque interface for passing options to the
 /// `TargetAttrInterface` methods. Users of this class must implement the
 /// `classof` method as well as using the macros `MLIR_*_EXPLICIT_TYPE_ID` to