Addressing review comments

Representing TargetSystemSpec as a set of key-value pairs where key is
the DeviceID (string) and the value is TargetDeviceSpec.

Author: nhasabni

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

@@ -116,24 +116,21 @@ def DLTI_TargetSystemSpecAttr :
   }];
   let description = [{
     A system specification describes the overall system containing
-    multiple devices, with each device having a unique ID
+    multiple devices, with each device having a unique ID (string)
     and its corresponding TargetDeviceSpec object.
 
     Example:
           dlti.target_system_spec =
-            #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.device_id", 0: ui32>,
-              #dlti.dl_entry<"dlti.device_type", "CPU">>,
-            #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.device_id", 1: ui32>,
-              #dlti.dl_entry<"dlti.device_type", "GPU">,
-              #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>,
-            #dlti.target_device_spec<
-              #dlti.dl_entry<"dlti.device_id", 2: ui32>,
-              #dlti.dl_entry<"dlti.device_type", "XPU">>>
+           #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>>>
   }];
   let parameters = (ins
-    ArrayRefParameter<"TargetDeviceSpecInterface", "">:$entries
+    ArrayRefParameter<"DeviceIDTargetDeviceSpecPair", "">:$entries
   );
   let mnemonic = "target_system_spec";
   let genVerifyDecl = 1;
@@ -142,15 +139,15 @@ def DLTI_TargetSystemSpecAttr :
     /// Return the device specification that matches the given device ID
     std::optional<TargetDeviceSpecInterface>
     getDeviceSpecForDeviceID(
-      TargetDeviceSpecInterface::DeviceID deviceID);
+      TargetSystemSpecInterface::DeviceID deviceID);
   }];
   let extraClassDefinition = [{
     std::optional<TargetDeviceSpecInterface>
     $cppClass::getDeviceSpecForDeviceID(
-        TargetDeviceSpecInterface::DeviceID deviceID) {
-      for (TargetDeviceSpecInterface entry : getEntries()) {
-        if (entry.getDeviceID() == deviceID)
-          return entry;
+        TargetSystemSpecInterface::DeviceID deviceID) {
+      for (const auto& entry : getEntries()) {
+        if (entry.first == deviceID)
+          return entry.second;
       }
       return std::nullopt;
     }
@@ -173,15 +170,12 @@ def DLTI_TargetDeviceSpecAttr :
   }];
   let description = [{
     Each device specification describes a single device and its
-    hardware properties. Each device specification must have a device_id
-    and a device_type. In addition, the specification can contain any number
+    hardware properties. Each device specification can contain any number
     of optional hardware properties (e.g., max_vector_op_width below).
 
     Example:
       #dlti.target_device_spec<
-            #dlti.dl_entry<"dlti.device_id", 1: ui32>,
-            #dlti.dl_entry<"dlti.device_type", "GPU">,
-            #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>
+        #dlti.dl_entry<"dlti.max_vector_op_width", 64 : ui32>>
   }];
   let parameters = (ins
     ArrayRefParameter<"DataLayoutEntryInterface", "">:$entries
@@ -190,28 +184,12 @@ def DLTI_TargetDeviceSpecAttr :
   let genVerifyDecl = 1;
   let assemblyFormat = "`<` $entries `>`";
   let extraClassDeclaration = [{
-    /// Returns the device ID identifier.
-    StringAttr getDeviceIDIdentifier();
-
-    /// Returns the device type identifier.
-    StringAttr getDeviceTypeIdentifier();
-
     /// Returns max vector op width identifier.
     StringAttr getMaxVectorOpWidthIdentifier();
 
     /// Returns L1 cache size identifier
     StringAttr getL1CacheSizeInBytesIdentifier();
 
-    /// Returns the interface spec for device ID
-    /// Since we verify that the spec contains device ID the function
-    /// will return a valid spec.
-    DataLayoutEntryInterface getSpecForDeviceID();
-
-    /// Returns the interface spec for device type
-    /// Since we verify that the spec contains device type the function
-    /// will return a valid spec.
-    DataLayoutEntryInterface getSpecForDeviceType();
-
     /// Returns the interface spec for max vector op width
     /// Since max vector op width is an optional property, this function will
     /// return a valid spec if the property is defined, otherwise it
@@ -223,9 +201,6 @@ def DLTI_TargetDeviceSpecAttr :
     /// return a valid spec if the property is defined, otherwise it
     /// will return an empty spec.
     DataLayoutEntryInterface getSpecForL1CacheSizeInBytes();
-
-    /// Return the value of device ID
-    TargetDeviceSpecInterface::DeviceID getDeviceID();
   }];
 }
 

mlir/include/mlir/Dialect/DLTI/DLTIBase.td

@@ -57,12 +57,6 @@ def DLTI_Dialect : Dialect {
     kDataLayoutStackAlignmentKey = "dlti.stack_alignment";
 
     // Constants used in target description part of DLTI.
-    constexpr const static ::llvm::StringLiteral
-    kTargetDeviceIDKey = "dlti.device_id";
-
-    constexpr const static ::llvm::StringLiteral
-    kTargetDeviceTypeKey = "dlti.device_type";
-
     constexpr const static ::llvm::StringLiteral
     kTargetDeviceMaxVectorOpWidthKey = "dlti.max_vector_op_width";
 

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

@@ -31,6 +31,10 @@ 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>;
 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
 class ModuleOp;
@@ -246,12 +250,12 @@ class DataLayout {
   /// Returns for max vector op width if the property is defined for the given
   /// device ID, otherwise return std::nullopt.
   std::optional<uint32_t>
-      getMaxVectorOpWidth(TargetDeviceSpecInterface::DeviceID) const;
+      getMaxVectorOpWidth(TargetSystemSpecInterface::DeviceID) const;
 
   /// Returns for L1 cache size if the property is defined for the given
   /// device ID, otherwise return std::nullopt.
   std::optional<uint32_t>
-      getL1CacheSizeInBytes(TargetDeviceSpecInterface::DeviceID) const;
+      getL1CacheSizeInBytes(TargetSystemSpecInterface::DeviceID) const;
 
 private:
   /// Combined layout spec at the given scope.

mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

@@ -238,20 +238,6 @@ def TargetDeviceSpecInterface : AttrInterface<"TargetDeviceSpecInterface"> {
       /*methodBody=*/"",
       /*defaultImplementation=*/[{ return ::mlir::success(); }]
     >,
-    InterfaceMethod<
-      /*description=*/"Returns the entry related to Device ID. The function"
-                      "will crash if the entry is missing.",
-      /*retTy=*/"::mlir::DataLayoutEntryInterface",
-      /*methodName=*/"getSpecForDeviceID",
-      /*args=*/(ins)
-    >,
-    InterfaceMethod<
-      /*description=*/"Returns the entry related to Device Type. "
-                      "The function will crash if the entry is missing.",
-      /*retTy=*/"::mlir::DataLayoutEntryInterface",
-      /*methodName=*/"getSpecForDeviceType",
-      /*args=*/(ins)
-    >,
     InterfaceMethod<
       /*description=*/"Returns the entry related to the given identifier, if "
                       "present. Otherwise, return empty spec.",
@@ -265,19 +251,8 @@ def TargetDeviceSpecInterface : AttrInterface<"TargetDeviceSpecInterface"> {
       /*retTy=*/"::mlir::DataLayoutEntryInterface",
       /*methodName=*/"getSpecForL1CacheSizeInBytes",
       /*args=*/(ins)
-    >,
-    InterfaceMethod<
-      /*description=*/"Returns the entry related to the given identifier, if "
-                      "present.",
-      /*retTy=*/"uint32_t",
-      /*methodName=*/"getDeviceID",
-      /*args=*/(ins)
-    >,
+    >
   ];
-
-  let extraClassDeclaration = [{
-    using DeviceID = uint32_t;
-  }];
 }
 
 def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface"> {
@@ -300,7 +275,7 @@ def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface"> {
   let methods = [
     InterfaceMethod<
       /*description=*/"Returns the list of layout entries.",
-      /*retTy=*/"llvm::ArrayRef<::mlir::TargetDeviceSpecInterface>",
+      /*retTy=*/"llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>",
       /*methodName=*/"getEntries",
       /*args=*/(ins)
     >,
@@ -309,7 +284,7 @@ def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface"> {
                       "ID",
       /*retTy=*/"std::optional<::mlir::TargetDeviceSpecInterface>",
       /*methodName=*/"getDeviceSpecForDeviceID",
-      /*args=*/(ins "TargetDeviceSpecInterface::DeviceID":$deviceID)
+      /*args=*/(ins "StringAttr":$deviceID)
     >,
     InterfaceMethod<
       /*description=*/"Verifies the validity of the specification and "
@@ -323,6 +298,10 @@ def TargetSystemSpecInterface : AttrInterface<"TargetSystemSpecInterface"> {
       }]
     >
   ];
+
+  let extraClassDeclaration = [{
+    using DeviceID = StringAttr;
+  }];
 }
 
 //===----------------------------------------------------------------------===//