[mlir] support memref of memref in standard-to-llvm conversion

[ftynse]

Summary

Now that memref supports arbitrary element types, add support for memref of
memref and make sure it is properly converted to the LLVM dialect. The type
support itself avoids adding the interface to the memref type itself similarly
to other built-in types. This allows the shape, and therefore byte size, of the
memref descriptor to remain a lowering aspect that is easier to customize and
evolve as opposed to sanctifying it in the data layout specification for the
memref type itself.

Factor out the code previously in a testing pass to live in a dedicated data
layout analysis and use that analysis in the conversion to compute the
allocation size for memref of memref. Other conversions will be ported
separately.

Depends On D103827

[ftynse]

cmake

[rriddle]

mlir/include/mlir/Analysis/DataLayoutAnalysis.h

28–35

Do we need two methods here? Can we just have the getAtOrAbove(albeit with a different name) variant and require the user to do op->getParentOp()?

39

What is the size of DataLayout? Should we store in a unique_ptr instead?

mlir/lib/Analysis/DataLayoutAnalysis.cpp

28

Why not store the default layout as a separate field? (Removing the need to look it up).

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1934–1938

Can you add a comment on where the size is coming from here?

4196

Can you remove the mlir:: here?

[ftynse]

Address review.

[ftynse]

mlir/include/mlir/Analysis/DataLayoutAnalysis.h

28–35

I expect getAbove to be used much more often than getAtOrAbove. The common request is "what is the data layout at the current operation, i.e. defined by its ancestors, not this op". Having to call getParentOp every time on it looks error-prone. getAtOrAbove is occasionally useful, and we can't assume that an op always has a child on which we could call getAbove.

39

Good idea!

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

1934–1938

I can do better and put this next to the code that defines what a memref descriptor is, the size is coming from there.

mlir/include/mlir/Analysis/DataLayoutAnalysis.h

+//===- DataLayoutAnalysis.h - API for Querying Nested Data Layout -*- C++ -*-=//
+//
+// 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_ANALYSIS_DATALAYOUTANALYSIS_H
+#define MLIR_ANALYSIS_DATALAYOUTANALYSIS_H
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/DenseMap.h"
+
+#include <memory>
+
+namespace mlir {
+
+class Operation;
+class DataLayout;
+
+/// Stores data layout objects for each operation that specifies the data layout
+/// above and below the given operation.
+class DataLayoutAnalysis {
+public:
+  /// Constructs the data layouts.
+  explicit DataLayoutAnalysis(Operation *root);
+
+  /// Returns the data layout active active at the given operation, that is the
+  /// data layout specified by the closest ancestor that can specify one, or the
+  /// default layout if there is no such ancestor.
+  const DataLayout &getAbove(Operation *operation) const;
+
+  /// Returns the data layout specified by the given operation or its closest
+  /// ancestor that can specify one.
+  const DataLayout &getAtOrAbove(Operation *operation) const;
+
+private:
+  /// Storage for individual data layouts.
+  DenseMap<Operation *, std::unique_ptr<DataLayout>> layouts;
+
+  /// Default data layout in case no operations specify one.
+  std::unique_ptr<DataLayout> defaultLayout;
+};
+
+} // namespace mlir
+
+#endif // MLIR_ANALYSIS_DATALAYOUTANALYSIS_H

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h

 class Module;
 class Type;
 } // namespace llvm
 
 namespace mlir {
 
 class BaseMemRefType;
 class ComplexType;
+class DataLayoutAnalysis;
 class LLVMTypeConverter;
 class UnrankedMemRefType;
 
 namespace LLVM {
 class LLVMDialect;
 class LLVMPointerType;
 } // namespace LLVM
 
class LLVMTypeConverter : public TypeConverter {
   friend LogicalResult
   structFuncArgTypeConverter(LLVMTypeConverter &converter, Type type,
                              SmallVectorImpl<Type> &result);
 
 public:
   using TypeConverter::convertType;
 
   /// Create an LLVMTypeConverter using the default LowerToLLVMOptions.
-  LLVMTypeConverter(MLIRContext *ctx);
-
-  /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
-  LLVMTypeConverter(MLIRContext *ctx, const LowerToLLVMOptions &options);
+  /// Optionally takes a data layout analysis to use in conversions.
+  LLVMTypeConverter(MLIRContext *ctx,
+                    const DataLayoutAnalysis *analysis = nullptr);
+
+  /// Create an LLVMTypeConverter using custom LowerToLLVMOptions. Optionally
+  /// takes a data layout analysis to use in conversions.
+  LLVMTypeConverter(MLIRContext *ctx, const LowerToLLVMOptions &options,
+                    const DataLayoutAnalysis *analysis = nullptr);
 
   /// Convert a function type.  The arguments and results are converted one by
   /// one and results are packed into a wrapped LLVM IR structure type. `result`
   /// is populated with argument mapping.
   Type convertFunctionSignature(FunctionType funcTy, bool isVariadic,
                                 SignatureConversion &result);
 
   /// Convert a non-empty list of types to be returned from a function into a
public:
   /// pointers to memref descriptors for arguments. Also converts the return
   /// type to a pointer argument if it is a struct. Returns true if this
   /// was the case.
   std::pair<Type, bool> convertFunctionTypeCWrapper(FunctionType type);
 
   /// Returns the data layout to use during and after conversion.
   const llvm::DataLayout &getDataLayout() { return options.dataLayout; }
 
+  /// Returns the data layout analysis to query during conversion.
+  const DataLayoutAnalysis *getDataLayoutAnalysis() const {
+    return dataLayoutAnalysis;
+  }
+
   /// Gets the LLVM representation of the index type. The returned type is an
   /// integer type with the size configured for this type converter.
   Type getIndexType();
 
   /// Gets the bitwidth of the index type when converted to LLVM.
   unsigned getIndexTypeBitwidth() { return options.getIndexBitwidth(); }
 
   /// Gets the pointer bitwidth.
   unsigned getPointerBitwidth(unsigned addressSpace = 0);
 
+  /// Returns the size of the memref descriptor object in bytes.
+  unsigned getMemRefDescriptorSize(MemRefType type, const DataLayout &layout);
+
+  /// Returns the size of the unranked memref descriptor object in bytes.
+  unsigned getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
+                                           const DataLayout &layout);
+
 protected:
   /// Pointer to the LLVM dialect.
   LLVM::LLVMDialect *llvmDialect;
 
 private:
   /// Convert a function type.  The arguments and results are converted one by
   /// one.  Additionally, if the function returns more than one value, pack the
   /// results into an LLVM IR structure type so that the converted function type
private:
 
   // Convert an unranked memref type to an LLVM type that captures the
   // runtime rank and a pointer to the static ranked memref desc
   Type convertUnrankedMemRefType(UnrankedMemRefType type);
 
   /// Convert a memref type to a bare pointer to the memref element type.
   Type convertMemRefToBarePtr(BaseMemRefType type);
 
-  // Convert a 1D vector type into an LLVM vector type.
+  /// Convert a 1D vector type into an LLVM vector type.
   Type convertVectorType(VectorType type);
 
   /// Options for customizing the llvm lowering.
   LowerToLLVMOptions options;
+
+  /// Data layout analysis mapping scopes to layouts active in them.
+  const DataLayoutAnalysis *dataLayoutAnalysis;
 };
 
 /// Helper class to produce LLVM dialect operations extracting or inserting
 /// values to a struct.
 class StructBuilder {
 public:
   /// Construct a helper for the given value.
   explicit StructBuilder(Value v);
protected:
   allocateBuffer(ConversionPatternRewriter &rewriter, Location loc,
                  Value sizeBytes, Operation *op) const = 0;
 
 private:
   static MemRefType getMemRefResultType(Operation *op) {
     return op->getResult(0).getType().cast<MemRefType>();
   }
 
-  LogicalResult match(Operation *op) const override {
-    MemRefType memRefType = getMemRefResultType(op);
-    return success(isConvertibleAndHasIdentityMaps(memRefType));
-  }
-
   // An `alloc` is converted into a definition of a memref descriptor value and
   // a call to `malloc` to allocate the underlying data buffer.  The memref
   // descriptor is of the LLVM structure type where:
   //   1. the first element is a pointer to the allocated (typed) data buffer,
   //   2. the second element is a pointer to the (typed) payload, aligned to the
   //      specified alignment,
   //   3. the remaining elements serve to store all the sizes and strides of the
   //      memref using LLVM-converted `index` type.
   //
   // Alignment is performed by allocating `alignment` more bytes than
   // requested and shifting the aligned pointer relative to the allocated
   // memory. Note: `alignment - <minimum malloc alignment>` would actually be
   // sufficient. If alignment is unspecified, the two pointers are equal.
 
   // An `alloca` is converted into a definition of a memref descriptor value and
   // an llvm.alloca to allocate the underlying data buffer.
-  void rewrite(Operation *op, ArrayRef<Value> operands,
+  LogicalResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 namespace LLVM {
 namespace detail {
 /// Replaces the given operation "op" with a new operation of type "targetOp"
 /// and given operands.
 LogicalResult oneToOneRewrite(Operation *op, StringRef targetOp,
                               ValueRange operands,

mlir/include/mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h

 /// to derive the bitwidth from the LLVM data layout.
 static constexpr unsigned kDeriveIndexBitwidthFromDataLayout = 0;
 
 /// Options to control the Standard dialect to LLVM lowering. The struct is used
 /// to share lowering options between passes, patterns, and type converter.
 class LowerToLLVMOptions {
 public:
   explicit LowerToLLVMOptions(MLIRContext *ctx);
-  explicit LowerToLLVMOptions(MLIRContext *ctx, DataLayout dl);
+  explicit LowerToLLVMOptions(MLIRContext *ctx, const DataLayout &dl);
 
   bool useBarePtrCallConv = false;
   bool emitCWrappers = false;
 
   enum class AllocLowering {
     /// Use malloc for for heap allocations.
     Malloc,
 

mlir/include/mlir/IR/BuiltinTypes.h

 // Deferred Method Definitions
 //===----------------------------------------------------------------------===//
 
 inline bool BaseMemRefType::classof(Type type) {
   return type.isa<MemRefType, UnrankedMemRefType>();
 }
 
 inline bool BaseMemRefType::isValidElementType(Type type) {
-  return type.isIntOrIndexOrFloat() || type.isa<ComplexType, VectorType>() ||
+  return type.isIntOrIndexOrFloat() ||
+         type.isa<ComplexType, MemRefType, VectorType, UnrankedMemRefType>() ||
          type.isa<MemRefElementTypeInterface>();
 }
 
 inline bool FloatType::classof(Type type) {
   return type.isa<BFloat16Type, Float16Type, Float32Type, Float64Type,
                   Float80Type, Float128Type>();
 }
 

mlir/include/mlir/IR/BuiltinTypes.td

let description = [{
     behavior.
 
     Are accepted as elements:
 
     - built-in integer types;
     - built-in index type;
     - built-in floating point types;
     - built-in vector types with elements of the above types;
+    - another memref type;
     - any other type implementing `MemRefElementTypeInterface`.
 
     ##### Codegen of Unranked Memref
 
     Using unranked memref in codegen besides the case mentioned above is highly
     discouraged. Codegen is concerned with generating loop nests and specialized
     instructions for high-performance, unranked memref is concerned with hiding
     the rank and thus, the number of enclosing loops required to iterate over

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   AffineAnalysis.cpp
   AffineStructures.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
+  DataLayoutAnalysis.cpp
   LinearTransform.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   NumberOfExecutions.cpp
   PresburgerSet.cpp
   SliceAnalysis.cpp
   Utils.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
   )
 
 add_mlir_library(MLIRAnalysis
   AliasAnalysis.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
+  DataLayoutAnalysis.cpp
   Liveness.cpp
   NumberOfExecutions.cpp
   SliceAnalysis.cpp
 
   AliasAnalysis/LocalAliasAnalysis.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis

mlir/lib/Analysis/DataLayoutAnalysis.cpp

+//===- DataLayoutAnalysis.cpp ---------------------------------------------===//
+//
+// 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/Analysis/DataLayoutAnalysis.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Interfaces/DataLayoutInterfaces.h"
+
+using namespace mlir;
+
+DataLayoutAnalysis::DataLayoutAnalysis(Operation *root)
+    : defaultLayout(std::make_unique<DataLayout>(DataLayoutOpInterface())) {
+  // Construct a DataLayout if possible from the op.
+  auto computeLayout = [this](Operation *op) {
+    if (auto iface = dyn_cast<DataLayoutOpInterface>(op))
+      layouts[op] = std::make_unique<DataLayout>(iface);
+    if (auto module = dyn_cast<ModuleOp>(op))
+      layouts[op] = std::make_unique<DataLayout>(module);
+  };
+
+  // Compute layouts for both ancestors and descendants.
+  root->walk(computeLayout);
+  for (Operation *ancestor = root->getParentOp(); ancestor != nullptr;
+       ancestor = ancestor->getParentOp()) {
+    computeLayout(ancestor);
+  }
+}
+
+const DataLayout &DataLayoutAnalysis::getAbove(Operation *operation) const {
+  for (Operation *ancestor = operation->getParentOp(); ancestor != nullptr;
+       ancestor = ancestor->getParentOp()) {
+    auto it = layouts.find(ancestor);
+    if (it != layouts.end())
+      return *it->getSecond();
+  }
+
+  // Fallback to the default layout.
+  return *defaultLayout;
+}
+
+const DataLayout &DataLayoutAnalysis::getAtOrAbove(Operation *operation) const {
+  auto it = layouts.find(operation);
+  if (it != layouts.end())
+    return *it->getSecond();
+  return getAbove(operation);
+}

mlir/lib/Conversion/StandardToLLVM/CMakeLists.txt

 add_mlir_conversion_library(MLIRStandardToLLVM
   StandardToLLVM.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Conversion/StandardToLLVM
 
   DEPENDS
   MLIRConversionPassIncGen
   intrinsics_gen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
+  MLIRAnalysis
   MLIRDataLayoutInterfaces
   MLIRLLVMIR
   MLIRMath
   MLIRMemRef
   MLIRTransforms
   )

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 //===- StandardToLLVM.cpp - Standard to LLVM dialect conversion -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements a pass to convert MLIR standard and builtin dialects
 // into the LLVM IR dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
+#include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
LogicalResult mlir::barePtrFuncArgTypeConverter(LLVMTypeConverter &converter,
   if (!llvmTy)
     return failure();
 
   result.push_back(llvmTy);
   return success();
 }
 
 /// Create an LLVMTypeConverter using default LowerToLLVMOptions.
-LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx)
-    : LLVMTypeConverter(ctx, LowerToLLVMOptions(ctx)) {}
+LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
+                                     const DataLayoutAnalysis *analysis)
+    : LLVMTypeConverter(ctx, LowerToLLVMOptions(ctx), analysis) {}
 
 /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
-                                     const LowerToLLVMOptions &options)
-    : llvmDialect(ctx->getOrLoadDialect<LLVM::LLVMDialect>()),
-      options(options) {
+                                     const LowerToLLVMOptions &options,
+                                     const DataLayoutAnalysis *analysis)
+    : llvmDialect(ctx->getOrLoadDialect<LLVM::LLVMDialect>()), options(options),
+      dataLayoutAnalysis(analysis) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
 
   // Register conversions for the builtin types.
   addConversion([&](ComplexType type) { return convertComplexType(type); });
   addConversion([&](FloatType type) { return convertFloatType(type); });
   addConversion([&](FunctionType type) { return convertFunctionType(type); });
   addConversion([&](IndexType type) { return convertIndexType(type); });
   addConversion([&](IntegerType type) { return convertIntegerType(type); });
LLVMTypeConverter::getMemRefDescriptorFields(MemRefType type,
 
   if (unpackAggregates)
     results.insert(results.end(), 2 * rank, indexTy);
   else
     results.insert(results.end(), 2, LLVM::LLVMArrayType::get(indexTy, rank));
   return results;
 }
 
+unsigned LLVMTypeConverter::getMemRefDescriptorSize(MemRefType type,
+                                                    const DataLayout &layout) {
+  // Compute the descriptor size given that of its components indicated above.
+  unsigned space = type.getMemorySpaceAsInt();
+  return 2 * llvm::divideCeil(getPointerBitwidth(space), 8) +
+         (1 + 2 * type.getRank()) * layout.getTypeSize(getIndexType());
+}
+
 /// Converts MemRefType to LLVMType. A MemRefType is converted to a struct that
 /// packs the descriptor fields as defined by `getMemRefDescriptorFields`.
 Type LLVMTypeConverter::convertMemRefType(MemRefType type) {
   // When converting a MemRefType to a struct with descriptor fields, do not
   // unpack the `sizes` and `strides` arrays.
   SmallVector<Type, 5> types =
       getMemRefDescriptorFields(type, /*unpackAggregates=*/false);
   if (types.empty())

 /// 2. void* ptr, pointer to the static ranked MemRef descriptor. This will be
 ///    stack allocated (alloca) copy of a MemRef descriptor that got casted to
 ///    be unranked.
 SmallVector<Type, 2> LLVMTypeConverter::getUnrankedMemRefDescriptorFields() {
   return {getIndexType(),
           LLVM::LLVMPointerType::get(IntegerType::get(&getContext(), 8))};
 }
 
+unsigned
+LLVMTypeConverter::getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
+                                                   const DataLayout &layout) {
+  // Compute the descriptor size given that of its components indicated above.
+  unsigned space = type.getMemorySpaceAsInt();
+  return layout.getTypeSize(getIndexType()) +
+         llvm::divideCeil(getPointerBitwidth(space), 8);
+}
+
 Type LLVMTypeConverter::convertUnrankedMemRefType(UnrankedMemRefType type) {
   if (!convertType(type.getElementType()))
     return {};
   return LLVM::LLVMStructType::getLiteral(&getContext(),
                                           getUnrankedMemRefDescriptorFields());
 }
 
 /// Convert a memref type to a bare pointer to the memref element type.
struct AllocOpLowering : public AllocLikeOpLLVMLowering {
   }
 };
 
 struct AlignedAllocOpLowering : public AllocLikeOpLLVMLowering {
   AlignedAllocOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLLVMLowering(memref::AllocOp::getOperationName(),
                                 converter) {}
 
-  /// Returns the memref's element size in bytes.
+  /// Returns the memref's element size in bytes using the data layout active at
+  /// `op`.
   // TODO: there are other places where this is used. Expose publicly?
-  static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
-    auto elementType = memRefType.getElementType();
-
-    unsigned sizeInBits;
-    if (elementType.isIntOrFloat()) {
-      sizeInBits = elementType.getIntOrFloatBitWidth();
-    } else {
-      auto vectorType = elementType.cast<VectorType>();
-      sizeInBits =
-          vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
-    }
-    return llvm::divideCeil(sizeInBits, 8);
+  unsigned getMemRefEltSizeInBytes(MemRefType memRefType, Operation *op) const {
+    const DataLayout *layout = &defaultLayout;
+    if (const DataLayoutAnalysis *analysis =
+            getTypeConverter()->getDataLayoutAnalysis()) {
+      layout = &analysis->getAbove(op);
+    }
+    Type elementType = memRefType.getElementType();
+    if (auto memRefElementType = elementType.dyn_cast<MemRefType>())
+      return getTypeConverter()->getMemRefDescriptorSize(memRefElementType,
+                                                         *layout);
+    if (auto memRefElementType = elementType.dyn_cast<UnrankedMemRefType>())
+      return getTypeConverter()->getUnrankedMemRefDescriptorSize(
+          memRefElementType, *layout);
+    return layout->getTypeSize(elementType);
   }
 
   /// Returns true if the memref size in bytes is known to be a multiple of
-  /// factor.
-  static bool isMemRefSizeMultipleOf(MemRefType type, uint64_t factor) {
-    uint64_t sizeDivisor = getMemRefEltSizeInBytes(type);
+  /// factor assuming the data layout active at `op`.
+  bool isMemRefSizeMultipleOf(MemRefType type, uint64_t factor,
+                              Operation *op) const {
+    uint64_t sizeDivisor = getMemRefEltSizeInBytes(type, op);
     for (unsigned i = 0, e = type.getRank(); i < e; i++) {
       if (type.isDynamic(type.getDimSize(i)))
         continue;
       sizeDivisor = sizeDivisor * type.getDimSize(i);
     }
     return sizeDivisor % factor == 0;
   }
 
   /// Returns the alignment to be used for the allocation call itself.
   /// aligned_alloc requires the allocation size to be a power of two, and the
   /// allocation size to be a multiple of alignment,
   int64_t getAllocationAlignment(memref::AllocOp allocOp) const {
     if (Optional<uint64_t> alignment = allocOp.alignment())
       return *alignment;
 
     // Whenever we don't have alignment set, we will use an alignment
     // consistent with the element type; since the allocation size has to be a
     // power of two, we will bump to the next power of two if it already isn't.
-    auto eltSizeBytes = getMemRefEltSizeInBytes(allocOp.getType());
+    auto eltSizeBytes = getMemRefEltSizeInBytes(allocOp.getType(), allocOp);
     return std::max(kMinAlignedAllocAlignment,
                     llvm::PowerOf2Ceil(eltSizeBytes));
   }
 
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
     // Heap allocations.
     memref::AllocOp allocOp = cast<memref::AllocOp>(op);
     MemRefType memRefType = allocOp.getType();
     int64_t alignment = getAllocationAlignment(allocOp);
     Value allocAlignment = createIndexConstant(rewriter, loc, alignment);
 
     // aligned_alloc requires size to be a multiple of alignment; we will pad
     // the size to the next multiple if necessary.
-    if (!isMemRefSizeMultipleOf(memRefType, alignment))
+    if (!isMemRefSizeMultipleOf(memRefType, alignment, op))
       sizeBytes = createAligned(rewriter, loc, sizeBytes, allocAlignment);
 
     Type elementPtrType = this->getElementPtrType(memRefType);
     auto allocFuncOp = LLVM::lookupOrCreateAlignedAllocFn(
         allocOp->getParentOfType<ModuleOp>(), getIndexType());
     auto results =
         createLLVMCall(rewriter, loc, allocFuncOp, {allocAlignment, sizeBytes},
                        getVoidPtrType());
     Value allocatedPtr =
         rewriter.create<LLVM::BitcastOp>(loc, elementPtrType, results[0]);
 
     return std::make_tuple(allocatedPtr, allocatedPtr);
   }
 
   /// The minimum alignment to use with aligned_alloc (has to be a power of 2).
   static constexpr uint64_t kMinAlignedAllocAlignment = 16UL;
+
+  /// Default layout to use in absence of the corresponding analysis.
+  DataLayout defaultLayout;
 };
 
 // Out of line definition, required till C++17.
 constexpr uint64_t AlignedAllocOpLowering::kMinAlignedAllocAlignment;
 
 struct AllocaOpLowering : public AllocLikeOpLLVMLowering {
   AllocaOpLowering(LLVMTypeConverter &converter)
       : AllocLikeOpLLVMLowering(memref::AllocaOp::getOperationName(),
if (failed(LLVM::LLVMDialect::verifyDataLayoutString(
             this->dataLayout, [this](const Twine &message) {
               getOperation().emitError() << message.str();
             }))) {
       signalPassFailure();
       return;
     }
 
     ModuleOp m = getOperation();
+    const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
 
-    LowerToLLVMOptions options(&getContext(), DataLayout(m));
+    LowerToLLVMOptions options(&getContext(),
+                               dataLayoutAnalysis.getAtOrAbove(m));
     options.useBarePtrCallConv = useBarePtrCallConv;
     options.emitCWrappers = emitCWrappers;
     if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout)
       options.overrideIndexBitwidth(indexBitwidth);
     options.allocLowering =
         (useAlignedAlloc ? LowerToLLVMOptions::AllocLowering::AlignedAlloc
                          : LowerToLLVMOptions::AllocLowering::Malloc);
     options.dataLayout = llvm::DataLayout(this->dataLayout);
-    LLVMTypeConverter typeConverter(&getContext(), options);
+
+    LLVMTypeConverter typeConverter(&getContext(), options,
+                                    &dataLayoutAnalysis);
 
     RewritePatternSet patterns(&getContext());
     populateStdToLLVMConversionPatterns(typeConverter, patterns);
 
     LLVMConversionTarget target(getContext());
     if (failed(applyPartialConversion(m, target, std::move(patterns))))
       signalPassFailure();
     m->setAttr(LLVM::LLVMDialect::getDataLayoutAttrName(),
                StringAttr::get(m.getContext(), this->dataLayout));
   }
 };
 } // end namespace
 
 Value AllocLikeOpLLVMLowering::createAligned(
     ConversionPatternRewriter &rewriter, Location loc, Value input,
     Value alignment) {
   Value one = createIndexAttrConstant(rewriter, loc, alignment.getType(), 1);
   Value bump = rewriter.create<LLVM::SubOp>(loc, alignment, one);
   Value bumped = rewriter.create<LLVM::AddOp>(loc, input, bump);
   Value mod = rewriter.create<LLVM::URemOp>(loc, bumped, alignment);
   return rewriter.create<LLVM::SubOp>(loc, bumped, mod);
 }
 
-void AllocLikeOpLLVMLowering::rewrite(
+LogicalResult AllocLikeOpLLVMLowering::matchAndRewrite(
     Operation *op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
   MemRefType memRefType = getMemRefResultType(op);
+  if (!isConvertibleAndHasIdentityMaps(memRefType))
+    return rewriter.notifyMatchFailure(op, "incompatible memref type");
   auto loc = op->getLoc();
 
   // Get actual sizes of the memref as values: static sizes are constant
   // values and dynamic sizes are passed to 'alloc' as operands.  In case of
   // zero-dimensional memref, assume a scalar (size 1).
   SmallVector<Value, 4> sizes;
   SmallVector<Value, 4> strides;
   Value sizeBytes;
   this->getMemRefDescriptorSizes(loc, memRefType, operands, rewriter, sizes,
                                  strides, sizeBytes);
 
   // Allocate the underlying buffer.
   Value allocatedPtr;
   Value alignedPtr;
   std::tie(allocatedPtr, alignedPtr) =
       this->allocateBuffer(rewriter, loc, sizeBytes, op);
 
   // Create the MemRef descriptor.
   auto memRefDescriptor = this->createMemRefDescriptor(
       loc, memRefType, allocatedPtr, alignedPtr, sizes, strides, rewriter);
 
   // Return the final value of the descriptor.
   rewriter.replaceOp(op, {memRefDescriptor});
+  return success();
 }
 
 mlir::LLVMConversionTarget::LLVMConversionTarget(MLIRContext &ctx)
     : ConversionTarget(ctx) {
   this->addLegalDialect<LLVM::LLVMDialect>();
   this->addIllegalOp<LLVM::DialectCastOp>();
 }
 
return std::make_unique<LLVMLoweringPass>(
       options.useBarePtrCallConv, options.emitCWrappers,
       options.getIndexBitwidth(), useAlignedAlloc, options.dataLayout);
 }
 
 mlir::LowerToLLVMOptions::LowerToLLVMOptions(MLIRContext *ctx)
     : LowerToLLVMOptions(ctx, DataLayout()) {}
 
 mlir::LowerToLLVMOptions::LowerToLLVMOptions(MLIRContext *ctx,
-                                             mlir::DataLayout dl) {
+                                             const DataLayout &dl) {
   indexBitwidth = dl.getTypeSizeInBits(IndexType::get(ctx));
 }

mlir/test/Conversion/StandardToLLVM/convert-dynamic-memref-ops.mlir

-// RUN: mlir-opt -convert-std-to-llvm %s | FileCheck %s
-// RUN: mlir-opt -convert-std-to-llvm='use-aligned-alloc=1' %s | FileCheck %s --check-prefix=ALIGNED-ALLOC
+// RUN: mlir-opt -split-input-file -convert-std-to-llvm %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -convert-std-to-llvm='use-aligned-alloc=1' %s | FileCheck %s --check-prefix=ALIGNED-ALLOC
 
 // CHECK-LABEL: func @check_strided_memref_arguments(
 // CHECK-COUNT-2: !llvm.ptr<f32>
 // CHECK-COUNT-5: i64
 // CHECK-COUNT-2: !llvm.ptr<f32>
 // CHECK-COUNT-5: i64
 // CHECK-COUNT-2: !llvm.ptr<f32>
 // CHECK-COUNT-5: i64

 // CHECK:   [[TARGET_STRIDE_PTR:%.*]] = llvm.getelementptr [[STRIDES_PTR]]{{\[}}[[DIM]]]
 // CHECK:   llvm.store [[CUR_STRIDE]], [[TARGET_STRIDE_PTR]] : !llvm.ptr<i64>
 // CHECK:   [[UPDATE_STRIDE:%.*]] = llvm.mul [[CUR_STRIDE]], [[SIZE]] : i64
 // CHECK:   [[STRIDE_COND:%.*]] = llvm.sub [[DIM]], [[C1_]] : i64
 // CHECK:   llvm.br ^bb1([[STRIDE_COND]], [[UPDATE_STRIDE]] : i64, i64)
 
 // CHECK: ^bb3:
 // CHECK:   llvm.return
+
+// -----
+
+// ALIGNED-ALLOC-LABEL: @memref_of_memref
+func @memref_of_memref() {
+  // Sizeof computation is as usual.
+  // ALIGNED-ALLOC: %[[NULL:.*]] = llvm.mlir.null
+  // ALIGNED-ALLOC: %[[PTR:.*]] = llvm.getelementptr
+  // ALIGNED-ALLOC: %[[SIZEOF:.*]] = llvm.ptrtoint
+
+  // Static alignment should be computed as ceilPowerOf2(2 * sizeof(pointer) +
+  // (1 + 2 * rank) * sizeof(index) = ceilPowerOf2(2 * 8 + 3 * 8) = 64.
+  // ALIGNED-ALLOC: llvm.mlir.constant(64 : index)
+
+  // Check that the types are converted as expected.
+  // ALIGNED-ALLOC: llvm.call @aligned_alloc
+  // ALIGNED-ALLOC: llvm.bitcast %{{.*}} : !llvm.ptr<i8> to
+  // ALIGNED-ALLOC-SAME: !llvm.
+  // ALIGNED-ALLOC-SAME: [[INNER:ptr<struct<\(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>\)>>]]
+  // ALIGNED-ALLOC: llvm.mlir.undef
+  // ALIGNED-ALLOC-SAME: !llvm.struct<([[INNER]], [[INNER]], i64, array<1 x i64>, array<1 x i64>)>
+  %0 = memref.alloc() : memref<1xmemref<1xf32>>
+  return
+}
+
+// -----
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 32>> } {
+  // ALIGNED-ALLOC-LABEL: @memref_of_memref_32
+  func @memref_of_memref_32() {
+    // Sizeof computation is as usual.
+    // ALIGNED-ALLOC: %[[NULL:.*]] = llvm.mlir.null
+    // ALIGNED-ALLOC: %[[PTR:.*]] = llvm.getelementptr
+    // ALIGNED-ALLOC: %[[SIZEOF:.*]] = llvm.ptrtoint
+
+    // Static alignment should be computed as ceilPowerOf2(2 * sizeof(pointer) +
+    // (1 + 2 * rank) * sizeof(index) = ceilPowerOf2(2 * 8 + 3 * 4) = 32.
+    // ALIGNED-ALLOC: llvm.mlir.constant(32 : index)
+
+    // Check that the types are converted as expected.
+    // ALIGNED-ALLOC: llvm.call @aligned_alloc
+    // ALIGNED-ALLOC: llvm.bitcast %{{.*}} : !llvm.ptr<i8> to
+    // ALIGNED-ALLOC-SAME: !llvm.
+    // ALIGNED-ALLOC-SAME: [[INNER:ptr<struct<\(ptr<f32>, ptr<f32>, i32, array<1 x i32>, array<1 x i32>\)>>]]
+    // ALIGNED-ALLOC: llvm.mlir.undef
+    // ALIGNED-ALLOC-SAME: !llvm.struct<([[INNER]], [[INNER]], i32, array<1 x i32>, array<1 x i32>)>
+    %0 = memref.alloc() : memref<1xmemref<1xf32>>
+    return
+  }
+}
+
+
+// -----
+
+// ALIGNED-ALLOC-LABEL: @memref_of_memref_of_memref
+func @memref_of_memref_of_memref() {
+  // Sizeof computation is as usual, also check the type.
+  // ALIGNED-ALLOC: %[[NULL:.*]] = llvm.mlir.null : !llvm.ptr<
+  // ALIGNED-ALLOC-SAME:   struct<(
+  // ALIGNED-ALLOC-SAME:     [[INNER:ptr<struct<\(ptr<f32>, ptr<f32>, i64, array<1 x i64>, array<1 x i64>\)>>]],
+  // ALIGNED-ALLOC-SAME:     [[INNER]],
+  // ALIGNED-ALLOC-SAME:     i64, array<1 x i64>, array<1 x i64>
+  // ALIGNED-ALLOC-SAME:   )>
+  // ALIGNED-ALLOC-SAME: >
+  // ALIGNED-ALLOC: %[[PTR:.*]] = llvm.getelementptr
+  // ALIGNED-ALLOC: %[[SIZEOF:.*]] = llvm.ptrtoint
+
+  // Static alignment should be computed as ceilPowerOf2(2 * sizeof(pointer) +
+  // (1 + 2 * rank) * sizeof(index) = ceilPowerOf2(2 * 8 + 3 * 8) = 64.
+  // ALIGNED-ALLOC: llvm.mlir.constant(64 : index)
+  // ALIGNED-ALLOC: llvm.call @aligned_alloc
+  %0 = memref.alloc() : memref<1 x memref<2 x memref<3 x f32>>>
+  return
+}
+
+// -----
+
+// ALIGNED-ALLOC-LABEL: @ranked_unranked
+func @ranked_unranked() {
+  // ALIGNED-ALLOC: llvm.mlir.null
+  // ALIGNED-ALLOC-SAME: !llvm.[[INNER:ptr<struct<\(i64, ptr<i8>\)>>]]
+  // ALIGNED-ALLOC: llvm.getelementptr
+  // ALIGNED-ALLOC: llvm.ptrtoint
+
+  // Static alignment should be computed as ceilPowerOf2(sizeof(index) +
+  // sizeof(pointer)) = 16.
+  // ALIGNED-ALLOC: llvm.mlir.constant(16 : index)
+  // ALIGNED-ALLOC: llvm.call @aligned_alloc
+  // ALIGNED-ALLOC: llvm.bitcast
+  // ALIGNED-ALLOC-SAME: !llvm.ptr<i8> to !llvm.[[INNER]]
+  %0 = memref.alloc() : memref<1 x memref<* x f32>>
+  memref.cast %0 : memref<1 x memref<* x f32>> to memref<* x memref<* x f32>>
+  return
+}
+

mlir/test/Conversion/StandardToLLVM/convert-static-memref-ops.mlir

 
 // CHECK: @unsupported_unranked_memref_element_type
 // CHECK-SAME: memref<
 // CHECK-NOT: !llvm.struct
 // BAREPTR: @unsupported_unranked_memref_element_type
 // BAREPTR-SAME: memref<
 // BAREPTR-NOT: !llvm.ptr
 func private @unsupported_unranked_memref_element_type() -> memref<* x !test.memref_element>
+

mlir/test/IR/parser.mlir

 func private @memref_with_complex_elems(memref<1x?xcomplex<f32>>)
 
 // CHECK: func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 func private @memref_with_vector_elems(memref<1x?xvector<10xf32>>)
 
 // CHECK: func private @memref_with_custom_elem(memref<1x?x!test.memref_element>)
 func private @memref_with_custom_elem(memref<1x?x!test.memref_element>)
 
+// CHECK: func private @memref_of_memref(memref<1xmemref<1xf64>>)
+func private @memref_of_memref(memref<1xmemref<1xf64>>)
+
+// CHECK: func private @memref_of_unranked_memref(memref<1xmemref<*xf32>>)
+func private @memref_of_unranked_memref(memref<1xmemref<*xf32>>)
+
+// CHECK: func private @unranked_memref_of_memref(memref<*xmemref<1xf32>>)
+func private @unranked_memref_of_memref(memref<*xmemref<1xf32>>)
+
+// CHECK: func private @unranked_memref_of_unranked_memref(memref<*xmemref<*xi32>>)
+func private @unranked_memref_of_unranked_memref(memref<*xmemref<*xi32>>)
+
 // CHECK: func private @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 func private @unranked_memref_with_complex_elems(memref<*xcomplex<f32>>)
 
 // CHECK: func private @unranked_memref_with_index_elems(memref<*xindex>)
 func private @unranked_memref_with_index_elems(memref<*xindex>)
 
 // CHECK: func private @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)
 func private @unranked_memref_with_vector_elems(memref<*xvector<10xf32>>)

mlir/test/lib/Dialect/DLTI/TestDataLayoutQuery.cpp

 //===- TestDataLayoutQuery.cpp - Test Data Layout Queries -----------------===//
 //
 // 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 "TestDialect.h"
+#include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Dialect/DLTI/DLTI.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 
 /// A pass that finds "test.data_layout_query" operations and attaches to them
 /// attributes containing the results of data layout queries for operation
 /// result types.
 struct TestDataLayoutQuery
     : public PassWrapper<TestDataLayoutQuery, FunctionPass> {
   void runOnFunction() override {
     FuncOp func = getFunction();
     Builder builder(func.getContext());
-    DenseMap<Operation *, DataLayout> layouts;
+    const DataLayoutAnalysis &layouts = getAnalysis<DataLayoutAnalysis>();
 
     func.walk([&](test::DataLayoutQueryOp op) {
       // Skip the ops with already processed in a deeper call.
       if (op->getAttr("size"))
         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(closest)->getSecond();
+      const DataLayout &layout = layouts.getAbove(op);
       unsigned size = layout.getTypeSize(op.getType());
       unsigned bitsize = layout.getTypeSizeInBits(op.getType());
       unsigned alignment = layout.getTypeABIAlignment(op.getType());
       unsigned preferred = layout.getTypePreferredAlignment(op.getType());
       op->setAttrs(
           {builder.getNamedAttr("size", builder.getIndexAttr(size)),
            builder.getNamedAttr("bitsize", builder.getIndexAttr(bitsize)),
            builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)),