[MLIR][DLTI] Pretty parsing and printing for DLTI attrs

Unifies parsing and printing for DLTI attributes. Introduces syntax of
`#dlti.attr<key1 = val1, ..., keyN = valN>` for all queryable DLTI
 attributes, while retaining support for specifying key-value entry
pairs with `#dlti.dl_entry` (whether to retain this is TBD).

As the new format does away with much of the boilerplate, it is much
easier on the eye. This makes an especially big difference for nested
attributes.

Updates the DLTI tests and includes fixes for misc error checking/
error messages.

Author: rolfmorel

mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td

@@ -93,6 +93,10 @@ def DLTI_DataLayoutSpecAttr :
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// MapAttr
+//===----------------------------------------------------------------------===//
+
 def DLTI_MapAttr : DLTIAttr<"Map", [DLTIQueryInterface]> {
   let summary = "A mapping of DLTI-information by way of key-value pairs";
   let description = [{
@@ -106,18 +110,16 @@ def DLTI_MapAttr : DLTIAttr<"Map", [DLTIQueryInterface]> {
 
     Consider the following flat encoding of a single-key dictionary
     ```
-    #dlti.map<#dlti.dl_entry<"CPU::cache::L1::size_in_bytes", 65536 : i32>>
+    #dlti.map<"CPU::cache::L1::size_in_bytes" = 65536 : i32>>
     ```
     versus nested maps, which make it possible to obtain sub-dictionaries of
     related information (with the following example making use of other
     attributes that also implement the `DLTIQueryInterface`):
     ```
-    #dlti.target_system_spec<"CPU":
-      #dlti.target_device_spec<#dlti.dl_entry<"cache",
-        #dlti.map<#dlti.dl_entry<"L1",
-          #dlti.map<#dlti.dl_entry<"size_in_bytes", 65536 : i32>>>,
-                  #dlti.dl_entry<"L1d",
-          #dlti.map<#dlti.dl_entry<"size_in_bytes", 32768 : i32>>> >>>>
+    #dlti.target_system_spec<"CPU" =
+      #dlti.target_device_spec<"cache" =
+        #dlti.map<"L1" = #dlti.map<"size_in_bytes" = 65536 : i32>,
+                  "L1d" = #dlti.map<"size_in_bytes" = 32768 : i32> >>>
     ```
 
     With the flat encoding, the implied structure of the key is ignored, that is
@@ -139,7 +141,7 @@ def DLTI_MapAttr : DLTIAttr<"Map", [DLTIQueryInterface]> {
   );
   let mnemonic = "map";
   let genVerifyDecl = 1;
-  let assemblyFormat = "`<` $entries `>`";
+  let hasCustomAssemblyFormat = 1;
   let extraClassDeclaration = [{
     /// Returns the attribute associated with the key.
     FailureOr<Attribute> query(DataLayoutEntryKey key) {
@@ -167,20 +169,23 @@ def DLTI_TargetSystemSpecAttr :
     ```
     dlti.target_system_spec =
      #dlti.target_system_spec<
-      "CPU": #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096: ui32>>,
-      "GPU": #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>,
-      "XPU": #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.max_vector_op_width", 4096 : ui32>>>
+      "CPU" = #dlti.target_device_spec<
+        "L1_cache_size_in_bytes" = 4096: ui32>,
+      "GPU" = #dlti.target_device_spec<
+        "max_vector_op_width" = 64 : ui32>,
+      "XPU" = #dlti.target_device_spec<
+        "max_vector_op_width" = 4096 : ui32>>
     ```
+
+    The verifier checks that keys are strings and pointed to values implement
+    DLTI's TargetDeviceSpecInterface.
   }];
   let parameters = (ins
-    ArrayRefParameter<"DeviceIDTargetDeviceSpecPair", "">:$entries
+    ArrayRefParameter<"DataLayoutEntryInterface">:$entries
   );
   let mnemonic = "target_system_spec";
   let genVerifyDecl = 1;
-  let assemblyFormat = "`<` $entries `>`";
+  let hasCustomAssemblyFormat = 1;
   let extraClassDeclaration = [{
     /// Return the device specification that matches the given device ID
     std::optional<TargetDeviceSpecInterface>
@@ -197,8 +202,10 @@ def DLTI_TargetSystemSpecAttr :
     $cppClass::getDeviceSpecForDeviceID(
         TargetSystemSpecInterface::DeviceID deviceID) {
       for (const auto& entry : getEntries()) {
-        if (entry.first == deviceID)
-          return entry.second;
+        if (entry.getKey() == DataLayoutEntryKey(deviceID))
+          if (auto deviceSpec =
+              llvm::dyn_cast<TargetDeviceSpecInterface>(entry.getValue()))
+            return deviceSpec;
       }
       return std::nullopt;
     }
@@ -219,16 +226,15 @@ def DLTI_TargetDeviceSpecAttr :
 
     Example:
     ```
-    #dlti.target_device_spec<
-      #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>
+    #dlti.target_device_spec<"max_vector_op_width" = 64 : ui32>
     ```
   }];
   let parameters = (ins
-    ArrayRefParameter<"DataLayoutEntryInterface", "">:$entries
+    ArrayRefParameter<"DataLayoutEntryInterface">:$entries
   );
   let mnemonic = "target_device_spec";
   let genVerifyDecl = 1;
-  let assemblyFormat = "`<` $entries `>`";
+  let hasCustomAssemblyFormat = 1;
 
   let extraClassDeclaration = [{
     /// Returns the attribute associated with the key.

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -15,6 +15,7 @@
 #ifndef MLIR_INTERFACES_DATALAYOUTINTERFACES_H
 #define MLIR_INTERFACES_DATALAYOUTINTERFACES_H
 
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/OpDefinition.h"
 #include "llvm/ADT/DenseMap.h"
@@ -32,10 +33,7 @@ using DataLayoutEntryKey = llvm::PointerUnion<Type, StringAttr>;
 using DataLayoutEntryList = llvm::SmallVector<DataLayoutEntryInterface, 4>;
 using DataLayoutEntryListRef = llvm::ArrayRef<DataLayoutEntryInterface>;
 using TargetDeviceSpecListRef = llvm::ArrayRef<TargetDeviceSpecInterface>;
-using DeviceIDTargetDeviceSpecPair =
-    std::pair<StringAttr, TargetDeviceSpecInterface>;
-using DeviceIDTargetDeviceSpecPairListRef =
-    llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>;
+using TargetDeviceSpecEntry = std::pair<StringAttr, TargetDeviceSpecInterface>;
 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
 class ModuleOp;

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -304,7 +304,7 @@ def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface", [DLTI
   let methods = [
     InterfaceMethod<
       /*description=*/"Returns the list of layout entries.",
-      /*retTy=*/"llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>",
+      /*retTy=*/"llvm::ArrayRef<DataLayoutEntryInterface>",
       /*methodName=*/"getEntries",
       /*args=*/(ins)
     >,