[Mlir-commits] [mlir] [mlir][gpu][target] Use promises to verify TargetAttrs IR correctness. (PR #65787)

[Fabian Mora]

https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/65787:

This patch employs the updated promise mechanism to enforce Target Attribute IR constraints. Due to this patch, TargetAttributes implementations no longer have to be registered before executing translation to LLVM IR in cases where they are not needed, like when translating `gpu.binary` operations.

>From 2d8539173659d3c9e73734c29c0fbc66cd23dbd2 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Fri, 8 Sep 2023 17:53:54 +0000
Subject: [PATCH] [mlir][gpu][target] Use promises to verify TargetAttrs IR
 correctness.

This patch employs the updated promise mechanism to enforce Target Attribute IR
constraints. Due to this patch, TargetAttributes implementations no longer have
to be registered before executing translation to LLVM IR in cases where they are
not needed, like when translating `gpu.binary` operations.
---
 .../GPU/IR/CompilationAttrInterfaces.td       | 13 ++++++++++--
 .../mlir/Dialect/GPU/IR/CompilationAttrs.td   |  3 ++-
 mlir/include/mlir/Target/LLVMIR/Dialect/All.h | 12 +++--------
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp        | 20 ++++++++++++++++---
 mlir/lib/Target/LLVMIR/CMakeLists.txt         |  2 --
 5 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
index f30f9d172b08ba9..5255286619e3bf2 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td
@@ -42,7 +42,15 @@ def GPUTargetAttrInterface : AttrInterface<"TargetAttrInterface"> {
   ];
 }
 
-def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttrInterface,
+def GPUTargetAttr :
+    ConfinedAttr<AnyAttr, [PromisedAttrInterface<GPUTargetAttrInterface>]> {
+  let description = [{
+    Generic GPU target attribute. These attributes must implement or promise
+    the `GPUTargetAttrInterface` interface.
+  }];
+}
+
+def GPUTargetArrayAttr : TypedArrayAttrBase<GPUTargetAttr,
   "array of GPU target attributes">;
 
 def GPUNonEmptyTargetArrayAttr :
@@ -51,6 +59,7 @@ def GPUNonEmptyTargetArrayAttr :
 //===----------------------------------------------------------------------===//
 // GPU offloading translation attribute trait.
 //===----------------------------------------------------------------------===//
+
 def OffloadingTranslationAttrTrait :
    NativeTrait<"OffloadingTranslationAttrTrait", ""> {
   let cppNamespace = "::mlir::gpu";
@@ -65,7 +74,7 @@ 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
+    implement an interface for handling the translation of GPU offloading
     operations like `gpu.binary` & `gpu.launch_func`. An example of such
     interface is the `OffloadingLLVMTranslationAttrInterface` interface.
   }];
diff --git a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
index 2e2a084413fa57c..9c1110d8e9a9463 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/CompilationAttrs.td
@@ -32,8 +32,9 @@ def GPU_ObjectAttr : GPU_Attr<"Object", "object"> {
       #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 :
diff --git a/mlir/include/mlir/Target/LLVMIR/Dialect/All.h b/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
index c25822c1c82bf1a..0563b9bf3d475a4 100644
--- a/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
+++ b/mlir/include/mlir/Target/LLVMIR/Dialect/All.h
@@ -14,8 +14,6 @@
 #ifndef MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 #define MLIR_TARGET_LLVMIR_DIALECT_ALL_H
 
-#include "mlir/Target/LLVM/NVVM/Target.h"
-#include "mlir/Target/LLVM/ROCDL/Target.h"
 #include "mlir/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/ArmSME/ArmSMEToLLVMIRTranslation.h"
@@ -51,13 +49,6 @@ static inline void registerAllToLLVMIRTranslations(DialectRegistry &registry) {
 
   // Extension required for translating GPU offloading Ops.
   gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(registry);
-
-  // GPU target attribute interfaces are not used during translation, however
-  // the IR fails to verify if they are not registered due to the promise
-  // mechanism.
-  // TODO: remove these.
-  NVVM::registerNVVMTargetInterfaceExternalModels(registry);
-  ROCDL::registerROCDLTargetInterfaceExternalModels(registry);
 }
 
 /// Registers all the translations to LLVM IR required by GPU passes.
@@ -70,6 +61,9 @@ registerAllGPUToLLVMIRTranslations(DialectRegistry &registry) {
   registerLLVMDialectTranslation(registry);
   registerNVVMDialectTranslation(registry);
   registerROCDLDialectTranslation(registry);
+
+  // Extension required for translating GPU offloading Ops.
+  gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(registry);
 }
 
 /// Registers all dialects that can be translated from LLVM IR and the
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index f417a083337fcaf..ad36b44763339f6 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -1954,6 +1954,20 @@ void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<SimplifyDimOfAllocOp>(context);
 }
 
+//===----------------------------------------------------------------------===//
+// GPU object attribute
+//===----------------------------------------------------------------------===//
+
+LogicalResult ObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
+                                 Attribute target, StringAttr object) {
+  if (!target)
+    return emitError() << "the target attribute cannot be null";
+  if (target.hasPromiseOrImplementsInterface<TargetAttrInterface>())
+    return success();
+  return emitError() << "the target attribute must implement or promise the "
+                        "`gpu::TargetAttrInterface`";
+}
+
 //===----------------------------------------------------------------------===//
 // GPU select object attribute
 //===----------------------------------------------------------------------===//
@@ -1965,11 +1979,11 @@ gpu::SelectObjectAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   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.hasPromiseOrImplementsInterface<TargetAttrInterface>()) {
       return emitError()
-             << "The target attribute must be a GPU Target attribute.";
+             << "the target attribute must be a GPU Target attribute";
     }
   }
   return success();
diff --git a/mlir/lib/Target/LLVMIR/CMakeLists.txt b/mlir/lib/Target/LLVMIR/CMakeLists.txt
index dd97ccf88688632..868ccbbb10620d9 100644
--- a/mlir/lib/Target/LLVMIR/CMakeLists.txt
+++ b/mlir/lib/Target/LLVMIR/CMakeLists.txt
@@ -57,8 +57,6 @@ add_mlir_translation_library(MLIRToLLVMIRTranslationRegistration
   MLIROpenACCToLLVMIRTranslation
   MLIROpenMPToLLVMIRTranslation
   MLIRROCDLToLLVMIRTranslation
-  MLIRNVVMTarget
-  MLIRROCDLTarget
   )
 
 add_mlir_translation_library(MLIRTargetLLVMIRImport