[mlir][gpu] Add the GPU Target Attr trait.

[fmorac]

Summary

The TargetAttrTrait indicates that an attribute implements or promises to implement
the TargetAttrInterface interface.

This trait is needed to avoid IR verification failures caused by not registering the
TargetAttrInterface in places where the implmentation is not needed. For example,
mlir-translate doesn't uses the interface, so there's no need to register them.

[fmorac]

This patch also solves the need to add registration calls for Target Attributes in Translation registration, see D157703.

[mehdi_amini]

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

58

How is it enforced? Is it just a "convention"?

[fmorac]

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

58

Right now it's just a convention, i think there's no proper way to enforce it yet.

I'm thinking proposing an extension to the promised interface mechanism, where this decoupling of the trait and the interfaces is made automatically by tablegen, with the trait verifying there's an implementation or a promise.

[mehdi_amini]

I'm not sure what does the trait provides right now? If you just remove the verifier for ObjectAttr and remove the trait, don't you get basically the same properties?

[fmorac]

In D158277#4600746, @mehdi_amini wrote:

I'm not sure what does the trait provides right now? If you just remove the verifier for ObjectAttr and remove the trait, don't you get basically the same properties?

If I remove the verifier and let #gpu.object hold any target attribute instead of constraining it to the trait or the interface, then yes, the issue is solved. The trait only removes the need for registering the interface in instances where is not needed, for example registering the NVVMTarget interface in mlir-translate.

[bondhugula]

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

23

The 'description` field needs to document what this is.

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1966

No full stops at the end of emit message. The diagnostic continues beyond this.

1985

Likewise.

[fmorac]

Addressed comments. If this patch D158464 gets through, this diff will be repurposed to use the changes in the aforementioned diff.

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

23

AFAIK NativeTrait doesn't have a description field, or at least is not in the class definition. Is there another way to add it?

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

1966

changed.

1985

changed.

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

 //===----------------------------------------------------------------------===//
 
 #ifndef GPU_COMPILATIONATTRINTERFACES
 #define GPU_COMPILATIONATTRINTERFACES
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 
+
 //===----------------------------------------------------------------------===//
-// GPU target attribute interface.
+// GPU target attribute.
 //===----------------------------------------------------------------------===//
+def GPUTargetAttrTrait : NativeTrait<"TargetAttrTrait", ""> {
+  let cppNamespace = "::mlir::gpu";
+}
+
 def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
   let description = [{
     Interface for GPU target attributes. Attributes implementing this interface
     compile GPU modules into binary objects, providing an opaque interface to
     hide implementation details.
   }];
   let cppNamespace = "::mlir::gpu";
   let methods = [
     InterfaceMethod<[{
         Serializes a GPU module to a string containing a representation of the
         module.
 
         If serialization fails then the method should return `std::nullopt`.
 
         The `module` argument must be a GPU Module Op. The `options` argument is
         meant to be used for passing additional options that are not in the
         attribute.
       }],
       "std::optional<SmallVector<char, 0>>", "serializeToObject",
       (ins "Operation*":$module, "const gpu::TargetOptions&":$options)>
   ];
 }
 
-def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttrInterface,
+def HasGPUTargetAttrTrait : AttrConstraint<
+    CPred<"$_self.hasTrait<::mlir::gpu::TargetAttrTrait>()">,
+  "with the GPU `TargetAttrTrait` trait."
+>;
+
+def GPUTargetAttr : ConfinedAttr<AnyAttr, [HasGPUTargetAttrTrait]> {
+  let description = [{
+    Generic GPU target attribute. These attributes must implement the GPU
+    `TargetAttrInterface` interface or promise the interface.
+  }];
+}
+
+def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttr,
   "array of GPU target attributes">;
 
 def GPUNonEmptyTargetArrayAttr :
   ConfinedAttr<GPUTargetArrayAttr, [ArrayMinCount<1>]>;
 
 //===----------------------------------------------------------------------===//
 // GPU offloading translation attribute trait.
 //===----------------------------------------------------------------------===//

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

let description = [{
     object itself.
 
     The target attribute must implement the `TargetAttrInterface` interface.
 
     ```
       #gpu.object<#nvvm.target, "...">
     ```
   }];
-  let parameters = (ins "TargetAttrInterface":$target, "StringAttr":$object);
+  let parameters = (ins "Attribute":$target, "StringAttr":$object);
   let assemblyFormat = [{`<` $target `,` $object `>`}];
+  let genVerifyDecl = 1;
 }
 
 def GPUObjectArrayAttr :
   TypedArrayAttrBase<GPU_ObjectAttr, "an array of GPU object attributes">;
 
 //===----------------------------------------------------------------------===//
 // GPU offloading LLVM translation handler attributes.
 //===----------------------------------------------------------------------===//

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

 }
 
 namespace mlir {
 namespace LLVM {
 class ModuleTranslation;
 }
 namespace gpu {
 /// This class indicates that the attribute associated with this trait is a GPU
+/// target attribute. These kinds of attributes must implement an interface for
+/// handling the serialization of GPU Modules into strings.
+template <typename ConcreteType>
+class TargetAttrTrait
+    : public AttributeTrait::TraitBase<ConcreteType, TargetAttrTrait> {
+  // TODO: Verify the attribute promises or implements the interface.
+};
+
+/// 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.

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

void getAsmValues(RewriterBase &rewriter,
         llvm::SmallVectorImpl<std::pair<mlir::Value, mlir::NVVM::PTXRegisterMod>> &asmValues);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // NVVM target attribute.
 //===----------------------------------------------------------------------===//
 
-def NVVM_TargettAttr : NVVM_Attr<"NVVMTarget", "target"> {
+def NVVM_TargettAttr : NVVM_Attr<"NVVMTarget", "target", [GPUTargetAttrTrait]> {
   let description = [{
     GPU target attribute for controlling compilation of NVIDIA targets. All
     parameters decay into default values if not present.
 
     Examples:
 
     1. Target with default values.
     ```

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

def ROCDL_RawBufferAtomicUMinOp :
   let hasCustomAssemblyFormat = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // ROCDL target attribute.
 //===----------------------------------------------------------------------===//
 
 def ROCDL_TargettAttr :
-    ROCDL_Attr<"ROCDLTarget", "target"> {
+    ROCDL_Attr<"ROCDLTarget", "target", [GPUTargetAttrTrait]> {
   let description = [{
     ROCDL target attribute for controlling compilation of AMDGPU targets. All
     parameters decay into default values if not present.
 
     Examples:
 
     1. Target with default values.
     ```

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 } // namespace
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
   results.add<SimplifyDimOfAllocOp>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// GPU object attribute
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+gpu::ObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                        Attribute target, StringAttr object) {
+  if (target && !target.hasTrait<TargetAttrTrait>())
+    return emitError()
+           << "attribute failed to have the `TargetAttrTrait` trait";
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // GPU select object attribute
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 gpu::SelectObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                               Attribute target) {
   // Check `target`, it can be null, an integer attr or a GPU Target attribute.
   if (target) {
     if (auto intAttr = mlir::dyn_cast<IntegerAttr>(target)) {
       if (intAttr.getInt() < 0) {
-        return emitError() << "The object index must be positive.";
+        return emitError() << "the object index must be positive";
       }
-    } else if (!(::mlir::isa<TargetAttrInterface>(target))) {
+    } else if (!target.hasTrait<TargetAttrTrait>()) {
       return emitError()
-             << "The target attribute must be a GPU Target attribute.";
+             << "the target attribute must be a GPU Target attribute";
     }
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GPU target options
 //===----------------------------------------------------------------------===//

mlir/lib/Target/LLVMIR/CMakeLists.txt

add_mlir_translation_library(MLIRToLLVMIRTranslationRegistration
   MLIRBuiltinToLLVMIRTranslation
   MLIRGPUToLLVMIRTranslation
   MLIRX86VectorToLLVMIRTranslation
   MLIRLLVMToLLVMIRTranslation
   MLIRNVVMToLLVMIRTranslation
   MLIROpenACCToLLVMIRTranslation
   MLIROpenMPToLLVMIRTranslation
   MLIRROCDLToLLVMIRTranslation
-  MLIRNVVMTarget
-  MLIRROCDLTarget
   )
 
 add_mlir_translation_library(MLIRTargetLLVMIRImport
   DataLayoutImporter.cpp
   DebugImporter.cpp
   LoopAnnotationImporter.cpp
   ModuleImport.cpp
   TypeFromLLVM.cpp

mlir/lib/Target/LLVMIR/ConvertToLLVMIR.cpp

 //===- ConvertToLLVMIR.cpp - MLIR to LLVM IR conversion -------------------===//
 //
 // 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 implements a translation between the MLIR LLVM dialect and LLVM IR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/Target/LLVM/NVVM/Target.h"
-#include "mlir/Target/LLVM/ROCDL/Target.h"
 #include "mlir/Target/LLVMIR/Dialect/All.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "mlir/Tools/mlir-translate/Translation.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 
 using namespace mlir;
 
 namespace mlir {
 void registerToLLVMIRTranslation() {
   TranslateFromMLIRRegistration registration(
       "mlir-to-llvmir", "Translate MLIR to LLVMIR",
       [](Operation *op, raw_ostream &output) {
         llvm::LLVMContext llvmContext;
         auto llvmModule = translateModuleToLLVMIR(op, llvmContext);
         if (!llvmModule)
           return failure();
 
         llvmModule->print(output, nullptr);
         return success();
       },
       [](DialectRegistry &registry) {
         registry.insert<DLTIDialect, func::FuncDialect>();
-        registerNVVMTarget(registry);
-        registerROCDLTarget(registry);
         registerAllToLLVMIRTranslations(registry);
       });
 }
 } // namespace mlir