[Mlir-commits] [mlir] [mlir][DLTI] Introduce an operation interface for scoped DLTI queries (PR #222776)

[Mehdi Amini]

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

>From 0660e74ac49640d6fad01d8624060f7bb1316771 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. Reject null and empty string keys before mutation.

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. Version func.func and gpu.module properties as well,
preserving their unversioned field order when reading older bytecode.
Add fixtures emitted by the parent revision for both operations.

Cover scoped lookup and precedence, nested queries, insertion,
replacement, removal, immutable values, invalid keys and attributes,
and bytecode round trips. Update the affected conversion and target
feature tests. Document the new inherent property and query behavior.

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

Assisted-by: Codex
---
 mlir/docs/DataLayout.md                       |  11 ++
 mlir/include/mlir/Dialect/DLTI/CMakeLists.txt |   5 +
 mlir/include/mlir/Dialect/DLTI/DLTI.h         |  16 ++-
 mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td   |  18 +--
 .../mlir/Dialect/DLTI/DLTIOpInterfaces.td     |  58 ++++++++
 .../DLTI/TransformOps/DLTITransformOps.td     |  20 +--
 mlir/include/mlir/Dialect/Func/IR/FuncOps.h   |   1 +
 mlir/include/mlir/Dialect/Func/IR/FuncOps.td  |  12 +-
 mlir/include/mlir/Dialect/GPU/IR/GPUDialect.h |   1 +
 mlir/include/mlir/Dialect/GPU/IR/GPUOps.td    |  15 ++-
 mlir/include/mlir/IR/BuiltinDialect.h         |   2 +-
 mlir/include/mlir/IR/BuiltinOps.td            |   9 +-
 .../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                | 107 +++++++++++----
 mlir/lib/Dialect/Func/IR/CMakeLists.txt       |   1 +
 mlir/lib/Dialect/Func/IR/FuncOps.cpp          |  53 ++++++++
 mlir/lib/Dialect/GPU/IR/GPUDialect.cpp        |  49 +++++++
 mlir/lib/IR/BuiltinDialect.cpp                |  31 +++++
 .../Transforms/TargetToTargetFeatures.cpp     |  13 +-
 .../dlti-query-op-interface-compat.mlir       |  27 ++++
 ...lti-query-op-interface-pre-dlti-gpu.mlirbc | Bin 0 -> 199 bytes
 .../dlti-query-op-interface-pre-dlti.mlirbc   | Bin 0 -> 290 bytes
 .../dlti-query-op-interface-roundtrip.mlir    |  12 ++
 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 |  32 +++++
 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  |  70 +++++++++-
 mlir/test/lib/Dialect/Test/TestOps.td         |   6 +-
 mlir/unittests/Bytecode/BytecodeTest.cpp      |   4 +
 38 files changed, 665 insertions(+), 139 deletions(-)
 create mode 100644 mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
 create mode 100644 mlir/test/Bytecode/dlti-query-op-interface-compat.mlir
 create mode 100644 mlir/test/Bytecode/dlti-query-op-interface-pre-dlti-gpu.mlirbc
 create mode 100644 mlir/test/Bytecode/dlti-query-op-interface-pre-dlti.mlirbc
 create mode 100644 mlir/test/Bytecode/dlti-query-op-interface-roundtrip.mlir
 create mode 100644 mlir/test/Dialect/DLTI/query-op-interface.mlir

diff --git a/mlir/docs/DataLayout.md b/mlir/docs/DataLayout.md
index 21430c44ca95fb..13608201b68677 100644
--- a/mlir/docs/DataLayout.md
+++ b/mlir/docs/DataLayout.md
@@ -294,3 +294,14 @@ The [DLTI](../Dialects/DLTIDialect/) dialect provides the attributes implementin
 attribute that can be used to attach the specification to a given operation. The
 verifier of this attribute triggers those of the specification and checks the
 compatibility of nested specifications.
+
+DLTI key-value queries use a separate interface from data layout scope
+calculation. `DLTIQueryOpInterface` does not require a single backing
+attribute: an operation can answer the key `"architecture"` using an inherent
+`architecture = "sm_90"` attribute, without storing a `dlti` attribute or a
+`#dlti.map`. Its default implementation answers keys from an operation's
+inherent `dlti` attribute.
+`func.func` and `gpu.module` implement this interface; `builtin.module`
+acquires it when the DLTI dialect is loaded. A query searches the target
+operation and then its ancestors until one answers the first key. Further keys
+are queried on the returned attribute through `DLTIQueryInterface`.
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..41e42cd940ce0c 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTI.h
+++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h
@@ -24,19 +24,25 @@ class DataLayoutEntryAttrStorage;
 } // namespace mlir
 namespace mlir {
 namespace dlti {
-/// Perform a DLTI-query at `op`, recursively querying each key of `keys` on
-/// query interface-implementing attrs, starting from attr obtained from `op`.
+namespace detail {
+FailureOr<Attribute> queryDlti(Operation *op, DataLayoutEntryKey key);
+LogicalResult setDlti(Operation *op, DataLayoutEntryKey key, Attribute value);
+} // namespace detail
+
+/// Query the first key at `op` or the nearest ancestor that can answer it.
+/// Query later keys recursively on the returned attributes, without searching
+/// operation ancestors again.
 FailureOr<Attribute> query(Operation *op, ArrayRef<DataLayoutEntryKey> keys,
                            bool emitError = false);
 
-/// 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`.
+/// As above, using each string in `keys` as a DLTI entry key.
 FailureOr<Attribute> query(Operation *op, ArrayRef<StringRef> keys,
                            bool emitError = false);
 } // 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/DLTIAttrs.td b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
index 56bd704d306560..8e16363adf8c33 100644
--- a/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIAttrs.td
@@ -115,14 +115,14 @@ def DLTI_MapAttr : DLTIAttr<"Map", [DLTIQueryInterface]> {
     A Data Layout and Target Information map is a list of entries effectively
     encoding a dictionary, mapping DLTI-related keys to DLTI-related values.
 
-    This attribute's main purpose is to facilate querying IR for arbitrary
-    key-value associations that encode DLTI. Facility functions exist to perform
-    recursive lookups on nested DLTI-map/query interface-implementing
-    attributes.
+    This attribute facilitates queries for arbitrary DLTI key-value
+    associations. An operation can store it in its inherent `dlti` attribute
+    and expose its entries through `DLTIQueryOpInterface`. Query results that
+    implement `DLTIQueryInterface` can be queried recursively.
 
     Consider the following flat encoding of a single-key dictionary
     ```
-    #dlti.map<"CPU::cache::L1::size_in_bytes" = 65536 : i32>>
+    #dlti.map<"CPU::cache::L1::size_in_bytes" = 65536 : i32>
     ```
     versus nested maps, which make it possible to obtain sub-dictionaries of
     related information (with the following example making use of other
@@ -137,12 +137,12 @@ def DLTI_MapAttr : DLTIAttr<"Map", [DLTIQueryInterface]> {
     With the flat encoding, the implied structure of the key is ignored, that is
     the only successful query (as expressed in the Transform Dialect) is:
     `transform.dlti.query ["CPU::cache::L1::size_in_bytes"] at %op`,
-    where `%op` is a handle to an operation which associates the flat-encoding
-    `#dlti.map` attribute.
+    where `%op` handles an operation with the flat `#dlti.map` stored as its
+    inherent `dlti` attribute.
 
     For querying nested dictionaries, the relevant keys need to be separately
-    provided. That is, if `%op` is an handle to an op which has the nesting
-    `#dlti.target_system_spec`-attribute from above attached, then
+    provided. If `%op` handles an operation with the nested
+    `#dlti.target_system_spec` above as its inherent `dlti` attribute, then
     `transform.dlti.query ["CPU","cache","L1","size_in_bytes"] at %op` gives
     back the first leaf value contained. To access the other leaf, we need to do
     `transform.dlti.query ["CPU","cache","L1d","size_in_bytes"] at %op`.
diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td b/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
new file mode 100644
index 00000000000000..074aff85c80b5a
--- /dev/null
+++ b/mlir/include/mlir/Dialect/DLTI/DLTIOpInterfaces.td
@@ -0,0 +1,58 @@
+//===- 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.
+
+    Queries and updates operate on one key at a time. The interface does not
+    prescribe how an operation stores values. For example, an operation with
+    an inherent `architecture = "sm_90"` attribute could return that string
+    attribute for the key `"architecture"`, without having a `dlti` attribute
+    or a `#dlti.map`.
+
+    The default implementations use the operation's inherent `dlti` attribute.
+    An update preserves other entries in a `#dlti.map`; setting a null value
+    removes an entry. Updating a non-map `dlti` attribute or using a null or
+    empty string key fails.
+  }];
+
+  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..748074a8e2c067 100644
--- a/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
+++ b/mlir/include/mlir/Dialect/DLTI/TransformOps/DLTITransformOps.td
@@ -19,25 +19,25 @@ def QueryOp : Op<Transform_Dialect, "dlti.query", [
   TransformOpInterface, TransformEachOpTrait,
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
 ]> {
-  let summary = "Return attribute (as param) associated to key via DTLI";
+  let summary = "Return attribute (as param) associated to key via DLTI";
   let description = [{
     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<"A" = #dlti.map<"B" = 42 : i32>>} {
       func.func private @f()
     }
     ```
-    and we have that `%func` is a Tranform handle to op `@f`, then
+    and we have that `%func` is a Transform handle to op `@f`, then
     `transform.dlti.query ["A", "B"] at %func` returns 42 as a param and
     `transform.dlti.query ["A"] at %func` returns the `#dlti.map` attribute
     containing just the key "B" and its value. Using `["B"]` or `["A","C"]` as
diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.h b/mlir/include/mlir/Dialect/Func/IR/FuncOps.h
index e6efeb5ee53f95..e18a7dc216574a 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/Dialect/Func/IR/FuncDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
diff --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index d7a0d5fc4f277a..3f4d6153c69dc7 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 = [{
@@ -272,7 +273,10 @@ def FuncOp : Func_Op<"func", [
     the region.
 
     Only dialect attribute names may be specified in the attribute dictionaries
-    for function arguments, results, or the function itself.
+    for function arguments, results, or the function itself. The function may
+    also have an inherent `dlti` attribute implementing `DLTIQueryInterface`.
+    Its values are available to scoped DLTI queries on the function and nested
+    operations, unless a closer scope answers the same key.
 
     Example:
 
@@ -303,7 +307,9 @@ 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 useCustomPropertiesEncoding = 1;
   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 038f7e8763ee0e..5978234f8f99d6 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/Dialect/GPU/IR/GPUDialectDecl.h"
diff --git a/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td b/mlir/include/mlir/Dialect/GPU/IR/GPUOps.td
index 2aa5219e904040..d3f48926d5dc30 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.";
@@ -1573,6 +1575,11 @@ def GPU_GPUModuleOp : GPU_Op<"module", [
     attribute will be used during the `gpu-module-to-binary` pass to specify the
     `OffloadingTranslationAttr` used when creating the `gpu.binary` operation.
 
+    A GPU module may also have an inherent `dlti` attribute implementing
+    `DLTIQueryInterface`. It provides values for scoped DLTI queries on the
+    module and its nested operations. In custom assembly, write it as
+    `dlti = #dlti.map<...>` before the optional attribute dictionary.
+
     ```
     gpu.module @symbol_name {
       gpu.func {}
@@ -1600,12 +1607,14 @@ 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 useCustomPropertiesEncoding = 1;
   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..2c2cf21c04e464 100644
--- a/mlir/include/mlir/IR/BuiltinOps.td
+++ b/mlir/include/mlir/IR/BuiltinOps.td
@@ -52,10 +52,16 @@ def ModuleOp : Builtin_Op<"module", [
       func.func @foo()
     }
     ```
+
+    A module may have an inherent `dlti` attribute implementing
+    `DLTIQueryInterface`, for example `dlti = #dlti.map<"chip" = "sm_90">`.
+    When the DLTI dialect is loaded, the module implements
+    `DLTIQueryOpInterface` and provides values for scoped DLTI queries.
   }];
 
   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 +96,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..021255099c284f 100644
--- a/mlir/lib/Dialect/DLTI/DLTI.cpp
+++ b/mlir/lib/Dialect/DLTI/DLTI.cpp
@@ -16,10 +16,13 @@
 #include "mlir/IR/DialectImplementation.h"
 
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
+#include "mlir/Dialect/DLTI/DLTIOpInterfaces.cpp.inc"
+
 #include "mlir/Dialect/DLTI/DLTIDialect.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
@@ -516,21 +519,61 @@ 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()));
+static constexpr llvm::StringLiteral kDltiAttrName = "dlti";
 
-  return std::pair(queryable, op);
+FailureOr<Attribute> mlir::dlti::detail::queryDlti(Operation *op,
+                                                   DataLayoutEntryKey key) {
+  auto queryable = op->getInherentAttrOfType<DLTIQueryInterface>(kDltiAttrName);
+  if (!queryable)
+    return failure();
+  return queryable.query(key);
+}
+
+LogicalResult mlir::dlti::detail::setDlti(Operation *op, DataLayoutEntryKey key,
+                                          Attribute value) {
+  if (key.isNull())
+    return failure();
+  if (auto stringKey = dyn_cast<StringAttr>(key);
+      stringKey && stringKey.getValue().empty())
+    return failure();
+
+  std::optional<Attribute> inherent = op->getInherentAttr(kDltiAttrName);
+  if (!inherent)
+    return failure();
+  auto current = dyn_cast_or_null<DLTIQueryInterface>(*inherent);
+  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();
+  }
+
+  StringAttr attrName = StringAttr::get(op->getContext(), kDltiAttrName);
+  if (entries.empty()) {
+    op->setInherentAttr(attrName, {});
+    return success();
+  }
+  op->setInherentAttr(attrName, MapAttr::get(op->getContext(), entries));
+  return success();
 }
 
 FailureOr<Attribute>
@@ -546,26 +589,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 +630,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 +661,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 +703,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/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index 6c41eeb36e9cb9..79816f84dd8825 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 
+#include "mlir/Bytecode/Encoding.h"
 #include "mlir/Conversion/ConvertToEmitC/ToEmitCInterface.h"
 #include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
@@ -30,11 +31,34 @@
 using namespace mlir;
 using namespace mlir::func;
 
+namespace {
+struct FuncDialectBytecodeInterface : BytecodeDialectInterface {
+  using BytecodeDialectInterface::BytecodeDialectInterface;
+
+  void writeVersion(DialectBytecodeWriter &writer) const override {
+    writer.writeVarInt(1);
+  }
+
+  std::unique_ptr<DialectVersion>
+  readVersion(DialectBytecodeReader &reader) const override {
+    uint64_t version;
+    if (failed(reader.readVarInt(version)))
+      return nullptr;
+    if (version != 1) {
+      reader.emitError("unsupported func dialect bytecode version");
+      return nullptr;
+    }
+    return std::make_unique<DialectVersion>();
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // FuncDialect
 //===----------------------------------------------------------------------===//
 
 void FuncDialect::initialize() {
+  addInterfaces<FuncDialectBytecodeInterface>();
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/Func/IR/FuncOps.cpp.inc"
@@ -139,6 +163,35 @@ bool ConstantOp::isBuildableWith(Attribute value, Type type) {
 // FuncOp
 //===----------------------------------------------------------------------===//
 
+LogicalResult FuncOp::readProperties(DialectBytecodeReader &reader,
+                                     OperationState &state) {
+  auto &properties = state.getOrAddProperties<Properties>();
+  // Preserve the pre-DLTI field order for unversioned bytecode.
+  if (failed(reader.readOptionalAttribute(properties.arg_attrs)) ||
+      failed(reader.readAttribute(properties.function_type)) ||
+      failed(reader.readOptionalAttribute(properties.no_inline)) ||
+      failed(reader.readOptionalAttribute(properties.res_attrs)) ||
+      failed(reader.readAttribute(properties.sym_name)) ||
+      failed(reader.readOptionalAttribute(properties.sym_visibility)))
+    return failure();
+
+  if (failed(reader.getDialectVersion<FuncDialect>()))
+    return success();
+  return reader.readOptionalAttribute(properties.dlti);
+}
+
+void FuncOp::writeProperties(DialectBytecodeWriter &writer) {
+  Properties &properties = getProperties();
+  writer.writeOptionalAttribute(properties.arg_attrs);
+  writer.writeAttribute(properties.function_type);
+  writer.writeOptionalAttribute(properties.no_inline);
+  writer.writeOptionalAttribute(properties.res_attrs);
+  writer.writeAttribute(properties.sym_name);
+  writer.writeOptionalAttribute(properties.sym_visibility);
+  if (writer.getBytecodeVersion() >= bytecode::kDialectVersioning)
+    writer.writeOptionalAttribute(properties.dlti);
+}
+
 FuncOp FuncOp::create(Location location, StringRef name, FunctionType type,
                       ArrayRef<NamedAttribute> attrs) {
   OpBuilder builder(location->getContext());
diff --git a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
index a39b43c9803eb9..15ec88ae66e9b3 100644
--- a/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
+++ b/mlir/lib/Dialect/GPU/IR/GPUDialect.cpp
@@ -12,6 +12,7 @@
 
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 
+#include "mlir/Bytecode/Encoding.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferDeallocationOpInterface.h"
 #include "mlir/Dialect/Math/IR/Math.h"
@@ -49,6 +50,28 @@ using namespace mlir::gpu;
 
 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.cpp.inc"
 
+namespace {
+struct GPUDialectBytecodeInterface : BytecodeDialectInterface {
+  using BytecodeDialectInterface::BytecodeDialectInterface;
+
+  void writeVersion(DialectBytecodeWriter &writer) const override {
+    writer.writeVarInt(1);
+  }
+
+  std::unique_ptr<DialectVersion>
+  readVersion(DialectBytecodeReader &reader) const override {
+    uint64_t version;
+    if (failed(reader.readVarInt(version)))
+      return nullptr;
+    if (version != 1) {
+      reader.emitError("unsupported gpu dialect bytecode version");
+      return nullptr;
+    }
+    return std::make_unique<DialectVersion>();
+  }
+};
+} // namespace
+
 //===----------------------------------------------------------------------===//
 // GPU Device Mapping Attributes
 //===----------------------------------------------------------------------===//
@@ -283,6 +306,7 @@ struct GPUInlinerInterface : public DialectInlinerInterface {
 } // namespace
 
 void GPUDialect::initialize() {
+  addInterfaces<GPUDialectBytecodeInterface>();
   addTypes<AsyncTokenType>();
   addTypes<MMAMatrixType>();
   addTypes<NamedBarrierType>();
@@ -1983,6 +2007,31 @@ LogicalResult gpu::ReturnOp::verify() {
 // GPUModuleOp
 //===----------------------------------------------------------------------===//
 
+LogicalResult GPUModuleOp::readProperties(DialectBytecodeReader &reader,
+                                          OperationState &state) {
+  auto &properties = state.getOrAddProperties<Properties>();
+  // Preserve the pre-DLTI field order for unversioned bytecode.
+  if (failed(reader.readOptionalAttribute(properties.offloadingHandler)) ||
+      failed(reader.readAttribute(properties.sym_name)) ||
+      failed(reader.readOptionalAttribute(properties.sym_visibility)) ||
+      failed(reader.readOptionalAttribute(properties.targets)))
+    return failure();
+
+  if (failed(reader.getDialectVersion<GPUDialect>()))
+    return success();
+  return reader.readOptionalAttribute(properties.dlti);
+}
+
+void GPUModuleOp::writeProperties(DialectBytecodeWriter &writer) {
+  Properties &properties = getProperties();
+  writer.writeOptionalAttribute(properties.offloadingHandler);
+  writer.writeAttribute(properties.sym_name);
+  writer.writeOptionalAttribute(properties.sym_visibility);
+  writer.writeOptionalAttribute(properties.targets);
+  if (writer.getBytecodeVersion() >= bytecode::kDialectVersioning)
+    writer.writeOptionalAttribute(properties.dlti);
+}
+
 void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
                         StringRef name, ArrayAttr targets,
                         Attribute offloadingHandler) {
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/Bytecode/dlti-query-op-interface-compat.mlir b/mlir/test/Bytecode/dlti-query-op-interface-compat.mlir
new file mode 100644
index 00000000000000..51b7791bb2651c
--- /dev/null
+++ b/mlir/test/Bytecode/dlti-query-op-interface-compat.mlir
@@ -0,0 +1,27 @@
+// Verify that bytecode written before func.func and gpu.module gained a dlti
+// property can still be read. The fixture was produced from this input by
+// mlir-opt built at fc696f791e06 (the PR's parent) using -emit-bytecode.
+// RUN: mlir-opt %S/dlti-query-op-interface-pre-dlti.mlirbc | FileCheck %s --check-prefix=OLD
+// RUN: mlir-opt %S/dlti-query-op-interface-pre-dlti-gpu.mlirbc | FileCheck %s --check-prefix=OLD-GPU
+// RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s --check-prefix=NEW
+
+// The GPU-only fixture uses the same parent revision and this input:
+//   module attributes {gpu.container_module} {
+//     gpu.module @legacy_gpu {}
+//   }
+
+// OLD: module attributes {gpu.container_module}
+// OLD: func.func @legacy()
+// OLD: gpu.module @legacy_gpu
+// OLD-GPU: module attributes {gpu.container_module}
+// OLD-GPU: gpu.module @legacy_gpu
+// NEW: module attributes {gpu.container_module}
+// NEW: func.func @legacy()
+// NEW: gpu.module @legacy_gpu
+module attributes {gpu.container_module} {
+  func.func @legacy() {
+    return
+  }
+  gpu.module @legacy_gpu {
+  }
+}
diff --git a/mlir/test/Bytecode/dlti-query-op-interface-pre-dlti-gpu.mlirbc b/mlir/test/Bytecode/dlti-query-op-interface-pre-dlti-gpu.mlirbc
new file mode 100644
index 0000000000000000000000000000000000000000..315113175c617079b01b4a9bd5dd75d1e73b6059
GIT binary patch
literal 199
zcmXX=OA5j;6wLVgj3tyJMY`|+E~Hh&E4UC8v=@+CLk*;9Dec0ec`C191Ys8Qd2g8W
z=~G1K)3Mm?^KHJZ%^<K~FoP6%q%jbyM9D-YlNC!MmK??@JW(lF1{vcT5 at fI>FkK>&
zfh1t+6~sGA*M~Tq-$!Ey<3!z!qG=za)nYR=-R9xltm}34EM2WLA7`zzue`CQ{~hO5
a>xR-e-QWIiR at Y_qy8Vh~iJyrO3itx<b19|(

literal 0
HcmV?d00001

diff --git a/mlir/test/Bytecode/dlti-query-op-interface-pre-dlti.mlirbc b/mlir/test/Bytecode/dlti-query-op-interface-pre-dlti.mlirbc
new file mode 100644
index 0000000000000000000000000000000000000000..3559fd2e57ff8f7a9749211cca905a63eff3bc48
GIT binary patch
literal 290
zcmXYr(MrTH5JfYa#u!2gB}g7b7ZIPN?IP|E2qMBRtos8?O|vD?rZpxH{c3-~6n(h|
z=FZ$Z``ynm-|ya!PtVn++H^J`SW_U30T!50NfI6u2E`?(k|&iiS*usVaq>wn6kOvR
zNa778g at SvkV9g0Y_D0}?VW8O-1H%V_X_jLk1YOaCBs8NK;35qG-L9k(kE0Nb7oXAg
z!8&w~PNQy$Xc*hr8{|!h-l3tlz7FO))Gu=|dTdSIrv36O8b6ogR9Y8|KR4QxdK{)E
mRLjK at RO!)$#yaCq|B=1v8a<;woR+T3JDNVvfiJ?K0QdzN*gHf3

literal 0
HcmV?d00001

diff --git a/mlir/test/Bytecode/dlti-query-op-interface-roundtrip.mlir b/mlir/test/Bytecode/dlti-query-op-interface-roundtrip.mlir
new file mode 100644
index 00000000000000..454e1c034e30cc
--- /dev/null
+++ b/mlir/test/Bytecode/dlti-query-op-interface-roundtrip.mlir
@@ -0,0 +1,12 @@
+// RUN: mlir-opt -emit-bytecode %s | mlir-opt | FileCheck %s
+
+// CHECK: module attributes {dlti = #dlti.map<"module" = 1 : i32>, gpu.container_module}
+module attributes {dlti = #dlti.map<"module" = 1 : i32>, gpu.container_module} {
+  // CHECK: func.func @new() attributes {dlti = #dlti.map<"func" = 2 : i32>}
+  func.func @new() attributes {dlti = #dlti.map<"func" = 2 : i32>} {
+    return
+  }
+  // CHECK: gpu.module @new_gpu dlti = #dlti.map<"gpu" = 3 : i32>
+  gpu.module @new_gpu dlti = #dlti.map<"gpu" = 3 : i32> {
+  }
+}
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 4f661b5a1d233336116b303958bc724529f88b5a..ee79226875ffae00eeceaf8a7bc78b06e147e216 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..74df84ba8a5458
--- /dev/null
+++ b/mlir/test/Dialect/DLTI/query-op-interface.mlir
@@ -0,0 +1,32 @@
+// 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"} : () -> ()
+
+  // CHECK: "test.op_with_data_layout"() <{dlti = #dlti.map<"kept" = 0 : i32>}> ({
+  // CHECK-NOT: test.dlti_action
+  "test.op_with_data_layout"() ({
+  }) {dlti = #dlti.map<"kept" = 0 : i32>,
+      test.dlti_action = "invalid"} : () -> ()
+}
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..e8e547cf7cf4ee 100644
--- a/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
+++ b/mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp
@@ -105,14 +105,82 @@ 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());
+      } else if (action.getValue() == "invalid") {
+        Attribute value = builder.getI32IntegerAttr(1);
+        if (succeeded(queryOp.setDlti(key(""), value)) ||
+            succeeded(queryOp.setDlti(DataLayoutEntryKey(), value)) ||
+            succeeded(queryOp.setDlti(key(""), Attribute()))) {
+          op->emitError("unexpectedly accepted an invalid DLTI key");
+          return WalkResult::interrupt();
+        }
+      }
+      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 c3e06dc968a270..769cc97c818faf 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"
@@ -3167,10 +3168,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;