[mlir][Memref] Add memref.memory_space_cast and its lowerings

[krzysz00]

Summary

Address space casts are present in common MLIR targets (LLVM, SPIRV).
Some planned rewrites (such as one of the potential fixes to the fact
that the AMDGPU backend requires alloca() to live in address space 5 /
the GPU private memory space) may require such casts to be inserted
into MLIR code, where those address spaces could be represented by
arbitrary memory space attributes.

Therefore, we define memref.memory_space_cast and its lowerings.

Depends on D141293

[krzysz00]

Split patch

[ftynse]

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

216

Nit

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

127–131

Nit: this is not needed when assemblyFormat is provided.

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

838

Nit

858

Nit: don't specify the explicit number of stack elements in small vectors unless there's a strong reason to.

872

Nit: please add trailing periods to all comments.

874–877

Memrefs allow for non-integer address spaces, but LLVM pointers do not. How about taking the address space from the converted type instead? Thus, the converted will have handled the address space conversion (or failed to produce a type, at which point we would have to stop anyway).

934–936

Are we sure that the pointers are tightly laid out, e.g. that if they are 32-bit values they cannot be aligned at 64 bit? I'd rather subtract the base address from the offset base address to be sure. This should normally be handled by constant folding at the LLVM IR level once it has full data layout.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

248–255

Aren't shape and element type guaranteed to be the same by the (trait) verifier? Equal shapes also require equal rank, so no need to check for both.

260–261

Nit: can't this be just return uaT.getElementType() == ubT.getElementType(); ? :)

[krzysz00]

Change operation name, address some review comments

[krzysz00]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

934–936

We've been making that assumption in all our unranked memref descriptor code up to this point - computeSizes() very much relies on the tight packing to compute the alloca() size we need for such a descriptor.

[ftynse]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

1203–1206

Nit: maybe use the GPU memory space attributes here instead of magic integers to promote good practices in the documentation.

[krzysz00]

Rebase for opaque pointers

mlir/include/mlir/Conversion/LLVMCommon/MemRefBuilder.h

static Value alignedPtr(OpBuilder &builder, Location loc,
                           LLVMTypeConverter &typeConverter, Value memRefDescPtr,
                           Type elemPtrPtrType);
   /// Builds IR inserting the aligned pointer into the descriptor.
   static void setAlignedPtr(OpBuilder &builder, Location loc,
                             LLVMTypeConverter &typeConverter,
                             Value memRefDescPtr, Type elemPtrPtrType,
                             Value alignedPtr);
 
+  /// Builds IR for getting the pointer to the offset's location.
+  /// Returns a pointer to a convertType(index), which points to the beggining
+  /// of a struct {index, index[rank], index[rank]}.
+  static Value offsetBasePtr(OpBuilder &builder, Location loc,
+                             LLVMTypeConverter &typeConverter,
+                             Value memRefDescPtr, Type elemPtrPtrType);
   /// Builds IR extracting the offset from the descriptor.
   static Value offset(OpBuilder &builder, Location loc,
                       LLVMTypeConverter &typeConverter, Value memRefDescPtr,
                       Type elemPtrPtrType);
   /// Builds IR inserting the offset into the descriptor.
   static void setOffset(OpBuilder &builder, Location loc,
                         LLVMTypeConverter &typeConverter, Value memRefDescPtr,
                         Type elemPtrPtrType, Value offset);

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let summary =
       "assertion that gives alignment information to the input memref";
   let description = [{
     The `assume_alignment` operation takes a memref and an integer of alignment
     value, and internally annotates the buffer with the given alignment. If
     the buffer isn't aligned to the given alignment, the behavior is undefined.
 
     This operation doesn't affect the semantics of a correct program. It's for
     optimization only, and the optimization is best-effort.
   }];
   let arguments = (ins AnyMemRef:$memref,
                        ConfinedAttr<I32Attr, [IntPositive]>:$alignment);
   let results = (outs);
 
   let assemblyFormat = "$memref `,` $alignment attr-dict `:` type($memref)";
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocOp
 //===----------------------------------------------------------------------===//
 
def LoadOp : MemRef_Op<"load",
 
   let hasFolder = 1;
   let hasVerifier = 1;
 
   let assemblyFormat = "$memref `[` $indices `]` attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
+// MemorySpaceCastOp
+//===----------------------------------------------------------------------===//
+def MemRef_MemorySpaceCastOp : MemRef_Op<"memory_space_cast", [
+      DeclareOpInterfaceMethods<CastOpInterface>,
+      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+      MemRefsNormalizable,
+      Pure,
+      SameOperandsAndResultElementType,
+      SameOperandsAndResultShape,
+      ViewLikeOpInterface
+    ]> {
+  let summary = "memref memory space cast operation";
+  let description = [{
+    This operation casts memref values between memory spaces.
+    The input and result will be memrefs of the same types and shape that alias
+    the same underlying memory, though, for some casts on some targets,
+    the underlying values of the pointer stored in the memref may be affected
+    by the cast.
+
+    The input and result must have the same shape, element type, rank, and layout.
+
+    If the source and target address spaces are the same, this operation is a noop.
+
+    Example:
+
+    ```mlir
+    // Cast a GPU private memory attribution into a generic pointer
+    %2 = memref.memory_space_cast %1 : memref<?xf32, 5> to memref<?xf32>
+    // Cast a generic pointer to workgroup-local memory
+    %4 = memref.memory_space_cast %3 : memref<5x4xi32> to memref<5x34xi32, 3>
+    // Cast between two non-default memory spaces
+    %6 = memref.memory_space_cast %5
+      : memref<*xmemref<?xf32>, 5> to memref<*xmemref<?xf32>, 3>
+    ```
+  }];
+
+  let arguments = (ins AnyRankedOrUnrankedMemRef:$source);
+  let results = (outs AnyRankedOrUnrankedMemRef:$dest);
+  let assemblyFormat = "$source attr-dict `:` type($source) `to` type($dest)";
+
+  let extraClassDeclaration = [{
+    Value getViewSource() { return getSource(); }
+  }];
+
+  let hasFolder = 1;
+}
+
+//===----------------------------------------------------------------------===//
 // PrefetchOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_PrefetchOp : MemRef_Op<"prefetch"> {
   let summary = "prefetch operation";
   let description = [{
     The "prefetch" op prefetches data from a memref location described with
     subscript indices similar to memref.load, and with three attributes: a

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

void UnrankedMemRefDescriptor::setAlignedPtr(OpBuilder &builder, Location loc,
   Value elementPtrPtr =
       builder.create<LLVM::BitcastOp>(loc, elemPtrPtrType, memRefDescPtr);
 
   Value alignedGep = builder.create<LLVM::GEPOp>(
       loc, elemPtrPtrType, elementPtrPtr, ArrayRef<LLVM::GEPArg>{1});
   builder.create<LLVM::StoreOp>(loc, alignedPtr, alignedGep);
 }
 
-Value UnrankedMemRefDescriptor::offset(OpBuilder &builder, Location loc,
+Value UnrankedMemRefDescriptor::offsetBasePtr(OpBuilder &builder, Location loc,
                                               LLVMTypeConverter &typeConverter,
                                               Value memRefDescPtr,
                                               Type elemPtrPtrType) {
   Value elementPtrPtr =
       builder.create<LLVM::BitcastOp>(loc, elemPtrPtrType, memRefDescPtr);
 
   Value offsetGep = builder.create<LLVM::GEPOp>(
       loc, elemPtrPtrType, elementPtrPtr, ArrayRef<LLVM::GEPArg>{2});
   offsetGep = builder.create<LLVM::BitcastOp>(
       loc, LLVM::LLVMPointerType::get(typeConverter.getIndexType()), offsetGep);
-  return builder.create<LLVM::LoadOp>(loc, offsetGep);
+  return offsetGep;
+}
+
+Value UnrankedMemRefDescriptor::offset(OpBuilder &builder, Location loc,
+                                       LLVMTypeConverter &typeConverter,
+                                       Value memRefDescPtr,
+                                       Type elemPtrPtrType) {
+  Value offsetPtr =
+      offsetBasePtr(builder, loc, typeConverter, memRefDescPtr, elemPtrPtrType);
+  return builder.create<LLVM::LoadOp>(loc, offsetPtr);
 }
 
 void UnrankedMemRefDescriptor::setOffset(OpBuilder &builder, Location loc,
                                          LLVMTypeConverter &typeConverter,
                                          Value memRefDescPtr,
                                          Type elemPtrPtrType, Value offset) {
-  Value elementPtrPtr =
-      builder.create<LLVM::BitcastOp>(loc, elemPtrPtrType, memRefDescPtr);
-
-  Value offsetGep = builder.create<LLVM::GEPOp>(
-      loc, elemPtrPtrType, elementPtrPtr, ArrayRef<LLVM::GEPArg>{2});
-  offsetGep = builder.create<LLVM::BitcastOp>(
-      loc, LLVM::LLVMPointerType::get(typeConverter.getIndexType()), offsetGep);
-  builder.create<LLVM::StoreOp>(loc, offset, offsetGep);
+  Value offsetPtr =
+      offsetBasePtr(builder, loc, typeConverter, memRefDescPtr, elemPtrPtrType);
+  builder.create<LLVM::StoreOp>(loc, offset, offsetPtr);
 }
 
 Value UnrankedMemRefDescriptor::sizeBasePtr(
     OpBuilder &builder, Location loc, LLVMTypeConverter &typeConverter,
     Value memRefDescPtr, LLVM::LLVMPointerType elemPtrPtrType) {
   Type elemPtrTy = elemPtrPtrType.getElementType();
   Type indexTy = typeConverter.getIndexType();
   Type structPtrTy =

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include <optional>
 
 namespace mlir {
 #define GEN_PASS_DEF_MEMREFTOLLVMCONVERSIONPASS
 #include "mlir/Conversion/Passes.h.inc"
 } // namespace mlir
 
LogicalResult match(memref::CastOp memRefCastOp) const override {
 
     // memref::CastOp reduce to bitcast in the ranked MemRef case and can be
     // used for type erasure. For now they must preserve underlying element type
     // and require source and result type to have the same rank. Therefore,
     // perform a sanity check that the underlying structs are the same. Once op
     // semantics are relaxed we can revisit.
     if (srcType.isa<MemRefType>() && dstType.isa<MemRefType>())
       return success(typeConverter->convertType(srcType) ==
                      typeConverter->convertType(dstType));
 
     // At least one of the operands is unranked type
     assert(srcType.isa<UnrankedMemRefType>() ||
            dstType.isa<UnrankedMemRefType>());
 
     // Unranked to unranked cast is disallowed
     return !(srcType.isa<UnrankedMemRefType>() &&
              dstType.isa<UnrankedMemRefType>())
                ? success()
                : failure();
   }
 
   void rewrite(memref::CastOp memRefCastOp, OpAdaptor adaptor,
                ConversionPatternRewriter &rewriter) const override {
     auto srcType = memRefCastOp.getOperand().getType();
     auto dstType = memRefCastOp.getType();
     auto targetStructType = typeConverter->convertType(memRefCastOp.getType());
     auto loc = memRefCastOp.getLoc();
 
     // For ranked/ranked case, just keep the original descriptor.
     if (srcType.isa<MemRefType>() && dstType.isa<MemRefType>())
       return rewriter.replaceOp(memRefCastOp, {adaptor.getSource()});
 
     if (srcType.isa<MemRefType>() && dstType.isa<UnrankedMemRefType>()) {
       // Casting ranked to unranked memref type
       // Set the rank in the destination from the memref type
       // Allocate space on the stack and copy the src memref descriptor
       // Set the ptr in the destination to the stack space
       auto srcMemRefType = srcType.cast<MemRefType>();
       int64_t rank = srcMemRefType.getRank();
       // ptr = AllocaOp sizeof(MemRefDescriptor)
       auto ptr = getTypeConverter()->promoteOneMemRefDescriptor(
           loc, adaptor.getSource(), rewriter);
       // voidptr = BitCastOp srcType* to void*
       auto voidPtr =
           rewriter.create<LLVM::BitcastOp>(loc, getVoidPtrType(), ptr)
               .getResult();
       // rank = ConstantOp srcRank
       auto rankVal = rewriter.create<LLVM::ConstantOp>(
           loc, getIndexType(), rewriter.getIndexAttr(rank));
       // undef = UndefOp
       UnrankedMemRefDescriptor memRefDesc =
           UnrankedMemRefDescriptor::undef(rewriter, loc, targetStructType);
       // d1 = InsertValueOp undef, rank, 0
       memRefDesc.setRank(rewriter, loc, rankVal);
       // d2 = InsertValueOp d1, voidptr, 1
       memRefDesc.setMemRefDescPtr(rewriter, loc, voidPtr);
lowerToMemCopyIntrinsic(memref::CopyOp op, OpAdaptor adaptor,
     // Compute number of elements.
     Value numElements = rewriter.create<LLVM::ConstantOp>(
         loc, getIndexType(), rewriter.getIndexAttr(1));
     for (int pos = 0; pos < srcType.getRank(); ++pos) {
       auto size = srcDesc.size(rewriter, loc, pos);
       numElements = rewriter.create<LLVM::MulOp>(loc, numElements, size);
     }
 
     // Get element size.
     auto sizeInBytes = getSizeInBytes(loc, srcType.getElementType(), rewriter);
     // Compute total.
     Value totalSize =
         rewriter.create<LLVM::MulOp>(loc, numElements, sizeInBytes);
 
     Value srcBasePtr = srcDesc.alignedPtr(rewriter, loc);
     Value srcOffset = srcDesc.offset(rewriter, loc);
     Value srcPtr = rewriter.create<LLVM::GEPOp>(loc, srcBasePtr.getType(),
                                                 srcBasePtr, srcOffset);
     MemRefDescriptor targetDesc(adaptor.getTarget());
matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor,
 
     if (isContiguousMemrefType(srcType) && isContiguousMemrefType(targetType))
       return lowerToMemCopyIntrinsic(op, adaptor, rewriter);
 
     return lowerToMemCopyFunctionCall(op, adaptor, rewriter);
   }
 };
 
+struct MemorySpaceCastOpLowering
+    : public ConvertOpToLLVMPattern<memref::MemorySpaceCastOp> {
+  using ConvertOpToLLVMPattern<
+      memref::MemorySpaceCastOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::MemorySpaceCastOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+
+    Type resultType = op.getDest().getType();
+    if (auto resultTypeR = resultType.dyn_cast<MemRefType>()) {
+      auto resultDescType =
+          typeConverter->convertType(resultTypeR).cast<LLVM::LLVMStructType>();
+      Type newPtrType = resultDescType.getBody()[0];
+
+      SmallVector<Value> descVals;
+      MemRefDescriptor::unpack(rewriter, loc, adaptor.getSource(), resultTypeR,
+                               descVals);
+      descVals[0] =
+          rewriter.create<LLVM::AddrSpaceCastOp>(loc, newPtrType, descVals[0]);
+      descVals[1] =
+          rewriter.create<LLVM::AddrSpaceCastOp>(loc, newPtrType, descVals[1]);
+      Value result = MemRefDescriptor::pack(rewriter, loc, *getTypeConverter(),
+                                            resultTypeR, descVals);
+      rewriter.replaceOp(op, result);
+      return success();
+    }
+    if (auto resultTypeU = resultType.dyn_cast<UnrankedMemRefType>()) {
+      // Since the type converter won't be doing this for us, get the address
+      // space.
+      auto sourceType = op.getSource().getType().cast<UnrankedMemRefType>();
+      FailureOr<unsigned> maybeSourceAddrSpace =
+          getTypeConverter()->getMemRefAddressSpace(sourceType);
+      if (failed(maybeSourceAddrSpace))
+        return rewriter.notifyMatchFailure(loc,
+                                           "non-integer source address space");
+      unsigned sourceAddrSpace = *maybeSourceAddrSpace;
+      FailureOr<unsigned> maybeResultAddrSpace =
+          getTypeConverter()->getMemRefAddressSpace(resultTypeU);
+      if (failed(maybeResultAddrSpace))
+        return rewriter.notifyMatchFailure(loc,
+                                           "non-integer result address space");
+      unsigned resultAddrSpace = *maybeResultAddrSpace;
+
+      UnrankedMemRefDescriptor sourceDesc(adaptor.getSource());
+      Value rank = sourceDesc.rank(rewriter, loc);
+      Value sourceUnderlyingDesc = sourceDesc.memRefDescPtr(rewriter, loc);
+
+      // Create and allocate storage for new memref descriptor.
+      auto result = UnrankedMemRefDescriptor::undef(
+          rewriter, loc, typeConverter->convertType(resultTypeU));
+      result.setRank(rewriter, loc, rank);
+      SmallVector<Value, 1> sizes;
+      UnrankedMemRefDescriptor::computeSizes(rewriter, loc, *getTypeConverter(),
+                                             result, resultAddrSpace, sizes);
+      Value resultUnderlyingSize = sizes.front();
+      Value resultUnderlyingDesc = rewriter.create<LLVM::AllocaOp>(
+          loc, getVoidPtrType(), resultUnderlyingSize, std::nullopt);
+      result.setMemRefDescPtr(rewriter, loc, resultUnderlyingDesc);
+
+      // Copy pointers, performing address space casts.
+      Type llvmElementType =
+          typeConverter->convertType(sourceType.getElementType());
+      auto sourceElemPtrPtrType = LLVM::LLVMPointerType::get(
+          LLVM::LLVMPointerType::get(llvmElementType, sourceAddrSpace));
+      auto resultElemPtrType =
+          LLVM::LLVMPointerType::get(llvmElementType, resultAddrSpace);
+      auto resultElemPtrPtrType = LLVM::LLVMPointerType::get(resultElemPtrType);
+
+      Value allocatedPtr = sourceDesc.allocatedPtr(
+          rewriter, loc, sourceUnderlyingDesc, sourceElemPtrPtrType);
+      Value alignedPtr =
+          sourceDesc.alignedPtr(rewriter, loc, *getTypeConverter(),
+                                sourceUnderlyingDesc, sourceElemPtrPtrType);
+      allocatedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
+          loc, resultElemPtrType, allocatedPtr);
+      alignedPtr = rewriter.create<LLVM::AddrSpaceCastOp>(
+          loc, resultElemPtrType, alignedPtr);
+
+      result.setAllocatedPtr(rewriter, loc, resultUnderlyingDesc,
+                             resultElemPtrPtrType, allocatedPtr);
+      result.setAlignedPtr(rewriter, loc, *getTypeConverter(),
+                           resultUnderlyingDesc, resultElemPtrPtrType,
+                           alignedPtr);
+
+      // Copy all the index-valued operands.
+      Value sourceIndexVals =
+          sourceDesc.offsetBasePtr(rewriter, loc, *getTypeConverter(),
+                                   sourceUnderlyingDesc, sourceElemPtrPtrType);
+      Value resultIndexVals =
+          result.offsetBasePtr(rewriter, loc, *getTypeConverter(),
+                               resultUnderlyingDesc, resultElemPtrPtrType);
+
+      int64_t bytesToSkip =
+          2 *
+          ceilDiv(getTypeConverter()->getPointerBitwidth(resultAddrSpace), 8);
+      Value bytesToSkipConst = rewriter.create<LLVM::ConstantOp>(
+          loc, getIndexType(), rewriter.getIndexAttr(bytesToSkip));
+      Value copySize = rewriter.create<LLVM::SubOp>(
+          loc, getIndexType(), resultUnderlyingSize, bytesToSkipConst);
+      Type llvmBool = typeConverter->convertType(rewriter.getI1Type());
+      Value nonVolatile = rewriter.create<LLVM::ConstantOp>(
+          loc, llvmBool, rewriter.getBoolAttr(false));
+      rewriter.create<LLVM::MemcpyOp>(loc, resultIndexVals, sourceIndexVals,
+                                      copySize, nonVolatile);
+
+      rewriter.replaceOp(op, ValueRange{result});
+      return success();
+    }
+    return rewriter.notifyMatchFailure(loc, "unexpected memref type");
+  }
+};
+
 /// Extracts allocated, aligned pointers and offset from a ranked or unranked
 /// memref type. In unranked case, the fields are extracted from the underlying
 /// ranked descriptor.
 static void extractPointersAndOffset(Location loc,
                                      ConversionPatternRewriter &rewriter,
                                      LLVMTypeConverter &typeConverter,
                                      Value originalOperand,
                                      Value convertedOperand,
patterns.add<
       DimOpLowering,
       ExtractStridedMetadataOpLowering,
       GenericAtomicRMWOpLowering,
       GlobalMemrefOpLowering,
       GetGlobalMemrefOpLowering,
       LoadOpLowering,
       MemRefCastOpLowering,
       MemRefCopyOpLowering,
+      MemorySpaceCastOpLowering,
       MemRefReinterpretCastOpLowering,
       MemRefReshapeOpLowering,
       PrefetchOpLowering,
       RankOpLowering,
       ReassociatingReshapeOpConversion<memref::ExpandShapeOp>,
       ReassociatingReshapeOpConversion<memref::CollapseShapeOp>,
       StoreOpLowering,
       SubViewOpLowering,

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 public:
   using OpConversionPattern<memref::StoreOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
+/// Converts memref.memory_space_cast to the appropriate spirv cast operations.
+class MemorySpaceCastOpPattern final
+    : public OpConversionPattern<memref::MemorySpaceCastOp> {
+public:
+  using OpConversionPattern<memref::MemorySpaceCastOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::MemorySpaceCastOp addrCastOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts memref.store to spirv.Store.
 class StoreOpPattern final : public OpConversionPattern<memref::StoreOp> {
 public:
   using OpConversionPattern<memref::StoreOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
IntStoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
   rewriter.eraseOp(storeOp);
 
   assert(accessChainOp.use_empty());
   rewriter.eraseOp(accessChainOp);
 
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// MemorySpaceCastOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult MemorySpaceCastOpPattern::matchAndRewrite(
+    memref::MemorySpaceCastOp addrCastOp, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  Location loc = addrCastOp.getLoc();
+  auto &typeConverter = *getTypeConverter<SPIRVTypeConverter>();
+  if (!typeConverter.allows(spirv::Capability::Kernel))
+    return rewriter.notifyMatchFailure(
+        loc, "address space casts require kernel capability");
+
+  auto sourceType = addrCastOp.getSource().getType().dyn_cast<MemRefType>();
+  if (!sourceType)
+    return rewriter.notifyMatchFailure(
+        loc, "SPIR-V lowering requires ranked memref types");
+  auto resultType = addrCastOp.getResult().getType().cast<MemRefType>();
+
+  auto sourceStorageClassAttr =
+      sourceType.getMemorySpace().dyn_cast_or_null<spirv::StorageClassAttr>();
+  if (!sourceStorageClassAttr)
+    return rewriter.notifyMatchFailure(loc, [sourceType](Diagnostic &diag) {
+      diag << "source address space " << sourceType.getMemorySpace()
+           << " must be a SPIR-V storage class";
+    });
+  auto resultStorageClassAttr =
+      resultType.getMemorySpace().dyn_cast_or_null<spirv::StorageClassAttr>();
+  if (!resultStorageClassAttr)
+    return rewriter.notifyMatchFailure(loc, [resultType](Diagnostic &diag) {
+      diag << "result address space " << resultType.getMemorySpace()
+           << " must be a SPIR-V storage class";
+    });
+
+  spirv::StorageClass sourceSc = sourceStorageClassAttr.getValue();
+  spirv::StorageClass resultSc = resultStorageClassAttr.getValue();
+
+  Value result = adaptor.getSource();
+  Type resultPtrType = typeConverter.convertType(resultType);
+  Type genericPtrType = resultPtrType;
+  // SPIR-V doesn't have a general address space cast operation. Instead, it has
+  // conversions to and from generic pointers. To implement the general case,
+  // we use specific-to-generic conversions when the source class is not
+  // generic. Then when the result storage class is not generic, we convert the
+  // generic pointer (either the input on ar intermediate result) to theat
+  // class. This also means that we'll need the intermediate generic pointer
+  // type if neither the source or destination have it.
+  if (sourceSc != spirv::StorageClass::Generic &&
+      resultSc != spirv::StorageClass::Generic) {
+    Type intermediateType =
+        MemRefType::get(sourceType.getShape(), sourceType.getElementType(),
+                        sourceType.getLayout(),
+                        rewriter.getAttr<spirv::StorageClassAttr>(
+                            spirv::StorageClass::Generic));
+    genericPtrType = typeConverter.convertType(intermediateType);
+  }
+  if (sourceSc != spirv::StorageClass::Generic) {
+    result =
+        rewriter.create<spirv::PtrCastToGenericOp>(loc, genericPtrType, result);
+  }
+  if (resultSc != spirv::StorageClass::Generic) {
+    result =
+        rewriter.create<spirv::GenericCastToPtrOp>(loc, resultPtrType, result);
+  }
+  rewriter.replaceOp(addrCastOp, result);
+  return success();
+}
+
 LogicalResult
 StoreOpPattern::matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
   auto memrefType = storeOp.getMemref().getType().cast<MemRefType>();
   if (memrefType.getElementType().isSignlessInteger())
     return failure();
   auto storePtr = spirv::getElementPtr(
       *getTypeConverter<SPIRVTypeConverter>(), memrefType, adaptor.getMemref(),

 
 //===----------------------------------------------------------------------===//
 // Pattern population
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 void populateMemRefToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                    RewritePatternSet &patterns) {
-  patterns
-      .add<AllocaOpPattern, AllocOpPattern, DeallocOpPattern, IntLoadOpPattern,
-           IntStoreOpPattern, LoadOpPattern, StoreOpPattern>(
-          typeConverter, patterns.getContext());
+  patterns.add<AllocaOpPattern, AllocOpPattern, DeallocOpPattern,
+               IntLoadOpPattern, IntStoreOpPattern, LoadOpPattern,
+               MemorySpaceCastOpPattern, StoreOpPattern>(typeConverter,
+                                                         patterns.getContext());
 }
 } // namespace mlir

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 
 template <typename AllocLikeOp>
 static LogicalResult verifyAllocLikeOp(AllocLikeOp op) {
   static_assert(llvm::is_one_of<AllocLikeOp, AllocOp, AllocaOp>::value,
                 "applies to only alloc or alloca");
   auto memRefType = op.getResult().getType().template dyn_cast<MemRefType>();
   if (!memRefType)
     return op.emitOpError("result must be a memref");
 
   if (static_cast<int64_t>(op.getDynamicSizes().size()) !=
       memRefType.getNumDynamicDims())
     return op.emitOpError("dimension operand count does not equal memref "
                           "dynamic dimension count");
 
   unsigned numSymbols = 0;
   if (!memRefType.getLayout().isIdentity())
     numSymbols = memRefType.getLayout().getAffineMap().getNumSymbols();
   if (op.getSymbolOperands().size() != numSymbols)
     return op.emitOpError("symbol operand count does not equal memref symbol "
                           "count: expected ")
            << numSymbols << ", got " << op.getSymbolOperands().size();
 
   return success();
 }
 
 LogicalResult AllocOp::verify() { return verifyAllocLikeOp(*this); }
 
 LogicalResult AllocaOp::verify() {
   // An alloca op needs to have an ancestor with an allocation scope trait.
   if (!(*this)->getParentWithTrait<OpTrait::AutomaticAllocationScope>())

 OpFoldResult LoadOp::fold(FoldAdaptor adaptor) {
   /// load(memrefcast) -> load
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
 //===----------------------------------------------------------------------===//
+// MemorySpaceCastOp
+//===----------------------------------------------------------------------===//
+
+void MemorySpaceCastOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  setNameFn(getResult(), "memspacecast");
+}
+
+bool MemorySpaceCastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {
+  if (inputs.size() != 1 || outputs.size() != 1)
+    return false;
+  Type a = inputs.front(), b = outputs.front();
+  auto aT = a.dyn_cast<MemRefType>();
+  auto bT = b.dyn_cast<MemRefType>();
+
+  auto uaT = a.dyn_cast<UnrankedMemRefType>();
+  auto ubT = b.dyn_cast<UnrankedMemRefType>();
+
+  if (aT && bT) {
+    if (aT.getElementType() != bT.getElementType())
+      return false;
+    if (aT.getLayout() != bT.getLayout())
+      return false;
+    if (aT.getShape() != bT.getShape())
+      return false;
+    return true;
+  }
+  if (uaT && ubT) {
+    return uaT.getElementType() == ubT.getElementType();
+  }
+  return false;
+}
+
+OpFoldResult MemorySpaceCastOp::fold(FoldAdaptor adaptor) {
+  // memory_space_cast(memory_space_cast(v, t1), t2) -> memory_space_cast(v,
+  // t2)
+  if (auto parentCast = getSource().getDefiningOp<MemorySpaceCastOp>()) {
+    getSourceMutable().assign(parentCast.getSource());
+    return getResult();
+  }
+  return Value{};
+}
+
+//===----------------------------------------------------------------------===//
 // PrefetchOp
 //===----------------------------------------------------------------------===//
 
 void PrefetchOp::print(OpAsmPrinter &p) {
   p << " " << getMemref() << '[';
   p.printOperands(getIndices());
   p << ']' << ", " << (getIsWrite() ? "write" : "read");
   p << ", locality<" << getLocalityHint();

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

 //  CHECK-NEXT:  %[[addr:.*]] = llvm.getelementptr %[[ptr]][%[[off1]]] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
 //  CHECK-NEXT:  llvm.store %{{.*}}, %[[addr]] : !llvm.ptr<f32>
   memref.store %val, %mixed[%i, %j] : memref<42x?xf32>
   return
 }
 
 // -----
 
+// FIXME: the *ToLLVM passes don't use information from data layouts
+// to set address spaces, so the constants below don't reflect the layout
+// Update this test once that data layout attribute works how we'd expect it to.
+module attributes { dlti.dl_spec = #dlti.dl_spec<
+  #dlti.dl_entry<!llvm.ptr, dense<[64, 64, 64]> : vector<3xi32>>,
+  #dlti.dl_entry<!llvm.ptr<1>, dense<[32, 32, 32]> : vector<3xi32>>> }  {
+  // CHECK-LABEL: @memref_memory_space_cast
+  func.func @memref_memory_space_cast(%input : memref<*xf32>) -> memref<*xf32, 1> {
+    %cast = memref.memory_space_cast %input : memref<*xf32> to memref<*xf32, 1>
+    return %cast : memref<*xf32, 1>
+  }
+}
+// CHECK: [[INPUT:%.*]] = builtin.unrealized_conversion_cast %{{.*}} to !llvm.struct<(i64, ptr<i8>)>
+// CHECK: [[RANK:%.*]] = llvm.extractvalue [[INPUT]][0] : !llvm.struct<(i64, ptr<i8>)>
+// CHECK: [[SOURCE_DESC:%.*]] = llvm.extractvalue [[INPUT]][1]
+// CHECK: [[RESULT_0:%.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
+// CHECK: [[RESULT_1:%.*]] = llvm.insertvalue [[RANK]], [[RESULT_0]][0] : !llvm.struct<(i64, ptr<i8>)>
+
+// Compute size in bytes to allocate result ranked descriptor
+// CHECK: [[C1:%.*]] = llvm.mlir.constant(1 : index) : i64
+// CHECK: [[C2:%.*]] = llvm.mlir.constant(2 : index) : i64
+// CHECK: [[INDEX_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
+// CHECK: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : index) : i64
+// CHECK: [[DOUBLE_PTR_SIZE:%.*]] = llvm.mul [[C2]], [[PTR_SIZE]]
+// CHECK: [[DOUBLE_RANK:%.*]] = llvm.mul [[C2]], %{{.*}}
+// CHECK: [[NUM_INDEX_VALS:%.*]] = llvm.add [[DOUBLE_RANK]], [[C1]]
+// CHECK: [[INDEX_VALS_SIZE:%.*]] = llvm.mul [[NUM_INDEX_VALS]], [[INDEX_SIZE]]
+// CHECK: [[DESC_ALLOC_SIZE:%.*]] = llvm.add [[DOUBLE_PTR_SIZE]], [[INDEX_VALS_SIZE]]
+// CHECK: [[RESULT_DESC:%.*]] = llvm.alloca [[DESC_ALLOC_SIZE]] x i8
+// CHECK: llvm.insertvalue [[RESULT_DESC]], [[RESULT_1]][1]
+
+// Cast pointers
+// CHECK: [[SOURCE_DESC_CAST_ALLOC:%.*]] = llvm.bitcast [[SOURCE_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32>>
+// CHECK: [[SOURCE_ALLOC:%.*]] = llvm.load [[SOURCE_DESC_CAST_ALLOC]]
+// CHECK: [[SOURCE_DESC_CAST_ALIGN:%.*]] = llvm.bitcast [[SOURCE_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32>>
+// CHECK: [[SOURCE_ALIGN_GEP:%.*]] = llvm.getelementptr [[SOURCE_DESC_CAST_ALIGN]][1]
+// CHECK: [[SOURCE_ALIGN:%.*]] = llvm.load [[SOURCE_ALIGN_GEP]] : !llvm.ptr<ptr<f32>>
+// CHECK: [[RESULT_ALLOC:%.*]] = llvm.addrspacecast [[SOURCE_ALLOC]] : !llvm.ptr<f32> to !llvm.ptr<f32, 1>
+// CHECK: [[RESULT_ALIGN:%.*]] = llvm.addrspacecast [[SOURCE_ALIGN]] : !llvm.ptr<f32> to !llvm.ptr<f32, 1>
+// CHECK: [[RESULT_DESC_CAST_ALLOC:%.*]] = llvm.bitcast [[RESULT_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32, 1>>
+// CHECK: llvm.store [[RESULT_ALLOC]], [[RESULT_DESC_CAST_ALLOC]] : !llvm.ptr<ptr<f32, 1>>
+// CHECK: [[RESULT_DESC_CAST_ALIGN:%.*]] = llvm.bitcast [[RESULT_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32, 1>>
+// CHECK: [[RESULT_ALIGN_GEP:%.*]] = llvm.getelementptr [[RESULT_DESC_CAST_ALIGN]][1]
+// CHECK: llvm.store [[RESULT_ALIGN]], [[RESULT_ALIGN_GEP]] : !llvm.ptr<ptr<f32, 1>>
+
+// Memcpy remaniing values
+
+// CHECK: [[SOURCE_DESC_CAST_OFFSET:%.*]] = llvm.bitcast [[SOURCE_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32>>
+// CHECK: [[SOURCE_OFFSET_GEP:%.*]] = llvm.getelementptr [[SOURCE_DESC_CAST_OFFSET]][2]
+// CHECK: [[SOURCE_OFFSET_IDX:%.*]] = llvm.bitcast [[SOURCE_OFFSET_GEP]] : !llvm.ptr<ptr<f32>> to !llvm.ptr<i64>
+// CHECK: [[RESULT_DESC_CAST_OFFSET:%.*]] = llvm.bitcast [[RESULT_DESC]] : !llvm.ptr<i8> to !llvm.ptr<ptr<f32, 1>>
+// CHECK: [[RESULT_OFFSET_GEP:%.*]] = llvm.getelementptr [[RESULT_DESC_CAST_OFFSET]][2]
+// CHECK: [[RESULT_OFFSET_IDX:%.*]] = llvm.bitcast [[RESULT_OFFSET_GEP]] : !llvm.ptr<ptr<f32, 1>> to !llvm.ptr<i64>
+// CHECK: [[SIZEOF_TWO_RESULT_PTRS:%.*]] = llvm.mlir.constant(16 : index) : i64
+// CHECK: [[COPY_SIZE:%.*]] = llvm.sub [[DESC_ALLOC_SIZE]], [[SIZEOF_TWO_RESULT_PTRS]]
+// CHECK: [[FALSE:%.*]] = llvm.mlir.constant(false) : i1
+// CHECK: "llvm.intr.memcpy"([[RESULT_OFFSET_IDX]], [[SOURCE_OFFSET_IDX]], [[COPY_SIZE]], [[FALSE]])
+
+// -----
+
 // CHECK-LABEL: func @memref_cast_static_to_dynamic
 func.func @memref_cast_static_to_dynamic(%static : memref<10x42xf32>) {
 // CHECK-NOT: llvm.bitcast
   %0 = memref.cast %static : memref<10x42xf32> to memref<?x?xf32>
   return
 }
 
 // -----

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

 // CHECK:           %[[descriptor_update4:.*]] = llvm.insertvalue %[[dst_dim]], %[[descriptor_update3]][3, 0]
 // CHECK:           %[[descriptor_update5:.*]] = builtin.unrealized_conversion_cast %[[descriptor_update4]]
 // CHECK:           return %[[descriptor_update5]] : memref<4xf32>
 
 
   %out = memref.realloc %in {alignment = 8} : memref<2xf32> to memref<4xf32>
   return %out : memref<4xf32>
 }
+
+// -----
+
+// CHECK-LABEL: @memref_memory_space_cast
+func.func @memref_memory_space_cast(%input : memref<?xf32>) -> memref<?xf32, 1> {
+  %cast = memref.memory_space_cast %input : memref<?xf32> to memref<?xf32, 1>
+  return %cast : memref<?xf32, 1>
+}
+// CHECK: [[INPUT:%.*]] = builtin.unrealized_conversion_cast %{{.*}}
+//  CHECK: [[ALLOC:%.*]] = llvm.extractvalue [[INPUT]][0]
+//  CHECK: [[ALIGN:%.*]] = llvm.extractvalue [[INPUT]][1]
+// CHECK: [[OFFSET:%.*]] = llvm.extractvalue [[INPUT]][2]
+//   CHECK: [[SIZE:%.*]] = llvm.extractvalue [[INPUT]][3, 0]
+// CHECK: [[STRIDE:%.*]] = llvm.extractvalue [[INPUT]][4, 0]
+// CHECK: [[CAST_ALLOC:%.*]] = llvm.addrspacecast [[ALLOC]] : !llvm.ptr<f32> to !llvm.ptr<f32, 1>
+// CHECK: [[CAST_ALIGN:%.*]] = llvm.addrspacecast [[ALIGN]] : !llvm.ptr<f32> to !llvm.ptr<f32, 1>
+// CHECK: [[RESULT_0:%.*]] = llvm.mlir.undef
+// CHECK: [[RESULT_1:%.*]] = llvm.insertvalue [[CAST_ALLOC]], [[RESULT_0]][0]
+// CHECK: [[RESULT_2:%.*]] = llvm.insertvalue [[CAST_ALIGN]], [[RESULT_1]][1]
+// CHECK: [[RESULT_3:%.*]] = llvm.insertvalue [[OFFSET]], [[RESULT_2]][2]
+// CHECK: [[RESULT_4:%.*]] = llvm.insertvalue [[SIZE]], [[RESULT_3]][3, 0]
+// CHECK: [[RESULT_5:%.*]] = llvm.insertvalue [[STRIDE]], [[RESULT_4]][4, 0]
+// CHECK: [[RESULT:%.*]] = builtin.unrealized_conversion_cast [[RESULT_5]] : {{.*}} to memref<?xf32, 1>
+// CHECK: return [[RESULT]]

mlir/test/Conversion/MemRefToSPIRV/memref-to-spirv.mlir

func.func @store_i1(%dst: memref<4xi1, #spirv.storage_class<CrossWorkgroup>>, %i: index) {
   memref.store %true, %dst[%i]: memref<4xi1, #spirv.storage_class<CrossWorkgroup>>
   return
 }
 
 } // end module
 
 // -----
 
+// Check address space casts
+
+module attributes {
+  spirv.target_env = #spirv.target_env<
+    #spirv.vce<v1.0,
+      [
+        Kernel, Addresses, GenericPointer], []>, #spirv.resource_limits<>>
+} {
+
+// CHECK-LABEL: func.func @memory_space_cast
+func.func @memory_space_cast(%arg: memref<4xf32, #spirv.storage_class<CrossWorkgroup>>)
+    -> memref<4xf32, #spirv.storage_class<Function>> {
+  // CHECK: %[[ARG_CAST:.+]] = builtin.unrealized_conversion_cast {{.*}} to !spirv.ptr<!spirv.array<4 x f32>, CrossWorkgroup>
+  // CHECK: %[[TO_GENERIC:.+]] = spirv.PtrCastToGeneric %[[ARG_CAST]] : !spirv.ptr<!spirv.array<4 x f32>, CrossWorkgroup> to !spirv.ptr<!spirv.array<4 x f32>, Generic>
+  // CHECK: %[[TO_PRIVATE:.+]] = spirv.GenericCastToPtr %[[TO_GENERIC]] : !spirv.ptr<!spirv.array<4 x f32>, Generic> to !spirv.ptr<!spirv.array<4 x f32>, Function>
+  // CHECK: %[[RET:.+]] = builtin.unrealized_conversion_cast %[[TO_PRIVATE]]
+  // CHECK: return %[[RET]]
+  %ret = memref.memory_space_cast %arg : memref<4xf32, #spirv.storage_class<CrossWorkgroup>>
+    to memref<4xf32, #spirv.storage_class<Function>>
+  return %ret : memref<4xf32, #spirv.storage_class<Function>>
+}
+
+} // end module
+
+// -----
+
 // Check that access chain indices are properly adjusted if non-32-bit types are
 // emulated via 32-bit types.
 // TODO: Test i64 types.
 module attributes {
   spirv.target_env = #spirv.target_env<
     #spirv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, #spirv.resource_limits<>>
 } {
 

mlir/test/Dialect/MemRef/canonicalize.mlir

func.func @fold_trivial_subviews(%m: memref<?xf32, strided<[?], offset: ?>>,
   %0 = memref.subview %m[5] [%sz] [1]
       : memref<?xf32, strided<[?], offset: ?>>
         to memref<?xf32, strided<[?], offset: ?>>
   %1 = memref.subview %0[0] [%sz] [1]
       : memref<?xf32, strided<[?], offset: ?>>
         to memref<?xf32, strided<[?], offset: ?>>
   return %1 : memref<?xf32, strided<[?], offset: ?>>
 }
+
+// -----
+
+// CHECK-LABEL: func @fold_trivial_memory_space_cast(
+//  CHECK-SAME:     %[[arg:.*]]: memref<?xf32>
+//       CHECK:   return %[[arg]]
+func.func @fold_trivial_memory_space_cast(%arg : memref<?xf32>) -> memref<?xf32> {
+  %0 = memref.memory_space_cast %arg : memref<?xf32> to memref<?xf32>
+  return %0 : memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_multiple_memory_space_cast(
+//  CHECK-SAME:     %[[arg:.*]]: memref<?xf32>
+//       CHECK:   %[[res:.*]] = memref.memory_space_cast %[[arg]] : memref<?xf32> to memref<?xf32, 2>
+//       CHECK:   return %[[res]]
+func.func @fold_multiple_memory_space_cast(%arg : memref<?xf32>) -> memref<?xf32, 2> {
+  %0 = memref.memory_space_cast %arg : memref<?xf32> to memref<?xf32, 1>
+  %1 = memref.memory_space_cast %0 : memref<?xf32, 1> to memref<?xf32, 2>
+  return %1 : memref<?xf32, 2>
+}

mlir/test/Dialect/MemRef/ops.mlir

func.func @memref_realloc_dd(%src : memref<?xf32>, %d: index)
   return %0 : memref<?xf32>
 }
 
 // CHECK-LABEL: func @memref_extract_aligned_pointer
 func.func @memref_extract_aligned_pointer(%src : memref<?xf32>) -> index {
   %0 = memref.extract_aligned_pointer_as_index %src : memref<?xf32> -> index
   return %0 : index
 }
+
+// CHECK-LABEL: func @memref_memory_space_cast
+func.func @memref_memory_space_cast(%src : memref<?xf32>) -> memref<?xf32, 1> {
+  %dst = memref.memory_space_cast %src : memref<?xf32> to memref<?xf32, 1>
+  return %dst : memref<?xf32, 1>
+}