[mlir] Propagate alignment annotations to LLVM loads and stores.

[timshen]

Backport.

[timshen]

Fix an arc mistake.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: diff.json, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62144 tests passed, 7 failed and 811 were skipped.

failed: MLIR.Conversion/StandardToLLVM/convert-memref-ops.mlir
failed: MLIR.Transforms/memref-normalize.mlir
failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: fail. 62144 tests passed, 7 failed and 811 were skipped.

failed: MLIR.Conversion/StandardToLLVM/convert-memref-ops.mlir
failed: MLIR.Transforms/memref-normalize.mlir
failed: libc++.std/language_support/cmp/cmp_partialord/partialord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongeq/cmp.strongeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_strongord/strongord.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakeq/cmp.weakeq.pass.cpp
failed: libc++.std/language_support/cmp/cmp_weakord/weakord.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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

static MemRefDescriptor fromStaticShape(OpBuilder &builder, Location loc,
                                           MemRefType type, Value memory);
 
   /// Builds IR extracting the allocated pointer from the descriptor.
   Value allocatedPtr(OpBuilder &builder, Location loc);
   /// Builds IR inserting the allocated pointer into the descriptor.
   void setAllocatedPtr(OpBuilder &builder, Location loc, Value ptr);
 
   /// Builds IR extracting the aligned pointer from the descriptor.
-  Value alignedPtr(OpBuilder &builder, Location loc);
+  Value alignedPtr(OpBuilder &builder, unsigned alignment, Location loc);
 
   /// Builds IR inserting the aligned pointer into the descriptor.
   void setAlignedPtr(OpBuilder &builder, Location loc, Value ptr);
 
   /// Builds IR extracting the offset from the descriptor.
   Value offset(OpBuilder &builder, Location loc);
 
   /// Builds IR inserting the offset into the descriptor.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

def LLVM_Prefetch : LLVM_ZeroResultOp<"intr.prefetch">,
   let llvmBuilder = [{
     llvm::Module *module = builder.GetInsertBlock()->getModule();
     llvm::Function *fn = llvm::Intrinsic::getDeclaration(
         module, llvm::Intrinsic::prefetch, $addr->getType());
     builder.CreateCall(fn, {$addr, $rw, $hint, $cache});
   }];
 }
 
+def LLVM_AssumeOp : LLVM_Op<"intr.assume", []>,
+                    Arguments<(ins LLVM_Type:$cond)> {
+  let llvmBuilder = [{
+    llvm::Module *module = builder.GetInsertBlock()->getModule();
+    llvm::Function *fn =
+        llvm::Intrinsic::getDeclaration(module, llvm::Intrinsic::assume, {});
+    builder.CreateCall(fn, {$cond});
+  }];
+}
+
 def AtomicBinOpXchg : I64EnumAttrCase<"xchg", 0>;
 def AtomicBinOpAdd  : I64EnumAttrCase<"add", 1>;
 def AtomicBinOpSub  : I64EnumAttrCase<"sub", 2>;
 def AtomicBinOpAnd  : I64EnumAttrCase<"_and", 3>;
 def AtomicBinOpNand : I64EnumAttrCase<"nand", 4>;
 def AtomicBinOpOr   : I64EnumAttrCase<"_or", 5>;
 def AtomicBinOpXor  : I64EnumAttrCase<"_xor", 6>;
 def AtomicBinOpMax  : I64EnumAttrCase<"max", 7>;

mlir/include/mlir/Dialect/StandardOps/Ops.td

let description = [{
     The optional list of symbol operands are bound to the symbols of the
     memrefs affine map. In the example below, the ssa value '%s' is bound to
     the symbol 's0' in the affine map specified in the allocs memref type.
 
       %0 = alloc()[%s] : memref<8x64xf32, (d0, d1)[s0] -> ((d0 + s0), d1), 1>
 
     This operation returns a single ssa value of memref type, which can be used
     by subsequent load and store operations.
-
-    The optional `alignment` attribute may be specified to ensure that the
-    region of memory that will be indexed is aligned at the specified byte
-    boundary. TODO(b/144281289) optional alignment attribute to MemRefType.
-
-      %0 = alloc()[%s] {alignment = 8} :
-        memref<8x64xf32, (d0, d1)[s0] -> ((d0 + s0), d1), 1>
   }];
 
-  let arguments = (ins Variadic<Index>:$value,
-                   Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment);
+  let arguments = (ins Variadic<Index>:$value);
   let results = (outs AnyMemRef);
 
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, MemRefType memrefType", [{
        result.types.push_back(memrefType);
      }]>,
     OpBuilder<
     "Builder *builder, OperationState &result, MemRefType memrefType, " #
-    "ArrayRef<Value> operands, IntegerAttr alignment = IntegerAttr()", [{
+    "ArrayRef<Value> operands", [{
        result.addOperands(operands);
        result.types.push_back(memrefType);
-       if (alignment)
-         result.addAttribute(getAlignmentAttrName(), alignment);
      }]>];
 
   let extraClassDeclaration = [{
-    static StringRef getAlignmentAttrName() { return "alignment"; }
-
     MemRefType getType() { return getResult().getType().cast<MemRefType>(); }
 
     /// Returns the number of symbolic operands (the ones in square brackets),
     /// which bind to the symbols of the memref's layout map.
     unsigned getNumSymbolicOperands() {
       return getNumOperands() - getType().getNumDynamicDims();
     }
 

mlir/include/mlir/IR/StandardTypes.h

public:
   using Base::Base;
 
   /// Get or create a new MemRefType based on shape, element type, affine
   /// map composition, and memory space.  Assumes the arguments define a
   /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
   /// construction failures.
   static MemRefType get(ArrayRef<int64_t> shape, Type elementType,
                         ArrayRef<AffineMap> affineMapComposition = {},
-                        unsigned memorySpace = 0);
+                        unsigned memorySpace = 0, unsigned alignment = 0);
 
   /// Get or create a new MemRefType based on shape, element type, affine
   /// map composition, and memory space declared at the given location.
   /// If the location is unknown, the last argument should be an instance of
   /// UnknownLoc.  If the MemRefType defined by the arguments would be
   /// ill-formed, emits errors (to the handler registered with the context or to
   /// the error stream) and returns nullptr.
   static MemRefType getChecked(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
-                               unsigned memorySpace, Location location);
+                               unsigned memorySpace, unsigned alignment,
+                               Location location);
 
   ArrayRef<int64_t> getShape() const;
 
   /// Returns an array of affine map pointers representing the memref affine
   /// map composition.
   ArrayRef<AffineMap> getAffineMaps() const;
 
   /// Returns the memory space in which data referred to by this memref resides.
   unsigned getMemorySpace() const;
 
+  /// Returns the alignment the underlying buffer is aligned to.
+  unsigned getAlignment() const;
+
   // TODO(ntv): merge these two special values in a single one used everywhere.
   // Unfortunately, uses of `-1` have crept deep into the codebase now and are
   // hard to track.
   static constexpr int64_t kDynamicSize = -1;
   static int64_t getDynamicStrideOrOffset() {
     return ShapedType::kDynamicStrideOrOffset;
   }
 
   static bool kindof(unsigned kind) { return kind == StandardTypes::MemRef; }
 
 private:
   /// Get or create a new MemRefType defined by the arguments.  If the resulting
   /// type would be ill-formed, return nullptr.  If the location is provided,
   /// emit detailed error messages.
   static MemRefType getImpl(ArrayRef<int64_t> shape, Type elementType,
                             ArrayRef<AffineMap> affineMapComposition,
-                            unsigned memorySpace, Optional<Location> location);
+                            unsigned memorySpace, unsigned alignment,
+                            Optional<Location> location);
   using Base::getImpl;
 };
 
+class MemRefTypeBuilder {
+  ArrayRef<int64_t> shape;
+  Type elementType;
+  ArrayRef<AffineMap> affineMaps;
+  unsigned memorySpace;
+  unsigned alignment;
+
+public:
+  // Builder that clones from another MemRefType.
+  explicit MemRefTypeBuilder(MemRefType other)
+      : shape(other.getShape()), elementType(other.getElementType()),
+        affineMaps(other.getAffineMaps()), memorySpace(other.getMemorySpace()),
+        alignment(other.getAlignment()) {}
+
+  MemRefTypeBuilder(ArrayRef<int64_t> shape, Type elementType)
+      : shape(shape), elementType(elementType), affineMaps(), memorySpace(0),
+        alignment(0) {}
+
+  MemRefTypeBuilder &setShape(ArrayRef<int64_t> newShape) {
+    shape = newShape;
+    return *this;
+  }
+
+  MemRefTypeBuilder &setElementType(Type newElementType) {
+    elementType = newElementType;
+    return *this;
+  }
+
+  MemRefTypeBuilder &setAffineMaps(ArrayRef<AffineMap> newAffineMaps) {
+    affineMaps = newAffineMaps;
+    return *this;
+  }
+
+  MemRefTypeBuilder &setMemorySpace(unsigned newMemorySpace) {
+    memorySpace = newMemorySpace;
+    return *this;
+  }
+
+  MemRefTypeBuilder &setAlignment(unsigned newAlignment) {
+    alignment = newAlignment;
+    return *this;
+  }
+
+  operator MemRefType() {
+    return MemRefType::get(shape, elementType, affineMaps, memorySpace,
+                           alignment);
+  }
+};
+
 /// Unranked MemRef type represent multi-dimensional MemRefs that
 /// have an unknown rank.
 class UnrankedMemRefType
     : public Type::TypeBase<UnrankedMemRefType, BaseMemRefType,
                             detail::UnrankedMemRefTypeStorage> {
 public:
   using Base::Base;
 

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

 class NVVMTypeConverter : public LLVMTypeConverter {
 public:
   using LLVMTypeConverter::LLVMTypeConverter;
 
   Type convertType(Type type) override {
     auto memref = type.dyn_cast<MemRefType>();
     if (memref &&
         memref.getMemorySpace() == gpu::GPUDialect::getPrivateAddressSpace()) {
-      type = MemRefType::get(memref.getShape(), memref.getElementType(),
-                             memref.getAffineMaps());
+      type = MemRefTypeBuilder(memref).setMemorySpace(0);
     }
 
     return LLVMTypeConverter::convertType(type);
   }
 };
 
 /// Converts all_reduce op to LLVM/NVVM ops.
 struct GPUAllReduceOpLowering : public LLVMOpLowering {

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

public:
   BaseViewConversionHelper(Type type)
       : d(MemRefDescriptor::undef(rewriter(), loc(), type)) {}
 
   BaseViewConversionHelper(Value v) : d(v) {}
 
   /// Wrappers around MemRefDescriptor that use EDSC builder and location.
   Value allocatedPtr() { return d.allocatedPtr(rewriter(), loc()); }
   void setAllocatedPtr(Value v) { d.setAllocatedPtr(rewriter(), loc(), v); }
-  Value alignedPtr() { return d.alignedPtr(rewriter(), loc()); }
+  Value alignedPtr() { return d.alignedPtr(rewriter(), 0, loc()); }
   void setAlignedPtr(Value v) { d.setAlignedPtr(rewriter(), loc(), v); }
   Value offset() { return d.offset(rewriter(), loc()); }
   void setOffset(Value v) { d.setOffset(rewriter(), loc(), v); }
   Value size(unsigned i) { return d.size(rewriter(), loc(), i); }
   void setSize(unsigned i, Value v) { d.setSize(rewriter(), loc(), i, v); }
   void setConstantSize(unsigned i, int64_t v) {
     d.setConstantSize(rewriter(), loc(), i, v);
   }

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

 }
 
 /// Builds IR inserting the allocated pointer into the descriptor.
 void MemRefDescriptor::setAllocatedPtr(OpBuilder &builder, Location loc,
                                        Value ptr) {
   setPtr(builder, loc, kAllocatedPtrPosInMemRefDescriptor, ptr);
 }
 
+// Creates a constant Op producing a value of `resultType` from an index-typed
+// integer attribute.
+static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
+                                     Type resultType, int64_t value) {
+  return builder.create<LLVM::ConstantOp>(
+      loc, resultType, builder.getIntegerAttr(builder.getIndexType(), value));
+}
+
 /// Builds IR extracting the aligned pointer from the descriptor.
-Value MemRefDescriptor::alignedPtr(OpBuilder &builder, Location loc) {
-  return extractPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor);
+Value MemRefDescriptor::alignedPtr(OpBuilder &builder, unsigned alignment,
+                                   Location loc) {
+  Value ptr = extractPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor);
+  if (alignment) {
+    assert(((alignment - 1) & alignment) == 0 &&
+           "Alignments must be power of 2");
+    builder.create<LLVM::AssumeOp>(
+        loc, builder.create<LLVM::ICmpOp>(
+                 loc, LLVM::ICmpPredicate::eq,
+                 builder.create<LLVM::AndOp>(
+                     loc, builder.create<LLVM::PtrToIntOp>(loc, indexType, ptr),
+                     createIndexAttrConstant(builder, loc, indexType,
+                                             alignment - 1)),
+                 createIndexAttrConstant(builder, loc, indexType, 0)));
+  }
+  return ptr;
 }
 
 /// Builds IR inserting the aligned pointer into the descriptor.
 void MemRefDescriptor::setAlignedPtr(OpBuilder &builder, Location loc,
                                      Value ptr) {
   setPtr(builder, loc, kAlignedPtrPosInMemRefDescriptor, ptr);
 }
 
-// Creates a constant Op producing a value of `resultType` from an index-typed
-// integer attribute.
-static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
-                                     Type resultType, int64_t value) {
-  return builder.create<LLVM::ConstantOp>(
-      loc, resultType, builder.getIntegerAttr(builder.getIndexType(), value));
-}
-
 /// Builds IR extracting the offset from the descriptor.
 Value MemRefDescriptor::offset(OpBuilder &builder, Location loc) {
   return builder.create<LLVM::ExtractValueOp>(
       loc, indexType, value,
       builder.getI64ArrayAttr(kOffsetPosInMemRefDescriptor));
 }
 
 /// Builds IR inserting the offset into the descriptor.
void rewrite(Operation *op, ArrayRef<Value> operands,
     auto elementSize =
         rewriter.create<LLVM::PtrToIntOp>(loc, getIndexType(), gep);
     cumulativeSize = rewriter.create<LLVM::MulOp>(
         loc, getIndexType(), ArrayRef<Value>{cumulativeSize, elementSize});
 
     // Allocate the underlying buffer and store a pointer to it in the MemRef
     // descriptor.
     Value allocated = nullptr;
-    int alignment = 0;
+    int alignment = type.getAlignment();
     Value alignmentValue = nullptr;
-    if (auto alignAttr = allocOp.alignment())
-      alignment = alignAttr.getValue().getSExtValue();
 
     if (useAlloca) {
       allocated = rewriter.create<LLVM::AllocaOp>(loc, getVoidPtrType(),
                                                   cumulativeSize, alignment);
     } else {
       // Insert the `malloc` declaration if it is not already present.
       auto module = op->getParentOfType<ModuleOp>();
       auto mallocFunc = module.lookupSymbol<LLVM::LLVMFuncOp>("malloc");
for (int i = 1, nSizes = allocSizes.size(); i < nSizes; ++i) {
           loc, this->getIndexType(), ArrayRef<Value>{linearized, indices[i]});
     }
     return linearized;
   }
 
   // This is a strided getElementPtr variant that linearizes subscripts as:
   //   `base_offset + index_0 * stride_0 + ... + index_n * stride_n`.
   Value getStridedElementPtr(Location loc, Type elementTypePtr,
-                             Value descriptor, ArrayRef<Value> indices,
-                             ArrayRef<int64_t> strides, int64_t offset,
+                             unsigned alignment, Value descriptor,
+                             ArrayRef<Value> indices, ArrayRef<int64_t> strides,
+                             int64_t offset,
                              ConversionPatternRewriter &rewriter) const {
     MemRefDescriptor memRefDescriptor(descriptor);
 
-    Value base = memRefDescriptor.alignedPtr(rewriter, loc);
+    Value base = memRefDescriptor.alignedPtr(rewriter, alignment, loc);
     Value offsetValue = offset == MemRefType::getDynamicStrideOrOffset()
                             ? memRefDescriptor.offset(rewriter, loc)
                             : this->createIndexConstant(rewriter, loc, offset);
 
     for (int i = 0, e = indices.size(); i < e; ++i) {
       Value stride = strides[i] == MemRefType::getDynamicStrideOrOffset()
                          ? memRefDescriptor.stride(rewriter, loc, i)
                          : this->createIndexConstant(rewriter, loc, strides[i]);
Value getDataPtr(Location loc, MemRefType type, Value memRefDesc,
                    ArrayRef<Value> indices, ConversionPatternRewriter &rewriter,
                    llvm::Module &module) const {
     LLVM::LLVMType ptrType = MemRefDescriptor(memRefDesc).getElementType();
     int64_t offset;
     SmallVector<int64_t, 4> strides;
     auto successStrides = getStridesAndOffset(type, strides, offset);
     assert(succeeded(successStrides) && "unexpected non-strided memref");
     (void)successStrides;
-    return getStridedElementPtr(loc, ptrType, memRefDesc, indices, strides,
-                                offset, rewriter);
+    return getStridedElementPtr(loc, ptrType, type.getAlignment(), memRefDesc,
+                                indices, strides, offset, rewriter);
   }
 };
 
 // Load operation is lowered to obtaining a pointer to the indexed element
 // and loading it.
 struct LoadOpLowering : public LoadStoreOpLowering<LoadOp> {
   using Base::Base;
 
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Copy the buffer pointer from the old descriptor to the new one.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
-    extracted = sourceMemRef.alignedPtr(rewriter, loc);
+    extracted = sourceMemRef.alignedPtr(rewriter, 0, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Extract strides needed to compute offset.
     SmallVector<Value, 4> strideValues;
     strideValues.reserve(viewMemRefType.getRank());
     for (int i = 0, e = viewMemRefType.getRank(); i < e; ++i)
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
 
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value extracted = sourceMemRef.allocatedPtr(rewriter, loc);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Field 2: Copy the actual aligned pointer to payload.
-    extracted = sourceMemRef.alignedPtr(rewriter, loc);
+    extracted = sourceMemRef.alignedPtr(rewriter, 0, loc);
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
         loc, targetElementTy.getPointerTo(), extracted);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Field 3: Copy the offset in aligned pointer.
     unsigned numDynamicSizes = llvm::size(viewOp.getDynamicSizes());
     (void)numDynamicSizes;
     bool hasDynamicOffset = offset == MemRefType::getDynamicStrideOrOffset();

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto desc = MemRefDescriptor::undef(rewriter, loc, llvmTargetDescriptorTy);
     Type llvmTargetElementTy = desc.getElementType();
     // Set allocated ptr.
     Value allocated = sourceMemRef.allocatedPtr(rewriter, loc);
     allocated =
         rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, allocated);
     desc.setAllocatedPtr(rewriter, loc, allocated);
     // Set aligned ptr.
-    Value ptr = sourceMemRef.alignedPtr(rewriter, loc);
+    Value ptr = sourceMemRef.alignedPtr(rewriter, 0, loc);
     ptr = rewriter.create<LLVM::BitcastOp>(loc, llvmTargetElementTy, ptr);
     desc.setAlignedPtr(rewriter, loc, ptr);
     // Fill offset 0.
     auto attr = rewriter.getIntegerAttr(rewriter.getIndexType(), 0);
     auto zero = rewriter.create<LLVM::ConstantOp>(loc, int64Ty, attr);
     desc.setOffset(rewriter, loc, zero);
 
     // Fill size and stride descriptors in memref.

mlir/lib/Dialect/FxpMathOps/Transforms/UniformKernelUtils.h

if (auto st = t.dyn_cast<ShapedType>()) {
     switch (st.getKind()) {
     case StandardTypes::Kind::Vector:
       return VectorType::get(st.getShape(), newElementType);
     case StandardTypes::Kind::RankedTensor:
       return RankedTensorType::get(st.getShape(), newElementType);
     case StandardTypes::Kind::UnrankedTensor:
       return UnrankedTensorType::get(newElementType);
     case StandardTypes::Kind::MemRef:
-      return MemRefType::get(st.getShape(), newElementType,
-                             st.cast<MemRefType>().getAffineMaps());
+      return MemRefTypeBuilder(st.cast<MemRefType>())
+          .setElementType(newElementType);
     }
   }
   assert(t.isIntOrFloat());
   return newElementType;
 }
 
 /// Creates an IntegerAttr with a type that matches the shape of 't' (which can
 /// be a scalar primitive or a shaped type).

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

for (AffineMap m : reassociation) {
     }
     newSizes.push_back(size);
     newStrides.push_back(stride);
     currentDim += dim;
   }
 
   // Early-exit: if `type` is contiguous, the result must be contiguous.
   if (canonicalizeStridedLayout(type).getAffineMaps().empty())
-    return MemRefType::get(newSizes, type.getElementType(), {});
+    return MemRefTypeBuilder(type).setShape(newSizes).setAffineMaps({});
 
   // Convert back to int64_t because we don't have enough information to create
   // new strided layouts from AffineExpr only. This corresponds to a case where
   // copies may be necessary.
   int64_t intOffset = ShapedType::kDynamicStrideOrOffset;
   if (auto o = offset.dyn_cast<AffineConstantExpr>())
     intOffset = o.getValue();
   SmallVector<int64_t, 4> intStrides;
   intStrides.reserve(strides.size());
   for (auto stride : newStrides) {
     if (auto cst = stride.dyn_cast_or_null<AffineConstantExpr>())
       intStrides.push_back(cst.getValue());
     else
       intStrides.push_back(ShapedType::kDynamicStrideOrOffset);
   }
   auto layout =
       makeStridedLinearLayoutMap(intStrides, intOffset, type.getContext());
   return canonicalizeStridedLayout(
-      MemRefType::get(newSizes, type.getElementType(), {layout}));
+      MemRefTypeBuilder(type).setShape(newSizes).setAffineMaps({layout}));
 }
 
 /// Helper functions assert Attribute of the proper type in attr and returns the
 /// corresponding vector.
 /// TODO(rridle,ntv) this should be evolved into a generic
 /// `getRangeOfType<AffineMap>(ArrayAttr attrs)` that does not copy.
 static SmallVector<AffineMap, 4> getAffineMaps(ArrayAttr attrs) {
   return functional::map(
void mlir::linalg::SliceOp::build(Builder *b, OperationState &result,
   SmallVector<int64_t, 4> strides;
   auto res = getStridesAndOffset(memRefType, strides, offset);
   assert(succeeded(res) && strides.size() == indexings.size());
   (void)res;
 
   unsigned rank = memRefType.getRank();
   // TODO(ntv): propagate static size and stride information when available.
   SmallVector<int64_t, 4> sizes(rank, -1); // -1 encodes dynamic size.
-  Type elementType = memRefType.getElementType();
-  result.addTypes({MemRefType::get(
-      sizes, elementType,
-      {makeStridedLinearLayoutMap(strides, offset, b->getContext())},
-      memRefType.getMemorySpace())});
+  result.addTypes({MemRefTypeBuilder(memRefType)
+                       .setShape(sizes)
+                       .setAffineMaps({makeStridedLinearLayoutMap(
+                           strides, offset, b->getContext())})});
 }
 
 static void print(OpAsmPrinter &p, SliceOp op) {
   auto indexings = op.indexings();
   p << SliceOp::getOperationName() << " " << op.view() << "[" << indexings
     << "] ";
   p.printOptionalAttrDict(op.getAttrs());
   p << " : " << op.getBaseViewType();
void mlir::linalg::TransposeOp::build(Builder *b, OperationState &result,
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   auto res = getStridesAndOffset(memRefType, strides, offset);
   assert(succeeded(res) && strides.size() == static_cast<unsigned>(rank));
   (void)res;
   auto map = makeStridedLinearLayoutMap(strides, offset, b->getContext());
   map = permutationMap ? map.compose(permutationMap) : map;
   // Compute result type.
-  auto resultType = MemRefType::get(sizes, memRefType.getElementType(), map,
-                                    memRefType.getMemorySpace());
+  MemRefType resultType =
+      MemRefTypeBuilder(memRefType).setShape(sizes).setAffineMaps(map);
 
   build(b, result, resultType, view, attrs);
   result.addAttribute(TransposeOp::getPermutationAttrName(), permutation);
 }
 
 static void print(OpAsmPrinter &p, TransposeOp op) {
   p << op.getOperationName() << " " << op.view() << " " << op.permutation();
   p.printOptionalAttrDict(op.getAttrs(),

mlir/lib/Dialect/StandardOps/Ops.cpp

for (unsigned dim = 0, e = memrefType.getRank(); dim < e; ++dim) {
         // Dynamic shape dimension not folded; copy operand from old memref.
         newShapeConstants.push_back(-1);
         newOperands.push_back(alloc.getOperand(dynamicDimPos));
       }
       dynamicDimPos++;
     }
 
     // Create new memref type (which will have fewer dynamic dimensions).
-    auto newMemRefType = MemRefType::get(
-        newShapeConstants, memrefType.getElementType(),
-        memrefType.getAffineMaps(), memrefType.getMemorySpace());
+    MemRefType newMemRefType =
+        MemRefTypeBuilder(memrefType).setShape(newShapeConstants);
     assert(static_cast<int64_t>(newOperands.size()) ==
            newMemRefType.getNumDynamicDims());
 
     // Create and insert the alloc op for the new memref.
-    auto newAlloc = rewriter.create<AllocOp>(alloc.getLoc(), newMemRefType,
-                                             newOperands, IntegerAttr());
+    auto newAlloc =
+        rewriter.create<AllocOp>(alloc.getLoc(), newMemRefType, newOperands);
     // Insert a cast so we have the same type as the old alloc.
     auto resultCast = rewriter.create<MemRefCastOp>(alloc.getLoc(), newAlloc,
                                                     alloc.getType());
 
     rewriter.replaceOp(alloc, {resultCast}, droppedOperands);
     return matchSuccess();
   }
 };
for (int i = rank - 2; i >= 0; --i) {
         newStrides[i] = newShapeConstants[i + 1] * newStrides[i + 1];
     }
 
     // Regenerate strided layout map with 'newStrides' and 'newOffset'.
     map = makeStridedLinearLayoutMap(newStrides, newOffset,
                                      rewriter.getContext());
 
     // Create new memref type with constant folded dims and/or offset/strides.
-    auto newMemRefType =
-        MemRefType::get(newShapeConstants, memrefType.getElementType(), {map},
-                        memrefType.getMemorySpace());
+    MemRefType newMemRefType = MemRefTypeBuilder(memrefType)
+                                   .setShape(newShapeConstants)
+                                   .setAffineMaps({map});
     (void)dynamicOffsetOperandCount; // unused in opt mode
     assert(static_cast<int64_t>(newOperands.size()) ==
            dynamicOffsetOperandCount + newMemRefType.getNumDynamicDims());
 
     // Create new ViewOp.
     auto newViewOp = rewriter.create<ViewOp>(viewOp.getLoc(), newMemRefType,
                                              viewOp.getOperand(0), newOperands);
     // Insert a cast so we have the same type as the old memref type.

 // Returns a MemRefType with dynamic sizes and offset and the same stride as the
 // `memRefType` passed as argument.
 // TODO(andydavis,ntv) Evolve to a more powerful inference that can also keep
 // sizes and offset static.
 static Type inferSubViewResultType(MemRefType memRefType) {
   auto rank = memRefType.getRank();
   int64_t offset;
   SmallVector<int64_t, 4> strides;
-  Type elementType = memRefType.getElementType();
   auto res = getStridesAndOffset(memRefType, strides, offset);
   assert(succeeded(res) && "SubViewOp expected strided memref type");
   (void)res;
 
   // Assume sizes and offset are fully dynamic for now until canonicalization
   // occurs on the ranges. Typed strides don't change though.
   offset = MemRefType::getDynamicStrideOrOffset();
   // Overwrite strides because verifier will not pass.
   // TODO(b/144419106): don't force degrade the strides to fully dynamic.
   for (auto &stride : strides)
     stride = MemRefType::getDynamicStrideOrOffset();
   auto stridedLayout =
       makeStridedLinearLayoutMap(strides, offset, memRefType.getContext());
   SmallVector<int64_t, 4> sizes(rank, ShapedType::kDynamicSize);
-  return MemRefType::get(sizes, elementType, stridedLayout,
-                         memRefType.getMemorySpace());
+  return MemRefTypeBuilder(memRefType)
+      .setShape(sizes)
+      .setAffineMaps(stridedLayout);
 }
 
 void mlir::SubViewOp::build(Builder *b, OperationState &result, Value source,
                             ValueRange offsets, ValueRange sizes,
                             ValueRange strides, Type resultType,
                             ArrayRef<NamedAttribute> attrs) {
   if (!resultType)
     resultType = inferSubViewResultType(source.getType().cast<MemRefType>());
if (subViewType.hasStaticShape() ||
       return matchFailure();
     }
     SmallVector<int64_t, 4> staticShape(subViewOp.getNumSizes());
     for (auto size : llvm::enumerate(subViewOp.sizes())) {
       auto defOp = size.value().getDefiningOp();
       assert(defOp);
       staticShape[size.index()] = cast<ConstantIndexOp>(defOp).getValue();
     }
-    MemRefType newMemRefType = MemRefType::get(
-        staticShape, subViewType.getElementType(), subViewType.getAffineMaps(),
-        subViewType.getMemorySpace());
+    MemRefType newMemRefType =
+        MemRefTypeBuilder(subViewType).setShape(staticShape);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), subViewOp.offsets(),
         ArrayRef<Value>(), subViewOp.strides(), newMemRefType);
     // Insert a memref_cast for compatibility of the uses of the op.
     rewriter.replaceOpWithNewOp<MemRefCastOp>(
         subViewOp.sizes(), subViewOp, newSubViewOp, subViewOp.getType());
     return matchSuccess();
   }
for (auto stride : llvm::enumerate(subViewOp.strides())) {
       assert(defOp);
       assert(baseStrides[stride.index()] > 0);
       staticStrides[stride.index()] =
           cast<ConstantIndexOp>(defOp).getValue() * baseStrides[stride.index()];
     }
     AffineMap layoutMap = makeStridedLinearLayoutMap(
         staticStrides, resultOffset, rewriter.getContext());
     MemRefType newMemRefType =
-        MemRefType::get(subViewType.getShape(), subViewType.getElementType(),
-                        layoutMap, subViewType.getMemorySpace());
+        MemRefTypeBuilder(subViewType).setAffineMaps(layoutMap);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), subViewOp.offsets(),
         subViewOp.sizes(), ArrayRef<Value>(), newMemRefType);
     // Insert a memref_cast for compatibility of the uses of the op.
     rewriter.replaceOpWithNewOp<MemRefCastOp>(
         subViewOp.strides(), subViewOp, newSubViewOp, subViewOp.getType());
     return matchSuccess();
   }
for (auto offset : llvm::enumerate(subViewOp.offsets())) {
       assert(baseStrides[offset.index()] > 0);
       staticOffset +=
           cast<ConstantIndexOp>(defOp).getValue() * baseStrides[offset.index()];
     }
 
     AffineMap layoutMap = makeStridedLinearLayoutMap(
         resultStrides, staticOffset, rewriter.getContext());
     MemRefType newMemRefType =
-        MemRefType::get(subViewType.getShape(), subViewType.getElementType(),
-                        layoutMap, subViewType.getMemorySpace());
+        MemRefTypeBuilder(subViewType).setAffineMaps(layoutMap);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), ArrayRef<Value>(),
         subViewOp.sizes(), subViewOp.strides(), newMemRefType);
     // Insert a memref_cast for compatibility of the uses of the op.
     rewriter.replaceOpWithNewOp<MemRefCastOp>(
         subViewOp.offsets(), subViewOp, newSubViewOp, subViewOp.getType());
     return matchSuccess();
   }

mlir/lib/IR/AsmPrinter.cpp

for (auto dim : v.getShape()) {
       os << 'x';
     }
     printType(v.getElementType());
     for (auto map : v.getAffineMaps()) {
       os << ", ";
       printAttribute(AffineMapAttr::get(map));
     }
     // Only print the memory space if it is the non-default one.
-    if (v.getMemorySpace())
+    if (v.getMemorySpace() || v.getAlignment())
       os << ", " << v.getMemorySpace();
+    if (v.getAlignment())
+      os << ", " << v.getAlignment();
     os << '>';
     return;
   }
   case StandardTypes::UnrankedMemRef: {
     auto v = type.cast<UnrankedMemRefType>();
     os << "memref<*x";
     printType(v.getElementType());
     os << '>';

mlir/lib/IR/StandardTypes.cpp

 //===----------------------------------------------------------------------===//
 
 /// Get or create a new MemRefType based on shape, element type, affine
 /// map composition, and memory space.  Assumes the arguments define a
 /// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
 /// construction failures.
 MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
                            ArrayRef<AffineMap> affineMapComposition,
-                           unsigned memorySpace) {
-  auto result = getImpl(shape, elementType, affineMapComposition, memorySpace,
-                        /*location=*/llvm::None);
+                           unsigned memorySpace, unsigned alignment) {
+  auto result =
+      getImpl(shape, elementType, affineMapComposition, memorySpace, alignment,
   assert(result && "Failed to construct instance of MemRefType.");
   return result;
 }
 
 /// Get or create a new MemRefType based on shape, element type, affine
 /// map composition, and memory space declared at the given location.
 /// If the location is unknown, the last argument should be an instance of
 /// UnknownLoc.  If the MemRefType defined by the arguments would be
 /// ill-formed, emits errors (to the handler registered with the context or to
 /// the error stream) and returns nullptr.
 MemRefType MemRefType::getChecked(ArrayRef<int64_t> shape, Type elementType,
                                   ArrayRef<AffineMap> affineMapComposition,
-                                  unsigned memorySpace, Location location) {
+                                  unsigned memorySpace, unsigned alignment,
+                                  Location location) {
   return getImpl(shape, elementType, affineMapComposition, memorySpace,
-                 location);
+                 alignment, location);
 }
 
 /// Get or create a new MemRefType defined by the arguments.  If the resulting
 /// type would be ill-formed, return nullptr.  If the location is provided,
 /// emit detailed error messages.  To emit errors when the location is unknown,
 /// pass in an instance of UnknownLoc.
 MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
-                               unsigned memorySpace,
+                               unsigned memorySpace, unsigned alignment,
                                Optional<Location> location) {
   auto *context = elementType.getContext();
 
   // Check that memref is formed from allowed types.
   if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
     return emitOptionalError(location, "invalid memref element type"),
            MemRefType();
 
MemRefType MemRefType::getImpl(ArrayRef<int64_t> shape, Type elementType,
   SmallVector<AffineMap, 2> cleanedAffineMapComposition;
   for (const auto &map : affineMapComposition) {
     if (map.isIdentity())
       continue;
     cleanedAffineMapComposition.push_back(map);
   }
 
   return Base::get(context, StandardTypes::MemRef, shape, elementType,
-                   cleanedAffineMapComposition, memorySpace);
+                   cleanedAffineMapComposition, memorySpace, alignment);
 }
 
 ArrayRef<int64_t> MemRefType::getShape() const { return getImpl()->getShape(); }
 
 ArrayRef<AffineMap> MemRefType::getAffineMaps() const {
   return getImpl()->getAffineMaps();
 }
 
 unsigned MemRefType::getMemorySpace() const { return getImpl()->memorySpace; }
 
+unsigned MemRefType::getAlignment() const { return getImpl()->alignment; }
+
 //===----------------------------------------------------------------------===//
 // UnrankedMemRefType
 //===----------------------------------------------------------------------===//
 
 UnrankedMemRefType UnrankedMemRefType::get(Type elementType,
                                            unsigned memorySpace) {
   return Base::get(elementType.getContext(), StandardTypes::UnrankedMemRef,
                    elementType, memorySpace);
MemRefType mlir::canonicalizeStridedLayout(MemRefType t) {
   // simplified layout of `t` we can just return an empty layout. Otherwise,
   // just simplify the existing layout.
   AffineExpr expr =
       makeCanonicalStridedLayoutExpr(t.getShape(), t.getContext());
   auto m = affineMaps[0];
   auto simplifiedLayoutExpr =
       simplifyAffineExpr(m.getResult(0), m.getNumDims(), m.getNumSymbols());
   if (expr != simplifiedLayoutExpr)
-    return MemRefType::get(t.getShape(), t.getElementType(),
-                           {AffineMap::get(m.getNumDims(), m.getNumSymbols(),
-                                           {simplifiedLayoutExpr})});
-
-  return MemRefType::get(t.getShape(), t.getElementType(), {});
+    return MemRefTypeBuilder(t).setAffineMaps({AffineMap::get(
+        m.getNumDims(), m.getNumSymbols(), {simplifiedLayoutExpr})});
+  return MemRefTypeBuilder(t).setAffineMaps({});
 }
 
 /// Return true if the layout for `t` is compatible with strided semantics.
 bool mlir::isStrided(MemRefType t) {
   int64_t offset;
   SmallVector<int64_t, 4> stridesAndOffset;
   auto res = getStridesAndOffset(t, stridesAndOffset, offset);
   return succeeded(res);
 }

mlir/lib/IR/TypeDetail.h

static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator,
     return new (allocator.allocate<UnrankedTensorTypeStorage>())
         UnrankedTensorTypeStorage(elementTy);
   }
 };
 
 struct MemRefTypeStorage : public ShapedTypeStorage {
   MemRefTypeStorage(unsigned shapeSize, Type elementType,
                     const int64_t *shapeElements, const unsigned numAffineMaps,
-                    AffineMap const *affineMapList, const unsigned memorySpace)
+                    AffineMap const *affineMapList, const unsigned memorySpace,
+                    const unsigned alignment)
       : ShapedTypeStorage(elementType, shapeSize), shapeElements(shapeElements),
         numAffineMaps(numAffineMaps), affineMapList(affineMapList),
-        memorySpace(memorySpace) {}
+        memorySpace(memorySpace), alignment(alignment) {}
 
   /// The hash key used for uniquing.
   // MemRefs are uniqued based on their shape, element type, affine map
   // composition, and memory space.
-  using KeyTy =
-      std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>, unsigned>;
+  using KeyTy = std::tuple<ArrayRef<int64_t>, Type, ArrayRef<AffineMap>,
+                           unsigned, unsigned>;
   bool operator==(const KeyTy &key) const {
-    return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace);
+    return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace,
+                        alignment);
   }
 
   /// Construction.
   static MemRefTypeStorage *construct(TypeStorageAllocator &allocator,
                                       const KeyTy &key) {
     // Copy the shape into the bump pointer.
     ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key));
 
     // Copy the affine map composition into the bump pointer.
     ArrayRef<AffineMap> affineMapComposition =
         allocator.copyInto(std::get<2>(key));
 
     // Initialize the memory using placement new.
-    return new (allocator.allocate<MemRefTypeStorage>())
-        MemRefTypeStorage(shape.size(), std::get<1>(key), shape.data(),
-                          affineMapComposition.size(),
-                          affineMapComposition.data(), std::get<3>(key));
+    return new (allocator.allocate<MemRefTypeStorage>()) MemRefTypeStorage(
+        shape.size(), std::get<1>(key), shape.data(),
+        affineMapComposition.size(), affineMapComposition.data(),
+        std::get<3>(key), std::get<4>(key));
   }
 
   ArrayRef<int64_t> getShape() const {
     return ArrayRef<int64_t>(shapeElements, getSubclassData());
   }
 
   ArrayRef<AffineMap> getAffineMaps() const {
     return ArrayRef<AffineMap>(affineMapList, numAffineMaps);
   }
 
   /// An array of integers which stores the shape dimension sizes.
   const int64_t *shapeElements;
   /// The number of affine maps in the 'affineMapList' array.
   const unsigned numAffineMaps;
   /// List of affine maps in the memref's layout/index map composition.
   AffineMap const *affineMapList;
   /// Memory space in which data referenced by memref resides.
   const unsigned memorySpace;
+  /// Alignment of the underlying pointer.
+  const unsigned alignment;
 };
 
 /// Unranked MemRef is a MemRef with unknown rank.
 /// Only element type and memory space are known
 struct UnrankedMemRefTypeStorage : public ShapedTypeStorage {
 
   UnrankedMemRefTypeStorage(Type elementTy, const unsigned memorySpace)
       : ShapedTypeStorage(elementTy), memorySpace(memorySpace) {}

mlir/lib/Parser/Parser.cpp

Type Parser::parseMemRefType() {
   auto typeLoc = getToken().getLoc();
   auto elementType = parseType();
   if (!elementType)
     return nullptr;
 
   // Parse semi-affine-map-composition.
   SmallVector<AffineMap, 2> affineMapComposition;
   unsigned memorySpace = 0;
+  unsigned alignment = 0;
   bool parsedMemorySpace = false;
+  bool parsedAlignment = false;
 
   auto parseElt = [&]() -> ParseResult {
     if (getToken().is(Token::integer)) {
       // Parse memory space.
-      if (parsedMemorySpace)
-        return emitError("multiple memory spaces specified in memref type");
+      if (!parsedMemorySpace) {
         auto v = getToken().getUnsignedIntegerValue();
         if (!v.hasValue())
           return emitError("invalid memory space in memref type");
         memorySpace = v.getValue();
         consumeToken(Token::integer);
         parsedMemorySpace = true;
+      } else if (!parsedAlignment) {
+        if (isUnranked)
+          return emitError("cannot have alignment for unranked memref type");
+        auto v = getToken().getUnsignedIntegerValue();
+        if (!v.hasValue())
+          return emitError("invalid alignment in memref type");
+        alignment = v.getValue();
+        consumeToken(Token::integer);
+        parsedAlignment = true;
+      } else {
+        return emitError("excessive integers specified in memref type");
+      }
     } else {
       if (isUnranked)
         return emitError("cannot have affine map for unranked memref type");
       if (parsedMemorySpace)
         return emitError("expected memory space to be last in memref type");
       if (getToken().is(Token::kw_offset)) {
         int64_t offset;
         SmallVector<int64_t, 4> strides;
if (parseToken(Token::greater, "expected ',' or '>' in memref type"))
       return nullptr;
   }
 
   if (isUnranked)
     return UnrankedMemRefType::getChecked(elementType, memorySpace,
                                           getEncodedSourceLocation(typeLoc));
 
   return MemRefType::getChecked(dimensions, elementType, affineMapComposition,
-                                memorySpace, getEncodedSourceLocation(typeLoc));
+                                memorySpace, alignment,
+                                getEncodedSourceLocation(typeLoc));
 }
 
 /// Parse any type except the function type.
 ///
 ///   non-function-type ::= integer-type
 ///                       | index-type
 ///                       | float-type
 ///                       | extended-type

mlir/lib/Transforms/PipelineDataTransfer.cpp

static bool doubleBuffer(Value oldMemRef, AffineForOp forOp) {
   // Doubles the shape with a leading dimension extent of 2.
   auto doubleShape = [&](MemRefType oldMemRefType) -> MemRefType {
     // Add the leading dimension in the shape for the double buffer.
     ArrayRef<int64_t> oldShape = oldMemRefType.getShape();
     SmallVector<int64_t, 4> newShape(1 + oldMemRefType.getRank());
     newShape[0] = 2;
     std::copy(oldShape.begin(), oldShape.end(), newShape.begin() + 1);
     auto newMemRefType =
-        MemRefType::get(newShape, oldMemRefType.getElementType(), {},
-                        oldMemRefType.getMemorySpace());
+        MemRefTypeBuilder(oldMemRefType).setShape(newShape).setAffineMaps({});
     return newMemRefType;
   };
 
   auto oldMemRefType = oldMemRef.getType().cast<MemRefType>();
   auto newMemRefType = doubleShape(oldMemRefType);
 
   // The double buffer is allocated right before 'forInst'.
   auto *forInst = forOp.getOperation();

mlir/lib/Transforms/Utils/Utils.cpp

if (ubConst.getValue() < 0)
       // This is due to an invalid map that maps to a negative space.
       return failure();
     newShape[d] = ubConst.getValue() + 1;
   }
 
   auto oldMemRef = allocOp.getResult();
   SmallVector<Value, 4> symbolOperands(allocOp.getSymbolicOperands());
 
-  auto newMemRefType = MemRefType::get(newShape, memrefType.getElementType(),
-                                       b.getMultiDimIdentityMap(newRank));
+  auto newMemRefType = MemRefTypeBuilder(memrefType)
+                           .setShape(newShape)
+                           .setAffineMaps(b.getMultiDimIdentityMap(newRank));
   auto newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType);
 
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
                                       /*extraOperands=*/{},
                                       /*symbolOperands=*/symbolOperands))) {

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

 // CHECK-NEXT:  %[[ptr:.*]] = llvm.extractvalue %[[ld]][0] : !llvm<"{ float*, float*, i64 }">
 // CHECK-NEXT:  %[[bc:.*]] = llvm.bitcast %[[ptr]] : !llvm<"float*"> to !llvm<"i8*">
 // CHECK-NEXT:  llvm.call @free(%[[bc]]) : (!llvm<"i8*">) -> ()
   dealloc %arg0 : memref<f32>
   return
 }
 
 // CHECK-LABEL: func @aligned_1d_alloc(
-func @aligned_1d_alloc() -> memref<42xf32> {
+func @aligned_1d_alloc() -> memref<42xf32, 0, 8> {
 // CHECK-NEXT:  llvm.mlir.constant(42 : index) : !llvm.i64
 // CHECK-NEXT:  %[[null:.*]] = llvm.mlir.null : !llvm<"float*">
 // CHECK-NEXT:  %[[one:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
 // CHECK-NEXT:  %[[gep:.*]] = llvm.getelementptr %[[null]][%[[one]]] : (!llvm<"float*">, !llvm.i64) -> !llvm<"float*">
 // CHECK-NEXT:  %[[sizeof:.*]] = llvm.ptrtoint %[[gep]] : !llvm<"float*"> to !llvm.i64
 // CHECK-NEXT:  llvm.mul %{{.*}}, %[[sizeof]] : !llvm.i64
 // CHECK-NEXT:  %[[alignment:.*]] = llvm.mlir.constant(8 : index) : !llvm.i64
 // CHECK-NEXT:  %[[alignmentMinus1:.*]] = llvm.add {{.*}}, %[[alignment]] : !llvm.i64
 // CHECK-NEXT:  %[[allocsize:.*]] = llvm.sub %[[alignmentMinus1]], %[[one]] : !llvm.i64
 // CHECK-NEXT:  %[[allocated:.*]] = llvm.call @malloc(%[[allocsize]]) : (!llvm.i64) -> !llvm<"i8*">
 // CHECK-NEXT:  %[[ptr:.*]] = llvm.bitcast %{{.*}} : !llvm<"i8*"> to !llvm<"float*">
 // CHECK-NEXT:  llvm.mlir.undef : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
 // CHECK-NEXT:  llvm.insertvalue %[[ptr]], %{{.*}}[0] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
 // CHECK-NEXT:  %[[allocatedAsInt:.*]] = llvm.ptrtoint %[[allocated]] : !llvm<"i8*"> to !llvm.i64
 // CHECK-NEXT:  %[[alignAdj1:.*]] = llvm.urem %[[allocatedAsInt]], %[[alignment]] : !llvm.i64
 // CHECK-NEXT:  %[[alignAdj2:.*]] = llvm.sub %[[alignment]], %[[alignAdj1]] : !llvm.i64
 // CHECK-NEXT:  %[[alignAdj3:.*]] = llvm.urem %[[alignAdj2]], %[[alignment]] : !llvm.i64
 // CHECK-NEXT:  %[[aligned:.*]] = llvm.getelementptr %9[%[[alignAdj3]]] : (!llvm<"i8*">, !llvm.i64) -> !llvm<"i8*">
 // CHECK-NEXT:  %[[alignedBitCast:.*]] = llvm.bitcast %[[aligned]] : !llvm<"i8*"> to !llvm<"float*">
 // CHECK-NEXT:  llvm.insertvalue %[[alignedBitCast]], %{{.*}}[1] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
 // CHECK-NEXT:  %[[c0:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
 // CHECK-NEXT:  llvm.insertvalue %[[c0]], %{{.*}}[2] : !llvm<"{ float*, float*, i64, [1 x i64], [1 x i64] }">
-  %0 = alloc() {alignment = 8} : memref<42xf32>
-  return %0 : memref<42xf32>
+  %0 = alloc() : memref<42xf32, 0, 8>
+  return %0 : memref<42xf32, 0, 8>
 }
 
 // CHECK-LABEL: func @mixed_alloc(
 //       CHECK:   %[[M:.*]]: !llvm.i64, %[[N:.*]]: !llvm.i64) -> !llvm<"{ float*, float*, i64, [3 x i64], [3 x i64] }"> {
 func @mixed_alloc(%arg0: index, %arg1: index) -> memref<?x42x?xf32> {
 //  CHECK-NEXT:  %[[c42:.*]] = llvm.mlir.constant(42 : index) : !llvm.i64
 //  CHECK-NEXT:  llvm.mul %[[M]], %[[c42]] : !llvm.i64
 //  CHECK-NEXT:  %[[sz:.*]] = llvm.mul %{{.*}}, %[[N]] : !llvm.i64
// CHECK-NEXT:  llvm.mlir.constant(15 : index) : !llvm.i64
   %2 = dim %static, 2 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(13 : index) : !llvm.i64
   %3 = dim %static, 3 : memref<42x32x15x13x27xf32>
 // CHECK-NEXT:  llvm.mlir.constant(27 : index) : !llvm.i64
   %4 = dim %static, 4 : memref<42x32x15x13x27xf32>
   return
 }
 
+// CHECK-LABEL: func @aligned_memref_load
+//       CHECK:   %[[A:.*]]: !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">, %[[I:.*]]: !llvm.i64
+func @aligned_memref_load(%arr : memref<4xf16, 0, 8>, %index : index) -> f16 {
+  // CHECK-NEXT: %[[ld:.*]] = llvm.load %{{.*}} : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">
+  // CHECK-NEXT: %[[ptr:.*]] = llvm.extractvalue %[[ld]][1] : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }">
+
+  // CHECK-NEXT: %[[zero:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[mask:.*]] = llvm.mlir.constant(7 : index) : !llvm.i64
+  // CHECK-NEXT: %[[intptr:.*]] = llvm.ptrtoint %1 : !llvm<"half*"> to !llvm.i64
+  // CHECK-NEXT: %[[masked:.*]] = llvm.and %[[intptr]], %[[mask]] : !llvm.i64
+  // CHECK-NEXT: %[[isaligned:.*]] = llvm.icmp "eq" %[[masked]], %[[zero]] : !llvm.i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[isaligned]]) : (!llvm.i1) -> ()
+
+  // CHECK-NEXT: %[[off:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[st0:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK-NEXT: %[[offI:.*]] = llvm.mul %[[I]], %[[st0]] : !llvm.i64
+  // CHECK-NEXT: %[[off0:.*]] = llvm.add %[[off]], %[[offI]] : !llvm.i64
+  // CHECK-NEXT: %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off0]]] : (!llvm<"half*">, !llvm.i64) -> !llvm<"half*">
+  // CHECK-NEXT: llvm.load %[[addr]] : !llvm<"half*">
+  %v = load %arr[%index] : memref<4xf16, 0, 8>
+  return %v : f16
+}
+
+// CHECK-LABEL: func @aligned_memref_store
+//       CHECK:   %[[A:.*]]: !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">, %[[V:.*]]: !llvm.half, %[[I:.*]]: !llvm.i64
+func @aligned_memref_store(%arr : memref<4xf16, 0, 8>, %v : f16, %index : index) {
+  // CHECK-NEXT: %[[ld:.*]] = llvm.load %{{.*}} : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }*">
+  // CHECK-NEXT: %[[ptr:.*]] = llvm.extractvalue %[[ld]][1] : !llvm<"{ half*, half*, i64, [1 x i64], [1 x i64] }">
+
+  // CHECK-NEXT: %[[zero:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[mask:.*]] = llvm.mlir.constant(7 : index) : !llvm.i64
+  // CHECK-NEXT: %[[intptr:.*]] = llvm.ptrtoint %1 : !llvm<"half*"> to !llvm.i64
+  // CHECK-NEXT: %[[masked:.*]] = llvm.and %[[intptr]], %[[mask]] : !llvm.i64
+  // CHECK-NEXT: %[[isaligned:.*]] = llvm.icmp "eq" %[[masked]], %[[zero]] : !llvm.i64
+  // CHECK-NEXT: "llvm.intr.assume"(%[[isaligned]]) : (!llvm.i1) -> ()
+
+  // CHECK-NEXT: %[[off:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
+  // CHECK-NEXT: %[[st0:.*]] = llvm.mlir.constant(1 : index) : !llvm.i64
+  // CHECK-NEXT: %[[offI:.*]] = llvm.mul %[[I]], %[[st0]] : !llvm.i64
+  // CHECK-NEXT: %[[off0:.*]] = llvm.add %[[off]], %[[offI]] : !llvm.i64
+  // CHECK-NEXT: %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off0]]] : (!llvm<"half*">, !llvm.i64) -> !llvm<"half*">
+  // CHECK-NEXT: llvm.store %[[V]], %[[addr]] : !llvm<"half*">
+  store %v, %arr[%index] : memref<4xf16, 0, 8>
+  return
+}

mlir/test/IR/core-ops.mlir

 // CHECK-LABEL: func @tensor_load_store
 func @tensor_load_store(%0 : memref<4x4xi32>) {
   // CHECK: %[[TENSOR:.*]] = tensor_load %[[MEMREF:.*]] : memref<4x4xi32>
   %1 = tensor_load %0 : memref<4x4xi32>
   // CHECK: tensor_store %[[TENSOR]], %[[MEMREF]] : memref<4x4xi32>
   tensor_store %1, %0 : memref<4x4xi32>
   return
 }
+
+// CHECK-LABEL: func @aligned_memref_load_store
+func @aligned_memref_load_store(%arr : memref<4xf16, 0, 8>) {
+  %c_0 = constant 0 : index
+  // CHECK: %[[VALUE:.*]] = load %[[MEMREF:.*]] : memref<4xf16, 0, 8>
+  %v = load %arr[%c_0] : memref<4xf16, 0, 8>
+  // CHECK: store %[[VALUE]], %[[MEMREF]] : memref<4xf16, 0, 8>
+  store %v, %arr[%c_0] : memref<4xf16, 0, 8>
+  return
+}

mlir/test/IR/invalid.mlir

 
 // -----
 // Test non-existent map in map composition of memref type.
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 
 func @memrefs(memref<2x4xi8, #map0, #map8>) // expected-error {{undefined symbol alias id 'map8'}}
 
 // -----
-// Test multiple memory space error.
+// Test multiple integers error.
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
-func @memrefs(memref<2x4xi8, #map0, 1, 2>) // expected-error {{multiple memory spaces specified in memref type}}
+func @memrefs(memref<2x4xi8, #map0, 1, 2, 3>) // expected-error {{excessive integers specified in memref type}}
+
+// -----
+// Test alignment for unranked memref error.
+func @memrefs(memref<*xi8, 1, 2>) // expected-error {{cannot have alignment for unranked memref type}}
 
 // -----
 // Test affine map after memory space.
 #map0 = affine_map<(d0, d1) -> (d0, d1)>
 #map1 = affine_map<(d0, d1) -> (d0, d1)>
 
 func @memrefs(memref<2x4xi8, #map0, 1, #map1>) // expected-error {{expected memory space to be last in memref type}}