[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

73

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?

89

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".

96

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

73

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

89

I'll change it.

96

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

73

Works for me.

96

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

89

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

70

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

70

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 target attributes & GPU object manager
+// attributes.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_COMPILATIONATTRINTERFACES
 #define GPU_COMPILATIONATTRINTERFACES
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"

 }
 
 def GPUTargetArrayAttr :
   TypedArrayAttrBase<GPUTargetAttr, "Array of GPU target attributes">;
 
 def GPUNonEmptyTargetArrayAttr :
   ConfinedAttr<GPUTargetArrayAttr, [ArrayMinCount<1>]>;
 
+//===----------------------------------------------------------------------===//
+// GPU object manager attribute interface.
+//===----------------------------------------------------------------------===//
+
+def GPUObjectManagerAttrInterface :
+    AttrInterface<"ObjectManagerAttrInterface"> {
+  let description = [{
+    Interface for GPU object manager attributes. Attributes implementing this
+    interface manage the interaction between GPU objects and host IR.
+  }];
+  let cppNamespace = "::mlir::gpu";
+  let methods = [
+    InterfaceMethod<[{
+        Embeds a GPU object into a host LLVM module. The operation expected by
+        this method must be a GPU BinaryOp.
+
+        All attributes implementing this interface must implement this method.
+        If the method fails then it must return `failure`.
+      }],
+      "LogicalResult", "embedBinary",
+      (ins "Operation*":$binaryOp, "llvm::IRBuilderBase&":$hostBuilder,
+           "LLVM::ModuleTranslation&":$hostModuleTranslation)
+    >,
+    InterfaceMethod<[{
+        Launches a kernel inside a binary. The first argument must be a GPU
+        LaunchFuncOp, while the second one a GPU BinaryOp.
+
+        All attributes implementing this interface must implement this method.
+        If the method fails then it must return `failure`.
+      }],
+      "LogicalResult", "launchKernel",
+      (ins "Operation*":$launchFunc, "Operation*":$binaryOp,
+           "llvm::IRBuilderBase&":$hostBuilder,
+           "LLVM::ModuleTranslation&":$hostModuleTranslation)
+    >
+  ];
+}
+
+def ImplementsObjectManagerAttrInterface : AttrConstraint<
+  CPred<"$_self.hasTrait<::mlir::gpu::ObjectManagerAttrInterface::Trait>()">,
+  "Attribute implementing the `ObjectManagerAttrInterface` interface."
+>;
+
+def GPUObjectManagerAttr :
+    ConfinedAttr<AnyAttr, [ImplementsObjectManagerAttrInterface]> {
+  let description = [{
+    Generic compilation attribute implementing the `ObjectManagerAttrInterface`
+    interface.
+  }];
+}
+
 #endif // GPU_COMPILATIONATTRINTERFACES

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

 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/InferIntRangeInterface.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/ADT/STLExtras.h"
 
+namespace llvm {
+class IRBuilderBase;
+}
+
 namespace mlir {
+namespace LLVM {
+class ModuleTranslation;
+}
 namespace gpu {
 
 /// Utility class for the GPU dialect to represent triples of `Value`s
 /// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
 struct KernelDim3 {
   Value x;
   Value y;
   Value z;