[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

822

Nit

842

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

856

Nit: please add trailing periods to all comments.

858–861

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).

918–920

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

918–920

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

1265–1268

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

public:
   static void unpack(OpBuilder &builder, Location loc, Value packed,
                      SmallVectorImpl<Value> &results);
 
   /// Returns the number of non-aggregate values that would be produced by
   /// `unpack`.
   static unsigned getNumUnpackedValues() { return 2; }
 
   /// Builds IR computing the sizes in bytes (suitable for opaque allocation)
-  /// and appends the corresponding values into `sizes`.
+  /// and appends the corresponding values into `sizes`. `addressSpaces`
+  /// which must have the same length as `values`, is needed to handle layouts
+  /// where sizeof(ptr addrspace(N)) != sizeof(ptr addrspace(0)).
   static void computeSizes(OpBuilder &builder, Location loc,
                            LLVMTypeConverter &typeConverter,
                            ArrayRef<UnrankedMemRefDescriptor> values,
+                           ArrayRef<unsigned> addressSpaces,
                            SmallVectorImpl<Value> &sizes);
 
   /// TODO: The following accessors don't take alignment rules between elements
   /// of the descriptor struct into account. For some architectures, it might be
   /// necessary to extend them and to use `llvm::DataLayout` contained in
   /// `LLVMTypeConverter`.
 
   /// Builds IR extracting the allocated pointer from the descriptor.
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 assemblyFormat = [{
     `(`$dynamicSizes`)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
+// AddressSpaceCastOp
+//===----------------------------------------------------------------------===//
+def MemRef_AddressSpaceCastOp : MemRef_Op<"address_space_cast", [
+      DeclareOpInterfaceMethods<CastOpInterface>,
+      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>,
+      MemRefsNormalizable,
+      Pure,
+      SameOperandsAndResultElementType,
+      SameOperandsAndResultShape,
+      ViewLikeOpInterface
+    ]> {
+  let summary = "memref address space cast operation";
+  let description = [{
+    Syntax:
+
+    ```
+    operation ::= ssa-id `=` `memref.address_space_cast` ssa-use `:` type `to` type
+    ```
+
+    This operation casts memref values between adress 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.address_space_cast %1 : memref<?xf32, 5> to memref<?xf32>
+    // Cast a generic pointer to workgroup-local memory
+    %4 = memref.address_space_cast %3 : memref<5x4xi32> to memref<5x34xi32, 3>
+    // Cast between two non-default memory spaces
+    %6 = memref.address_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;
+}
+
+//===----------------------------------------------------------------------===//
 // AssumeAlignmentOp
 //===----------------------------------------------------------------------===//
 
 def AssumeAlignmentOp : MemRef_Op<"assume_alignment"> {
   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
let description = [{
     instruction cache.
   }];
 
   let arguments = (ins AnyMemRef:$memref, Variadic<Index>:$indices,
                        BoolAttr:$isWrite,
                        ConfinedAttr<I32Attr, [IntMinValue<0>,
                                           IntMaxValue<3>]>:$localityHint,
                        BoolAttr:$isDataCache);
 
   let extraClassDeclaration = [{
     MemRefType getMemRefType() {
       return getMemref().getType().cast<MemRefType>();
     }
     static StringRef getLocalityHintAttrStrName() { return "localityHint"; }
     static StringRef getIsWriteAttrStrName() { return "isWrite"; }
     static StringRef getIsDataCacheAttrStrName() { return "isDataCache"; }
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasFolder = 1;

mlir/lib/Conversion/LLVMCommon/MemRefBuilder.cpp

void UnrankedMemRefDescriptor::unpack(OpBuilder &builder, Location loc,
   UnrankedMemRefDescriptor d(packed);
   results.reserve(results.size() + 2);
   results.push_back(d.rank(builder, loc));
   results.push_back(d.memRefDescPtr(builder, loc));
 }
 
 void UnrankedMemRefDescriptor::computeSizes(
     OpBuilder &builder, Location loc, LLVMTypeConverter &typeConverter,
-    ArrayRef<UnrankedMemRefDescriptor> values, SmallVectorImpl<Value> &sizes) {
+    ArrayRef<UnrankedMemRefDescriptor> values, ArrayRef<unsigned> addressSpaces,
+    SmallVectorImpl<Value> &sizes) {
   if (values.empty())
     return;
-
+  assert(values.size() == addressSpaces.size() &&
+         "must provide address space for each descriptor");
   // Cache the index type.
   Type indexType = typeConverter.getIndexType();
 
   // Initialize shared constants.
   Value one = createIndexAttrConstant(builder, loc, indexType, 1);
   Value two = createIndexAttrConstant(builder, loc, indexType, 2);
-  Value pointerSize = createIndexAttrConstant(
-      builder, loc, indexType, ceilDiv(typeConverter.getPointerBitwidth(), 8));
   Value indexSize =
       createIndexAttrConstant(builder, loc, indexType,
                               ceilDiv(typeConverter.getIndexTypeBitwidth(), 8));
 
   sizes.reserve(sizes.size() + values.size());
-  for (UnrankedMemRefDescriptor desc : values) {
+  for (auto pair : llvm::zip(values, addressSpaces)) {
+    UnrankedMemRefDescriptor desc = std::get<0>(pair);
+    unsigned addressSpace = std::get<1>(pair);
     // Emit IR computing the memory necessary to store the descriptor. This
     // assumes the descriptor to be
     //   { type*, type*, index, index[rank], index[rank] }
     // and densely packed, so the total size is
     //   2 * sizeof(pointer) + (1 + 2 * rank) * sizeof(index).
     // TODO: consider including the actual size (including eventual padding due
     // to data layout) into the unranked descriptor.
+    Value pointerSize = createIndexAttrConstant(
+        builder, loc, indexType,
+        ceilDiv(typeConverter.getPointerBitwidth(addressSpace), 8));
     Value doublePointerSize =
         builder.create<LLVM::MulOp>(loc, indexType, two, pointerSize);
 
     // (1 + 2 * rank) * sizeof(index)
     Value rank = desc.rank(builder, loc);
     Value doubleRank = builder.create<LLVM::MulOp>(loc, indexType, two, rank);
     Value doubleRankIncremented =
         builder.create<LLVM::AddOp>(loc, indexType, doubleRank, one);
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/LLVMCommon/Pattern.cpp

 LogicalResult ConvertToLLVMPattern::copyUnrankedDescriptors(
     OpBuilder &builder, Location loc, TypeRange origTypes,
     SmallVectorImpl<Value> &operands, bool toDynamic) const {
   assert(origTypes.size() == operands.size() &&
          "expected as may original types as operands");
 
   // Find operands of unranked memref type and store them.
   SmallVector<UnrankedMemRefDescriptor, 4> unrankedMemrefs;
-  for (unsigned i = 0, e = operands.size(); i < e; ++i)
-    if (origTypes[i].isa<UnrankedMemRefType>())
+  SmallVector<unsigned, 4> unrankedAddressSpaces;
+  for (unsigned i = 0, e = operands.size(); i < e; ++i) {
+    if (auto memRefType = origTypes[i].dyn_cast<UnrankedMemRefType>()) {
       unrankedMemrefs.emplace_back(operands[i]);
+      unrankedAddressSpaces.emplace_back(memRefType.getMemorySpaceAsInt());
+    }
+  }
 
   if (unrankedMemrefs.empty())
     return success();
 
   // Compute allocation sizes.
   SmallVector<Value, 4> sizes;
   UnrankedMemRefDescriptor::computeSizes(builder, loc, *getTypeConverter(),
-                                         unrankedMemrefs, sizes);
+                                         unrankedMemrefs, unrankedAddressSpaces,
+                                         sizes);
 
   // Get frequently used types.
   MLIRContext *context = builder.getContext();
   Type voidPtrType = LLVM::LLVMPointerType::get(IntegerType::get(context, 8));
   auto i1Type = IntegerType::get(context, 1);
   Type indexType = getTypeConverter()->getIndexType();
 
   // Find the malloc and free, or declare them if necessary.

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/BlockAndValueMapping.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
 
if (auto rankedMemRefType = operandType.dyn_cast<MemRefType>()) {
       rewriter.replaceOp(
           op, {createIndexConstant(rewriter, loc, rankedMemRefType.getRank())});
       return success();
     }
     return failure();
   }
 };
 
+struct AddressSpareCastOpLowering
+    : public ConvertOpToLLVMPattern<memref::AddressSpaceCastOp> {
+  using ConvertOpToLLVMPattern<
+      memref::AddressSpaceCastOp>::ConvertOpToLLVMPattern;
+
+  FailureOr<unsigned> getMemorySpace(Type type) const {
+    Attribute memorySpace;
+    if (auto mR = type.dyn_cast<MemRefType>())
+      memorySpace = mR.getMemorySpace();
+    else if (auto mU = type.dyn_cast<UnrankedMemRefType>())
+      memorySpace = mU.getMemorySpace();
+    else
+      return failure();
+
+    // Default memory space is 0
+    if (!memorySpace)
+      return 0u;
+    auto memorySpaceInt = memorySpace.dyn_cast<IntegerAttr>();
+    if (!memorySpaceInt)
+      return failure();
+    return static_cast<unsigned>(memorySpaceInt.getInt());
+  }
+
+  LogicalResult
+  matchAndRewrite(memref::AddressSpaceCastOp 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, 8> 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 = getMemorySpace(sourceType);
+      if (failed(maybeSourceAddrSpace))
+        return rewriter.notifyMatchFailure(loc,
+                                           "non-integer source address space");
+      unsigned sourceAddrSpace = *maybeSourceAddrSpace;
+      FailureOr<unsigned> maybeResultAddrSpace = getMemorySpace(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");
+  }
+};
+
 struct MemRefCastOpLowering : public ConvertOpToLLVMPattern<memref::CastOp> {
   using ConvertOpToLLVMPattern<memref::CastOp>::ConvertOpToLLVMPattern;
 
   LogicalResult match(memref::CastOp memRefCastOp) const override {
     Type srcType = memRefCastOp.getOperand().getType();
     Type dstType = memRefCastOp.getType();
 
     // memref::CastOp reduce to bitcast in the ranked MemRef case and can be
convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
 
     // Create the unranked memref descriptor that holds the ranked one. The
     // inner descriptor is allocated on stack.
     auto targetDesc = UnrankedMemRefDescriptor::undef(
         rewriter, loc, typeConverter->convertType(targetType));
     targetDesc.setRank(rewriter, loc, resultRank);
     SmallVector<Value, 4> sizes;
     UnrankedMemRefDescriptor::computeSizes(rewriter, loc, *getTypeConverter(),
-                                           targetDesc, sizes);
+                                           targetDesc, addressSpace, sizes);
     Value underlyingDescPtr = rewriter.create<LLVM::AllocaOp>(
         loc, getVoidPtrType(), sizes.front(), std::nullopt);
     targetDesc.setMemRefDescPtr(rewriter, loc, underlyingDescPtr);
 
     // Extract pointers and offset from the source memref.
     Value allocatedPtr, alignedPtr, offset;
     extractPointersAndOffset(loc, rewriter, *getTypeConverter(),
                              reshapeOp.getSource(), adaptor.getSource(),

 };
 
 } // namespace
 
 void mlir::populateMemRefToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                   RewritePatternSet &patterns) {
   // clang-format off
   patterns.add<
+      AddressSpareCastOpLowering,
       AllocaOpLowering,
       AllocaScopeOpLowering,
       AtomicRMWOpLowering,
       AssumeAlignmentOpLowering,
       ConvertExtractAlignedPointerAsIndex,
       DimOpLowering,
       ExtractStridedMetadataOpLowering,
       GenericAtomicRMWOpLowering,

mlir/lib/Conversion/MemRefToSPIRV/MemRefToSPIRV.cpp

 //===----------------------------------------------------------------------===//
 
 // Note that DRR cannot be used for the patterns in this file: we may need to
 // convert type along the way, which requires ConversionPattern. DRR generates
 // normal RewritePattern.
 
 namespace {
 
+/// Converts memref.address_space_cast to the appropriate spirv cast operations.
+class AddressSpaceCastOpPattern final
+    : public OpConversionPattern<memref::AddressSpaceCastOp> {
+public:
+  using OpConversionPattern<memref::AddressSpaceCastOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::AddressSpaceCastOp addrCastOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+
 /// Converts memref.alloca to SPIR-V Function variables.
 class AllocaOpPattern final : public OpConversionPattern<memref::AllocaOp> {
 public:
   using OpConversionPattern<memref::AllocaOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
public:
   LogicalResult
   matchAndRewrite(memref::StoreOp storeOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// AddressSpaceCastOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult AddressSpaceCastOpPattern::matchAndRewrite(
+    memref::AddressSpaceCastOp 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();
+}
+
+//===----------------------------------------------------------------------===//
 // AllocaOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 AllocaOpPattern::matchAndRewrite(memref::AllocaOp allocaOp, OpAdaptor adaptor,
                                  ConversionPatternRewriter &rewriter) const {
   MemRefType allocType = allocaOp.getType();
   if (!isAllocationSupported(allocaOp, allocType))

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

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

static SmallVector<int64_t> getConstantStrides(MemRefType memrefType) {
   LogicalResult hasStaticInformation =
       getStridesAndOffset(memrefType, strides, offset);
   if (failed(hasStaticInformation))
     return SmallVector<int64_t>();
   return strides;
 }
 
 //===----------------------------------------------------------------------===//
+// AddressSpaceCastOp
+//===----------------------------------------------------------------------===//
+
+void AddressSpaceCastOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+  setNameFn(getResult(), "addrcast");
+}
+
+bool AddressSpaceCastOp::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.getRank() != bT.getRank())
+      return false;
+    if (aT.getShape() != bT.getShape())
+      return false;
+    return true;
+  }
+  if (uaT && ubT) {
+    // The style from above led to a lint
+    bool areCompatible = uaT.getElementType() == ubT.getElementType();
+    return areCompatible;
+  }
+  return false;
+}
+
+OpFoldResult AddressSpaceCastOp::fold(ArrayRef<Attribute> operands) {
+  // address_space_cast(address_space_cast(v, t1), t2) -> address_space_cast(v,
+  // t2)
+  if (auto parentCast = getSource().getDefiningOp<AddressSpaceCastOp>()) {
+    getSourceMutable().assign(parentCast.getSource());
+    return getResult();
+  }
+  return Value{};
+}
+
+//===----------------------------------------------------------------------===//
 // AllocOp / AllocaOp
 //===----------------------------------------------------------------------===//
 
 void AllocOp::getAsmResultNames(
     function_ref<void(Value, StringRef)> setNameFn) {
   setNameFn(getResult(), "alloc");
 }
 

mlir/test/Conversion/FuncToLLVM/calling-convention.mlir

 // CHECK-LABEL: llvm.func @return_var_memref_caller
 func.func @return_var_memref_caller(%arg0: memref<4x3xf32>) {
   // CHECK: %[[CALL_RES:.*]] = llvm.call @return_var_memref
   %0 = call @return_var_memref(%arg0) : (memref<4x3xf32>) -> memref<*xf32>
 
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : index)
   // CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : index)
   // These sizes may depend on the data layout, not matching specific values.
-  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[IDX_SIZE:.*]] = llvm.mlir.constant
 
+  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[DOUBLE_PTR_SIZE:.*]] = llvm.mul %[[TWO]], %[[PTR_SIZE]]
   // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[CALL_RES]][0] : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[DOUBLE_RANK:.*]] = llvm.mul %[[TWO]], %[[RANK]]
   // CHECK: %[[DOUBLE_RANK_INC:.*]] = llvm.add %[[DOUBLE_RANK]], %[[ONE]]
   // CHECK: %[[TABLES_SIZE:.*]] = llvm.mul %[[DOUBLE_RANK_INC]], %[[IDX_SIZE]]
   // CHECK: %[[ALLOC_SIZE:.*]] = llvm.add %[[DOUBLE_PTR_SIZE]], %[[TABLES_SIZE]]
   // CHECK: %[[FALSE:.*]] = llvm.mlir.constant(false)
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca %[[ALLOC_SIZE]] x i8
func.func @return_var_memref(%arg0: memref<4x3xf32>) -> memref<*xf32> attributes { llvm.emit_c_interface } {
   // CHECK: %[[ALLOCA:.*]] = llvm.alloca
   // CHECK: %[[MEMORY:.*]] = llvm.bitcast %[[ALLOCA]]
   // CHECK: %[[RANK:.*]] = llvm.mlir.constant(2 : index)
   // CHECK: %[[DESC_0:.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
   // CHECK: %[[DESC_1:.*]] = llvm.insertvalue %[[RANK]], %[[DESC_0]][0]
   // CHECK: %[[DESC_2:.*]] = llvm.insertvalue %[[MEMORY]], %[[DESC_1]][1]
   %0 = memref.cast %arg0: memref<4x3xf32> to memref<*xf32>
 
-
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1 : index)
   // CHECK: %[[TWO:.*]] = llvm.mlir.constant(2 : index)
   // These sizes may depend on the data layout, not matching specific values.
-  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[IDX_SIZE:.*]] = llvm.mlir.constant
 
+  // CHECK: %[[PTR_SIZE:.*]] = llvm.mlir.constant
   // CHECK: %[[DOUBLE_PTR_SIZE:.*]] = llvm.mul %[[TWO]], %[[PTR_SIZE]]
   // CHECK: %[[DOUBLE_RANK:.*]] = llvm.mul %[[TWO]], %[[RANK]]
   // CHECK: %[[DOUBLE_RANK_INC:.*]] = llvm.add %[[DOUBLE_RANK]], %[[ONE]]
   // CHECK: %[[TABLES_SIZE:.*]] = llvm.mul %[[DOUBLE_RANK_INC]], %[[IDX_SIZE]]
   // CHECK: %[[ALLOC_SIZE:.*]] = llvm.add %[[DOUBLE_PTR_SIZE]], %[[TABLES_SIZE]]
   // CHECK: %[[FALSE:.*]] = llvm.mlir.constant(false)
   // CHECK: %[[ALLOCATED:.*]] = llvm.call @malloc(%[[ALLOC_SIZE]])
   // CHECK: "llvm.intr.memcpy"(%[[ALLOCATED]], %[[MEMORY]], %[[ALLOC_SIZE]], %[[FALSE]])

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_address_space_cast
+  func.func @memref_address_space_cast(%input : memref<*xf32>) -> memref<*xf32, 1> {
+    %cast = memref.address_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
 }
 
 // -----

 // CHECK: [[SHAPE:%.*]] = builtin.unrealized_conversion_cast %{{.*}} to [[SHAPE_TY:!.*]]
 // CHECK: [[RANK:%.*]] = llvm.extractvalue [[SHAPE]][3, 0] : [[SHAPE_TY]]
 // CHECK: [[UNRANKED_OUT_O:%.*]] = llvm.mlir.undef : !llvm.struct<(i64, ptr<i8>)>
 // CHECK: [[UNRANKED_OUT_1:%.*]] = llvm.insertvalue [[RANK]], [[UNRANKED_OUT_O]][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: [[PTR_SIZE:%.*]] = llvm.mlir.constant(8 : 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]] : i64
 // CHECK: [[DESC_ALLOC_SIZE:%.*]] = llvm.add [[DOUBLE_PTR_SIZE]], %{{.*}}
 // CHECK: [[UNDERLYING_DESC:%.*]] = llvm.alloca [[DESC_ALLOC_SIZE]] x i8
 // CHECK: llvm.insertvalue [[UNDERLYING_DESC]], [[UNRANKED_OUT_1]][1]
 
 // Set allocated, aligned pointers and offset.
 // CHECK: [[ALLOC_PTR:%.*]] = llvm.extractvalue [[INPUT]][0] : [[INPUT_TY]]
 // CHECK: [[ALIGN_PTR:%.*]] = llvm.extractvalue [[INPUT]][1] : [[INPUT_TY]]

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_address_space_cast
+func.func @memref_address_space_cast(%input : memref<?xf32>) -> memref<?xf32, 1> {
+  %cast = memref.address_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 @address_space_cast
+func.func @address_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.address_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_address_space_cast(
+//  CHECK-SAME:     %[[arg:.*]]: memref<?xf32>
+//       CHECK:   return %[[arg]]
+func.func @fold_trivial_address_space_cast(%arg : memref<?xf32>) -> memref<?xf32> {
+  %0 = memref.address_space_cast %arg : memref<?xf32> to memref<?xf32>
+  return %0 : memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @fold_multiple_address_space_cast(
+//  CHECK-SAME:     %[[arg:.*]]: memref<?xf32>
+//       CHECK:   %[[res:.*]] = memref.address_space_cast %[[arg]] : memref<?xf32> to memref<?xf32, 2>
+//       CHECK:   return %[[res]]
+func.func @fold_multiple_address_space_cast(%arg : memref<?xf32>) -> memref<?xf32, 2> {
+  %0 = memref.address_space_cast %arg : memref<?xf32> to memref<?xf32, 1>
+  %1 = memref.address_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_address_space_cast
+func.func @memref_address_space_cast(%src : memref<?xf32>) -> memref<?xf32, 1> {
+  %dst = memref.address_space_cast %src : memref<?xf32> to memref<?xf32, 1>
+  return %dst : memref<?xf32, 1>
+}