[mlir] support data layout specs on ModuleOp

[ftynse]

Summary

ModuleOp is a natural place to provide scoped data layout information. However,
it is undesirable for ModuleOp to implement the entirety of
DataLayoutOpInterface because that would require either pushing the interface
inside the IR library instead of a separate library, or putting the default
implementation of the interface as inline functions in headers leading to
binary bloat. Instead, ModuleOp accepts an arbitrary data layout spec attribute
and has a dedicated hook to extract it, and DataLayout is modified to know
about ModuleOp particularities.

[ftynse]

a bit more doc

[ftynse]

rebase

[ftynse]

forgot to git-add a test

[herhut]

mlir/docs/DataLayout.md

30

194

I don't under stand the otherwise here. Is it "otherwise if it is a..." ?

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

136

For skip this would need to be a continue. Should this not skip the module op if it is not top-level? So a nested ModuleOp that has no data layout needs to be skipped.

[ftynse]

Address review.

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

136

No, it shouldn't. We can have null layouts in the middle and should handle them. It only skips the top-level module because it is most likely the one implicitly added when parsing the IR and using its location is very confusing. And we cannot continue because we are in the lambda. Extended comment to explain this.

[herhut]

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

136

Thank you for clarifying this.

mlir/docs/DataLayout.md

*   attribute interfaces that can be implemented by concrete data layout
     specifications;
 *   type interfaces that should be implemented by types subject to data layout;
 *   operation interfaces that must be implemented by operations that can serve
     as data layout scopes (e.g., modules);
 *   and dialect interfaces for data layout properties unrelated to specific
     types.
 
 Built-in types are handled specially to decrease the overall query cost.
+Similarly, built-in `ModuleOp` supports data layouts without going through the
+interface.
 
 ## Usage
 
 ### Scoping
 
 Following MLIR's nested structure, data layout properties are _scoped_ to
-regions belonging to specific operations that implement the
-`DataLayoutOpInterface`. Such scoping operations partially control the data
-layout properties and may have attributes that affect them, typically organized
-in a data layout specification.
+regions belonging to either operations that implement the
+`DataLayoutOpInterface` or `ModuleOp` operations. Such scoping operations
+partially control the data layout properties and may have attributes that affect
+them, typically organized in a data layout specification.
 
 Types may have a different data layout in different scopes, including scopes
 that are nested in other scopes such as modules contained in other modules. At
 the same time, within the given scope excluding any nested scope, a given type
 has fixed data layout properties. Types are also expected to have a default,
 "natural" data layout in case they are used outside of any operation that
-provides data layout scope for them. This ensure data layout queries always have
-a valid result.
+provides data layout scope for them. This ensures that data layout queries
+always have a valid result.
 
 ### Compatibility and Transformations
 
 The information necessary to compute layout properties can be combined from
 nested scopes. For example, an outer scope can define layout properties for a
 subset of types while inner scopes define them for a disjoint subset, or scopes
 can progressively relax alignment requirements on a type. This mechanism is
 supported by the notion of data layout _compatibility_: the layout defined in a

 dialect is expected to implement the `DataLayoutDialectInterface`. This dialect
 provides hooks for verifying the validity of the entry value attributes and for
 and the compatibility of nested entries.
 
 ### Query Dispatch
 
 The overall flow of a data layout property query is as follows.
 
--   The user constructs a `DataLayout` at the given scope. The constructor
+1.  The user constructs a `DataLayout` at the given scope. The constructor
     fetches the data layout specification and combines it with those of
     enclosing scopes (layouts are expected to be compatible).
--   The user calls `DataLayout::query(Type ty)`.
--   If `DataLayout` has a cached response, this response is returned
+2.  The user calls `DataLayout::query(Type ty)`.
+3.  If `DataLayout` has a cached response, this response is returned
     immediately.
--   Otherwise, the query is handed down by `DataLayout` to
-    `DataLayoutOpInterface::query(ty, *this, relevantEntries)` where the
-    relevant entries are computed as described above.
--   Unless the `query` hook is reimplemented by the op interface, the query is
+4.  Otherwise, the query is handed down by `DataLayout` to the closest layout
+    scoping operation. If it implements `DataLayoutOpInterface`, then the query
+    is forwarded to`DataLayoutOpInterface::query(ty, *this, relevantEntries)`
+    where the relevant entries are computed as described above. If it does not
+    implement `DataLayoutOpInterface`, it must be a `ModuleOp`, and the query is
+    forwarded to `DataLayoutTypeInterface::query(dataLayout, relevantEntries)`
+    after casting `ty` to the type interface.
+5.  Unless the `query` hook is reimplemented by the op interface, the query is
     handled further down to `DataLayoutTypeInterface::query(dataLayout,
     relevantEntries)` after casting `ty` to the type interface. If the type does
     not implement the interface, an unrecoverable fatal error is produced.
--   The type is expected to always provide the response, which is returned up
+6.  The type is expected to always provide the response, which is returned up
     the call stack and cached by the `DataLayout.`
 
 ## Default Implementation
 
 The default implementation of the data layout interfaces directly handles
 queries for a subset of built-in types.
 
+### Built-in Modules
+
+Built-in `ModuleOp` allows at most one attribute that implements
+`DataLayoutSpecInterface`. It does not implement the entire interface for
+efficiency and layering reasons. Instead, `DataLayout` can be constructed for
+`ModuleOp` and handles modules transparently alongside other operations that
+implement the interface.
+
 ### Built-in Types
 
 The following describes the default properties of built-in types.
 
 The size of built-in integers and floats in bytes is computed as
 `ceildiv(bitwidth, 8)`. The ABI alignment of integer types with bitwidth below
 64 and of the float types is the closest from above power-of-two number of
 bytes. The ABI alignment of integer types with bitwidth 64 and above is 4 bytes

mlir/include/mlir/IR/BuiltinOps.h

 #ifndef MLIR_IR_BUILTINOPS_H_
 #define MLIR_IR_BUILTINOPS_H_
 
 #include "mlir/IR/FunctionSupport.h"
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 //===----------------------------------------------------------------------===//
 // Dialect Operations
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES

mlir/include/mlir/IR/BuiltinOps.td

 
 #ifndef BUILTIN_OPS
 #define BUILTIN_OPS
 
 include "mlir/IR/BuiltinDialect.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/CastInterfaces.td"
+include "mlir/Interfaces/DataLayoutInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Base class for Builtin dialect ops.
 class Builtin_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Builtin_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // FuncOp
let extraClassDeclaration = [{
     Optional<StringRef> getName() { return sym_name(); }
 
     //===------------------------------------------------------------------===//
     // SymbolOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// A ModuleOp may optionally define a symbol.
     bool isOptionalSymbol() { return true; }
+
+    //===------------------------------------------------------------------===//
+    // DataLayoutOpInterface Methods
+    //===------------------------------------------------------------------===//
+
+    DataLayoutSpecInterface getDataLayoutSpec();
   }];
   let verifier = [{ return ::verify(*this); }];
 
   // We need to ensure the block inside the region is properly terminated;
   // the auto-generated builders do not guarantee that.
   let skipDefaultBuilders = 1;
 }
 

mlir/include/mlir/Interfaces/DataLayoutInterfaces.h

 class DataLayoutEntryInterface;
 using DataLayoutEntryKey = llvm::PointerUnion<Type, Identifier>;
 // Using explicit SmallVector size because we cannot infer the size from the
 // forward declaration, and we need the typedef in the actual declaration.
 using DataLayoutEntryList = llvm::SmallVector<DataLayoutEntryInterface, 4>;
 using DataLayoutEntryListRef = llvm::ArrayRef<DataLayoutEntryInterface>;
 class DataLayoutOpInterface;
 class DataLayoutSpecInterface;
+class ModuleOp;
 
 namespace detail {
 /// Default handler for the type size request. Computes results for built-in
 /// types and dispatches to the DataLayoutTypeInterface for other types.
 unsigned getDefaultTypeSize(Type type, const DataLayout &dataLayout,
                             ArrayRef<DataLayoutEntryInterface> params);
 
 /// Default handler for the required alignemnt request. Computes results for

 DataLayoutEntryList filterEntriesForType(DataLayoutEntryListRef entries,
                                          TypeID typeID);
 
 /// Given a list of data layout entries, returns the entry that has the given
 /// identifier as key, if such an entry exists in the list.
 DataLayoutEntryInterface
 filterEntryForIdentifier(DataLayoutEntryListRef entries, Identifier id);
 
-/// Verifies that the operation implementing the data layout interface is valid.
-/// This calls the verifier of the spec attribute and checks if the layout is
-/// compatible with specs attached to the enclosing operations.
-LogicalResult verifyDataLayoutOp(DataLayoutOpInterface op);
+/// Verifies that the operation implementing the data layout interface, or a
+/// module operation, is valid. This calls the verifier of the spec attribute
+/// and checks if the layout is compatible with specs attached to the enclosing
+/// operations.
+LogicalResult verifyDataLayoutOp(Operation *op);
 
 /// Verifies that a data layout spec is valid. This dispatches to individual
 /// entry verifiers, and then to the verifiers implemented by the relevant type
 /// and dialect interfaces for type and identifier keys respectively.
 LogicalResult verifyDataLayoutSpec(DataLayoutSpecInterface spec, Location loc);
 } // namespace detail
 } // namespace mlir
 

 /// efficiency, results to all requests will be cached in this object.
 /// Therefore, if the data layout spec for the scoping operation, or any of the
 /// enclosing operations, changes, the cache is no longer valid. The user is
 /// responsible creating a new DataLayout object after any spec change. In debug
 /// mode, the cache validity is being checked in every request.
 class DataLayout {
 public:
   explicit DataLayout(DataLayoutOpInterface op);
+  explicit DataLayout(ModuleOp op);
 
   /// Returns the size of the given type in the current scope.
   unsigned getTypeSize(Type t) const;
 
   /// Returns the required alignment of the given type in the current scope.
   unsigned getTypeABIAlignment(Type t) const;
 
   /// Returns the preferred of the given type in the current scope.
#endif
 
   /// Asserts that the cache is still valid. Expensive in debug mode. No-op in
   /// release mode.
   void checkValid() const;
 
   /// Operation defining the scope of requests.
   // TODO: this is mutable because the generated interface method are not const.
   // Update the generator to support const methods and change this to const.
-  mutable DataLayoutOpInterface scope;
+  mutable Operation *scope;
 
   /// Caches for individual requests.
   mutable DenseMap<Type, unsigned> sizes;
   mutable DenseMap<Type, unsigned> abiAlignments;
   mutable DenseMap<Type, unsigned> preferredAlignments;
 };
 
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H

mlir/lib/Dialect/DLTI/DLTI.cpp

 //===- DLTI.cpp - Data Layout And Target Info MLIR Dialect Implementation -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // DataLayoutEntryAttr

 
 LogicalResult DLTIDialect::verifyOperationAttribute(Operation *op,
                                                     NamedAttribute attr) {
   if (attr.first == DLTIDialect::kDataLayoutAttrName) {
     if (!attr.second.isa<DataLayoutSpecAttr>()) {
       return op->emitError() << "'" << DLTIDialect::kDataLayoutAttrName
                              << "' is expected to be a #dlti.dl_spec attribute";
     }
+    if (isa<ModuleOp>(op))
+      return detail::verifyDataLayoutOp(op);
     return success();
   }
 
   return op->emitError() << "attribute '" << attr.first
                          << "' not supported by dialect";
 }

mlir/lib/IR/BuiltinDialect.cpp

 }
 
 /// Construct a module from the given context.
 ModuleOp ModuleOp::create(Location loc, Optional<StringRef> name) {
   OpBuilder builder(loc->getContext());
   return builder.create<ModuleOp>(loc, name);
 }
 
+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
+  // layout object construction that is used for repeated queries.
+  for (Attribute attr : llvm::make_second_range(getOperation()->getAttrs())) {
+    if (auto spec = attr.dyn_cast<DataLayoutSpecInterface>())
+      return spec;
+  }
+  return {};
+}
+
 static LogicalResult verify(ModuleOp op) {
   // Check that none of the attributes are non-dialect attributes, except for
   // the symbol related attributes.
   for (auto attr : op->getAttrs()) {
     if (!attr.first.strref().contains('.') &&
         !llvm::is_contained(
             ArrayRef<StringRef>{mlir::SymbolTable::getSymbolAttrName(),
                                 mlir::SymbolTable::getVisibilityAttrName()},
             attr.first.strref()))
       return op.emitOpError() << "can only contain attributes with "
                                  "dialect-prefixed names, found: '"
                               << attr.first << "'";
   }
 
+  // Check that there is at most one data layout spec attribute.
+  StringRef layoutSpecAttrName;
+  DataLayoutSpecInterface layoutSpec;
+  for (const NamedAttribute &na : op->getAttrs()) {
+    if (auto spec = na.second.dyn_cast<DataLayoutSpecInterface>()) {
+      if (layoutSpec) {
+        InFlightDiagnostic diag =
+            op.emitOpError() << "expects at most one data layout attribute";
+        diag.attachNote() << "'" << layoutSpecAttrName
+                          << "' is a data layout attribute";
+        diag.attachNote() << "'" << na.first << "' is a data layout attribute";
+      }
+      layoutSpecAttrName = na.first.strref();
+      layoutSpec = spec;
+    }
+  }
+
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // UnrealizedConversionCastOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult

mlir/lib/IR/CMakeLists.txt

add_mlir_library(MLIRIR
   DEPENDS
   MLIRBuiltinAttributesIncGen
   MLIRBuiltinDialectIncGen
   MLIRBuiltinLocationAttributesIncGen
   MLIRBuiltinOpsIncGen
   MLIRBuiltinTypesIncGen
   MLIRCallInterfacesIncGen
   MLIRCastInterfacesIncGen
+  MLIRDataLayoutInterfacesIncGen
   MLIROpAsmInterfaceIncGen
   MLIRRegionKindInterfaceIncGen
   MLIRSideEffectInterfacesIncGen
   MLIRSymbolInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

mlir/lib/Interfaces/DataLayoutInterfaces.cpp

 //===- DataLayoutInterfaces.cpp - Data Layout Interface Implementation ----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
 #include "mlir/IR/BuiltinDialect.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Operation.h"
 
+#include "llvm/ADT/TypeSwitch.h"
+
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Default implementations
 //===----------------------------------------------------------------------===//
 
 /// Reports that the given type is missing the data layout information and
 /// exits.
mlir::detail::filterEntryForIdentifier(DataLayoutEntryListRef entries,
   const auto *it = llvm::find_if(entries, [id](DataLayoutEntryInterface entry) {
     if (!entry.getKey().is<Identifier>())
       return false;
     return entry.getKey().get<Identifier>() == id;
   });
   return it == entries.end() ? DataLayoutEntryInterface() : *it;
 }
 
+static DataLayoutSpecInterface getSpec(Operation *operation) {
+  return llvm::TypeSwitch<Operation *, DataLayoutSpecInterface>(operation)
+      .Case<ModuleOp, DataLayoutOpInterface>(
+          [&](auto op) { return op.getDataLayoutSpec(); })
+      .Default([](Operation *) {
+        llvm_unreachable("expected an op with data layout spec");
+        return DataLayoutSpecInterface();
+      });
+}
+
 /// Populates `opsWithLayout` with the list of proper ancestors of `leaf` that
-/// implement the `DataLayoutOpInterface`.
-static void findProperAscendantsWithLayout(
-    Operation *leaf, SmallVectorImpl<DataLayoutOpInterface> &opsWithLayout) {
+/// are either modules or implement the `DataLayoutOpInterface`.
+static void
+collectParentLayouts(Operation *leaf,
+                     SmallVectorImpl<DataLayoutSpecInterface> &specs,
+                     SmallVectorImpl<Location> *opLocations = nullptr) {
   if (!leaf)
     return;
 
-  while (auto opLayout = leaf->getParentOfType<DataLayoutOpInterface>()) {
-    opsWithLayout.push_back(opLayout);
-    leaf = opLayout;
+  for (Operation *parent = leaf->getParentOp(); parent != nullptr;
+       parent = parent->getParentOp()) {
+    llvm::TypeSwitch<Operation *>(parent)
+        .Case<ModuleOp>([&](ModuleOp op) {
+          // Skip top-level module op unless it has a layout. Top-level module
+          // without layout is most likely the one implicitly added by the
+          // parser and it doesn't have location. Top-level null specification
+          // would have had the same effect as not having a specification at all
+          // (using type defaults).
+          if (!op->getParentOp() && !op.getDataLayoutSpec())
+            return;
+          specs.push_back(op.getDataLayoutSpec());
+          if (opLocations)
+            opLocations->push_back(op.getLoc());
+        })
+        .Case<DataLayoutOpInterface>([&](DataLayoutOpInterface op) {
+          specs.push_back(op.getDataLayoutSpec());
+          if (opLocations)
+            opLocations->push_back(op.getLoc());
+        });
   }
 }
 
 /// Returns a layout spec that is a combination of the layout specs attached
 /// to the given operation and all its ancestors.
-static DataLayoutSpecInterface
-getCombinedDataLayout(DataLayoutOpInterface leaf) {
+static DataLayoutSpecInterface getCombinedDataLayout(Operation *leaf) {
   if (!leaf)
     return {};
 
+  assert((isa<ModuleOp, DataLayoutOpInterface>(leaf)) &&
+         "expected an op with data layout spec");
+
   SmallVector<DataLayoutOpInterface> opsWithLayout;
-  findProperAscendantsWithLayout(leaf, opsWithLayout);
+  SmallVector<DataLayoutSpecInterface> specs;
+  collectParentLayouts(leaf, specs);
 
   // Fast track if there are no ancestors.
-  if (opsWithLayout.empty())
-    return leaf.getDataLayoutSpec();
+  if (specs.empty())
+    return getSpec(leaf);
 
   // Create the list of non-null specs (null/missing specs can be safely
   // ignored) from the outermost to the innermost.
-  SmallVector<DataLayoutSpecInterface> specs;
-  specs.reserve(opsWithLayout.size());
-  for (DataLayoutOpInterface op : llvm::reverse(opsWithLayout))
-    if (DataLayoutSpecInterface current = op.getDataLayoutSpec())
-      specs.push_back(current);
+  auto nonNullSpecs = llvm::to_vector<2>(llvm::make_filter_range(
+      llvm::reverse(specs),
+      [](DataLayoutSpecInterface iface) { return iface != nullptr; }));
 
   // Combine the specs using the innermost as anchor.
-  if (DataLayoutSpecInterface current = leaf.getDataLayoutSpec())
-    return current.combineWith(specs);
-  if (specs.empty())
+  if (DataLayoutSpecInterface current = getSpec(leaf))
+    return current.combineWith(nonNullSpecs);
+  if (nonNullSpecs.empty())
     return {};
-  return specs.back().combineWith(llvm::makeArrayRef(specs).drop_back());
+  return nonNullSpecs.back().combineWith(
+      llvm::makeArrayRef(nonNullSpecs).drop_back());
 }
 
-LogicalResult mlir::detail::verifyDataLayoutOp(DataLayoutOpInterface op) {
-  DataLayoutSpecInterface spec = op.getDataLayoutSpec();
+LogicalResult mlir::detail::verifyDataLayoutOp(Operation *op) {
+  DataLayoutSpecInterface spec = getSpec(op);
   // The layout specification may be missing and it's fine.
   if (!spec)
     return success();
 
-  if (failed(spec.verifySpec(op.getLoc())))
+  if (failed(spec.verifySpec(op->getLoc())))
     return failure();
   if (!getCombinedDataLayout(op)) {
     InFlightDiagnostic diag =
-        op.emitError()
-        << "data layout is not a refinement of the layouts in enclosing ops";
-    SmallVector<DataLayoutOpInterface> opsWithLayout;
-    findProperAscendantsWithLayout(op, opsWithLayout);
-    for (DataLayoutOpInterface parent : opsWithLayout)
-      diag.attachNote(parent.getLoc()) << "enclosing op with data layout";
+        op->emitError()
+        << "data layout does not combine with layouts of enclosing ops";
+    SmallVector<DataLayoutSpecInterface> specs;
+    SmallVector<Location> opLocations;
+    collectParentLayouts(op, specs, &opLocations);
+    for (Location loc : opLocations)
+      diag.attachNote(loc) << "enclosing op with data layout";
     return diag;
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // DataLayout
 //===----------------------------------------------------------------------===//
 
-mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
-    : originalLayout(getCombinedDataLayout(op)), scope(op) {
+template <typename OpTy>
+void checkMissingLayout(DataLayoutSpecInterface originalLayout, OpTy op) {
   if (!originalLayout) {
     assert((!op || !op.getDataLayoutSpec()) &&
            "could not compute layout information for an op (failed to "
            "combine attributes?)");
   }
+}
 
+mlir::DataLayout::DataLayout(DataLayoutOpInterface op)
+    : originalLayout(getCombinedDataLayout(op)), scope(op) {
 #ifndef NDEBUG
-  SmallVector<DataLayoutOpInterface> opsWithLayout;
-  findProperAscendantsWithLayout(op, opsWithLayout);
-  layoutStack = llvm::to_vector<2>(
-      llvm::map_range(opsWithLayout, [](DataLayoutOpInterface iface) {
-        return iface.getDataLayoutSpec();
-      }));
+  checkMissingLayout(originalLayout, op);
+  collectParentLayouts(op, layoutStack);
+#endif
+}
+
+mlir::DataLayout::DataLayout(ModuleOp op)
+    : originalLayout(getCombinedDataLayout(op)), scope(op) {
+#ifndef NDEBUG
+  checkMissingLayout(originalLayout, op);
+  collectParentLayouts(op, layoutStack);
 #endif
 }
 
 void mlir::DataLayout::checkValid() const {
 #ifndef NDEBUG
-  SmallVector<DataLayoutOpInterface> opsWithLayout;
-  findProperAscendantsWithLayout(scope, opsWithLayout);
-  assert(opsWithLayout.size() == layoutStack.size() &&
+  SmallVector<DataLayoutSpecInterface> specs;
+  collectParentLayouts(scope, specs);
+  assert(specs.size() == layoutStack.size() &&
          "data layout object used, but no longer valid due to the change in "
          "number of nested layouts");
-  for (auto pair : llvm::zip(opsWithLayout, layoutStack)) {
-    Attribute newLayout = std::get<0>(pair).getDataLayoutSpec();
+  for (auto pair : llvm::zip(specs, layoutStack)) {
+    Attribute newLayout = std::get<0>(pair);
     Attribute origLayout = std::get<1>(pair);
     assert(newLayout == origLayout &&
            "data layout object used, but no longer valid "
            "due to the change in layout attributes");
   }
 #endif
   assert(((!scope && !this->originalLayout) ||
           (scope && this->originalLayout == getCombinedDataLayout(scope))) &&
static unsigned cachedLookup(Type t, DenseMap<Type, unsigned> &cache,
 
   auto result = cache.try_emplace(t, compute(t));
   return result.first->second;
 }
 
 unsigned mlir::DataLayout::getTypeSize(Type t) const {
   checkValid();
   return cachedLookup(t, sizes, [&](Type ty) {
-    return (scope && originalLayout)
-               ? scope.getTypeSize(
-                     ty, *this, originalLayout.getSpecForType(ty.getTypeID()))
-               : detail::getDefaultTypeSize(ty, *this, {});
+    if (originalLayout) {
+      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
+      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
+        return iface.getTypeSize(ty, *this, list);
+      return detail::getDefaultTypeSize(ty, *this, list);
+    }
+    return detail::getDefaultTypeSize(ty, *this, {});
   });
 }
 
 unsigned mlir::DataLayout::getTypeABIAlignment(Type t) const {
   checkValid();
   return cachedLookup(t, abiAlignments, [&](Type ty) {
-    return (scope && originalLayout)
-               ? scope.getTypeABIAlignment(
-                     ty, *this, originalLayout.getSpecForType(ty.getTypeID()))
-               : detail::getDefaultABIAlignment(ty, *this, {});
+    if (originalLayout) {
+      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
+      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
+        return iface.getTypeABIAlignment(ty, *this, list);
+      return detail::getDefaultABIAlignment(ty, *this, list);
+    }
+    return detail::getDefaultABIAlignment(ty, *this, {});
   });
 }
 
 unsigned mlir::DataLayout::getTypePreferredAlignment(Type t) const {
   checkValid();
   return cachedLookup(t, preferredAlignments, [&](Type ty) {
-    return (scope && originalLayout)
-               ? scope.getTypePreferredAlignment(
-                     ty, *this, originalLayout.getSpecForType(ty.getTypeID()))
-               : detail::getDefaultPreferredAlignment(ty, *this, {});
+    if (originalLayout) {
+      DataLayoutEntryList list = originalLayout.getSpecForType(ty.getTypeID());
+      if (auto iface = dyn_cast<DataLayoutOpInterface>(scope))
+        return iface.getTypePreferredAlignment(ty, *this, list);
+      return detail::getDefaultPreferredAlignment(ty, *this, list);
+    }
+    return detail::getDefaultPreferredAlignment(ty, *this, {});
   });
 }
 
 //===----------------------------------------------------------------------===//
 // DataLayoutSpecInterface
 //===----------------------------------------------------------------------===//
 
 void DataLayoutSpecInterface::bucketEntriesByType(

mlir/test/Dialect/DLTI/invalid.mlir

 // expected-error@below {{'dlti.endianness' data layout entry is expected to be either 'big' or 'little'}}
 "test.op_with_data_layout"() ({
 }) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.endianness", "some">> } : () -> ()
 
 // -----
 
 // Mismatching entries don't combine.
 "test.op_with_data_layout"() ({
-  // expected-error@below {{data layout is not a refinement of the layouts in enclosing ops}}
+  // expected-error@below {{data layout does not combine with layouts of enclosing ops}}
   // expected-note@above {{enclosing op with data layout}}
   "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> ()
   "test.maybe_terminator_op"() : () -> ()
 }) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 33>> } : () -> ()
 
 // -----
 
 // Layout not supported for built-in types.
 // expected-error@below {{unexpected data layout for a built-in type}}
 "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<i32, 32>> } : () -> ()
 
 // -----
 
 // expected-error@below {{data layout specified for a type that does not support it}}
 "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<!test.test_type, 32>> } : () -> ()
+
+// -----
+
+// Mismatching entries are checked on module ops as well.
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 33>>} {
+  // expected-error@below {{data layout does not combine with layouts of enclosing ops}}
+  // expected-note@above {{enclosing op with data layout}}
+  module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>>} {
+  }
+}
+
+// -----
+
+// Mismatching entries are checked on a combination of modules and other ops.
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 33>>} {
+  // expected-error@below {{data layout does not combine with layouts of enclosing ops}}
+  // expected-note@above {{enclosing op with data layout}}
+  "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>>} : () -> ()
+}

mlir/test/IR/module-op.mlir

 module @foo {
   // CHECK: module
   module {
     // CHECK: module @bar attributes
     module @bar attributes {foo.bar} {
     }
   }
 }
+
+// -----
+
+// expected-error@below {{expects at most one data layout attribute}}
+// expected-note@below {{'test.another_attribute' is a data layout attribute}}
+// expected-note@below {{'test.random_attribute' is a data layout attribute}}
+module attributes { test.random_attribute = #dlti.dl_spec<>,
+                    test.another_attribute = #dlti.dl_spec<>} {
+}

mlir/test/Interfaces/DataLayoutInterfaces/module.mlir

+// RUN: mlir-opt --test-data-layout-query %s | FileCheck %s
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+      #dlti.dl_entry<!test.test_type_with_layout<10>, ["size", 12]>,
+      #dlti.dl_entry<!test.test_type_with_layout<20>, ["alignment", 32]>>} {
+  // CHECK-LABEL: @module_level_layout
+  func @module_level_layout() {
+     // CHECK: alignment = 32
+     // CHECK: preferred = 1
+     // CHECK: size = 12
+    "test.data_layout_query"() : () -> !test.test_type_with_layout<10>
+    return
+  }
+}

mlir/test/lib/Transforms/TestDataLayoutQuery.cpp

func.walk([&](test::DataLayoutQueryOp op) {
         return;
 
       auto scope = op->getParentOfType<test::OpWithDataLayoutOp>();
       if (!layouts.count(scope)) {
         layouts.try_emplace(
             scope, scope ? cast<DataLayoutOpInterface>(scope.getOperation())
                          : nullptr);
       }
+      auto module = op->getParentOfType<ModuleOp>();
+      if (!layouts.count(module))
+        layouts.try_emplace(module, module);
+
+      Operation *closest = (scope && module && module->isProperAncestor(scope))
+                               ? scope.getOperation()
+                               : module.getOperation();
 
-      const DataLayout &layout = layouts.find(scope)->getSecond();
+      const DataLayout &layout = layouts.find(closest)->getSecond();
       unsigned size = layout.getTypeSize(op.getType());
       unsigned alignment = layout.getTypeABIAlignment(op.getType());
       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
       op->setAttrs(
           {builder.getNamedAttr("size", builder.getIndexAttr(size)),
            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),
            builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))});
     });