[Mlir-commits] [mlir] [MLIR][DLTI] Introduce DLTIQueryInterface and impl for DLTI attrs (PR #104595)

[Rolf Morel]

================
@@ -393,39 +407,50 @@ TargetSystemSpecAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 // DLTIDialect
 //===----------------------------------------------------------------------===//
 
-DataLayoutSpecInterface dlti::getDataLayoutSpec(Operation *op) {
-  DataLayoutSpecInterface dlSpec = nullptr;
-
-  for (Operation *cur = op; cur && !dlSpec; cur = cur->getParentOp()) {
-    if (auto dataLayoutOp = dyn_cast<DataLayoutOpInterface>(cur))
-      dlSpec = dataLayoutOp.getDataLayoutSpec();
-    else if (auto moduleOp = dyn_cast<ModuleOp>(cur))
-      dlSpec = moduleOp.getDataLayoutSpec();
-    else
-      for (NamedAttribute attr : cur->getAttrs())
-        if ((dlSpec = llvm::dyn_cast<DataLayoutSpecInterface>(attr.getValue())))
-          break;
-  }
+/// Retrieve the first `DLTIQueryInterface`-implementing attribute that is
+/// attached to `op` or such an attr on as close as possible an ancestor. The
+/// op the attribute is attached to is returned as well.
+static std::pair<DLTIQueryInterface, Operation *>
+getClosestQueryable(Operation *op) {
+  DLTIQueryInterface queryable = {};
+
+  // Search op and its ancestors for the first attached DLTIQueryInterface attr.
+  do {
+    for (NamedAttribute attr : op->getAttrs())
+      if ((queryable = llvm::dyn_cast<DLTIQueryInterface>(attr.getValue())))
+        break;
+  } while (!queryable && (op = op->getParentOp()));
 
-  return dlSpec;
+  return std::pair(queryable, op);
 }
 
-TargetSystemSpecInterface dlti::getTargetSystemSpec(Operation *op) {
-  TargetSystemSpecInterface sysSpec = nullptr;
-
-  for (Operation *cur = op; cur && !sysSpec; cur = cur->getParentOp()) {
-    if (auto dataLayoutOp = dyn_cast<DataLayoutOpInterface>(cur))
-      sysSpec = dataLayoutOp.getTargetSystemSpec();
-    else if (auto moduleOp = dyn_cast<ModuleOp>(cur))
-      sysSpec = moduleOp.getTargetSystemSpec();
-    else
-      for (NamedAttribute attr : cur->getAttrs())
-        if ((sysSpec =
-                 llvm::dyn_cast<TargetSystemSpecInterface>(attr.getValue())))
-          break;
+FailureOr<Attribute> dlti::query(Operation *op, ArrayRef<StringAttr> keys,
+                                 InFlightDiagnostic *diag) {
+  auto [queryable, queryOp] = getClosestQueryable(op);
+
+#define FAIL(message)                                                          \
----------------
rolfmorel wrote:

Have gotten rid of the macro magic and simplified `dlti::query()`'s interface.

https://github.com/llvm/llvm-project/pull/104595