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

[christopherbate]

Summary

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.

[christopherbate]

Fix missing newline

[ThomasRaoux]

nice :)

mlir/include/mlir/Dialect/NVGPU/IR/NVGPU.td

 // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef NVGPU
 #define NVGPU
 
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/OpBase.td"
 
 def NVGPU_Dialect : Dialect {
   let name = "nvgpu";
   let cppNamespace = "::mlir::nvgpu";
   let description = [{
     This `NVGPU` dialect provides a bridge between the target agnostic GPU and
     Vector dialects and the lower level LLVM IR based NVVM dialect. This allow
     representing PTX specific operations while using MLIR high level concepts
     like memref and 2-D vector.
   }];
 
   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 = []> :
   Op<NVGPU_Dialect, mnemonic, traits> {}
 
 def NVGPU_LdMatrixOp : NVGPU_Op<"ldmatrix", [
                                 MemoryEffects<[MemRead]>,
                                 PredOpTrait<"srcMemref and res have same element type", 

mlir/include/mlir/Dialect/NVGPU/IR/NVGPUDialect.h

 #ifndef MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
 #define MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #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"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/NVGPU/IR/NVGPU.h.inc"
 
 #endif // MLIR_DIALECT_NVGPU_NVGPUDIALECT_H_

mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 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
 //===----------------------------------------------------------------------===//
 
 /// Return true if the last dimension of the MemRefType has unit stride. Also
 /// return true for memrefs with no strides.
 static bool isLastMemrefDimUnitStride(MemRefType type) {
   int64_t offset;
return emitError() << "expected vector register shape[1] = "
                        << numElementsPer32b;
   if (!(resShape[0] == numTiles))
     return emitError()
            << "expected vector register shape[0] and numTiles to match";
 
   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"
 No newline at end of file