Skip to content

Commit bdeee9b

Browse files
nhasabninhasabni
authored
DLTI: Simplifying getDevicePropertyValue API by returning Attribute type value (#96706)
**Rationale** - With the current flexibility of supporting any type of value, we will need to offer type-specific APIs to fetch a value (e.g., `getDevicePropertyValueAsInt` for integer type, `getDevicePropertyValueAsFloat` for float type, etc.) A single type of value will eliminate this need. - Current flexibility can also lead to typing errors when a user fetches the value of a property using an API that is not consistent with the type of the value. **What is the change** For following system description, ``` module attributes { dlti.target_system_spec = #dlti.target_system_spec< "CPU": #dlti.target_device_spec< #dlti.dl_entry<"max_vector_op_width", 64.0 : f32>>, "GPU": #dlti.target_device_spec< #dlti.dl_entry<"max_vector_op_width", 128 : ui32>> >} {} ``` a user no longer needs to use `getDevicePropertyValueAsInt` for retrieving GPU's `max_vector_op_width` and `getDevicePropertyValueAsFloat` for retrieving CPU's `max_vector_op_width`. Instead it can be done with a uniform API of `getDevicePropertyValue`.
1 parent 7a969ec commit bdeee9b

File tree

8 files changed

+79
-40
lines changed

8 files changed

+79
-40
lines changed

‎mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -93,8 +93,7 @@ uint64_t getDefaultStackAlignment(DataLayoutEntryInterface entry);
9393

9494
/// Returns the value of the property from the specified DataLayoutEntry. If the
9595
/// property is missing from the entry, returns std::nullopt.
96-
std::optional<int64_t>
97-
getDevicePropertyValueAsInt(DataLayoutEntryInterface entry);
96+
std::optional<Attribute> getDevicePropertyValue(DataLayoutEntryInterface entry);
9897

9998
/// Given a list of data layout entries, returns a new list containing the
10099
/// entries with keys having the given type ID, i.e. belonging to the same type
@@ -246,9 +245,9 @@ class DataLayout {
246245

247246
/// Returns the value of the specified property if the property is defined for
248247
/// the given device ID, otherwise returns std::nullopt.
249-
std::optional<int64_t>
250-
getDevicePropertyValueAsInt(TargetSystemSpecInterface::DeviceID,
251-
StringAttr propertyName) const;
248+
std::optional<Attribute>
249+
getDevicePropertyValue(TargetSystemSpecInterface::DeviceID,
250+
StringAttr propertyName) const;
252251

253252
private:
254253
/// Combined layout spec at the given scope.

‎mlir/include/mlir/Interfaces/DataLayoutInterfaces.td

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -468,12 +468,12 @@ def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> {
468468
StaticInterfaceMethod<
469469
/*description=*/"Returns the value of the property, if the property is "
470470
"defined. Otherwise, it returns std::nullopt.",
471-
/*retTy=*/"std::optional<int64_t>",
472-
/*methodName=*/"getDevicePropertyValueAsInt",
471+
/*retTy=*/"std::optional<Attribute>",
472+
/*methodName=*/"getDevicePropertyValue",
473473
/*args=*/(ins "::mlir::DataLayoutEntryInterface":$entry),
474474
/*methodBody=*/"",
475475
/*defaultImplementation=*/[{
476-
return ::mlir::detail::getDevicePropertyValueAsInt(entry);
476+
return ::mlir::detail::getDevicePropertyValue(entry);
477477
}]
478478
>
479479
];

‎mlir/lib/Dialect/DLTI/DLTI.cpp

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -352,6 +352,7 @@ TargetDeviceSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
352352
<< "dlti.target_device_spec does not allow type as a key: "
353353
<< type;
354354
} else {
355+
// Check that keys in a target device spec are unique.
355356
auto id = entry.getKey().get<StringAttr>();
356357
if (!ids.insert(id).second)
357358
return emitError() << "repeated layout entry key: " << id.getValue();

‎mlir/lib/Interfaces/DataLayoutInterfaces.cpp

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -293,13 +293,12 @@ mlir::detail::getDefaultStackAlignment(DataLayoutEntryInterface entry) {
293293
return value.getValue().getZExtValue();
294294
}
295295

296-
std::optional<int64_t>
297-
mlir::detail::getDevicePropertyValueAsInt(DataLayoutEntryInterface entry) {
296+
std::optional<Attribute>
297+
mlir::detail::getDevicePropertyValue(DataLayoutEntryInterface entry) {
298298
if (entry == DataLayoutEntryInterface())
299299
return std::nullopt;
300300

301-
auto value = cast<IntegerAttr>(entry.getValue());
302-
return value.getValue().getZExtValue();
301+
return entry.getValue();
303302
}
304303

305304
DataLayoutEntryList
@@ -661,7 +660,7 @@ uint64_t mlir::DataLayout::getStackAlignment() const {
661660
return *stackAlignment;
662661
}
663662

664-
std::optional<int64_t> mlir::DataLayout::getDevicePropertyValueAsInt(
663+
std::optional<Attribute> mlir::DataLayout::getDevicePropertyValue(
665664
TargetSystemSpecInterface::DeviceID deviceID,
666665
StringAttr propertyName) const {
667666
checkValid();
@@ -676,9 +675,9 @@ std::optional<int64_t> mlir::DataLayout::getDevicePropertyValueAsInt(
676675
// missing, we return std::nullopt so that the users can resort to
677676
// the default value however they want.
678677
if (auto iface = dyn_cast_or_null<DataLayoutOpInterface>(scope))
679-
return iface.getDevicePropertyValueAsInt(entry);
678+
return iface.getDevicePropertyValue(entry);
680679
else
681-
return detail::getDevicePropertyValueAsInt(entry);
680+
return detail::getDevicePropertyValue(entry);
682681
}
683682

684683
//===----------------------------------------------------------------------===//

‎mlir/test/Dialect/DLTI/invalid.mlir

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ module attributes {
126126
// expected-error@+2 {{failed to parse DLTI_TargetSystemSpecAttr parameter 'entries' which is to be a `::llvm::ArrayRef<DeviceIDTargetDeviceSpecPair>`}}
127127
dlti.target_system_spec = #dlti.target_system_spec<
128128
0: #dlti.target_device_spec<
129-
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>>
129+
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>
130130
>} {}
131131

132132
// -----

‎mlir/test/Dialect/DLTI/roundtrip.mlir

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,8 +66,8 @@
6666
module attributes {
6767
dlti.target_system_spec = #dlti.target_system_spec<
6868
"CPU": #dlti.target_device_spec<
69-
#dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096: ui32>>,
69+
#dlti.dl_entry<"dlti.L1_cache_size_in_bytes", 4096 : ui32>>,
7070
"GPU": #dlti.target_device_spec<
71-
#dlti.dl_entry<"dlti.max_vector_op_width", 128: ui32>>
71+
#dlti.dl_entry<"dlti.max_vector_op_width", 128 : ui32>>
7272
>} {}
7373

‎mlir/test/Dialect/DLTI/valid.mlir

Lines changed: 52 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@
1313
module attributes {
1414
dlti.target_system_spec = #dlti.target_system_spec<
1515
"CPU": #dlti.target_device_spec<
16-
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>>,
16+
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
1717
"GPU": #dlti.target_device_spec<
18-
#dlti.dl_entry<"max_vector_op_width", 128: i32>>
18+
#dlti.dl_entry<"max_vector_op_width", 128 : i32>>
1919
>} {}
2020

2121
// -----
@@ -31,9 +31,9 @@ module attributes {
3131
module attributes {
3232
dlti.target_system_spec = #dlti.target_system_spec<
3333
"CPU": #dlti.target_device_spec<
34-
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i32>>,
34+
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i32>>,
3535
"GPU": #dlti.target_device_spec<
36-
#dlti.dl_entry<"L1_cache_size_in_bytes", 8192: i32>>
36+
#dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i32>>
3737
>} {}
3838

3939
// -----
@@ -49,9 +49,9 @@ module attributes {
4949
module attributes {
5050
dlti.target_system_spec = #dlti.target_system_spec<
5151
"CPU": #dlti.target_device_spec<
52-
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096: i64>>,
52+
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : i64>>,
5353
"GPU": #dlti.target_device_spec<
54-
#dlti.dl_entry<"L1_cache_size_in_bytes", 8192: i64>>
54+
#dlti.dl_entry<"L1_cache_size_in_bytes", 8192 : i64>>
5555
>} {}
5656

5757
// -----
@@ -67,9 +67,9 @@ module attributes {
6767
module attributes {
6868
dlti.target_system_spec = #dlti.target_system_spec<
6969
"CPU": #dlti.target_device_spec<
70-
#dlti.dl_entry<"max_vector_op_width", 64: i32>>,
70+
#dlti.dl_entry<"max_vector_op_width", 64 : i32>>,
7171
"GPU": #dlti.target_device_spec<
72-
#dlti.dl_entry<"max_vector_op_width", 128: i32>>
72+
#dlti.dl_entry<"max_vector_op_width", 128 : i32>>
7373
>} {}
7474

7575
// -----
@@ -85,9 +85,9 @@ module attributes {
8585
module attributes {
8686
dlti.target_system_spec = #dlti.target_system_spec<
8787
"CPU": #dlti.target_device_spec<
88-
#dlti.dl_entry<"max_vector_op_width", 64: i64>>,
88+
#dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
8989
"GPU": #dlti.target_device_spec<
90-
#dlti.dl_entry<"max_vector_op_width", 128: i64>>
90+
#dlti.dl_entry<"max_vector_op_width", 128 : i64>>
9191
>} {}
9292

9393
// -----
@@ -103,7 +103,47 @@ module attributes {
103103
module attributes {
104104
dlti.target_system_spec = #dlti.target_system_spec<
105105
"CPU": #dlti.target_device_spec<
106-
#dlti.dl_entry<"max_vector_op_width", 64>>,
106+
#dlti.dl_entry<"max_vector_op_width", 64 : i64>>,
107107
"GPU": #dlti.target_device_spec<
108-
#dlti.dl_entry<"max_vector_op_width", 128>>
108+
#dlti.dl_entry<"max_vector_op_width", 128 : i64>>
109+
>} {}
110+
111+
// -----
112+
113+
// Check values of mixed type
114+
//
115+
// CHECK: module attributes {
116+
// CHECK-SAME: dlti.target_system_spec = #dlti.target_system_spec<
117+
// CHECK-SAME: "CPU" : #dlti.target_device_spec<
118+
// CHECK-SAME: #dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>>,
119+
// CHECK-SAME: "GPU" : #dlti.target_device_spec<
120+
// CHECK-SAME: #dlti.dl_entry<"max_vector_op_width", "128">>
121+
// CHECK-SAME: >} {
122+
// CHECK: }
123+
module attributes {
124+
dlti.target_system_spec = #dlti.target_system_spec<
125+
"CPU": #dlti.target_device_spec<
126+
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>>,
127+
"GPU": #dlti.target_device_spec<
128+
#dlti.dl_entry<"max_vector_op_width", "128">>
129+
>} {}
130+
131+
// -----
132+
133+
// Check values of mixed type
134+
//
135+
// CHECK: module attributes {
136+
// CHECK-SAME: dlti.target_system_spec = #dlti.target_system_spec<
137+
// CHECK-SAME: "CPU" : #dlti.target_device_spec<
138+
// CHECK-SAME: #dlti.dl_entry<"max_vector_op_width", 4.096000e+03 : f32>>,
139+
// CHECK-SAME: "GPU" : #dlti.target_device_spec<
140+
// CHECK-SAME: #dlti.dl_entry<"L1_cache_size_in_bytes", "128">>
141+
// CHECK-SAME: >} {
142+
// CHECK: }
143+
module attributes {
144+
dlti.target_system_spec = #dlti.target_system_spec<
145+
"CPU": #dlti.target_device_spec<
146+
#dlti.dl_entry<"max_vector_op_width", 4096.0 : f32>>,
147+
"GPU": #dlti.target_device_spec<
148+
#dlti.dl_entry<"L1_cache_size_in_bytes", "128">>
109149
>} {}

‎mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -495,11 +495,11 @@ TEST(DataLayout, NullSpec) {
495495
EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
496496
EXPECT_EQ(layout.getStackAlignment(), 0u);
497497

498-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
498+
EXPECT_EQ(layout.getDevicePropertyValue(
499499
Builder(&ctx).getStringAttr("CPU" /* device ID*/),
500500
Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
501501
std::nullopt);
502-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
502+
EXPECT_EQ(layout.getDevicePropertyValue(
503503
Builder(&ctx).getStringAttr("CPU" /* device ID*/),
504504
Builder(&ctx).getStringAttr("max_vector_width")),
505505
std::nullopt);
@@ -535,11 +535,11 @@ TEST(DataLayout, EmptySpec) {
535535
EXPECT_EQ(layout.getGlobalMemorySpace(), Attribute());
536536
EXPECT_EQ(layout.getStackAlignment(), 0u);
537537

538-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
538+
EXPECT_EQ(layout.getDevicePropertyValue(
539539
Builder(&ctx).getStringAttr("CPU" /* device ID*/),
540540
Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
541541
std::nullopt);
542-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
542+
EXPECT_EQ(layout.getDevicePropertyValue(
543543
Builder(&ctx).getStringAttr("CPU" /* device ID*/),
544544
Builder(&ctx).getStringAttr("max_vector_width")),
545545
std::nullopt);
@@ -599,8 +599,8 @@ TEST(DataLayout, SpecWithTargetSystemDescEntries) {
599599
module attributes { dl_target_sys_desc_test.target_system_spec =
600600
#dl_target_sys_desc_test.target_system_spec<
601601
"CPU": #dlti.target_device_spec<
602-
#dlti.dl_entry<"L1_cache_size_in_bytes", 4096 : ui32>,
603-
#dlti.dl_entry<"max_vector_op_width", 128 : ui32>>
602+
#dlti.dl_entry<"L1_cache_size_in_bytes", "4096">,
603+
#dlti.dl_entry<"max_vector_op_width", "128">>
604604
> } {}
605605
)MLIR";
606606

@@ -610,14 +610,14 @@ TEST(DataLayout, SpecWithTargetSystemDescEntries) {
610610

611611
OwningOpRef<ModuleOp> module = parseSourceString<ModuleOp>(ir, &ctx);
612612
DataLayout layout(*module);
613-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
613+
EXPECT_EQ(layout.getDevicePropertyValue(
614614
Builder(&ctx).getStringAttr("CPU") /* device ID*/,
615615
Builder(&ctx).getStringAttr("L1_cache_size_in_bytes")),
616-
std::optional<int64_t>(4096));
617-
EXPECT_EQ(layout.getDevicePropertyValueAsInt(
616+
std::optional<Attribute>(Builder(&ctx).getStringAttr("4096")));
617+
EXPECT_EQ(layout.getDevicePropertyValue(
618618
Builder(&ctx).getStringAttr("CPU") /* device ID*/,
619619
Builder(&ctx).getStringAttr("max_vector_op_width")),
620-
std::optional<int64_t>(128));
620+
std::optional<Attribute>(Builder(&ctx).getStringAttr("128")));
621621
}
622622

623623
TEST(DataLayout, Caching) {

0 commit comments

Comments
 (0)