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

[fmorac]

Summary

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

[fmorac]

Rebasing.

[krzysz00]

I'm going to assume this interface is sufficient.

[fmorac]

Rebasing.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

59

I think more context would be appreciated here about the role of the "manager" and the context where it is used (translation)

I'm not convinced that "ObjectManager" is expressive enough in terms of what this provides.

The name should likely include "Translation" somehow, GPUKernelOffloadingLLVMTranslationAttrInterface would the first thing I can come up with right now, thoughts?

75

The description say that this method launches a kernel: shall we say instead "Translate a gpu launch func operation to a LLVM IR target specific instructions sequence".

82

Shall we take a GPULaunchFuncOp instead of Operation* here? Same for the binaryOp being a GPUBinaryOp?

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

59

I'll add it.

I settled in ObjectManager because it can be used to implement many concepts for managing GPU objects, like fat binaries & fat bin launch. However I agree the name needs to reflect it's for LLVM translation, what about GPUBinaryHandlerLLVMTranslationAttrInterface

75

I'll change it.

82

The reason these are like this is I couldn't break C++ dependencies and have nvptx_target & amdgpu_target in gpu, as

nvptx_target depends on TargetAttrInterface that depends on the definitions of those operations that depends on attributes (fwd declarations didn't work).

If we move, gpu.nvptx_target to nvvm.target, then there are no issues and the above methods can use GPULaunchFuncOp & GPUBinaryOp.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

59

Works for me.

82

Is this the same discussion we're having in the parent about the location of the NVPTXTarget attribute?

[fmorac]

Switched name from GPUObjectManagerAttrInterface to GPUBinaryHandlerLLVMTranslationAttrInterface.

[fmorac]

Rebasing.

[fmorac]

Remove remnants of the ObjectManager name.

[mehdi_amini]

LG!

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

75

I don't think this one was done? It still says Launches a kernel

[fmorac]

Improved docs & removed unnecessary tablegen code. Also switched the name from Binary Handler to Offloading Handler.

[mehdi_amini]

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

56

Can we keep the TableGen name identical to the C++ name here? It got me confused on subsequence patches to read one and the other.

[fmorac]

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

56

Yes, I'll change I just used the GPU prefix because I saw it was customary with other Ops & Attributes.

[fmorac]

Add the OffloadingTranslationAttrTrait trait. This traits decouples gpu.binary op
from LLVM translation, making gpu.binary usable by other translation mechanisms like SPIRV.

mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td

 //===-- CompilationAttrInterfaces.td - GPU compilation interfaces ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
-// This file defines interfaces for GPU target attributes.
+// This file defines interfaces for GPU compilation attributes.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_COMPILATIONATTRINTERFACES
 #define GPU_COMPILATIONATTRINTERFACES
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"

 }
 
 def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttrInterface,
   "array of GPU target attributes">;
 
 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

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

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_GPU_IR_COMPILATIONINTERFACES_H
 #define MLIR_DIALECT_GPU_IR_COMPILATIONINTERFACES_H
 
 #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
 /// ensure type safeness. Targets are free to ignore these options.
 class TargetOptions {
 public:
   /// The target representation of the compilation process.
   typedef enum {