diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td --- 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 GPUNonEmptyTargetArrayAttr : ConfinedAttr]>; +//===----------------------------------------------------------------------===// +// 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 { + 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 --- 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 +class OffloadingTranslationAttrTrait + : public AttributeTrait::TraitBase { + // 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