[mlir][nvgpu] Use TableGen TypeDef for NVGPU dialect types
authorChristopher Bate <cbate@nvidia.com>
Fri, 23 Sep 2022 20:37:38 +0000 (14:37 -0600)
committerChristopher Bate <cbate@nvidia.com>
Sat, 24 Sep 2022 01:46:03 +0000 (19:46 -0600)
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

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

index 8858b6c..d4b1c84 100644 (file)
@@ -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 = []> :
index e642fc7..bdc887f 100644 (file)
 #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"
 
index 3a2999e..9ed04b4 100644 (file)
 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"