[Mlir-commits] [mlir] DLTI: Simplifying getDevicePropertyValue API by returning Attribute type value (PR #96706)

Niranjan Hasabnis llvmlistbot at llvm.org
Wed Jun 26 10:18:51 PDT 2024


https://github.com/nhasabni updated https://github.com/llvm/llvm-project/pull/96706

>From e02797697aa57edabd9ad5ca9cd12315fe27f341 Mon Sep 17 00:00:00 2001
From: Niranjan Hasabnis <niranjan.hasabnis at intel.com>
Date: Tue, 25 Jun 2024 15:07:34 -0700
Subject: [PATCH 1/3] DLTI: Restricting value type to StringAttr in a target
 device spec

---
 .../mlir/Interfaces/DataLayoutInterfaces.h    | 10 ++--
 .../mlir/Interfaces/DataLayoutInterfaces.td   |  6 +--
 mlir/lib/Dialect/DLTI/DLTI.cpp                |  6 +++
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp  | 13 +++--
 mlir/test/Dialect/DLTI/invalid.mlir           | 25 +++++++---
 mlir/test/Dialect/DLTI/roundtrip.mlir         |  8 ++--
 mlir/test/Dialect/DLTI/valid.mlir             | 48 +++++++++----------
 .../Interfaces/DataLayoutInterfacesTest.cpp   | 20 ++++----
 8 files changed, 77 insertions(+), 59 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index 4cbad38df42ba..ec90c82131246 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -93,8 +93,8 @@ uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 
 /// Returns the value of the property from the specified DataLayoutEntry. If the
 /// property is missing from the entry, returns std::nullopt.
-std::optional<int64_t>
-getDevicePropertyValueAsInt(DataLayoutEntryInterface entry);
+std::optional<StringAttr>
+getDevicePropertyValue(DataLayoutEntryInterface entry);
 
 /// Given a list of data layout entries, returns a new list containing the
 /// entries with keys having the given type ID, i.e. belonging to the same type
@@ -246,9 +246,9 @@ class DataLayout {
 
   /// Returns the value of the specified property if the property is defined for
   /// the given device ID, otherwise returns std::nullopt.
-  std::optional<int64_t>
-  getDevicePropertyValueAsInt(TargetSystemSpecInterface::DeviceID,
-                              StringAttr propertyName) const;
+  std::optional<StringAttr>
+  getDevicePropertyValue(TargetSystemSpecInterface::DeviceID,
+                         StringAttr propertyName) const;
 
 private:
   /// Combined layout spec at the given scope.
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 8ff0e3f5f863b..4975e36551a42 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -468,12 +468,12 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
     StaticInterfaceMethod<
       /*description=*/"Returns the value of the property, if the property is "
                       "defined. Otherwise, it returns std::nullopt.",
-      /*retTy=*/"std::optional<int64_t>",
-      /*methodName=*/"getDevicePropertyValueAsInt",
+      /*retTy=*/"std::optional<StringAttr>",
+      /*methodName=*/"getDevicePropertyValue",
       /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        return ::mlir::detail::getDevicePropertyValueAsInt(entry);
+        return ::mlir::detail::getDevicePropertyValue(entry);
       }]
     >
   ];
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 592fced0b2abd..080cba397b549 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -352,9 +352,15 @@ TargetDeviceSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
              << "dlti.target_device_spec does not allow type as a key: "
              << type;
     } else {
+      // Check that keys in a target device spec are unique.
       auto id = entry.getKey().get<StringAttr>();
       if (!ids.insert(id).second)
         return emitError() << "repeated layout entry key: " << id.getValue();
+
+      // Check that values in a target device spec are of StringAttr type.
+      if (!llvm::isa<StringAttr>(entry.getValue()))
+        return emitError() << "dlti.target_device_spec supports values of "
+                              "StringAttr type only.";
     }
   }
 
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index df86bd757b628..b8bbe27fa15dd 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -293,13 +293,12 @@ mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
   return value.getValue().getZExtValue();
 }
 
-std::optional<int64_t>
-mlir::detail::getDevicePropertyValueAsInt(DataLayoutEntryInterface entry) {
+std::optional<StringAttr>
+mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
   if (entry == DataLayoutEntryInterface())
     return std::nullopt;
 
-  auto value = cast<IntegerAttr>(entry.getValue());
-  return value.getValue().getZExtValue();
+  return cast<StringAttr>(entry.getValue());
 }
 
 DataLayoutEntryList
@@ -661,7 +660,7 @@ uint64_t mlir::DataLayout::getStackAlignment() const {
   return *stackAlignment;
 }
 
-std::optional<int64_t> mlir::DataLayout::getDevicePropertyValueAsInt(
+std::optional<StringAttr> mlir::DataLayout::getDevicePropertyValue(
     TargetSystemSpecInterface::DeviceID deviceID,
     StringAttr propertyName) const {
   checkValid();
@@ -676,9 +675,9 @@ std::optional<int64_t> mlir::DataLayout::getDevicePropertyValueAsInt(
   // missing, we return std::nullopt so that the users can resort to
   // the default value however they want.
   if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
-    return iface.getDevicePropertyValueAsInt(entry);
+    return iface.getDevicePropertyValue(entry);
   else
-    return detail::getDevicePropertyValueAsInt(entry);
+    return detail::getDevicePropertyValue(entry);
 }
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/DLTI/invalid.mlir b/mlir/test/Dialect/DLTI/invalid.mlir
index d732fb462e68e..5eb9559e5d415 100644
--- a/mlir/test/Dialect/DLTI/invalid.mlir
+++ b/mlir/test/Dialect/DLTI/invalid.mlir
@@ -113,7 +113,7 @@ module attributes {
   // expected-error at +2 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     : #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>
   >} {}
 
 // -----
@@ -126,7 +126,7 @@ module attributes {
   // expected-error at +2 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     0: #dlti.target_device_spec<
-        #dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>> 
+        #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>
   >} {}
 
 // -----
@@ -137,9 +137,9 @@ module attributes {
   // expected-error at below {{repeated Device ID in dlti.target_system_spec: "CPU"}}
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>>,
+            #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", 8192>>
+            #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
   >} {}
 
 // -----
@@ -152,6 +152,19 @@ module attributes {
   // expected-error at +5 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>,
-            #dlti.dl_entry<"L1_cache_size_in_bytes", 8192>>
+            #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
+            #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+  >} {}
+
+// -----
+
+module attributes {
+  // Repeated DLTI entry
+  //
+  // expected-error at +4 {{dlti.target_device_spec supports values of StringAttr type only.}}
+  // expected-error at +5 {{Error in parsing target device spec}}
+  // expected-error at +4 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
+  dlti.target_system_spec = #dlti.target_system_spec<
+    "CPU": #dlti.target_device_spec<
+            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>>
   >} {}
diff --git a/mlir/test/Dialect/DLTI/roundtrip.mlir b/mlir/test/Dialect/DLTI/roundtrip.mlir
index 7b8255ad71d4d..c5be3bc805330 100644
--- a/mlir/test/Dialect/DLTI/roundtrip.mlir
+++ b/mlir/test/Dialect/DLTI/roundtrip.mlir
@@ -58,16 +58,16 @@
 // CHECK: module attributes {
 // CHECK: dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK:  "CPU" : #dlti.target_device_spec<
-// CHECK:   #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096 : ui32>>,
+// CHECK:   #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", "4096">>,
 // CHECK: "GPU" : #dlti.target_device_spec<
-// CHECK:   #dlti.dl_entry<"dlti.max_vector_op_width", 128 : ui32>>
+// CHECK:   #dlti.dl_entry<"dlti.max_vector_op_width", "128">>
 // CHECK: >} {
 // CHECK: }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096: ui32>>,
+      #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", "4096">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"dlti.max_vector_op_width", 128: ui32>>
+      #dlti.dl_entry<"dlti.max_vector_op_width", "128">>
   >} {}
  
diff --git a/mlir/test/Dialect/DLTI/valid.mlir b/mlir/test/Dialect/DLTI/valid.mlir
index 98e9d0b8de491..2084168d865df 100644
--- a/mlir/test/Dialect/DLTI/valid.mlir
+++ b/mlir/test/Dialect/DLTI/valid.mlir
@@ -5,17 +5,17 @@
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 128: i32>>
+      #dlti.dl_entry<"max_vector_op_width", "128">>
   >} {}
 
 // -----
@@ -23,17 +23,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i32>>
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192: i32>>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
   >} {}
 
 // -----
@@ -41,17 +41,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i64>>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i64>>
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i64>>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192: i64>>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
   >} {}
 
 // -----
@@ -59,17 +59,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i32>>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 64: i32>>,
+      #dlti.dl_entry<"max_vector_op_width", "64">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 128: i32>>
+      #dlti.dl_entry<"max_vector_op_width", "128">>
   >} {}
 
 // -----
@@ -77,17 +77,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 64: i64>>,
+      #dlti.dl_entry<"max_vector_op_width", "64">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 128: i64>>
+      #dlti.dl_entry<"max_vector_op_width", "128">>
   >} {}
 
 // -----
@@ -95,15 +95,15 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 64>>,
+      #dlti.dl_entry<"max_vector_op_width", "64">>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", 128>>
+      #dlti.dl_entry<"max_vector_op_width", "128">>
   >} {}
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index bcbfa88c5d815..2f1ea548bea90 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -495,11 +495,11 @@ TEST(DataLayout, NullSpec) {
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
 
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
                 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
             std::nullopt);
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
                 Builder(&ctx).getStringAttr("max_vector_width")),
             std::nullopt);
@@ -535,11 +535,11 @@ TEST(DataLayout, EmptySpec) {
   EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
   EXPECT_EQ(layout.getStackAlignment(), 0u);
 
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
                 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
             std::nullopt);
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU" /* device ID*/),
                 Builder(&ctx).getStringAttr("max_vector_width")),
             std::nullopt);
@@ -599,8 +599,8 @@ TEST(DataLayout, SpecWithTargetSystemDescEntries) {
   module attributes { dl_target_sys_desc_test.target_system_spec =
     #dl_target_sys_desc_test.target_system_spec<
       "CPU": #dlti.target_device_spec<
-              #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>,
-              #dlti.dl_entry<"max_vector_op_width", 128 : ui32>>
+              #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
+              #dlti.dl_entry<"max_vector_op_width", "128">>
     > } {}
   )MLIR";
 
@@ -610,14 +610,14 @@ TEST(DataLayout, SpecWithTargetSystemDescEntries) {
 
   OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
   DataLayout layout(*module);
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
                 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
-            std::optional<int64_t>(4096));
-  EXPECT_EQ(layout.getDevicePropertyValueAsInt(
+            std::optional<StringAttr>(Builder(&ctx).getStringAttr("4096")));
+  EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
                 Builder(&ctx).getStringAttr("max_vector_op_width")),
-            std::optional<int64_t>(128));
+            std::optional<StringAttr>(Builder(&ctx).getStringAttr("128")));
 }
 
 TEST(DataLayout, Caching) {

>From 970e7168f1ef4fb69897bae948ee41d505cdc536 Mon Sep 17 00:00:00 2001
From: Niranjan Hasabnis <niranjan.hasabnis at intel.com>
Date: Wed, 26 Jun 2024 10:00:30 -0700
Subject: [PATCH 2/3] Using Attribute as the return type of value; removing
 restriction of StringAttr value type

---
 .../mlir/Interfaces/DataLayoutInterfaces.h    |  5 ++-
 .../mlir/Interfaces/DataLayoutInterfaces.td   |  2 +-
 mlir/lib/Dialect/DLTI/DLTI.cpp                |  5 ---
 mlir/lib/Interfaces/DataLayoutInterfaces.cpp  |  6 ++--
 mlir/test/Dialect/DLTI/invalid.mlir           | 13 -------
 mlir/test/Dialect/DLTI/valid.mlir             | 36 +++++++++++++++++++
 .../Interfaces/DataLayoutInterfacesTest.cpp   |  4 +--
 7 files changed, 44 insertions(+), 27 deletions(-)

diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
index ec90c82131246..ab65f92820a6a 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h
@@ -93,8 +93,7 @@ uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
 
 /// Returns the value of the property from the specified DataLayoutEntry. If the
 /// property is missing from the entry, returns std::nullopt.
-std::optional<StringAttr>
-getDevicePropertyValue(DataLayoutEntryInterface entry);
+std::optional<Attribute> getDevicePropertyValue(DataLayoutEntryInterface entry);
 
 /// Given a list of data layout entries, returns a new list containing the
 /// entries with keys having the given type ID, i.e. belonging to the same type
@@ -246,7 +245,7 @@ class DataLayout {
 
   /// Returns the value of the specified property if the property is defined for
   /// the given device ID, otherwise returns std::nullopt.
-  std::optional<StringAttr>
+  std::optional<Attribute>
   getDevicePropertyValue(TargetSystemSpecInterface::DeviceID,
                          StringAttr propertyName) const;
 
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 4975e36551a42..d4b89997d9c86 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -468,7 +468,7 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
     StaticInterfaceMethod<
       /*description=*/"Returns the value of the property, if the property is "
                       "defined. Otherwise, it returns std::nullopt.",
-      /*retTy=*/"std::optional<StringAttr>",
+      /*retTy=*/"std::optional<Attribute>",
       /*methodName=*/"getDevicePropertyValue",
       /*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
       /*methodBody=*/"",
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 080cba397b549..420c605d1a19b 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -356,11 +356,6 @@ TargetDeviceSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
       auto id = entry.getKey().get<StringAttr>();
       if (!ids.insert(id).second)
         return emitError() << "repeated layout entry key: " << id.getValue();
-
-      // Check that values in a target device spec are of StringAttr type.
-      if (!llvm::isa<StringAttr>(entry.getValue()))
-        return emitError() << "dlti.target_device_spec supports values of "
-                              "StringAttr type only.";
     }
   }
 
diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
index b8bbe27fa15dd..2634245a4b7b1 100644
--- a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
+++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp
@@ -293,12 +293,12 @@ mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
   return value.getValue().getZExtValue();
 }
 
-std::optional<StringAttr>
+std::optional<Attribute>
 mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
   if (entry == DataLayoutEntryInterface())
     return std::nullopt;
 
-  return cast<StringAttr>(entry.getValue());
+  return entry.getValue();
 }
 
 DataLayoutEntryList
@@ -660,7 +660,7 @@ uint64_t mlir::DataLayout::getStackAlignment() const {
   return *stackAlignment;
 }
 
-std::optional<StringAttr> mlir::DataLayout::getDevicePropertyValue(
+std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
     TargetSystemSpecInterface::DeviceID deviceID,
     StringAttr propertyName) const {
   checkValid();
diff --git a/mlir/test/Dialect/DLTI/invalid.mlir b/mlir/test/Dialect/DLTI/invalid.mlir
index 5eb9559e5d415..8d7d7e92315d2 100644
--- a/mlir/test/Dialect/DLTI/invalid.mlir
+++ b/mlir/test/Dialect/DLTI/invalid.mlir
@@ -155,16 +155,3 @@ module attributes {
             #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
             #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
   >} {}
-
-// -----
-
-module attributes {
-  // Repeated DLTI entry
-  //
-  // expected-error at +4 {{dlti.target_device_spec supports values of StringAttr type only.}}
-  // expected-error at +5 {{Error in parsing target device spec}}
-  // expected-error at +4 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
-  dlti.target_system_spec = #dlti.target_system_spec<
-    "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>>
-  >} {}
diff --git a/mlir/test/Dialect/DLTI/valid.mlir b/mlir/test/Dialect/DLTI/valid.mlir
index 2084168d865df..8f171faf23eb6 100644
--- a/mlir/test/Dialect/DLTI/valid.mlir
+++ b/mlir/test/Dialect/DLTI/valid.mlir
@@ -107,3 +107,39 @@ module attributes {
     "GPU": #dlti.target_device_spec<
       #dlti.dl_entry<"max_vector_op_width", "128">>
   >} {}
+
+// -----
+
+// CHECK:      module attributes {
+// CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
+// CHECK-SAME:    "CPU" : #dlti.target_device_spec<
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>>,
+// CHECK-SAME:    "GPU" : #dlti.target_device_spec<
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
+// CHECK-SAME:  >} {
+// CHECK:        }
+module attributes {
+  dlti.target_system_spec = #dlti.target_system_spec<
+    "CPU": #dlti.target_device_spec<
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>>,
+    "GPU": #dlti.target_device_spec<
+      #dlti.dl_entry<"max_vector_op_width", "128">>
+  >} {}
+
+// -----
+
+// CHECK:      module attributes {
+// CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
+// CHECK-SAME:    "CPU" : #dlti.target_device_spec<
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 4.096000e+03 : f32>>,
+// CHECK-SAME:    "GPU" : #dlti.target_device_spec<
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "128">>
+// CHECK-SAME:  >} {
+// CHECK:        }
+module attributes {
+  dlti.target_system_spec = #dlti.target_system_spec<
+    "CPU": #dlti.target_device_spec<
+      #dlti.dl_entry<"max_vector_op_width", 4096.0 : f32>>,
+    "GPU": #dlti.target_device_spec<
+      #dlti.dl_entry<"L1_cache_size_in_bytes", "128">>
+  >} {}
diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
index 2f1ea548bea90..d1227b045d4ed 100644
--- a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
+++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp
@@ -613,11 +613,11 @@ TEST(DataLayout, SpecWithTargetSystemDescEntries) {
   EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
                 Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
-            std::optional<StringAttr>(Builder(&ctx).getStringAttr("4096")));
+            std::optional<Attribute>(Builder(&ctx).getStringAttr("4096")));
   EXPECT_EQ(layout.getDevicePropertyValue(
                 Builder(&ctx).getStringAttr("CPU") /* device ID*/,
                 Builder(&ctx).getStringAttr("max_vector_op_width")),
-            std::optional<StringAttr>(Builder(&ctx).getStringAttr("128")));
+            std::optional<Attribute>(Builder(&ctx).getStringAttr("128")));
 }
 
 TEST(DataLayout, Caching) {

>From 43352c1e3a8c2fa297678efa7ead455c9f5eafba Mon Sep 17 00:00:00 2001
From: Niranjan Hasabnis <niranjan.hasabnis at intel.com>
Date: Wed, 26 Jun 2024 10:18:28 -0700
Subject: [PATCH 3/3] Removing spurious changes introduced by first commit

---
 mlir/test/Dialect/DLTI/invalid.mlir   | 12 +++----
 mlir/test/Dialect/DLTI/roundtrip.mlir |  8 ++---
 mlir/test/Dialect/DLTI/valid.mlir     | 52 ++++++++++++++-------------
 3 files changed, 38 insertions(+), 34 deletions(-)

diff --git a/mlir/test/Dialect/DLTI/invalid.mlir b/mlir/test/Dialect/DLTI/invalid.mlir
index 8d7d7e92315d2..36bf6088fab7e 100644
--- a/mlir/test/Dialect/DLTI/invalid.mlir
+++ b/mlir/test/Dialect/DLTI/invalid.mlir
@@ -113,7 +113,7 @@ module attributes {
   // expected-error at +2 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     : #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>
   >} {}
 
 // -----
@@ -126,7 +126,7 @@ module attributes {
   // expected-error at +2 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     0: #dlti.target_device_spec<
-        #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>
+        #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>
   >} {}
 
 // -----
@@ -137,9 +137,9 @@ module attributes {
   // expected-error at below {{repeated Device ID in dlti.target_system_spec: "CPU"}}
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>>,
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+            #dlti.dl_entry<"L1_cache_size_in_bytes", 8192>>
   >} {}
 
 // -----
@@ -152,6 +152,6 @@ module attributes {
   // expected-error at +5 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-            #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
-            #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+            #dlti.dl_entry<"L1_cache_size_in_bytes", 4096>,
+            #dlti.dl_entry<"L1_cache_size_in_bytes", 8192>>
   >} {}
diff --git a/mlir/test/Dialect/DLTI/roundtrip.mlir b/mlir/test/Dialect/DLTI/roundtrip.mlir
index c5be3bc805330..43188aad595a7 100644
--- a/mlir/test/Dialect/DLTI/roundtrip.mlir
+++ b/mlir/test/Dialect/DLTI/roundtrip.mlir
@@ -58,16 +58,16 @@
 // CHECK: module attributes {
 // CHECK: dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK:  "CPU" : #dlti.target_device_spec<
-// CHECK:   #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", "4096">>,
+// CHECK:   #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096 : ui32>>,
 // CHECK: "GPU" : #dlti.target_device_spec<
-// CHECK:   #dlti.dl_entry<"dlti.max_vector_op_width", "128">>
+// CHECK:   #dlti.dl_entry<"dlti.max_vector_op_width", 128 : ui32>>
 // CHECK: >} {
 // CHECK: }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", "4096">>,
+      #dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096 : ui32>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"dlti.max_vector_op_width", "128">>
+      #dlti.dl_entry<"dlti.max_vector_op_width", 128 : ui32>>
   >} {}
  
diff --git a/mlir/test/Dialect/DLTI/valid.mlir b/mlir/test/Dialect/DLTI/valid.mlir
index 8f171faf23eb6..8ee7a7eaa069d 100644
--- a/mlir/test/Dialect/DLTI/valid.mlir
+++ b/mlir/test/Dialect/DLTI/valid.mlir
@@ -5,17 +5,17 @@
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "128">>
+      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
   >} {}
 
 // -----
@@ -23,17 +23,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i32>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i32>>
   >} {}
 
 // -----
@@ -41,17 +41,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i64>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+// CHECK-SAME:      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i64>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "4096">>,
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i64>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"L1_cache_size_in_bytes", "8192">>
+      #dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i64>>
   >} {}
 
 // -----
@@ -59,17 +59,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i32>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "64">>,
+      #dlti.dl_entry<"max_vector_op_width", 64 : i32>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "128">>
+      #dlti.dl_entry<"max_vector_op_width", 128 : i32>>
   >} {}
 
 // -----
@@ -77,17 +77,17 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "64">>,
+      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "128">>
+      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
   >} {}
 
 // -----
@@ -95,21 +95,23 @@ module attributes {
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "64">>,
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
 // CHECK-SAME:    "GPU" : #dlti.target_device_spec<
-// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", "128">>
+// CHECK-SAME:      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
 // CHECK-SAME:  >} {
 // CHECK:        }
 module attributes {
   dlti.target_system_spec = #dlti.target_system_spec<
     "CPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "64">>,
+      #dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
     "GPU": #dlti.target_device_spec<
-      #dlti.dl_entry<"max_vector_op_width", "128">>
+      #dlti.dl_entry<"max_vector_op_width", 128 : i64>>
   >} {}
 
 // -----
 
+// Check values of mixed type
+//
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<
@@ -128,6 +130,8 @@ module attributes {
 
 // -----
 
+// Check values of mixed type
+//
 // CHECK:      module attributes {
 // CHECK-SAME:  dlti.target_system_spec = #dlti.target_system_spec<
 // CHECK-SAME:    "CPU" : #dlti.target_device_spec<



More information about the Mlir-commits mailing list