[Mlir-commits] [mlir] afb0b21 - [mlir][nvgpu] Use TableGen TypeDef for NVGPU dialect types

[Christopher Bate]

Author: Christopher Bate
Date: 2022-09-23T19:46:03-06:00
New Revision: afb0b21f24adecfa0b302c79366980453b15ba7b

URL: https://github.com/llvm/llvm-project/commit/afb0b21f24adecfa0b302c79366980453b15ba7b
DIFF: https://github.com/llvm/llvm-project/commit/afb0b21f24adecfa0b302c79366980453b15ba7b.diff

LOG: [mlir][nvgpu] Use TableGen TypeDef for NVGPU dialect types

Moves definition of DeviceAsyncToken to use the declarative Tablegen
TypeDef since the type is trivial. This also allows for removing the
current code for parsing/printing types by using the auto-generated
functions.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D134564

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
    mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h
    mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
index 8858b6cfbf7a..d4b1c84fa253 100644
--- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
+++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td
@@ -21,6 +21,7 @@
 #define NVGPU
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 
 def NVGPU_Dialect : Dialect {
@@ -36,16 +37,29 @@ def NVGPU_Dialect : Dialect {
   let useDefaultTypePrinterParser = 1;
 }
 
-/// Device-side synchronization token.
-def NVGPU_DeviceAsyncToken : DialectType<
-  NVGPU_Dialect, CPred<"$_self.isa<::mlir::nvgpu::DeviceAsyncTokenType>()">,
-   "device async token type">,
-   BuildableType<
-     "mlir::nvgpu::DeviceAsyncTokenType::get($_builder.getContext())">;
+//===----------------------------------------------------------------------===//
+// NVGPU Type Definitions
+//===----------------------------------------------------------------------===//
+
+class NVGPU_Type<string name, string typeMnemonic,
+        list<Trait> traits = []> : TypeDef<NVGPU_Dialect, name, traits> {
+  let mnemonic = typeMnemonic;
+}
 
+def NVGPU_DeviceAsyncToken : NVGPU_Type<"DeviceAsyncToken",  
+                                        "device.async.token", []> {
+  let summary = "device async token type";
+  let description = [{
+    `nvgpu.device.async.token` is a type returned by an asynchronous operation
+    that runs on the GPU (device). It is used to establish an SSA-based link
+    between the async operation (e.g. DeviceAsyncCopy) and operations that
+    group or synchronize the async operations (e.g. DeviceAsyncCreateGroupOp,
+    DeviceAsyncWaitOp).
+  }];
+}
 
 //===----------------------------------------------------------------------===//
-// NVGPU Op definitions
+// NVGPU Op Definitions
 //===----------------------------------------------------------------------===//
 
 class NVGPU_Op<string mnemonic, list<Trait> traits = []> :

diff  --git a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h
index e642fc7126f5..bdc887f68c19 100644
--- a/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h
+++ b/mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h
@@ -18,19 +18,8 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
-namespace mlir {
-namespace nvgpu {
-
-/// Device-side token storage type. There is only one type of device-side token.
-class DeviceAsyncTokenType
-    : public Type::TypeBase<DeviceAsyncTokenType, Type, TypeStorage> {
-public:
-  // Used for generic hooks in TypeBase.
-  using Base::Base;
-};
-
-} // namespace nvgpu
-} // namespace mlir
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/NVGPU/IR/NVGPUTypes.h.inc"
 
 #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h.inc"
 

diff  --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index 3a2999ed5e43..9ed04b45aa1c 100644
--- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
+++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
@@ -21,35 +21,17 @@
 using namespace mlir;
 using namespace mlir::nvgpu;
 
-#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.cpp.inc"
-
 void nvgpu::NVGPUDialect::initialize() {
-  addTypes<DeviceAsyncTokenType>();
+  addTypes<
+#define GET_TYPEDEF_LIST
+#include "mlir/Dialect/NVGPU/IR/NVGPUTypes.cpp.inc"
+      >();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"
       >();
 }
 
-Type NVGPUDialect::parseType(DialectAsmParser &parser) const {
-  // Parse the main keyword for the type.
-  StringRef keyword;
-  if (parser.parseKeyword(&keyword))
-    return Type();
-  MLIRContext *context = getContext();
-  // Handle 'device async token' types.
-  if (keyword == "device.async.token")
-    return DeviceAsyncTokenType::get(context);
-
-  parser.emitError(parser.getNameLoc(), "unknown nvgpu type: " + keyword);
-  return Type();
-}
-
-void NVGPUDialect::printType(Type type, DialectAsmPrinter &os) const {
-  TypeSwitch<Type>(type)
-      .Case<DeviceAsyncTokenType>([&](Type) { os << "device.async.token"; })
-      .Default([](Type) { llvm_unreachable("unexpected 'nvgpu' type kind"); });
-}
 //===----------------------------------------------------------------------===//
 // NVGPU_DeviceAsyncCopyOp
 //===----------------------------------------------------------------------===//
@@ -254,5 +236,14 @@ LogicalResult LdMatrixOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TableGen'd dialect, type, and op definitions
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.cpp.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.cpp.inc"
+
+#define GET_TYPEDEF_CLASSES
+#include "mlir/Dialect/NVGPU/IR/NVGPUTypes.cpp.inc"