[mlir] Add alignment to MemRefType and remove alignment from AllocOp.

[timshen]

Backport the actual code of removing alignment from AllocOp from the child patch.

[merge_guards_bot]

Unit tests: fail. 62145 tests passed, 6 failed and 811 were skipped.

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. 62145 tests passed, 6 failed and 811 were skipped.

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.

[timshen]

Add a test for power of 2 alignment.

[merge_guards_bot]

Unit tests: fail. 62145 tests passed, 6 failed and 811 were skipped.

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.

[timshen]

I'm holding on the review since the discussion in the mailing list is still ongoing.

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()) {
-  }
+        affineMaps(other.getAffineMaps()), memorySpace(other.getMemorySpace()),
+        alignment(other.getAlignment()) {}
 
   MemRefTypeBuilder(ArrayRef<int64_t> shape, Type elementType)
-      : shape(shape), elementType(elementType), affineMaps(), memorySpace(0) {}
+      : 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);
+    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> {

mlir/lib/Conversion/StandardToLLVM/ConvertStandardToLLVM.cpp

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");

mlir/lib/Dialect/StandardOps/Ops.cpp

PatternMatchResult matchAndRewrite(AllocOp alloc,
 
     // Create new memref type (which will have fewer dynamic dimensions).
     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();
   }
 };

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,
   // implicit identity.
   SmallVector<AffineMap, 2> cleanedAffineMapComposition;
   for (const auto &map : affineMapComposition) {
     if (map.isIdentity())
       continue;
     cleanedAffineMapComposition.push_back(map);
   }
 
+  if ((alignment & (alignment - 1)) != 0) {
+    return emitOptionalError(location, "alignment must be power of 2"),
+           MemRefType();
+  }
+
   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);

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/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

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 invalid alignment.
+#map0 = affine_map<(d0, d1) -> (d0, d1)>
+func @memrefs(memref<2x4xi8, #map0, 1, 3>) // expected-error {{alignment must be power of 2}}
+
+// -----
+// 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}}