From bf0aa8aab4b6cd08567874930a3b1702555c2060 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk@intel.com>
Date: Tue, 11 Feb 2025 12:20:45 +0100
Subject: [PATCH 1/2] [mlir][dlti] Query by strings

Adds DLTI utility to query using strings directly as keys.
---
 mlir/include/mlir/Dialect/DLTI/DLTI.h |  6 ++++++
 mlir/lib/Dialect/DLTI/DLTI.cpp        | 16 ++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index f268fea340a6f..4004c2f64e6da 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -28,6 +28,12 @@ namespace dlti {
 /// query interface-implementing attrs, starting from attr obtained from `op`.
 FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
                            bool emitError = false);
+
+/// Perform a DLTI-query at `op` using string `keys` as DLTI entry keys,
+/// recursively querying on query interface-implementing attrs, starting from
+/// attr obtained from `op`.
+FailureOr<Attribute> query(Operation *op, ArrayRef<StringRef> keys,
+                           bool emitError = false);
 } // namespace dlti
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 2510e774f2b2a..86bee1933268b 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -508,6 +508,9 @@ getClosestQueryable(Operation *op) {
 
 FailureOr<Attribute>
 dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
+  if (!op)
+    return failure();
+
   if (keys.empty()) {
     if (emitError) {
       auto diag = op->emitError() << "target op of failed DLTI query";
@@ -562,6 +565,19 @@ dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
   return currentAttr;
 }
 
+FailureOr<Attribute> dlti::query(Operation *op, ArrayRef<StringRef> keys,
+                                 bool emitError) {
+  if (!op)
+    return failure();
+
+  MLIRContext *ctx = op->getContext();
+  SmallVector<DataLayoutEntryKey> entryKeys;
+  for (StringRef key : keys)
+    entryKeys.push_back(StringAttr::get(ctx, key));
+
+  return dlti::query(op, entryKeys, emitError);
+}
+
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName;
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey;
 constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig;

From ac3ad83e2e74bfde130119a5a5a878754665ee95 Mon Sep 17 00:00:00 2001
From: Adam Siemieniuk <adam.siemieniuk@intel.com>
Date: Tue, 11 Feb 2025 13:57:57 +0100
Subject: [PATCH 2/2] Address comments

---
 mlir/include/mlir/Dialect/DLTI/DLTI.h | 6 +++---
 mlir/lib/Dialect/DLTI/DLTI.cpp        | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 4004c2f64e6da..84cf6eabc5fa9 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -29,9 +29,9 @@ namespace dlti {
 FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
                            bool emitError = false);
 
-/// Perform a DLTI-query at `op` using string `keys` as DLTI entry keys,
-/// recursively querying on query interface-implementing attrs, starting from
-/// attr obtained from `op`.
+/// Perform a DLTI-query at `op` using each string in `keys` as a separate DLTI
+/// entry key, recursively querying on query interface-implementing attrs,
+/// starting from attr obtained from `op`.
 FailureOr<Attribute> query(Operation *op, ArrayRef<StringRef> keys,
                            bool emitError = false);
 } // namespace dlti
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index 86bee1933268b..b057554c40d8c 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -571,7 +571,7 @@ FailureOr<Attribute> dlti::query(Operation *op, ArrayRef<StringRef> keys,
     return failure();
 
   MLIRContext *ctx = op->getContext();
-  SmallVector<DataLayoutEntryKey> entryKeys;
+  SmallVector<DataLayoutEntryKey> entryKeys(keys.size());
   for (StringRef key : keys)
     entryKeys.push_back(StringAttr::get(ctx, key));