[Mlir-commits] [mlir] [mlir][DLTI] Stop using Operation::getAttrs (NFC) (PR #222776)

Mehdi Amini llvmlistbot at llvm.org
Fri Sep 25 02:52:58 PDT 2026


https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/222776

>From d688c28c19796b954d254ef19658d351531ff7fc Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 25 Sep 2026 02:52:32 -0700
Subject: [PATCH] [mlir][DLTI] Introduce an operation interface for scoped DLTI
 queries

So far, DLTI queries have relied on scanning discardable attributes
for values implementing DLTIQueryInterface. This is problematic
when that data is used to determine operation semantics: semantic
information belongs in an operation's inherent attributes/properties,
and should have an explicit contract rather than depend on arbitrary
discardable attribute attachments.

Introduce DLTIQueryOpInterface with symmetric queryDlti and setDlti
methods for individual keys. Provide default implementations backed
by an inherent dlti attribute. Updates preserve unrelated map
entries, a null value removes a key, and attempts to mutate non-map
representations fail.

Resolve the first query key by walking the operation and its ancestors
that implement the new interface, continuing upward when a scope
cannot answer that key. Resolve subsequent keys recursively through the
returned attribute's DLTIQueryInterface without restarting the ancestor
search.  Keep this query contract separate from DataLayoutOpInterface,
which models formal data layout scopes.

Add inherent dlti storage to builtin.module, func.func, gpu.module,
and the data layout test operation. Attach the builtin module interface
externally when the DLTI dialect is loaded. Migrate DLTI, MPI, Shard,
and LLVM target feature attachments to the new storage convention,
and update target feature propagation through setDlti so unrelated
DLTI keys are retained.  Remove the generic discardable-attribute
scanning fallback.

Version builtin module property bytecode to preserve the historical
sym_name/sym_visibility encoding for version 0 and include dlti in
version 1. Update the bytecode fixture and alignment test for the
new encoding.

Add coverage for scoped lookup and precedence, nested queries,
insertion, replacement, removal, immutable values, invalid attributes,
and bytecode round trips. Update the affected conversion and target
feature tests.

Validation: a fresh Release build with assertions enabled passed
ninja check-mlir (3,989 passed, 623 unsupported, 1 expected failure).
git diff --check also passed.

Assisted-by: Codex
---
 mlir/include/mlir/Dialect/DLTI/CMakeLists.txt |   5 +
 mlir/include/mlir/Dialect/DLTI/DLTI.h         |   7 +
 .../mlir/Dialect/DLTI/DLTIOpInterfaces.td     |  51 +++++++
 .../DLTI/TransformOps/DLTITransformOps.td     |  17 +--
 mlir/include/mlir/Dialect/Func/IR/FuncOps.h   |   1 +
 mlir/include/mlir/Dialect/Func/IR/FuncOps.td  |   6 +-
 mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h |   1 +
 mlir/include/mlir/Dialect/GPU/IR/GPUOps.td    |   9 +-
 mlir/include/mlir/IR/BuiltinDialect.h         |   2 +-
 mlir/include/mlir/IR/BuiltinOps.td            |   4 +-
 .../mlir/Interfaces/DataLayoutInterfaces.td   |   6 +
 .../mlir/Target/LLVMIR/Transforms/Passes.td   |   9 +-
 mlir/lib/Dialect/DLTI/CMakeLists.txt          |   1 +
 mlir/lib/Dialect/DLTI/DLTI.cpp                |  94 +++++++++----
 mlir/lib/Dialect/Func/IR/CMakeLists.txt       |   1 +
 mlir/lib/IR/BuiltinDialect.cpp                |  31 +++++
 .../Transforms/TargetToTargetFeatures.cpp     |  13 +-
 mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir |  16 +--
 .../ShardToMPI/convert-shard-to-mpi.mlir      |   8 +-
 .../ShardToMPI/convert-shardshape-to-mpi.mlir |   2 +-
 .../Builtin/Bytecode/builtin_fixed_0.mlirbc   | Bin 5470 -> 5508 bytes
 mlir/test/Dialect/DLTI/invalid.mlir           |   5 +
 .../test/Dialect/DLTI/query-op-interface.mlir |  26 ++++
 mlir/test/Dialect/DLTI/query.mlir             | 125 +++++++++++++-----
 mlir/test/Dialect/MPI/canonicalize.mlir       |   2 +-
 .../Shard/forward-sharding-propagation.mlir   |   2 +-
 ...et-to-data-layout-and-target-features.mlir |  55 +++++---
 .../lib/Dialect/DLTI/TestDataLayoutQuery.cpp  |  62 ++++++++-
 mlir/test/lib/Dialect/Test/TestOps.td         |   6 +-
 mlir/unittests/Bytecode/BytecodeTest.cpp      |   4 +
 30 files changed, 449 insertions(+), 122 deletions(-)
 create mode 100644 mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
 create mode 100644 mlir/test/Dialect/DLTI/query-op-interface.mlir

diff --git a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
index a7695e67e2d8a5..d42f4663b39902 100644
--- a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt
@@ -7,3 +7,8 @@ set(LLVM_TARGET_DEFINITIONS DLTIAttrs.td)
 mlir_tablegen(DLTIAttrs.h.inc -gen-attrdef-decls -attrdefs-dialect=dlti)
 mlir_tablegen(DLTIAttrs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=dlti)
 add_mlir_dialect_tablegen_target(MLIRDLTIAttrsIncGen)
+
+set(LLVM_TARGET_DEFINITIONS DLTIOpInterfaces.td)
+mlir_tablegen(DLTIOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(DLTIOpInterfaces.cpp.inc -gen-op-interface-defs)
+add_mlir_generic_tablegen_target(MLIRDLTIOpInterfacesIncGen)
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h
index 84cf6eabc5fa93..9272a369292645 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -24,6 +24,11 @@ class DataLayoutEntryAttrStorage;
 } // namespace mlir
 namespace mlir {
 namespace dlti {
+namespace detail {
+FailureOr<Attribute> queryDlti(Operation *op, DataLayoutEntryKey key);
+LogicalResult setDlti(Operation *op, DataLayoutEntryKey key, Attribute value);
+} // namespace detail
+
 /// Perform a DLTI-query at `op`, recursively querying each key of `keys` on
 /// query interface-implementing attrs, starting from attr obtained from `op`.
 FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
@@ -37,6 +42,8 @@ FailureOr<Attribute> query(Operation *op, ArrayRef<StringRef> keys,
 } // namespace dlti
 } // namespace mlir
 
+#include "mlir/Dialect/DLTI/DLTIOpInterfaces.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/DLTI/DLTIAttrs.h.inc"
 #include "mlir/Dialect/DLTI/DLTIDialect.h.inc"
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td b/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
new file mode 100644
index 00000000000000..fd7b8d3aa5644e
--- /dev/null
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
@@ -0,0 +1,51 @@
+//===- DLTIOpInterfaces.td - DLTI operation interfaces -----*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_DLTI_DLTIOPINTERFACES
+#define MLIR_DIALECT_DLTI_DLTIOPINTERFACES
+
+include "mlir/Interfaces/DataLayoutInterfaces.td"
+
+def DLTIQueryOpInterface : OpInterface<"DLTIQueryOpInterface"> {
+  let cppNamespace = "::mlir";
+
+  let description = [{
+    Operation interface for scoped DLTI key-value associations.
+
+    Implementing operations expose DLTI data stored in their inherent `dlti`
+    attribute. Queries and updates operate on one key at a time. Setting a null
+    attribute removes the association.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*description=*/"Returns the DLTI value associated with the key.",
+      /*retTy=*/"::mlir::FailureOr<::mlir::Attribute>",
+      /*methodName=*/"queryDlti",
+      /*args=*/(ins "::mlir::DataLayoutEntryKey":$key),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::dlti::detail::queryDlti($_op.getOperation(), key);
+      }]
+    >,
+    InterfaceMethod<
+      /*description=*/"Creates, replaces, or removes a DLTI key-value "
+                      "association.",
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"setDlti",
+      /*args=*/(ins "::mlir::DataLayoutEntryKey":$key,
+                    "::mlir::Attribute":$value),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        return ::mlir::dlti::detail::setDlti($_op.getOperation(), key, value);
+      }]
+    >
+  ];
+}
+
+#endif // MLIR_DIALECT_DLTI_DLTIOPINTERFACES
diff --git a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
index f25bb383912d45..66ed2116f2d11a 100644
--- a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
@@ -24,16 +24,17 @@ def QueryOp : Op<Transform_Dialect, "dlti.query", [
     This op queries data layout and target information associated to payload
     IR by way of the DLTI dialect.
 
-    A lookup is performed for the given `keys` at `target` op - or its closest
-    interface-implementing ancestor - by way of the `DLTIQueryInterface`, which
-    returns an attribute for a key. Each key should be either a (quoted) string
-    or a type. If more than one key is provided, the lookup continues
-    recursively, now on the returned attributes, with the condition that these
-    implement the above interface. For example if the payload IR is
+    A lookup of the first key is performed at `target` or its ancestors by way
+    of `DLTIQueryOpInterface`. Ancestors that do not implement the interface or
+    cannot answer that key are skipped. Each key should be either a (quoted)
+    string or a type. If more than one key is provided, the lookup continues
+    recursively on the returned attributes by way of `DLTIQueryInterface`.
+    Operation ancestor lookup is not restarted for these later keys. For
+    example if the payload IR is
 
     ```
-    module attributes {#dlti.map = #dlti.map<#dlti.dl_entry<"A",
-                                     #dlti.map<#dlti.dl_entry<"B", 42: int>>>} {
+    module attributes {dlti = #dlti.map<#dlti.dl_entry<"A",
+                                 #dlti.map<#dlti.dl_entry<"B", 42 : i32>>>} {
       func.func private @f()
     }
     ```
diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.h b/mlir/include/mlir/Dialect/Func/IR/FuncOps.h
index e12c31855e9181..0bcdc4ff532cca 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.h
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.h
@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_FUNC_IR_OPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index d7a0d5fc4f277a..231b9235091997 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -10,6 +10,7 @@
 #define MLIR_DIALECT_FUNC_IR_FUNCOPS_TD
 
 include "mlir/IR/EnumAttr.td"
+include "mlir/Dialect/DLTI/DLTIOpInterfaces.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
@@ -257,7 +258,7 @@ def ConstantOp : Func_Op<"constant",
 def FuncOp : Func_Op<"func", [
   AffineScope, AutomaticAllocationScope,
   SymbolName, SymbolVisibility, FunctionOpInterface, IsolatedFromAbove,
-  OpAsmOpInterface
+  OpAsmOpInterface, DLTIQueryOpInterface
 ]> {
   let summary = "An operation with a name containing a single `SSACFG` region";
   let description = [{
@@ -303,7 +304,8 @@ def FuncOp : Func_Op<"func", [
                        OptionalAttr<StrAttr>:$sym_visibility,
                        OptionalAttr<DictArrayAttr>:$arg_attrs,
                        OptionalAttr<DictArrayAttr>:$res_attrs,
-                       UnitAttr:$no_inline);
+                       UnitAttr:$no_inline,
+                       OptionalAttr<DLTIQueryAttr>:$dlti);
   let regions = (region AnyRegion:$body);
 
   let builders = [OpBuilder<(ins
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
index 5880c78136566a..9b94e2c6bd8781 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h
@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_GPU_IR_GPUDIALECT_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
+#include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/Dialect/DLTI/Traits.h"
 #include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
 #include "mlir/IR/Builders.h"
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 2aa5219e904040..2d2281c969d6c0 100644
--- a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
+++ b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
@@ -14,6 +14,7 @@
 #define GPU_OPS
 
 include "mlir/Dialect/DLTI/DLTIBase.td"
+include "mlir/Dialect/DLTI/DLTIOpInterfaces.td"
 include "mlir/Dialect/GPU/IR/GPUBase.td"
 include "mlir/Dialect/GPU/IR/CompilationAttrInterfaces.td"
 include "mlir/Dialect/GPU/IR/CompilationAttrs.td"
@@ -1546,7 +1547,8 @@ def GPU_InitializeNamedBarrierOp
 }
 
 def GPU_GPUModuleOp : GPU_Op<"module", [
-      IsolatedFromAbove, DataLayoutOpInterface, HasDefaultDLTIDataLayout,
+      IsolatedFromAbove, DataLayoutOpInterface, DLTIQueryOpInterface,
+      HasDefaultDLTIDataLayout,
       NoRegionArguments, SymbolTable, SymbolName, SymbolVisibility,
       Symbol] # GraphRegionNoTerminator.traits> {
   let summary = "A top level compilation unit containing code to be run on a GPU.";
@@ -1600,12 +1602,13 @@ def GPU_GPUModuleOp : GPU_Op<"module", [
       SymbolNameAttr:$sym_name,
       OptionalAttr<StrAttr>:$sym_visibility,
       OptionalAttr<GPUNonEmptyTargetArrayAttr>:$targets,
-      OptionalAttr<OffloadingTranslationAttr>:$offloadingHandler);
+      OptionalAttr<OffloadingTranslationAttr>:$offloadingHandler,
+      OptionalAttr<DLTIQueryAttr>:$dlti);
   let regions = (region SizedRegion<1>:$bodyRegion);
   let assemblyFormat = [{
     ($sym_visibility^)? $sym_name
     (`<` $offloadingHandler^ `>`)?
-    ($targets^)? attr-dict-with-keyword $bodyRegion
+    ($targets^)? (`dlti` `=` $dlti^)? attr-dict-with-keyword $bodyRegion
   }];
 
   // We need to ensure the block inside the region is properly terminated;
diff --git a/mlir/include/mlir/IR/BuiltinDialect.h b/mlir/include/mlir/IR/BuiltinDialect.h
index d06b63da6b892d..b7c7a271b07fcf 100644
--- a/mlir/include/mlir/IR/BuiltinDialect.h
+++ b/mlir/include/mlir/IR/BuiltinDialect.h
@@ -26,7 +26,7 @@ struct BuiltinDialectVersion : public mlir::DialectVersion {
 
   int64_t getVersion() const { return version; }
 
-  static BuiltinDialectVersion getCurrentVersion() { return {0}; }
+  static BuiltinDialectVersion getCurrentVersion() { return {1}; }
 
   bool operator<(const BuiltinDialectVersion &other) const {
     return version < other.version;
diff --git a/mlir/include/mlir/IR/BuiltinOps.td b/mlir/include/mlir/IR/BuiltinOps.td
index f16b686762b48f..aa53ffbd2ca3dd 100644
--- a/mlir/include/mlir/IR/BuiltinOps.td
+++ b/mlir/include/mlir/IR/BuiltinOps.td
@@ -55,7 +55,8 @@ def ModuleOp : Builtin_Op<"module", [
   }];
 
   let arguments = (ins OptionalAttr<SymbolNameAttr>:$sym_name,
-                       OptionalAttr<StrAttr>:$sym_visibility);
+                       OptionalAttr<StrAttr>:$sym_visibility,
+                       OptionalAttr<DLTIQueryAttr>:$dlti);
   let regions = (region SizedRegion<1>:$bodyRegion);
 
   let assemblyFormat = "($sym_name^)? attr-dict-with-keyword $bodyRegion";
@@ -90,6 +91,7 @@ def ModuleOp : Builtin_Op<"module", [
     }
   }];
   let hasVerifier = 1;
+  let useCustomPropertiesEncoding = 1;
 
   // We need to ensure the block inside the region is properly terminated;
   // the auto-generated builders do not guarantee that.
diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
index 9d35f78b218f8f..a4520dc3df0576 100644
--- a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
+++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td
@@ -43,6 +43,12 @@ def DLTIQueryInterface : AttrInterface<"DLTIQueryInterface"> {
   ];
 }
 
+def ImplementsDLTIQueryInterface : AttrConstraint<
+    CPred<"::llvm::isa<::mlir::DLTIQueryInterface>($_self)">,
+    "implementing the `DLTIQueryInterface` attr interface">;
+
+def DLTIQueryAttr : ConfinedAttr<AnyAttr, [ImplementsDLTIQueryInterface]>;
+
 def DataLayoutEntryInterface : AttrInterface<"DataLayoutEntryInterface"> {
   let cppNamespace = "::mlir";
 
diff --git a/mlir/include/mlir/Target/LLVMIR/Transforms/Passes.td b/mlir/include/mlir/Target/LLVMIR/Transforms/Passes.td
index c858e9f4c34cb0..7002333b99976c 100644
--- a/mlir/include/mlir/Target/LLVMIR/Transforms/Passes.td
+++ b/mlir/include/mlir/Target/LLVMIR/Transforms/Passes.td
@@ -28,12 +28,13 @@ def LLVMTargetToDataLayout : Pass<"llvm-target-to-data-layout"> {
 }
 
 def LLVMTargetToTargetFeatures : Pass<"llvm-target-to-target-features"> {
-  let summary = "Update attached #llvm.target's features per the described target";
+  let summary = "Derive and store the target's DLTI features";
+  let dependentDialects = ["mlir::DLTIDialect"];
   let description = [{
     Obtain the TargetMachine specified by the attached #llvm.target's attributes
-    and obtain from it the full list of features of the selected target. Updates
-    the attached #llvm.target so that its features reflect the full list of
-    features.
+    and obtain from it the full list of features of the selected target. Stores
+    the resulting features under the `features` key of the operation's DLTI
+    map.
   }];
   let options = [
     Option<"initializeLLVMTargets", "initialize-llvm-targets", "bool",
diff --git a/mlir/lib/Dialect/DLTI/CMakeLists.txt b/mlir/lib/Dialect/DLTI/CMakeLists.txt
index 7ae80f9e913bb3..5c366122f25a1b 100644
--- a/mlir/lib/Dialect/DLTI/CMakeLists.txt
+++ b/mlir/lib/Dialect/DLTI/CMakeLists.txt
@@ -6,6 +6,7 @@ add_mlir_dialect_library(MLIRDLTIDialect
   DEPENDS
   MLIRDLTIIncGen
   MLIRDLTIAttrsIncGen
+  MLIRDLTIOpInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRIR
diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp
index da572f123fdc15..8df062aaa075de 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -20,6 +20,8 @@
 
 using namespace mlir;
 
+#include "mlir/Dialect/DLTI/DLTIOpInterfaces.cpp.inc"
+
 #include "mlir/Dialect/DLTI/DLTIDialect.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
@@ -516,21 +518,49 @@ void TargetSystemSpecAttr::print(AsmPrinter &printer) const {
 // DLTIDialect
 //===----------------------------------------------------------------------===//
 
-/// 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 = dyn_cast<DLTIQueryInterface>(attr.getValue())))
-        break;
-  } while (!queryable && (op = op->getParentOp()));
+FailureOr<Attribute> mlir::dlti::detail::queryDlti(Operation *op,
+                                                   DataLayoutEntryKey key) {
+  auto queryable = dyn_cast_or_null<DLTIQueryInterface>(op->getAttr("dlti"));
+  if (!queryable)
+    return failure();
+  return queryable.query(key);
+}
+
+LogicalResult mlir::dlti::detail::setDlti(Operation *op, DataLayoutEntryKey key,
+                                          Attribute value) {
+  auto current = op->getAttrOfType<DLTIQueryInterface>("dlti");
+  if (current && !isa<MapAttr>(current))
+    return failure();
+
+  SmallVector<DataLayoutEntryInterface> entries;
+  if (auto map = dyn_cast_or_null<MapAttr>(current))
+    llvm::append_range(entries, map.getEntries());
+
+  auto makeEntry = [&]() -> DataLayoutEntryInterface {
+    if (auto type = dyn_cast<Type>(key))
+      return DataLayoutEntryAttr::get(type, value);
+    return DataLayoutEntryAttr::get(cast<StringAttr>(key), value);
+  };
+  auto entry = llvm::find_if(entries, [&](DataLayoutEntryInterface entry) {
+    return entry.getKey() == key;
+  });
+  if (entry != entries.end()) {
+    if (value)
+      *entry = makeEntry();
+    else
+      entries.erase(entry);
+  } else if (value) {
+    entries.push_back(makeEntry());
+  } else {
+    return success();
+  }
 
-  return std::pair(queryable, op);
+  if (entries.empty()) {
+    op->removeAttr("dlti");
+    return success();
+  }
+  op->setAttr("dlti", MapAttr::get(op->getContext(), entries));
+  return success();
 }
 
 FailureOr<Attribute>
@@ -546,26 +576,37 @@ dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
     return failure();
   }
 
-  auto [queryable, queryOp] = getClosestQueryable(op);
-  Operation *reportOp = (queryOp ? queryOp : op);
+  Attribute currentAttr;
+  Operation *queryOp = op;
+  for (; queryOp; queryOp = queryOp->getParentOp()) {
+    auto queryable = dyn_cast<DLTIQueryOpInterface>(queryOp);
+    if (!queryable)
+      continue;
+    FailureOr<Attribute> result = queryable.queryDlti(keys.front());
+    if (succeeded(result)) {
+      currentAttr = *result;
+      break;
+    }
+  }
 
-  if (!queryable) {
+  if (!currentAttr) {
     if (emitError) {
       auto diag = op->emitError() << "target op of failed DLTI query";
-      diag.attachNote(reportOp->getLoc())
-          << "no DLTI-queryable attrs on target op or any of its ancestors";
+      diag.attachNote(op->getLoc())
+          << "no DLTI-queryable operation on the target or any of its "
+             "ancestors could answer key "
+          << keyToStr(keys.front());
     }
     return failure();
   }
 
-  Attribute currentAttr = queryable;
-  for (auto &&[idx, key] : llvm::enumerate(keys)) {
+  for (auto &&[idx, key] : llvm::enumerate(keys.drop_front())) {
     if (auto map = dyn_cast<DLTIQueryInterface>(currentAttr)) {
       auto maybeAttr = map.query(key);
       if (failed(maybeAttr)) {
         if (emitError) {
           auto diag = op->emitError() << "target op of failed DLTI query";
-          diag.attachNote(reportOp->getLoc())
+          diag.attachNote(queryOp->getLoc())
               << "key " << keyToStr(key)
               << " has no DLTI-mapping per attr: " << map;
         }
@@ -576,12 +617,12 @@ dlti::query(Operation *op, ArrayRef<DataLayoutEntryKey> keys, bool emitError) {
       if (emitError) {
         std::string commaSeparatedKeys;
         llvm::interleave(
-            keys.take_front(idx), // All prior keys.
+            keys.take_front(idx + 1), // All prior keys.
             [&](auto key) { commaSeparatedKeys += keyToStr(key); },
             [&]() { commaSeparatedKeys += ","; });
 
         auto diag = op->emitError() << "target op of failed DLTI query";
-        diag.attachNote(reportOp->getLoc())
+        diag.attachNote(queryOp->getLoc())
             << "got non-DLTI-queryable attribute upon looking up keys ["
             << commaSeparatedKeys << "] at op";
       }
@@ -607,6 +648,10 @@ FailureOr<Attribute> dlti::query(Operation *op, ArrayRef<StringRef> keys,
 }
 
 namespace {
+struct ModuleDLTIQueryOpInterface
+    : DLTIQueryOpInterface::ExternalModel<ModuleDLTIQueryOpInterface,
+                                          ModuleOp> {};
+
 class TargetDataLayoutInterface : public DataLayoutDialectInterface {
 public:
   using DataLayoutDialectInterface::DataLayoutDialectInterface;
@@ -645,6 +690,7 @@ void DLTIDialect::initialize() {
 #include "mlir/Dialect/DLTI/DLTIAttrs.cpp.inc"
       >();
   addInterfaces<TargetDataLayoutInterface>();
+  ModuleOp::attachInterface<ModuleDLTIQueryOpInterface>(*getContext());
 }
 
 LogicalResult DLTIDialect::verifyOperationAttribute(Operation *op,
diff --git a/mlir/lib/Dialect/Func/IR/CMakeLists.txt b/mlir/lib/Dialect/Func/IR/CMakeLists.txt
index 329301c6fbafd2..987c5954c0a488 100644
--- a/mlir/lib/Dialect/Func/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Func/IR/CMakeLists.txt
@@ -10,6 +10,7 @@ add_mlir_dialect_library(MLIRFuncDialect
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
+  MLIRDLTIDialect
   MLIRFunctionInterfaces
   MLIRInferTypeOpInterface
   MLIRIR
diff --git a/mlir/lib/IR/BuiltinDialect.cpp b/mlir/lib/IR/BuiltinDialect.cpp
index c692a4dda518cd..080b79c36a2e41 100644
--- a/mlir/lib/IR/BuiltinDialect.cpp
+++ b/mlir/lib/IR/BuiltinDialect.cpp
@@ -103,6 +103,7 @@ struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface {
 } // namespace
 
 void BuiltinDialect::initialize() {
+  setVersion(BuiltinDialectVersion::getCurrentVersion());
   registerTypes();
   registerAttributes();
   registerLocationAttributes();
@@ -135,6 +136,36 @@ ModuleOp ModuleOp::create(Location loc, std::optional<StringRef> name) {
   return ModuleOp::create(builder, loc, name);
 }
 
+LogicalResult ModuleOp::readProperties(DialectBytecodeReader &reader,
+                                       OperationState &state) {
+  auto &properties = state.getOrAddProperties<Properties>();
+  if (failed(reader.readOptionalAttribute(properties.sym_name)) ||
+      failed(reader.readOptionalAttribute(properties.sym_visibility)))
+    return failure();
+
+  FailureOr<const DialectVersion *> maybeVersion =
+      reader.getDialectVersion<BuiltinDialect>();
+  if (failed(maybeVersion) ||
+      static_cast<const BuiltinDialectVersion *>(*maybeVersion)->getVersion() <
+          1)
+    return success();
+  return reader.readOptionalAttribute(properties.dlti);
+}
+
+void ModuleOp::writeProperties(DialectBytecodeWriter &writer) {
+  Properties &properties = getProperties();
+  writer.writeOptionalAttribute(properties.sym_name);
+  writer.writeOptionalAttribute(properties.sym_visibility);
+
+  FailureOr<const DialectVersion *> maybeVersion =
+      writer.getDialectVersion<BuiltinDialect>();
+  if (succeeded(maybeVersion) &&
+      static_cast<const BuiltinDialectVersion *>(*maybeVersion)->getVersion() <
+          1)
+    return;
+  writer.writeOptionalAttribute(properties.dlti);
+}
+
 DataLayoutSpecInterface ModuleOp::getDataLayoutSpec() {
   // Take the first and only (if present) attribute that implements the
   // interface. This needs a linear search, but is called only once per data
diff --git a/mlir/lib/Target/LLVMIR/Transforms/TargetToTargetFeatures.cpp b/mlir/lib/Target/LLVMIR/Transforms/TargetToTargetFeatures.cpp
index 0cc5f5aabb6308..0d4e994478be61 100644
--- a/mlir/lib/Target/LLVMIR/Transforms/TargetToTargetFeatures.cpp
+++ b/mlir/lib/Target/LLVMIR/Transforms/TargetToTargetFeatures.cpp
@@ -70,11 +70,12 @@ struct TargetToTargetFeaturesPass
     auto fullTargetFeaturesAttr =
         LLVM::TargetFeaturesAttr::get(&getContext(), plussedFeaturesRefs);
 
-    auto updatedTargetAttr =
-        LLVM::TargetAttr::get(&getContext(), targetAttr.getTriple(),
-                              targetAttr.getChip(), fullTargetFeaturesAttr);
-
-    op->setDiscardableAttr(LLVM::LLVMDialect::getTargetAttrName(),
-                           updatedTargetAttr);
+    auto queryOp = dyn_cast<DLTIQueryOpInterface>(op);
+    if (!queryOp ||
+        failed(queryOp.setDlti(StringAttr::get(&getContext(), "features"),
+                               fullTargetFeaturesAttr))) {
+      op->emitError() << "failed to update DLTI target features";
+      return signalPassFailure();
+    }
   }
 };
diff --git a/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir b/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
index 6445f5afd80002..5beb280ef760b9 100644
--- a/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
+++ b/mlir/test/Conversion/MPIToLLVM/mpitollvm.mlir
@@ -3,7 +3,7 @@
 // RUN: mlir-opt -split-input-file -convert-to-llvm="dynamic=true" %s | FileCheck %s
 
 // COM: Test MPICH ABI
-// CHECK-LABEL: module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
+// CHECK-LABEL: module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH">} {
 // CHECK-DAG: llvm.func @MPI_Finalize() -> i32
 // CHECK-DAG: llvm.func @MPI_Reduce_scatter_block(!llvm.ptr, !llvm.ptr, i32, i32, i32, i32) -> i32
 // CHECK-DAG: llvm.func @MPI_Allreduce(!llvm.ptr, !llvm.ptr, i32, i32, i32, i32) -> i32
@@ -14,7 +14,7 @@
 // CHECK-DAG: llvm.func @MPI_Send(!llvm.ptr, i32, i32, i32, i32, i32) -> i32
 // CHECK-DAG: llvm.func @MPI_Comm_rank(i32, !llvm.ptr) -> i32
 // CHECK-DAG: llvm.func @MPI_Init(!llvm.ptr, !llvm.ptr) -> i32
-module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH">} {
 
   // CHECK-LABEL: llvm.func @test_init_finalize_mpich
   func.func @test_init_finalize_mpich() {
@@ -155,7 +155,7 @@ module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
 // -----
 
 // COM: Test OpenMPI ABI
-// CHECK-LABEL: module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "OpenMPI">} {
+// CHECK-LABEL: module attributes {dlti = #dlti.map<"MPI:Implementation" = "OpenMPI">} {
 // CHECK-DAG: llvm.func @MPI_Finalize() -> i32
 // CHECK-DAG: llvm.func @MPI_Comm_split(!llvm.ptr, i32, i32, !llvm.ptr) -> i32
 // CHECK-DAG: llvm.func @MPI_Reduce_scatter_block(!llvm.ptr, !llvm.ptr, i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> i32
@@ -169,7 +169,7 @@ module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
 // CHECK-DAG: llvm.func @MPI_Comm_rank(!llvm.ptr, !llvm.ptr) -> i32
 // CHECK-DAG: llvm.mlir.global external @ompi_mpi_comm_world() {addr_space = 0 : i32} : !llvm.struct<"ompi_communicator_t", opaque>
 // CHECK-DAG: llvm.func @MPI_Init(!llvm.ptr, !llvm.ptr) -> i32
-module attributes { dlti.map = #dlti.map<"MPI:Implementation" = "OpenMPI"> } {
+module attributes { dlti = #dlti.map<"MPI:Implementation" = "OpenMPI"> } {
 
   // CHECK-LABEL: llvm.func @test_init_finalize_openmpi
   func.func @test_init_finalize_openmpi() {
@@ -315,7 +315,7 @@ module attributes { dlti.map = #dlti.map<"MPI:Implementation" = "OpenMPI"> } {
 
 // -----
 
-module attributes {mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH", "MPI:comm_world_size" = 4, "MPI:comm_world_rank" = 1> } {
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH", "MPI:comm_world_size" = 4, "MPI:comm_world_rank" = 1> } {
   // CHECK-LABEL: llvm.func @test_fold
   func.func @test_fold(%arg0: memref<100xf32>) {
     // CHECK: [[v0:%.*]] = llvm.mlir.constant(1140850688 : i64) : i64
@@ -332,7 +332,7 @@ module attributes {mpi.dlti = #dlti.map<"MPI:Implementation" = "MPICH", "MPI:com
 
 // COM: Test that an index type that already matches the MPI element count width
 // COM: is used as is, both for the offset and for the extents.
-module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH">,
                    dlti.dl_spec = #dlti.dl_spec<index = 32 : i32>} {
   // CHECK-LABEL: llvm.func @test_send_index32
   func.func @test_send_index32(%arg0: memref<100xf32>, %rank: i32) {
@@ -357,7 +357,7 @@ module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
 
 // COM: Test that an index type narrower than the MPI element count is zero
 // COM: extended.
-module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH">,
                    dlti.dl_spec = #dlti.dl_spec<index = 16 : i32>} {
   // CHECK-LABEL: llvm.func @test_send_index16
   func.func @test_send_index16(%arg0: memref<100xf32>, %rank: i32) {
@@ -380,7 +380,7 @@ module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">,
 
 // COM: Test that a rank-zero memref, whose descriptor carries no extents, uses
 // COM: the element count of one directly.
-module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "MPICH">} {
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "MPICH">} {
   // CHECK-LABEL: llvm.func @test_send_rank_zero
   func.func @test_send_rank_zero(%arg0: memref<f32>, %rank: i32) {
     // CHECK: [[v0:%.*]] = llvm.insertvalue {{.*}}[2] : !llvm.struct<(ptr, ptr, i64)>
diff --git a/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi.mlir b/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi.mlir
index 062f05b5c5e131..02612061676d5e 100644
--- a/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi.mlir
+++ b/mlir/test/Conversion/ShardToMPI/convert-shard-to-mpi.mlir
@@ -69,7 +69,7 @@ func.func @neighbors_dim2(%arg0 : tensor<120x120x120xi8>) -> (index, index) {
 
 // -----
 // CHECK: shard.grid @grid0
-module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
+module attributes { dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
   shard.grid @grid0(shape = 3x4x5)
   func.func @process_multi_index() -> (index, index, index) {
     // CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index
@@ -107,7 +107,7 @@ module {
 }
 
 // -----
-module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 7> } {
+module attributes { dlti = #dlti.map<"MPI:comm_world_rank" = 7> } {
   shard.grid @grid0(shape = 3x4x5)
   // CHECK-LABEL: func.func @allreduce_tensor(
   func.func @allreduce_tensor(
@@ -291,7 +291,7 @@ func.func @update_halo_1d_first(
 }
 
 // -----
-module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 1> } {
+module attributes { dlti = #dlti.map<"MPI:comm_world_rank" = 1> } {
   shard.grid @grid0(shape = 4)
   // CHECK-LABEL: func @update_halo_1d_with_zero
   func.func @update_halo_1d_with_zero (
@@ -316,7 +316,7 @@ module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 1> } {
 }
 
 // -----
-module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
+module attributes { dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
   shard.grid @grid0(shape = 3x4x5)
   // CHECK-LABEL: func @update_halo_3d
   func.func @update_halo_3d(
diff --git a/mlir/test/Conversion/ShardToMPI/convert-shardshape-to-mpi.mlir b/mlir/test/Conversion/ShardToMPI/convert-shardshape-to-mpi.mlir
index 9729d2bfb384e2..6337a1f0a178b3 100644
--- a/mlir/test/Conversion/ShardToMPI/convert-shardshape-to-mpi.mlir
+++ b/mlir/test/Conversion/ShardToMPI/convert-shardshape-to-mpi.mlir
@@ -1,6 +1,6 @@
 // RUN: mlir-opt %s --convert-shard-to-mpi -canonicalize | FileCheck %s
 
-module attributes { mpi.dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
+module attributes { dlti = #dlti.map<"MPI:comm_world_rank" = 24> } {
 
   // CHECK: shard.grid @grid0
   shard.grid @grid0(shape = 3x4x5)
diff --git a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc b/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc
index 95be4de8fd1ab202b51658ce69f8037276dd74e8..e1da0c883cc5e0c97d770ebdbb6a444fbe78afda 100644
GIT binary patch
delta 184
zcmcbo)uJur>+?Q{H!ZUwHHCpul$DvCnRz3>Bp2&x5ZJ82b%d+lipiRtk&%&oI}l$7
z;?InXoOVn=A`V29frx(~VjeS)&|?7-`#=OAE0C}Q5n*gV;v0zQ0~*A64n+I|5o#Pj
jA`L{efrw=w;u(k#;{+0AAmSLOAV?6z6ypLCen0{MjB_5M

delta 144
zcmZqCzNf|S>+?Q{H!ZUwHHCqZpOtZ=kR;b;1Fj=nMUSmn85vo(GqPT1Wc|#@ZpXyP
z9tXr_K>QDg=P@&~>#;Dh?*n2!R-il(hp{oTe*@w^c1HGdK>QDg)i at a0(}1`Qh?fEJ
ZGawe at 1eycH$AI`55Q}j!vior{0s!B*8p{9x

diff --git a/mlir/test/Dialect/DLTI/invalid.mlir b/mlir/test/Dialect/DLTI/invalid.mlir
index 2436e4f7484f5a..c0d93ce686354a 100644
--- a/mlir/test/Dialect/DLTI/invalid.mlir
+++ b/mlir/test/Dialect/DLTI/invalid.mlir
@@ -10,6 +10,11 @@
 
 // -----
 
+// expected-error at below {{attribute 'dlti' failed to satisfy constraint}}
+module attributes {dlti = 42 : i32} {}
+
+// -----
+
 // expected-error at below {{'dlti.dl_spec' is expected to be a #dlti.dl_spec attribute}}
 "test.unknown_op"() { dlti.dl_spec = 42 } : () -> ()
 
diff --git a/mlir/test/Dialect/DLTI/query-op-interface.mlir b/mlir/test/Dialect/DLTI/query-op-interface.mlir
new file mode 100644
index 00000000000000..6bb9dd3d285364
--- /dev/null
+++ b/mlir/test/Dialect/DLTI/query-op-interface.mlir
@@ -0,0 +1,26 @@
+// RUN: mlir-opt %s -test-dlti-query-op-interface --mlir-print-op-generic | FileCheck %s
+// RUN: mlir-opt %s -emit-bytecode | mlir-opt --mlir-print-op-generic | FileCheck --check-prefix=BYTECODE %s
+
+// BYTECODE: "builtin.module"() <{dlti = #dlti.map<"module" = 1 : i32>}> ({
+module attributes {dlti = #dlti.map<"module" = 1 : i32>} {
+  // CHECK: "test.op_with_data_layout"() <{dlti = #dlti.map<"replaced" = 2 : i32, i32 = 32 : i32, "inserted" = 1 : i32>}> ({
+  // CHECK-NOT: test.dlti_action
+  "test.op_with_data_layout"() ({
+  }) {dlti = #dlti.map<"replaced" = 0 : i32,
+                       "removed" = 0 : i32,
+                       i32 = 0 : i32>,
+      test.dlti_action = "update"} : () -> ()
+
+  // CHECK: "test.op_with_data_layout"() ({
+  // CHECK-NOT: dlti =
+  // CHECK-NOT: test.dlti_action
+  "test.op_with_data_layout"() ({
+  }) {dlti = #dlti.map<"only" = 0 : i32>,
+      test.dlti_action = "clear"} : () -> ()
+
+  // CHECK: "test.op_with_data_layout"() <{dlti = #dlti.dl_spec<"immutable" = 0 : i32>}> ({
+  // CHECK: }) {test.set_dlti_failed}
+  "test.op_with_data_layout"() ({
+  }) {dlti = #dlti.dl_spec<"immutable" = 0 : i32>,
+      test.dlti_action = "fail"} : () -> ()
+}
diff --git a/mlir/test/Dialect/DLTI/query.mlir b/mlir/test/Dialect/DLTI/query.mlir
index 3825cee6f16164..08aa97e7cefaed 100644
--- a/mlir/test/Dialect/DLTI/query.mlir
+++ b/mlir/test/Dialect/DLTI/query.mlir
@@ -1,7 +1,7 @@
 // RUN: mlir-opt -transform-interpreter -canonicalize -split-input-file -verify-diagnostics %s | FileCheck %s
 
 // expected-remark @below {{attr associated to "test.id" = 42 : i32}}
-module attributes { test.dlti = #dlti.map<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.map<"test.id" = 42 : i32> } {
   func.func private @f()
 }
 
@@ -18,7 +18,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-remark @below {{i32 present in set : unit}}
-module attributes { test.dlti = #dlti.map<i32 = unit> } {
+module attributes { dlti = #dlti.map<i32 = unit> } {
   func.func private @f()
 }
 
@@ -34,8 +34,65 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+module {
+  // expected-remark @below {{associated attr 7 : i32}}
+  gpu.module @gpu dlti = #dlti.map<"test.id" = 7 : i32> {
+  }
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %gpu_module = transform.structured.match ops{["gpu.module"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query ["test.id"] at %gpu_module : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %gpu_module : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// A nearer interface-bearing operation that cannot answer the first key does
+// not hide an association on an ancestor.
+module attributes { dlti = #dlti.map<"test.id" = 42 : i32> } {
+  // expected-remark @below {{associated attr 42 : i32}}
+  func.func private @f() attributes {
+    dlti = #dlti.map<"other" = 24 : i32>
+  }
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    %param = transform.dlti.query ["test.id"] at %func : (!transform.any_op) -> !transform.any_param
+    transform.debug.emit_param_as_remark %param, "associated attr" at %func : !transform.any_param, !transform.any_op
+    transform.yield
+  }
+}
+
+// -----
+
+// Once the first key succeeds, later keys query only the returned attribute.
+module attributes { dlti = #dlti.map<"root" = #dlti.map<"leaf" = 42 : i32>> } {
+  // expected-error @below {{target op of failed DLTI query}}
+  // expected-note @below {{key "leaf" has no DLTI-mapping per attr: #dlti.map<"other" = 24 : i32>}}
+  func.func private @f() attributes {
+    dlti = #dlti.map<"root" = #dlti.map<"other" = 24 : i32>>
+  }
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg: !transform.any_op) {
+    %func = transform.structured.match ops{["func.func"]} in %arg : (!transform.any_op) -> !transform.any_op
+    // expected-error @below {{'transform.dlti.query' op failed to apply}}
+    %param = transform.dlti.query ["root", "leaf"] at %func : (!transform.any_op) -> !transform.any_param
+    transform.yield
+  }
+}
+
+// -----
+
 // expected-remark @below {{attr associated to i32's "width_in_bits" = 32 : i32}}
-module attributes { test.dlti = #dlti.map<i32 = #dlti.map<"width_in_bits" = 32 : i32>> } {
+module attributes { dlti = #dlti.map<i32 = #dlti.map<"width_in_bits" = 32 : i32>> } {
   func.func private @f()
 }
 
@@ -53,7 +110,7 @@ module attributes {transform.with_named_sequence} {
 
 // expected-remark @below {{width in bits of i32 = 32 : i64}}
 // expected-remark @below {{width in bits of f64 = 64 : i64}}
-module attributes { test.dlti = #dlti.map<"width_in_bits" = #dlti.map<i32 = 32, f64 = 64>> } {
+module attributes { dlti = #dlti.map<"width_in_bits" = #dlti.map<i32 = 32, f64 = 64>> } {
   func.func private @f()
 }
 
@@ -72,7 +129,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-remark @below {{attr associated to "test.id" = 42 : i32}}
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
   func.func private @f()
 }
 
@@ -88,9 +145,9 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
   // expected-remark @below {{attr associated to "test.id" = 24 : i32}}
-  func.func private @f() attributes { test.dlti = #dlti.dl_spec<"test.id" = 24 : i32>}
+  func.func private @f() attributes { dlti = #dlti.dl_spec<"test.id" = 24 : i32>}
 }
 
 module attributes {transform.with_named_sequence} {
@@ -105,8 +162,8 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-remark @below {{attr associated to "test.id" = 42 : i32}}
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
-  func.func private @f() attributes { test.dlti = #dlti.dl_spec<"test.id" = 24 : i32> }
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+  func.func private @f() attributes { dlti = #dlti.dl_spec<"test.id" = 24 : i32> }
 }
 
 module attributes {transform.with_named_sequence} {
@@ -121,7 +178,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
   func.func @matmul_tensors(
     %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
       -> tensor<?x?xf32> {
@@ -144,10 +201,10 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
   func.func @matmul_tensors(
     %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
-      -> tensor<?x?xf32> attributes {test.dlti = #dlti.dl_spec<"test.id" = 24 : i32> } {
+      -> tensor<?x?xf32> attributes {dlti = #dlti.dl_spec<"test.id" = 24 : i32> } {
     // expected-remark @below {{associated attr 24 : i32}}
     %0 = linalg.matmul  ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
                        outs(%arg2: tensor<?x?xf32>)
@@ -168,7 +225,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-remark @below {{associated attr 42 : i32}}
-module attributes { test.dlti =
+module attributes { dlti =
   #dlti.target_system_spec<"CPU" =
     #dlti.target_device_spec<"test.id" = 42 : i32>> } {
   func.func private @f()
@@ -186,7 +243,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
+module attributes { dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
                                                          "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
   // expected-remark @below {{associated attr 43 : i32}}
   func.func private @f()
@@ -203,10 +260,10 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
+module attributes { dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
                                                          "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
   // expected-remark @below {{associated attr 24 : i32}}
-  func.func private @f() attributes { test.dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 24 : i32>> }
+  func.func private @f() attributes { dlti = #dlti.target_system_spec<"CPU" = #dlti.target_device_spec<"test.id" = 24 : i32>> }
 }
 
 module attributes {transform.with_named_sequence} {
@@ -220,7 +277,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.target_system_spec<
+module attributes { dlti = #dlti.target_system_spec<
   "CPU" = #dlti.target_device_spec<
     "cache::L1::size_in_bytes" = 65536 : i32,
     "cache::L1d::size_in_bytes" = 32768 : i32>> } {
@@ -244,7 +301,7 @@ module attributes {transform.with_named_sequence} {
 
 #l1_size = #dlti.map<"size_in_bytes" = 65536 : i32>
 #l1d_size = #dlti.map<"size_in_bytes" = 32768 : i32>
-module attributes { test.dlti =
+module attributes { dlti =
   #dlti.target_system_spec<"CPU" =
     #dlti.target_device_spec<"cache" =
       #dlti.map<"L1" = #l1_size,
@@ -267,7 +324,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.target_system_spec<
+module attributes { dlti = #dlti.target_system_spec<
   "CPU" = #dlti.target_device_spec<"inner_most_tile_size" = 42 : i32>> } {
   // CHECK-LABEL: func @matmul_tensors
   func.func @matmul_tensors(
@@ -298,11 +355,11 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// expected-note @below {{key "NPU" has no DLTI-mapping per attr: #dlti.target_system_spec}}
-module attributes { test.dlti = #dlti.target_system_spec<
+module attributes { dlti = #dlti.target_system_spec<
     "CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
     "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
   // expected-error @below {{target op of failed DLTI query}}
+  // expected-note @below {{no DLTI-queryable operation on the target or any of its ancestors could answer key "NPU"}}
   func.func private @f()
 }
 
@@ -318,7 +375,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-note @below {{key "unspecified" has no DLTI-mapping per attr: #dlti.target_device_spec}}
-module attributes { test.dlti = #dlti.target_system_spec<
+module attributes { dlti = #dlti.target_system_spec<
     "CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
     "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
   // expected-error @below {{target op of failed DLTI query}}
@@ -336,11 +393,11 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// expected-note @below {{key "test.id" has no DLTI-mapping per attr: #dlti.target_system_spec}}
-module attributes { test.dlti = #dlti.target_system_spec<
+module attributes { dlti = #dlti.target_system_spec<
   "CPU" = #dlti.target_device_spec<"test.id" = 42 : i32>,
   "GPU" = #dlti.target_device_spec<"test.id" = 43 : i32>> } {
   // expected-error @below {{target op of failed DLTI query}}
+  // expected-note @below {{no DLTI-queryable operation on the target or any of its ancestors could answer key "test.id"}}
   func.func private @f()
 }
 
@@ -355,9 +412,9 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-// expected-note @below {{key "CPU" has no DLTI-mapping per attr: #dlti.dl_spec}}
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32> } {
   // expected-error @below {{target op of failed DLTI query}}
+  // expected-note @below {{no DLTI-queryable operation on the target or any of its ancestors could answer key "CPU"}}
   func.func private @f()
 }
 
@@ -373,7 +430,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-note @below {{got non-DLTI-queryable attribute upon looking up keys ["CPU"]}}
-module attributes { test.dlti = #dlti.dl_spec<"CPU" = 42 : i32> } {
+module attributes { dlti = #dlti.dl_spec<"CPU" = 42 : i32> } {
   // expected-error @below {{target op of failed DLTI query}}
   func.func private @f()
 }
@@ -390,7 +447,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-note @below {{got non-DLTI-queryable attribute upon looking up keys [i32]}}
-module attributes { test.dlti = #dlti.dl_spec<i32 = 32 : i32> } {
+module attributes { dlti = #dlti.dl_spec<i32 = 32 : i32> } {
   // expected-error @below {{target op of failed DLTI query}}
   func.func private @f()
 }
@@ -408,7 +465,7 @@ module attributes {transform.with_named_sequence} {
 
 module {
   // expected-error @below {{target op of failed DLTI query}}
-  // expected-note @below {{no DLTI-queryable attrs on target op or any of its ancestors}}
+  // expected-note @below {{no DLTI-queryable operation on the target or any of its ancestors could answer key "CPU"}}
   func.func private @f()
 }
 
@@ -424,7 +481,7 @@ module attributes {transform.with_named_sequence} {
 // -----
 
 // expected-note @below {{key i64 has no DLTI-mapping per attr: #dlti.map<i32 = 32 : i64>}}
-module attributes { test.dlti = #dlti.map<"width_in_bits" = #dlti.map<i32 = 32>>} {
+module attributes { dlti = #dlti.map<"width_in_bits" = #dlti.map<i32 = 32>>} {
   // expected-error @below {{target op of failed DLTI query}}
   func.func private @f()
 }
@@ -440,7 +497,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.dl_spec<"test.id" = 42 : i32>} {
+module attributes { dlti = #dlti.dl_spec<"test.id" = 42 : i32>} {
   func.func private @f()
 }
 
@@ -455,7 +512,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.map<#dlti.dl_entry<"test.id", 42 : i32>>} {
+module attributes { dlti = #dlti.map<#dlti.dl_entry<"test.id", 42 : i32>>} {
   // expected-error @below {{target op of failed DLTI query}}
   // expected-note @below {{no keys provided to attempt query with}}
   func.func private @f()
@@ -472,7 +529,7 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
-module attributes { test.dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
+module attributes { dlti = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>} {
   func.func private @f()
 }
 
@@ -483,4 +540,4 @@ module attributes {transform.with_named_sequence} {
     %param = transform.dlti.query ["test.id"] at %funcs : (!transform.any_op) -> !transform.param<i64>
     transform.yield
   }
-}
\ No newline at end of file
+}
diff --git a/mlir/test/Dialect/MPI/canonicalize.mlir b/mlir/test/Dialect/MPI/canonicalize.mlir
index 3523d46e212196..76e9398a2c6083 100644
--- a/mlir/test/Dialect/MPI/canonicalize.mlir
+++ b/mlir/test/Dialect/MPI/canonicalize.mlir
@@ -1,6 +1,6 @@
 // RUN: mlir-opt %s -canonicalize | mlir-opt | FileCheck %s
 
-module attributes {mpi.dlti = #dlti.map<"MPI:comm_world_size" = 12, "MPI:comm_world_rank" = 5> } {
+module attributes {dlti = #dlti.map<"MPI:comm_world_size" = 12, "MPI:comm_world_rank" = 5> } {
   // CHECK-LABEL: func.func @mpi_test
   func.func @mpi_test(%ref : memref<100xf32>) -> (i32, i32) {
     %comm = mpi.comm_world : !mpi.comm
diff --git a/mlir/test/Dialect/Shard/forward-sharding-propagation.mlir b/mlir/test/Dialect/Shard/forward-sharding-propagation.mlir
index a5faafb438da0e..b6f6732fd7d9d8 100644
--- a/mlir/test/Dialect/Shard/forward-sharding-propagation.mlir
+++ b/mlir/test/Dialect/Shard/forward-sharding-propagation.mlir
@@ -1,7 +1,7 @@
 // RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward}))" %s | FileCheck %s
 
 #map = affine_map<(d0, d1) -> (d0, d1)>
-module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "mpich", "MPI:comm_world_rank" = 0 : i32>} {
+module attributes {dlti = #dlti.map<"MPI:Implementation" = "mpich", "MPI:comm_world_rank" = 0 : i32>} {
   shard.grid "private" @grid(shape = 1)
   func.func @test_forward() -> (tensor<6x6xi32>, tensor<6x6xi32>, tensor<i32>) attributes {llvm.emit_c_interface} {
     %c1_i32 = arith.constant 1 : i32
diff --git a/mlir/test/Target/LLVMIR/target-to-data-layout-and-target-features.mlir b/mlir/test/Target/LLVMIR/target-to-data-layout-and-target-features.mlir
index b6b29768605416..5b042a80269c36 100644
--- a/mlir/test/Target/LLVMIR/target-to-data-layout-and-target-features.mlir
+++ b/mlir/test/Target/LLVMIR/target-to-data-layout-and-target-features.mlir
@@ -13,14 +13,15 @@
 
 // TARGET-FEATURES: module attributes
 // TARGET-FEATURES-NOT:  dlti.dl_spec
-// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
-// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
-// TARGET-FEATURES-SAME:   chip = ""
-// TARGET-FEATURES-SAME:   features = <[
+// TARGET-FEATURES-SAME: dlti = #dlti.map<"features" = #llvm.target_features<[
 // TARGET-FEATURES-SAME:     +64bit
 // TARGET-FEATURES-NOT:      +avx
 // TARGET-FEATURES-SAME:     +sse
 // TARGET-FEATURES-NOT:      +mmx
+// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
+// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
+// TARGET-FEATURES-SAME:   chip = ""
+// TARGET-FEATURES-NOT:    features =
 
 module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                                chip = ""> } {
@@ -38,14 +39,15 @@ module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
 
 // TARGET-FEATURES: module attributes
 // TARGET-FEATURES-NOT:  dlti.dl_spec
-// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
-// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
-// TARGET-FEATURES-SAME:   chip = ""
-// TARGET-FEATURES-SAME:   features = <[
+// TARGET-FEATURES-SAME: dlti = #dlti.map<"features" = #llvm.target_features<[
 // TARGET-FEATURES-SAME:     +64bit
 // TARGET-FEATURES-NOT:      +avx
 // TARGET-FEATURES-SAME:     +mmx
 // TARGET-FEATURES-SAME:     +sse
+// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
+// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
+// TARGET-FEATURES-SAME:   chip = ""
+// TARGET-FEATURES-SAME:   features = <["+mmx", "+sse"]>
 
 module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                                chip = "",
@@ -64,16 +66,17 @@ module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
 
 // TARGET-FEATURES: module attributes
 // TARGET-FEATURES-NOT:  dlti.dl_spec
-// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
-// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
-// TARGET-FEATURES-SAME:   chip = "skylake"
-// TARGET-FEATURES-SAME:   features = <[
+// TARGET-FEATURES-SAME: dlti = #dlti.map<"features" = #llvm.target_features<[
 // TARGET-FEATURES-SAME:     +64bit
 // TARGET-FEATURES-SAME:     +avx
 // TARGET-FEATURES-SAME:     +avx2
 // TARGET-FEATURES-NOT:      +avx512f
 // TARGET-FEATURES-SAME:     +mmx
 // TARGET-FEATURES-SAME:     +sse
+// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
+// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
+// TARGET-FEATURES-SAME:   chip = "skylake"
+// TARGET-FEATURES-NOT:    features =
 
 module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                                chip = "skylake"> } {
@@ -91,15 +94,16 @@ module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
 
 // TARGET-FEATURES: module attributes
 // TARGET-FEATURES-NOT:  dlti.dl_spec
-// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
-// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
-// TARGET-FEATURES-SAME:   chip = "skylake"
-// TARGET-FEATURES-SAME:   features = <[
+// TARGET-FEATURES-SAME: dlti = #dlti.map<"features" = #llvm.target_features<[
 // TARGET-FEATURES-SAME:     +64bit
 // TARGET-FEATURES-NOT:      +avx
 // TARGET-FEATURES-NOT:      +avx2
 // TARGET-FEATURES-SAME:     +mmx
 // TARGET-FEATURES-NOT:      +sse
+// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
+// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
+// TARGET-FEATURES-SAME:   chip = "skylake"
+// TARGET-FEATURES-SAME:   features = <["-sse", "-avx"]>
 
 module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                                chip = "skylake",
@@ -118,19 +122,26 @@ module attributes { llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
 // DATA-LAYOUT-SAME:   features = <["-mmx", "+avx512f"]>
 
 // TARGET-FEATURES: module attributes
-// TARGET-FEATURES-SAME: #dlti.dl_spec<index = 32 : i64>
-// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
-// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
-// TARGET-FEATURES-SAME:   chip = "skylake"
-// TARGET-FEATURES-SAME:   features = <[
+// TARGET-FEATURES-SAME: dlti = #dlti.map<
+// TARGET-FEATURES-SAME:   "MPI:comm_world_size" = 4 : i64,
+// TARGET-FEATURES-SAME:   "features" = #llvm.target_features<[
 // TARGET-FEATURES-SAME:     +64bit
 // TARGET-FEATURES-SAME:     +avx
 // TARGET-FEATURES-SAME:     +avx2
 // TARGET-FEATURES-SAME:     +avx512f
 // TARGET-FEATURES-NOT:      +mmx
 // TARGET-FEATURES-SAME:     +sse
+// TARGET-FEATURES-SAME:   "triple" = "preserved"
+// TARGET-FEATURES-SAME: #dlti.dl_spec<index = 32 : i64>
+// TARGET-FEATURES-SAME: llvm.target = #llvm.target<
+// TARGET-FEATURES-SAME:   triple = "x86_64-unknown-linux"
+// TARGET-FEATURES-SAME:   chip = "skylake"
+// TARGET-FEATURES-SAME:   features = <["-mmx", "+avx512f"]>
 
-module attributes { dlti.dl_spec = #dlti.dl_spec<index = 32>,
+module attributes { dlti = #dlti.map<"MPI:comm_world_size" = 4,
+                                     "features" = #llvm.target_features<["+old"]>,
+                                     "triple" = "preserved">,
+                    dlti.dl_spec = #dlti.dl_spec<index = 32>,
                     llvm.target = #llvm.target<triple = "x86_64-unknown-linux",
                                                chip = "skylake",
                                                features = <["-mmx", "+avx512f"]>> } {
diff --git a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
index 58c77cb94b47f8..f941d47d5da8da 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -105,14 +105,74 @@ struct TestDataLayoutQuery
                                     : legalIntWidths)
 
           });
+    });
+  }
+};
+
+struct TestDLTIQueryOpInterface
+    : public PassWrapper<TestDLTIQueryOpInterface, OperationPass<>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestDLTIQueryOpInterface)
+
+  StringRef getArgument() const final { return "test-dlti-query-op-interface"; }
+  StringRef getDescription() const final {
+    return "Test the DLTI query operation interface";
+  }
+
+  void runOnOperation() override {
+    WalkResult result = getOperation()->walk([&](Operation *op) {
+      auto action =
+          op->getDiscardableAttrOfType<StringAttr>("test.dlti_action");
+      if (!action)
+        return WalkResult::advance();
+
+      auto queryOp = dyn_cast<DLTIQueryOpInterface>(op);
+      if (!queryOp) {
+        op->emitError("expected DLTIQueryOpInterface");
+        return WalkResult::interrupt();
+      }
 
+      Builder builder(op->getContext());
+      auto key = [&](StringRef value) { return builder.getStringAttr(value); };
+      if (action.getValue() == "update") {
+        if (failed(queryOp.setDlti(key("inserted"),
+                                   builder.getI32IntegerAttr(1))) ||
+            failed(queryOp.setDlti(key("replaced"),
+                                   builder.getI32IntegerAttr(2))) ||
+            failed(queryOp.setDlti(builder.getI32Type(),
+                                   builder.getI32IntegerAttr(32))) ||
+            failed(queryOp.setDlti(key("removed"), Attribute())) ||
+            failed(queryOp.setDlti(key("absent"), Attribute()))) {
+          op->emitError("failed to update DLTI map");
+          return WalkResult::interrupt();
+        }
+      } else if (action.getValue() == "clear") {
+        if (failed(queryOp.setDlti(key("only"), Attribute())) ||
+            failed(queryOp.setDlti(key("absent"), Attribute()))) {
+          op->emitError("failed to clear DLTI map");
+          return WalkResult::interrupt();
+        }
+      } else if (action.getValue() == "fail") {
+        if (succeeded(
+                queryOp.setDlti(key("new"), builder.getI32IntegerAttr(1)))) {
+          op->emitError("unexpectedly updated immutable DLTI representation");
+          return WalkResult::interrupt();
+        }
+        op->setDiscardableAttr("test.set_dlti_failed", builder.getUnitAttr());
+      }
+      op->removeDiscardableAttr("test.dlti_action");
+      return WalkResult::advance();
     });
+    if (result.wasInterrupted())
+      signalPassFailure();
   }
 };
 } // namespace
 
 namespace mlir {
 namespace test {
-void registerTestDataLayoutQuery() { PassRegistration<TestDataLayoutQuery>(); }
+void registerTestDataLayoutQuery() {
+  PassRegistration<TestDataLayoutQuery>();
+  PassRegistration<TestDLTIQueryOpInterface>();
+}
 } // namespace test
 } // namespace mlir
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 492cd448333801..defff8baf9f17b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -24,6 +24,7 @@ include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"
+include "mlir/Dialect/DLTI/DLTIOpInterfaces.td"
 include "mlir/Interfaces/DestinationStyleOpInterface.td"
 include "mlir/Interfaces/InferIntRangeInterface.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
@@ -3119,10 +3120,13 @@ def MakeTupleOp: TEST_Op<"make_tuple"> {
 //===----------------------------------------------------------------------===//
 
 def OpWithDataLayoutOp : TEST_Op<"op_with_data_layout",
-                                 [IsolatedFromAbove, HasDefaultDLTIDataLayout, DataLayoutOpInterface]> {
+                                 [IsolatedFromAbove, HasDefaultDLTIDataLayout,
+                                  DataLayoutOpInterface,
+                                  DLTIQueryOpInterface]> {
   let summary =
       "An op that uses DataLayout implementation from the Target dialect";
   let regions = (region VariadicRegion<AnyRegion>:$regions);
+  let arguments = (ins OptionalAttr<DLTIQueryAttr>:$dlti);
 }
 
 def DataLayoutQueryOp : TEST_Op<"data_layout_query"> {
diff --git a/mlir/unittests/Bytecode/BytecodeTest.cpp b/mlir/unittests/Bytecode/BytecodeTest.cpp
index 7e11aa14fda441..b6eb84fb4e60b3 100644
--- a/mlir/unittests/Bytecode/BytecodeTest.cpp
+++ b/mlir/unittests/Bytecode/BytecodeTest.cpp
@@ -129,6 +129,10 @@ TEST(Bytecode, AlignmentFailure) {
   OwningOpRef<Operation *> module =
       parseSourceString<Operation *>(irWithResources, parseConfig);
   ASSERT_TRUE(module);
+  // Ensure that the resource section needs an explicit alignment marker. The
+  // test below exercises the reader's handling of that marker.
+  module.get()->setDiscardableAttr("bytecode.padding",
+                                   builder.getStringAttr("padding"));
 
   // Write the module to bytecode.
   std::string serializedBytecode;



More information about the Mlir-commits mailing list