Add generic type attribute mapping infrastructure, use it in GpuToX

[krzysz00]

Summary

Remapping memory spaces is a function often needed in type
conversions, most often when going to LLVM or to/from SPIR-V (a future
commit), and it is possible that such remappings may become more
common in the future as dialects take advantage of the more generic
memory space infrastructure.

Currently, memory space remappings are handled by running a
special-purpose conversion pass before the main conversion that
changes the address space attributes. In this commit, this approach is
replaced by adding a notion of type attribute conversions
TypeConverter, which is then used to convert memory space attributes.

Then, we use this infrastructure throughout the *ToLLVM conversions.
This has the advantage of loosing the requirements on the inputs to
those passes from "all address spaces must be integers" to "all
memory spaces must be convertible to integer spaces", a looser
requirement that reduces the coupling between portions of MLIR.

ON top of that, this change leads to the removal of most of the calls
to getMemorySpaceAsInt(), bringing us closer to removing it.

(A rework of the SPIR-V conversions to use this new system will be in
a folowup commit.)

As a note, one long-term motivation for this change is that I would
eventually like to add an allocaMemorySpace key to MLIR data layouts
and then call getMemRefAddressSpace(allocaMemorySpace) in the
relevant *ToLLVM in order to ensure all alloca()s, whether incoming or
produces during the LLVM lowering, have the correct address space for
a given target.

I expect that the type attribute conversion system may be useful in
other contexts.

[ftynse]

The overall approach looks good to me. Two design comments:

1. There is no reason to limit this to memory space attributes only and no good reason to introduce the notion of memory space to the top-level TypeConverter. Instead, this can be a facility to convert attributes. (Renaming TypeConverter to AttrAndTypeConverter is probably too much churn for one use case, so one can argue that it intends to support attributes present in types for now).
2. It looks strange that the attribute conversion must always succeed. IMO, it should have an ability to fail, especially given the extension beyond memory spaces. When set up for conversions where there is a default, it is possible to add a low-priority catch-all conversion rule, e.g., for all attributes implementing the memory space interface. The code, inherited from type converters, already does a loop around conversions to support this.

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

135

mlir/include/mlir/Transforms/DialectConversion.h

163

167

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

316

We usually want a test for user-visible error messages.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

689

Nit: I see it was there before this change, but could we avoid the magic number here?

[ftynse]

(requesting changes so it shows up in my review thread again)

[krzysz00]

@ftynse I appreciate the design feedback!

I can agree that having a more general facility for attributes in types would be quite nice, and that having the conversion being able to fail would be good (and would make for a nice symmetry with the actual type converter).

My one concern is that attributes are used in different context, and the conversions could differ depending on those contexts. For example, we could imagine that, while memory spaces need to be converted to integers when lowering to LLVM (which makes IntegerAttr a trivial conversion), there could be some other type and attribute where you want integers to be converted to some kind of backend-specific enum, and so you'd need a different set of conversions for that context.

So the general signature might be Optional<Attribute> convertTypeAttribute(SomethingIDontKnowYet context, Attribute attribute);

One option for SomethingIDontKnowYet is the type the attribute came from.
Then, the signature of the conversion functions would be addTypeAttributeConversion([Type, Attribute] -> FailureOr<Attribute>) and you could, for example, add conversions for BaseMemRefType and so on. This also means attribute converters get some extra context to peek in to.

Another thought would be some sort of key, but then we have a stringly-typed hash table of conversion groups.

.... yeah, having written that out, convertTypeAttribute(Type type, Attribute attribute) is the better general signature.

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

689

I'd leave it there for backwards compatibility.

[krzysz00]

To amend, attribute conversions will need some way to return three states

1. Failed, abort conversion attempt
2. Not applicable
3. Here's your result

which we can't do the way that type conversions can because nullptr is a meaningful attribute value for a lot of contexts.

Or we could not have the early-stop system.

[krzysz00]

Make the mapping system more generic and allow it to fail.

[ftynse]

Originally, I was thinking that it was okay to just have context-free attribute conversion here. It's the same for types nested in other types: one may eventually need to convert types differently based on what they are nested in, but there is no support for context. Same for attributes, there's currently no need for context. But I also like the approach adopted here.

Regarding the early stopping and return values, it is possible to define a custom return type that steals two bits from the pointer to the attribute internal storage, using AttrConvresionResult = llvm::PointerIntPair<Attribute, 2> and interpret those bits as: 00 here's your result and it's null, 01 here's your result and it's a concrete attribute, 10 keep trying, 11 stop.

mlir/include/mlir/Transforms/DialectConversion.h

177–180

Nit: would putting this in a using declaration inside the function work instead?

263–264

Who determines the semantics of returning nullptr here, the caller?

403–405

This can grow quite fast... Not sure if caching with pair as index is worth it, especially given that the logic differs per type class (e.g, the same callback works for instances of MemRefType, but we will store for each instance).

[krzysz00]

@ftynse Thanks for the feedback!

On the PointerIntPair method, would the solution to be making a wrapper around PointerIntPair<Attribute, 2>? Something along the lines of

enum TypeAttrConversionResult {
   Result(Attribute),
   NullAttribute,
    Continue,
   Abort
}

(that there being how I'd phrase it in Rust - the actual implementation would, I think, involve PointerIntPair and such)

The question I have is whether this struct would go inside TypeConverter or be free-standing?

mlir/include/mlir/Transforms/DialectConversion.h

177–180

We need it in the type of the vector of converters

263–264

Yeah. Good question. I should clarify.

403–405

Agreed, should probably remove.

[ftynse]

On the PointerIntPair method, would the solution to be making a wrapper around PointerIntPair<Attribute, 2>? Something along the lines of <...>

SGTM, good for a separate patch.

[krzysz00]

Use AttributeConversionResult instead of std::optional<Attribute>

[krzysz00]

Rebase away a merge conflict

[mehdi_amini]

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

97

You're dereferencing an optional here: this is crashing behavior, the error should be handled and propagated instead. Just hit this in https://github.com/llvm/llvm-project/issues/70160#issuecomment-1780644011

mlir/include/mlir/Conversion/LLVMCommon/TypeConverter.h

public:
 
   /// Returns the size of the memref descriptor object in bytes.
   unsigned getMemRefDescriptorSize(MemRefType type, const DataLayout &layout);
 
   /// Returns the size of the unranked memref descriptor object in bytes.
   unsigned getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
                                            const DataLayout &layout);
 
+  /// Return the LLLVM address space corresponding to the memory space attribute
+  /// `memorySpace` or failure if the space cannot be converted to an integer.
+  FailureOr<unsigned> getMemRefAddressSpace(Attribute memorySpace);
+
+  /// Return the LLVM address space corresponding to the memory space of the
+  /// memref type `type` or failure if the memory space cannot be converted to
+  /// an integer.
+  FailureOr<unsigned> getMemRefAddressSpace(BaseMemRefType type);
+
   /// Check if a memref type can be converted to a bare pointer.
   static bool canConvertToBarePtr(BaseMemRefType type);
 
 protected:
   /// Pointer to the LLVM dialect.
   LLVM::LLVMDialect *llvmDialect;
 
 private:

mlir/include/mlir/Dialect/AMDGPU/AMDGPU.td

let description = [{
     The `AMDGPU` dialect provides wrappers around AMD-specific functionality
     and LLVM intrinsics. These wrappers should be used in conjunction with
     more generic dialects, such as `gpu` and `vector`, when generating LLVM IR
     that will eventually be executed on AMD hardware.
   }];
 
 
   let dependentDialects = [
-    "arith::ArithDialect"
+    "arith::ArithDialect",
+    "gpu::GPUDialect"
   ];
   let useDefaultAttributePrinterParser = 1;
   let useFoldAPI = kEmitFoldAdaptorFolder;
 }
 
 //===----------------------------------------------------------------------===//
 // AMDGPU Op definitions
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/GPU/Transforms/Passes.h

 void populateGpuAllReducePatterns(RewritePatternSet &patterns);
 
 /// Collect all patterns to rewrite ops within the GPU dialect.
 inline void populateGpuRewritePatterns(RewritePatternSet &patterns) {
   populateGpuAllReducePatterns(patterns);
 }
 
 namespace gpu {
-/// A function that maps a MemorySpace enum to a target-specific integer value.
-using MemorySpaceMapping =
-    std::function<unsigned(gpu::AddressSpace gpuAddressSpace)>;
-
-/// Populates type conversion rules for lowering memory space attributes to
-/// numeric values.
-void populateMemorySpaceAttributeTypeConversions(
-    TypeConverter &typeConverter, const MemorySpaceMapping &mapping);
-
-/// Populates patterns to lower memory space attributes to numeric values.
-void populateMemorySpaceLoweringPatterns(TypeConverter &typeConverter,
-                                         RewritePatternSet &patterns);
-
-/// Populates legality rules for lowering memory space attriutes to numeric
-/// values.
-void populateLowerMemorySpaceOpLegality(ConversionTarget &target);
-
 /// Returns the default annotation name for GPU binary blobs.
 std::string getDefaultGpuBinaryAnnotation();
 
 /// Base pass class to serialize kernel functions through LLVM into
 /// user-specified IR and add the resulting blob as module attribute.
 class SerializeToBlobPass : public OperationPass<gpu::GPUModuleOp> {
 public:
   SerializeToBlobPass(TypeID passID);

mlir/include/mlir/Dialect/GPU/Transforms/Passes.td

 def GpuMapParallelLoopsPass
     : Pass<"gpu-map-parallel-loops", "mlir::func::FuncOp"> {
   let summary = "Greedily maps loops to GPU hardware dimensions.";
   let constructor = "mlir::createGpuMapParallelLoopsPass()";
   let description = "Greedily maps loops to GPU hardware dimensions.";
   let dependentDialects = ["mlir::gpu::GPUDialect"];
 }
 
-def GPULowerMemorySpaceAttributesPass
-    : Pass<"gpu-lower-memory-space-attributes"> {
-  let summary = "Assign numeric values to memref memory space symbolic placeholders";
-  let description = [{
-    Updates all memref types that have a memory space attribute
-    that is a `gpu::AddressSpaceAttr`. These attributes are
-    changed to `IntegerAttr`'s using a mapping that is given in the
-    options.
-  }];
-  let options = [
-    Option<"privateAddrSpace", "private", "unsigned", "5",
-      "private address space numeric value">,
-    Option<"workgroupAddrSpace", "workgroup", "unsigned", "3",
-      "workgroup address space numeric value">,
-    Option<"globalAddrSpace", "global", "unsigned", "1",
-      "global address space numeric value">
-  ];
-}
-
 #endif // MLIR_DIALECT_GPU_PASSES

mlir/include/mlir/Transforms/DialectConversion.h

 #include "mlir/Rewrite/FrozenRewritePatternSet.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/StringMap.h"
 #include <type_traits>
 
 namespace mlir {
 
 // Forward declarations.
+class Attribute;
 class Block;
 class ConversionPatternRewriter;
 class MLIRContext;
 class Operation;
 class Type;
 class Value;
 
 //===----------------------------------------------------------------------===//
public:
   /// converting type from an illegal, or source, type to a legal type.
   template <typename FnT, typename T = typename llvm::function_traits<
                               std::decay_t<FnT>>::template arg_t<1>>
   void addTargetMaterialization(FnT &&callback) {
     targetMaterializations.emplace_back(
         wrapMaterialization<T>(std::forward<FnT>(callback)));
   }
 
+  /// Register a memory space conversion function. Memory space conversion
+  /// functions take a non-null Attribute or a subtype of Attribute (in which
+  /// case they are only called on attributes of that type) and return
+  /// std::optional<Attribute> Returning std::nullopt from the function
+  /// indicates that it did not apply to the input argument, while returning the
+  /// null Attribute indicates conversion to the default memory space.
+  ///
+  /// Registered conversion functions are callled in the reverse of the order in
+  /// which they were registered.
+  template <typename FnT, typename T = typename llvm::function_traits<
+                              std::decay_t<FnT>>::template arg_t<0>>
+  void addMemorySpaceConversion(FnT &&callback) {
+    registerMemorySpaceConversion(
+        wrapMemorySpaceConversion<T>(std::forward<FnT>(callback)));
+  }
+
+  /// Get the memory space that the default (null attribute) memory space will
+  /// be converted to, which is the default memory space by default.
+  Attribute getDefaultMemorySpaceConversion() const {
+    return defaultMemorySpaceConversion;
+  }
+
+  /// Set the memory space to which the default (null attribute) memory space
+  /// will be converted.
+  void setDefaultMemorySpaceConversion(Attribute memorySpace) {
+    defaultMemorySpaceConversion = memorySpace;
+    // Clear type caches.
+    cachedDirectConversions.clear();
+    cachedMultiConversions.clear();
+  }
+
   /// Convert the given type. This function should return failure if no valid
   /// conversion exists, success otherwise. If the new set of types is empty,
   /// the type is removed and any usages of the existing value are expected to
   /// be removed during conversion.
   LogicalResult convertType(Type t, SmallVectorImpl<Type> &results);
 
   /// This hook simplifies defining 1-1 type conversions. This function returns
   /// the type to convert to on success, and a null type on failure.
return materializeConversion(sourceMaterializations, builder, loc,
                                  resultType, inputs);
   }
   Value materializeTargetConversion(OpBuilder &builder, Location loc,
                                     Type resultType, ValueRange inputs) {
     return materializeConversion(targetMaterializations, builder, loc,
                                  resultType, inputs);
   }
 
+  /// Convert a memory space attribute to one that may be legal after
+  /// conversion. Memory space conversion cannot fail: if no applicable
+  /// conversion callback has been provided, the provided memory space attribute
+  /// is returned. Note that the null Attribute is not a failure condition but
+  /// instead represents the default memory space. The null memory space itself
+  /// can be converted using setDefaultMemorySpace()
+  Attribute convertMemorySpace(Attribute memorySpace);
+
 private:
   /// The signature of the callback used to convert a type. If the new set of
   /// types is empty, the type is removed and any usages of the existing value
   /// are expected to be removed during conversion.
   using ConversionCallbackFn = std::function<std::optional<LogicalResult>(
       Type, SmallVectorImpl<Type> &, ArrayRef<Type>)>;
 
   /// The signature of the callback used to materialize a conversion.
   using MaterializationCallbackFn = std::function<std::optional<Value>(
       OpBuilder &, Type, ValueRange, Location)>;
 
+  /// The signature of the callback used to convert a memory space.
+  using MemorySpaceConversionCallbackFn =
+      std::function<std::optional<Attribute>(Attribute)>;
+
   /// Attempt to materialize a conversion using one of the provided
   /// materialization functions.
   Value materializeConversion(
       MutableArrayRef<MaterializationCallbackFn> materializations,
       OpBuilder &builder, Location loc, Type resultType, ValueRange inputs);
 
   /// Generate a wrapper for the given callback. This allows for accepting
   /// different callback forms, that all compose into a single version.
return [callback = std::forward<FnT>(callback)](
                OpBuilder &builder, Type resultType, ValueRange inputs,
                Location loc) -> std::optional<Value> {
       if (T derivedType = resultType.dyn_cast<T>())
         return callback(builder, derivedType, inputs, loc);
       return std::nullopt;
     };
   }
 
+  /// Generate a wrapper for the given memory space conversion callback. The
+  /// callback may take any subclass of `Attribute` and the wrapper will check
+  /// for the target attribute to be of the expected class before calling the
+  /// callback.
+  template <typename T, typename FnT>
+  MemorySpaceConversionCallbackFn wrapMemorySpaceConversion(FnT &&callback) {
+    return [callback = std::forward<FnT>(callback)](
+               Attribute memorySpace) -> std::optional<Attribute> {
+      if (T derivedMemorySpace = memorySpace.dyn_cast<T>())
+        return callback(derivedMemorySpace);
+      return std::nullopt;
+    };
+  }
+
+  /// Register a memory space conversion, clearing caches.
+  void registerMemorySpaceConversion(MemorySpaceConversionCallbackFn callback) {
+    memorySpaceConversions.emplace_back(std::move(callback));
+    cachedMemorySpaceConversions.clear();
+    // Clear type conversions in case a memory space is lingering inside.
+    cachedDirectConversions.clear();
+    cachedMultiConversions.clear();
+  }
+
   /// The set of registered conversion functions.
   SmallVector<ConversionCallbackFn, 4> conversions;
 
   /// The list of registered materialization functions.
   SmallVector<MaterializationCallbackFn, 2> argumentMaterializations;
   SmallVector<MaterializationCallbackFn, 2> sourceMaterializations;
   SmallVector<MaterializationCallbackFn, 2> targetMaterializations;
 
+  /// The conversion for the default memory space.
+  Attribute defaultMemorySpaceConversion;
+
+  /// The list of registered memory space conversion functions.
+  SmallVector<MemorySpaceConversionCallbackFn, 2> memorySpaceConversions;
+
   /// A set of cached conversions to avoid recomputing in the common case.
   /// Direct 1-1 conversions are the most common, so this cache stores the
   /// successful 1-1 conversions as well as all failed conversions.
   DenseMap<Type, Type> cachedDirectConversions;
   /// This cache stores the successful 1->N conversions, where N != 1.
   DenseMap<Type, SmallVector<Type, 2>> cachedMultiConversions;
+  /// This cache stores the successful memory space conversions.
+  DenseMap<Attribute, Attribute> cachedMemorySpaceConversions;
 
   /// Stores the types that are being converted in the case when convertType
   /// is being called recursively to convert nested types.
   SmallVector<Type, 2> conversionCallStack;
 };
 
 //===----------------------------------------------------------------------===//
 // Conversion Patterns
 //===----------------------------------------------------------------------===//

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h

public:
   LogicalResult
   matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     return impl::scalarizeVectorOp(op, adaptor.getOperands(), rewriter,
                                    *this->getTypeConverter());
   }
 };
 
+/// A function that maps a MemorySpace enum to a target-specific integer value.
+using MemorySpaceMapping =
+    std::function<unsigned(gpu::AddressSpace gpuAddressSpace)>;
+
+/// Populates memory space attribute conversion rules for lowering
+/// gpu.address_space to integer values.
+void populateGpuMemorySpaceAttributeConversions(
+    TypeConverter &typeConverter, const MemorySpaceMapping &mapping);
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_GPUCOMMON_GPUOPSLOWERING_H_

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

Operation *scalarOp =
         rewriter.create(loc, name, scalarOperands, elementType, op->getAttrs());
     rewriter.create<LLVM::InsertElementOp>(loc, result, scalarOp->getResult(0),
                                            index);
   }
 
   rewriter.replaceOp(op, result);
   return success();
 }
+
+static IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
+  return IntegerAttr::get(IntegerType::get(ctx, 64), space);
+}
+
+void mlir::populateGpuMemorySpaceAttributeConversions(
+    TypeConverter &typeConverter, const MemorySpaceMapping &mapping) {
+  typeConverter.addMemorySpaceConversion(
+      [mapping](gpu::AddressSpaceAttr memorySpaceAttr) -> Attribute {
+        gpu::AddressSpace memorySpace = memorySpaceAttr.getValue();
+        unsigned addressSpace = mapping(memorySpace);
+        return wrapNumericMemorySpace(memorySpaceAttr.getContext(),
+                                      addressSpace);
+      });
+}

mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

void runOnOperation() override {
     // single conversion pass.
     {
       RewritePatternSet patterns(m.getContext());
       populateGpuRewritePatterns(patterns);
       if (failed(applyPatternsAndFoldGreedily(m, std::move(patterns))))
         return signalPassFailure();
     }
 
-    // MemRef conversion for GPU to NVVM lowering.
-    {
-      RewritePatternSet patterns(m.getContext());
-      TypeConverter typeConverter;
-      typeConverter.addConversion([](Type t) { return t; });
-      // NVVM uses alloca in the default address space to represent private
-      // memory allocations, so drop private annotations. NVVM uses address
-      // space 3 for shared memory. NVVM uses the default address space to
-      // represent global memory.
-      gpu::populateMemorySpaceAttributeTypeConversions(
-          typeConverter, [](gpu::AddressSpace space) -> unsigned {
-            switch (space) {
-            case gpu::AddressSpace::Global:
-              return static_cast<unsigned>(
-                  NVVM::NVVMMemorySpace::kGlobalMemorySpace);
-            case gpu::AddressSpace::Workgroup:
-              return static_cast<unsigned>(
-                  NVVM::NVVMMemorySpace::kSharedMemorySpace);
-            case gpu::AddressSpace::Private:
-              return 0;
-            }
-            llvm_unreachable("unknown address space enum value");
-            return 0;
-          });
-      gpu::populateMemorySpaceLoweringPatterns(typeConverter, patterns);
-      ConversionTarget target(getContext());
-      gpu::populateLowerMemorySpaceOpLegality(target);
-      if (failed(applyFullConversion(m, target, std::move(patterns))))
-        return signalPassFailure();
-    }
-
-    LLVMTypeConverter converter(m.getContext(), options);
+    LLVMTypeConverter converter(m.getContext(), options);
+    populateGpuMemorySpaceAttributeConversions(
+        converter, [](gpu::AddressSpace space) -> unsigned {
     // Lowering for MMAMatrixType.
     converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
       return convertMMAToLLVMType(type);
     });
     RewritePatternSet llvmPatterns(m.getContext());
 
     arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

void runOnOperation() override {
     // ops which need to be lowered further, which is not supported by a
     // single conversion pass.
     {
       RewritePatternSet patterns(ctx);
       populateGpuRewritePatterns(patterns);
       (void)applyPatternsAndFoldGreedily(m, std::move(patterns));
     }
 
-    // Apply memory space lowering. The target uses 3 for workgroup memory and 5
-    // for private memory.
-    {
-      RewritePatternSet patterns(ctx);
-      TypeConverter typeConverter;
-      typeConverter.addConversion([](Type t) { return t; });
-      gpu::populateMemorySpaceAttributeTypeConversions(
-          typeConverter, [](gpu::AddressSpace space) {
-            switch (space) {
-            case gpu::AddressSpace::Global:
-              return 1;
-            case gpu::AddressSpace::Workgroup:
-              return 3;
-            case gpu::AddressSpace::Private:
-              return 5;
-            }
-            llvm_unreachable("unknown address space enum value");
-            return 0;
-          });
-      ConversionTarget target(getContext());
-      gpu::populateLowerMemorySpaceOpLegality(target);
-      gpu::populateMemorySpaceLoweringPatterns(typeConverter, patterns);
-      if (failed(applyFullConversion(m, target, std::move(patterns))))
-        return signalPassFailure();
-    }
+    LLVMTypeConverter converter(ctx, options);
+    populateGpuMemorySpaceAttributeConversions(
+        converter, [](gpu::AddressSpace space) {
 
-    LLVMTypeConverter converter(ctx, options);
     RewritePatternSet llvmPatterns(ctx);
 
     mlir::arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);
     populateAMDGPUToROCDLConversionPatterns(converter, llvmPatterns,
                                             *maybeChipset);
     populateVectorToLLVMConversionPatterns(converter, llvmPatterns);
     cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns);
     populateFuncToLLVMConversionPatterns(converter, llvmPatterns);

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

bool ConvertToLLVMPattern::isConvertibleAndHasIdentityMaps(
   if (!typeConverter->convertType(type.getElementType()))
     return false;
   return type.getLayout().isIdentity();
 }
 
 Type ConvertToLLVMPattern::getElementPtrType(MemRefType type) const {
   auto elementType = type.getElementType();
   auto structElementType = typeConverter->convertType(elementType);
-  return LLVM::LLVMPointerType::get(structElementType,
-                                    type.getMemorySpaceAsInt());
+  auto addressSpace = getTypeConverter()->getMemRefAddressSpace(type);
+  if (failed(addressSpace))
+    return {};
+  return LLVM::LLVMPointerType::get(structElementType, *addressSpace);
 }
 
 void ConvertToLLVMPattern::getMemRefDescriptorSizes(
     Location loc, MemRefType memRefType, ValueRange dynamicSizes,
     ConversionPatternRewriter &rewriter, SmallVectorImpl<Value> &sizes,
     SmallVectorImpl<Value> &strides, Value &sizeBytes) const {
   assert(isConvertibleAndHasIdentityMaps(memRefType) &&
          "layout maps must have been normalized away");

mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp

emitError(
         "conversion to strided form failed either due to non-strided layout "
         "maps (which should have been normalized away) or other reasons");
     return {};
   }
 
   Type elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
-  auto ptrTy =
-      LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
+  FailureOr<unsigned> addressSpace = getMemRefAddressSpace(type);
+  if (failed(addressSpace)) {
+    emitError(UnknownLoc::get(type.getContext()),
+              "conversion of memref memory space to integer address space "
+              "failed. Consider adding memory space conversions");
+  }
+  auto ptrTy = LLVM::LLVMPointerType::get(elementType, *addressSpace);
   auto indexTy = getIndexType();
 
   SmallVector<Type, 5> results = {ptrTy, ptrTy, indexTy};
   auto rank = type.getRank();
   if (rank == 0)
     return results;
 
   if (unpackAggregates)
     results.insert(results.end(), 2 * rank, indexTy);
   else
     results.insert(results.end(), 2, LLVM::LLVMArrayType::get(indexTy, rank));
   return results;
 }
 
 unsigned LLVMTypeConverter::getMemRefDescriptorSize(MemRefType type,
                                                     const DataLayout &layout) {
   // Compute the descriptor size given that of its components indicated above.
-  unsigned space = type.getMemorySpaceAsInt();
+  unsigned space = *getMemRefAddressSpace(type);
   return 2 * llvm::divideCeil(getPointerBitwidth(space), 8) +
          (1 + 2 * type.getRank()) * layout.getTypeSize(getIndexType());
 }
 
 /// Converts MemRefType to LLVMType. A MemRefType is converted to a struct that
 /// packs the descriptor fields as defined by `getMemRefDescriptorFields`.
 Type LLVMTypeConverter::convertMemRefType(MemRefType type) {
   // When converting a MemRefType to a struct with descriptor fields, do not
SmallVector<Type, 2> LLVMTypeConverter::getUnrankedMemRefDescriptorFields() {
   return {getIndexType(),
           LLVM::LLVMPointerType::get(IntegerType::get(&getContext(), 8))};
 }
 
 unsigned
 LLVMTypeConverter::getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
                                                    const DataLayout &layout) {
   // Compute the descriptor size given that of its components indicated above.
-  unsigned space = type.getMemorySpaceAsInt();
+  unsigned space = *getMemRefAddressSpace(type);
   return layout.getTypeSize(getIndexType()) +
          llvm::divideCeil(getPointerBitwidth(space), 8);
 }
 
 Type LLVMTypeConverter::convertUnrankedMemRefType(UnrankedMemRefType type) {
   if (!convertType(type.getElementType()))
     return {};
   return LLVM::LLVMStructType::getLiteral(&getContext(),
                                           getUnrankedMemRefDescriptorFields());
 }
 
+FailureOr<unsigned>
+LLVMTypeConverter::getMemRefAddressSpace(Attribute memorySpace) {
+  Attribute converted = convertMemorySpace(memorySpace);
+  if (!converted) // default memory space -> 0 in LLVM
+    return 0;
+  if (auto explicitSpace = converted.dyn_cast<IntegerAttr>())
+    return explicitSpace.getInt();
+  return failure();
+}
+
+FailureOr<unsigned>
+LLVMTypeConverter::getMemRefAddressSpace(BaseMemRefType type) {
+  return getMemRefAddressSpace(type.getMemorySpace());
+}
+
 // Check if a memref type can be converted to a bare pointer.
 bool LLVMTypeConverter::canConvertToBarePtr(BaseMemRefType type) {
   if (type.isa<UnrankedMemRefType>())
     // Unranked memref is not supported in the bare pointer calling convention.
     return false;
 
   // Check that the memref has static shape, strides and offset. Otherwise, it
   // cannot be lowered to a bare pointer.

 
 /// Convert a memref type to a bare pointer to the memref element type.
 Type LLVMTypeConverter::convertMemRefToBarePtr(BaseMemRefType type) {
   if (!canConvertToBarePtr(type))
     return {};
   Type elementType = convertType(type.getElementType());
   if (!elementType)
     return {};
-  return LLVM::LLVMPointerType::get(elementType, type.getMemorySpaceAsInt());
+  FailureOr<unsigned> addressSpace = getMemRefAddressSpace(type);
+  if (failed(addressSpace))
+    return {};
+  return LLVM::LLVMPointerType::get(elementType, *addressSpace);
 }
 
 /// Convert an n-D vector type to an LLVM vector type:
 ///  * 0-D `vector<T>` are converted to vector<1xT>
 ///  * 1-D `vector<axT>` remains as is while,
 ///  * n>1 `vector<ax...xkxT>` convert via an (n-1)-D array type to
 ///    `!llvm.array<ax...array<jxvector<kxT>>>`.
 Type LLVMTypeConverter::convertVectorType(VectorType type) {

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

struct AllocaOpLowering : public AllocLikeOpLLVMLowering {
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
 
     // With alloca, one gets a pointer to the element type right away.
     // For stack allocations.
     auto allocaOp = cast<memref::AllocaOp>(op);
     auto elementPtrType = this->getElementPtrType(allocaOp.getType());
 
     auto allocatedElementPtr = rewriter.create<LLVM::AllocaOp>(
         loc, elementPtrType, sizeBytes, allocaOp.getAlignment().value_or(0));
 
     return std::make_tuple(allocatedElementPtr, allocatedElementPtr);
   }
 };
 
 /// The base class for lowering realloc op, to support the implementation of

 struct DimOpLowering : public ConvertOpToLLVMPattern<memref::DimOp> {
   using ConvertOpToLLVMPattern<memref::DimOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(memref::DimOp dimOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type operandType = dimOp.getSource().getType();
     if (operandType.isa<UnrankedMemRefType>()) {
-      rewriter.replaceOp(
-          dimOp, {extractSizeOfUnrankedMemRef(
-                     operandType, dimOp, adaptor.getOperands(), rewriter)});
-
+      FailureOr<Value> extractedSize = extractSizeOfUnrankedMemRef(
+          operandType, dimOp, adaptor.getOperands(), rewriter);
+      if (failed(extractedSize))
+        return failure();
+      rewriter.replaceOp(dimOp, {*extractedSize});
       return success();
     }
     if (operandType.isa<MemRefType>()) {
       rewriter.replaceOp(
           dimOp, {extractSizeOfRankedMemRef(operandType, dimOp,
                                             adaptor.getOperands(), rewriter)});
       return success();
     }
     llvm_unreachable("expected MemRefType or UnrankedMemRefType");
   }
 
 private:
-  Value extractSizeOfUnrankedMemRef(Type operandType, memref::DimOp dimOp,
-                                    OpAdaptor adaptor,
-                                    ConversionPatternRewriter &rewriter) const {
+  FailureOr<Value>
+  extractSizeOfUnrankedMemRef(Type operandType, memref::DimOp dimOp,
     Location loc = dimOp.getLoc();
 
     auto unrankedMemRefType = operandType.cast<UnrankedMemRefType>();
     auto scalarMemRefType =
         MemRefType::get({}, unrankedMemRefType.getElementType());
-    unsigned addressSpace = unrankedMemRefType.getMemorySpaceAsInt();
+    FailureOr<unsigned> maybeAddressSpace =
+        getTypeConverter()->getMemRefAddressSpace(unrankedMemRefType);
+    if (failed(maybeAddressSpace)) {
+      dimOp.emitOpError("memref memory space must be convertible to an integer "
+                        "address space");
+      return failure();
+    }
+    unsigned addressSpace = *maybeAddressSpace;
 
     // Extract pointer to the underlying ranked descriptor and bitcast it to a
     // memref<element_type> descriptor pointer to minimize the number of GEP
     // operations.
     UnrankedMemRefDescriptor unrankedDesc(adaptor.getSource());
     Value underlyingRankedDesc = unrankedDesc.memRefDescPtr(rewriter, loc);
     Value scalarMemRefDescPtr = rewriter.create<LLVM::BitcastOp>(
         loc,
         LLVM::LLVMPointerType::get(typeConverter->convertType(scalarMemRefType),
                                    addressSpace),
         underlyingRankedDesc);
 
     // Get pointer to offset field of memref<element_type> descriptor.
     Type indexPtrTy = LLVM::LLVMPointerType::get(
         getTypeConverter()->getIndexType(), addressSpace);
     Value offsetPtr = rewriter.create<LLVM::GEPOp>(
         loc, indexPtrTy, scalarMemRefDescPtr, ArrayRef<LLVM::GEPArg>{0, 2});
 
     // The size value that we have to extract can be obtained using GEPop with
     // `dimOp.index() + 1` index argument.
     Value idxPlusOne = rewriter.create<LLVM::AddOp>(
         loc, createIndexConstant(rewriter, loc, 1), adaptor.getIndex());
     Value sizePtr =
         rewriter.create<LLVM::GEPOp>(loc, indexPtrTy, offsetPtr, idxPlusOne);
-    return rewriter.create<LLVM::LoadOp>(loc, sizePtr);
+    return rewriter.create<LLVM::LoadOp>(loc, sizePtr).getResult();
   }
 
   std::optional<int64_t> getConstantDimIndex(memref::DimOp dimOp) const {
     if (auto idx = dimOp.getConstantIndex())
       return idx;
 
     if (auto constantOp = dimOp.getIndex().getDefiningOp<LLVM::ConstantOp>())
       return constantOp.getValue()
if (!global.isExternal() && !global.isUninitialized()) {
 
       // For scalar memrefs, the global variable created is of the element type,
       // so unpack the elements attribute to extract the value.
       if (type.getRank() == 0)
         initialValue = elementsAttr.getSplatValue<Attribute>();
     }
 
     uint64_t alignment = global.getAlignment().value_or(0);
-
+    FailureOr<unsigned> addressSpace =
+        getTypeConverter()->getMemRefAddressSpace(type);
+    if (failed(addressSpace))
+      return global.emitOpError(
+          "memory space cannot be converted to an integer address space");
     auto newGlobal = rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         global, arrayTy, global.getConstant(), linkage, global.getSymName(),
-        initialValue, alignment, type.getMemorySpaceAsInt());
+        initialValue, alignment, *addressSpace);
     if (!global.isExternal() && global.isUninitialized()) {
       Block *blk = new Block();
       newGlobal.getInitializerRegion().push_back(blk);
       rewriter.setInsertionPointToStart(blk);
       Value undef[] = {
           rewriter.create<LLVM::UndefOp>(global.getLoc(), arrayTy)};
       rewriter.create<LLVM::ReturnOp>(global.getLoc(), undef);
     }
struct GetGlobalMemrefOpLowering : public AllocLikeOpLLVMLowering {
 
   /// Buffer "allocation" for memref.get_global op is getting the address of
   /// the global variable referenced.
   std::tuple<Value, Value> allocateBuffer(ConversionPatternRewriter &rewriter,
                                           Location loc, Value sizeBytes,
                                           Operation *op) const override {
     auto getGlobalOp = cast<memref::GetGlobalOp>(op);
     MemRefType type = getGlobalOp.getResult().getType().cast<MemRefType>();
-    unsigned memSpace = type.getMemorySpaceAsInt();
+
+    // This is called after a type conversion, which would have failed if this
+    // call fails.
+    unsigned memSpace = *getTypeConverter()->getMemRefAddressSpace(type);
 
     Type arrayTy = convertGlobalMemrefTypeToLLVM(type, *getTypeConverter());
     auto addressOf = rewriter.create<LLVM::AddressOfOp>(
         loc, LLVM::LLVMPointerType::get(arrayTy, memSpace),
         getGlobalOp.getName());
 
     // Get the address of the first element in the array by creating a GEP with
     // the address of the GV as the base, and (rank + 1) number of 0 indices.
if (operandType.isa<MemRefType>()) {
     MemRefDescriptor desc(convertedOperand);
     *allocatedPtr = desc.allocatedPtr(rewriter, loc);
     *alignedPtr = desc.alignedPtr(rewriter, loc);
     if (offset != nullptr)
       *offset = desc.offset(rewriter, loc);
     return;
   }
 
-  unsigned memorySpace =
-      operandType.cast<UnrankedMemRefType>().getMemorySpaceAsInt();
+  // These will all cause assert()s on unconvertible types.
+  unsigned memorySpace = *typeConverter.getMemRefAddressSpace(
+      operandType.cast<UnrankedMemRefType>());
   Type elementType = operandType.cast<UnrankedMemRefType>().getElementType();
   Type llvmElementType = typeConverter.convertType(elementType);
   Type elementPtrPtrType = LLVM::LLVMPointerType::get(
       LLVM::LLVMPointerType::get(llvmElementType, memorySpace));
 
   // Extract pointer to the underlying ranked memref descriptor and cast it to
   // ElemType**.
   UnrankedMemRefDescriptor unrankedDesc(convertedOperand);
convertSourceMemRefToDescriptor(ConversionPatternRewriter &rewriter,
     // The shape is a rank-1 tensor with unknown length.
     Location loc = reshapeOp.getLoc();
     MemRefDescriptor shapeDesc(adaptor.getShape());
     Value resultRank = shapeDesc.size(rewriter, loc, 0);
 
     // Extract address space and element type.
     auto targetType =
         reshapeOp.getResult().getType().cast<UnrankedMemRefType>();
-    unsigned addressSpace = targetType.getMemorySpaceAsInt();
+    unsigned addressSpace =
+        *getTypeConverter()->getMemRefAddressSpace(targetType);
     Type elementType = targetType.getElementType();
 
     // 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;
matchAndRewrite(memref::ViewOp viewOp, OpAdaptor adaptor,
 
     // Create the descriptor.
     MemRefDescriptor sourceMemRef(adaptor.getSource());
     auto targetMemRef = MemRefDescriptor::undef(rewriter, loc, targetDescTy);
 
     // Field 1: Copy the allocated pointer, used for malloc/free.
     Value allocatedPtr = sourceMemRef.allocatedPtr(rewriter, loc);
     auto srcMemRefType = viewOp.getSource().getType().cast<MemRefType>();
+    unsigned sourceMemorySpace =
+        *getTypeConverter()->getMemRefAddressSpace(srcMemRefType);
     Value bitcastPtr = rewriter.create<LLVM::BitcastOp>(
-        loc,
-        LLVM::LLVMPointerType::get(targetElementTy,
-                                   srcMemRefType.getMemorySpaceAsInt()),
+        loc, LLVM::LLVMPointerType::get(targetElementTy, sourceMemorySpace),
         allocatedPtr);
     targetMemRef.setAllocatedPtr(rewriter, loc, bitcastPtr);
 
     // Field 2: Copy the actual aligned pointer to payload.
     Value alignedPtr = sourceMemRef.alignedPtr(rewriter, loc);
     alignedPtr = rewriter.create<LLVM::GEPOp>(
         loc, alignedPtr.getType(), alignedPtr, adaptor.getByteShift());
     bitcastPtr = rewriter.create<LLVM::BitcastOp>(
-        loc,
-        LLVM::LLVMPointerType::get(targetElementTy,
-                                   srcMemRefType.getMemorySpaceAsInt()),
+        loc, LLVM::LLVMPointerType::get(targetElementTy, sourceMemorySpace),
         alignedPtr);
     targetMemRef.setAlignedPtr(rewriter, loc, bitcastPtr);
 
     // Field 3: The offset in the resulting type must be 0. This is because of
     // the type change: an offset on srcType* may not be expressible as an
     // offset on dstType*.
     targetMemRef.setOffset(rewriter, loc,
                            createIndexConstant(rewriter, loc, offset));

mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp

struct NVGPUAsyncCopyLowering
   LogicalResult
   matchAndRewrite(nvgpu::DeviceAsyncCopyOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     auto dstMemrefType = op.getDst().getType().cast<MemRefType>();
     Value dstPtr = getStridedElementPtr(loc, dstMemrefType, adaptor.getDst(),
                                         adaptor.getDstIndices(), rewriter);
     auto i8Ty = IntegerType::get(op.getContext(), 8);
-    auto dstPointerType =
-        LLVM::LLVMPointerType::get(i8Ty, dstMemrefType.getMemorySpaceAsInt());
+    FailureOr<unsigned> dstAddressSpace =
+        getTypeConverter()->getMemRefAddressSpace(dstMemrefType);
+    if (failed(dstAddressSpace))
+      return rewriter.notifyMatchFailure(
+          loc, "destination memref address space not convertible to integer");
+    auto dstPointerType = LLVM::LLVMPointerType::get(i8Ty, *dstAddressSpace);
     dstPtr = rewriter.create<LLVM::BitcastOp>(loc, dstPointerType, dstPtr);
 
     auto srcMemrefType = op.getSrc().getType().cast<MemRefType>();
+    FailureOr<unsigned> srcAddressSpace =
+        getTypeConverter()->getMemRefAddressSpace(srcMemrefType);
+    if (failed(srcAddressSpace))
+      return rewriter.notifyMatchFailure(
+          loc, "source memref address space not convertible to integer");
 
     Value scrPtr = getStridedElementPtr(loc, srcMemrefType, adaptor.getSrc(),
                                         adaptor.getSrcIndices(), rewriter);
-    auto srcPointerType =
-        LLVM::LLVMPointerType::get(i8Ty, srcMemrefType.getMemorySpaceAsInt());
+    auto srcPointerType = LLVM::LLVMPointerType::get(i8Ty, *srcAddressSpace);
     scrPtr = rewriter.create<LLVM::BitcastOp>(loc, srcPointerType, scrPtr);
     // Intrinsics takes a global pointer so we need an address space cast.
     auto srcPointerGlobalType = LLVM::LLVMPointerType::get(
         i8Ty, NVVM::NVVMMemorySpace::kGlobalMemorySpace);
     scrPtr = rewriter.create<LLVM::AddrSpaceCastOp>(loc, srcPointerGlobalType,
                                                     scrPtr);
     int64_t dstElements = adaptor.getDstElements().getZExtValue();
     int64_t sizeInBytes =

mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp

convertTransferReadToLoads(vector::TransferReadOp op,
                            llvm::DenseMap<Value, Value> &valueMapping) {
   OpBuilder b(op);
 
   FailureOr<nvgpu::WarpMatrixInfo> warpMatrixInfo =
       nvgpu::getWarpMatrixInfo(op);
   if (failed(warpMatrixInfo))
     return failure();
 
+  Attribute memorySpace =
+      op.getSource().getType().cast<MemRefType>().getMemorySpace();
+  bool isSourceWorkgroupShared = false;
+  if (auto integerAttr = memorySpace.dyn_cast_or_null<IntegerAttr>())
+    isSourceWorkgroupShared = integerAttr.getInt() == 3;
+  else if (auto gpuAddrSpaceAttr =
+               memorySpace.dyn_cast_or_null<gpu::AddressSpaceAttr>())
+    isSourceWorkgroupShared =
+        gpuAddrSpaceAttr.getValue() == gpu::AddressSpace::Workgroup;
   bool isLdMatrixCompatible =
-      op.getSource().getType().cast<MemRefType>().getMemorySpaceAsInt() == 3 &&
+      isSourceWorkgroupShared &&
       nvgpu::inferTileWidthInBits(*warpMatrixInfo) == 128;
 
   VectorType vecTy = op.getVectorType();
   int64_t bitWidth = vecTy.getElementType().getIntOrFloatBitWidth();
 
   // When we are transposing the B operand, ldmatrix will only work if we have
   // at least 8 rows to read and the width to read for the transpose is 128
   // bits.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 //===- VectorToLLVM.cpp - Conversion from Vector to the LLVM dialect ------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 
+#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/LLVMCommon/VectorPattern.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Vector/Transforms/VectorTransforms.h"
 #include "mlir/IR/BuiltinTypes.h"
LogicalResult getMemRefAlignment(LLVMTypeConverter &typeConverter,
   // stop depending on translation.
   llvm::LLVMContext llvmContext;
   align = LLVM::TypeToLLVMIRTranslator(llvmContext)
               .getPreferredAlignment(elementTy, typeConverter.getDataLayout());
   return success();
 }
 
 // Check if the last stride is non-unit or the memory space is not zero.
-static LogicalResult isMemRefTypeSupported(MemRefType memRefType) {
+static LogicalResult isMemRefTypeSupported(MemRefType memRefType,
+                                           LLVMTypeConverter &converter) {
   int64_t offset;
   SmallVector<int64_t, 4> strides;
   auto successStrides = getStridesAndOffset(memRefType, strides, offset);
-  if (failed(successStrides) || strides.back() != 1 ||
-      memRefType.getMemorySpaceAsInt() != 0)
+  FailureOr<unsigned> addressSpace =
+      converter.getMemRefAddressSpace(memRefType);
+  if (failed(successStrides) || strides.back() != 1 || failed(addressSpace) ||
+      *addressSpace != 0)
     return failure();
   return success();
 }
 
 // Add an index vector component to a base pointer.
 static Value getIndexedPtrs(ConversionPatternRewriter &rewriter, Location loc,
+                            LLVMTypeConverter &typeConverter,
                             MemRefType memRefType, Value llvmMemref, Value base,
                             Value index, uint64_t vLen) {
-  assert(succeeded(isMemRefTypeSupported(memRefType)) &&
+  assert(succeeded(isMemRefTypeSupported(memRefType, typeConverter)) &&
          "unsupported memref type");
   auto pType = MemRefDescriptor(llvmMemref).getElementPtrType();
   auto ptrsType = LLVM::getFixedVectorType(pType, vLen);
   return rewriter.create<LLVM::GEPOp>(loc, ptrsType, base, index);
 }
 
 // Casts a strided element pointer to a vector pointer.  The vector pointer
 // will be in the same address space as the incoming memref type.
 static Value castDataPtr(ConversionPatternRewriter &rewriter, Location loc,
-                         Value ptr, MemRefType memRefType, Type vt) {
-  auto pType = LLVM::LLVMPointerType::get(vt, memRefType.getMemorySpaceAsInt());
+                         Value ptr, MemRefType memRefType, Type vt,
+                         LLVMTypeConverter &converter) {
+  unsigned addressSpace = *converter.getMemRefAddressSpace(memRefType);
+  auto pType = LLVM::LLVMPointerType::get(vt, addressSpace);
   return rewriter.create<LLVM::BitcastOp>(loc, pType, ptr);
 }
 
 namespace {
 
 /// Trivial Vector to LLVM conversions
 using VectorScaleOpConversion =
     OneToOneConvertToLLVMPattern<vector::VectorScaleOp, LLVM::vscale>;
matchAndRewrite(LoadOrStoreOp loadOrStoreOp,
     if (failed(getMemRefAlignment(*this->getTypeConverter(), memRefTy, align)))
       return failure();
 
     // Resolve address.
     auto vtype = this->typeConverter->convertType(loadOrStoreOp.getVectorType())
                      .template cast<VectorType>();
     Value dataPtr = this->getStridedElementPtr(loc, memRefTy, adaptor.getBase(),
                                                adaptor.getIndices(), rewriter);
-    Value ptr = castDataPtr(rewriter, loc, dataPtr, memRefTy, vtype);
+    Value ptr = castDataPtr(rewriter, loc, dataPtr, memRefTy, vtype,
+                            *this->getTypeConverter());
 
     replaceLoadOrStoreOp(loadOrStoreOp, adaptor, vtype, ptr, align, rewriter);
     return success();
   }
 };
 
 /// Conversion pattern for a vector.gather.
 class VectorGatherOpConversion
     : public ConvertOpToLLVMPattern<vector::GatherOp> {
 public:
   using ConvertOpToLLVMPattern<vector::GatherOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(vector::GatherOp gather, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     MemRefType memRefType = gather.getBaseType().dyn_cast<MemRefType>();
     assert(memRefType && "The base should be bufferized");
 
-    if (failed(isMemRefTypeSupported(memRefType)))
+    if (failed(isMemRefTypeSupported(memRefType, *this->getTypeConverter())))
       return failure();
 
     auto loc = gather->getLoc();
 
     // Resolve alignment.
     unsigned align;
     if (failed(getMemRefAlignment(*getTypeConverter(), memRefType, align)))
       return failure();
 
     Value ptr = getStridedElementPtr(loc, memRefType, adaptor.getBase(),
                                      adaptor.getIndices(), rewriter);
     Value base = adaptor.getBase();
 
     auto llvmNDVectorTy = adaptor.getIndexVec().getType();
     // Handle the simple case of 1-D vector.
     if (!llvmNDVectorTy.isa<LLVM::LLVMArrayType>()) {
       auto vType = gather.getVectorType();
       // Resolve address.
-      Value ptrs = getIndexedPtrs(rewriter, loc, memRefType, base, ptr,
-                                  adaptor.getIndexVec(),
+      Value ptrs = getIndexedPtrs(rewriter, loc, *this->getTypeConverter(),
+                                  memRefType, base, ptr, adaptor.getIndexVec(),
                                   /*vLen=*/vType.getDimSize(0));
       // Replace with the gather intrinsic.
       rewriter.replaceOpWithNewOp<LLVM::masked_gather>(
           gather, typeConverter->convertType(vType), ptrs, adaptor.getMask(),
           adaptor.getPassThru(), rewriter.getI32IntegerAttr(align));
       return success();
     }
 
-    auto callback = [align, memRefType, base, ptr, loc, &rewriter](
-                        Type llvm1DVectorTy, ValueRange vectorOperands) {
+    LLVMTypeConverter &typeConverter = *this->getTypeConverter();
+    auto callback = [align, memRefType, base, ptr, loc, &rewriter,
+                     &typeConverter](Type llvm1DVectorTy,
+                                     ValueRange vectorOperands) {
       // Resolve address.
       Value ptrs = getIndexedPtrs(
-          rewriter, loc, memRefType, base, ptr, /*index=*/vectorOperands[0],
+          rewriter, loc, typeConverter, memRefType, base, ptr,
+          /*index=*/vectorOperands[0],
           LLVM::getVectorNumElements(llvm1DVectorTy).getFixedValue());
       // Create the gather intrinsic.
       return rewriter.create<LLVM::masked_gather>(
           loc, llvm1DVectorTy, ptrs, /*mask=*/vectorOperands[1],
           /*passThru=*/vectorOperands[2], rewriter.getI32IntegerAttr(align));
     };
     SmallVector<Value> vectorOperands = {
         adaptor.getIndexVec(), adaptor.getMask(), adaptor.getPassThru()};
public:
   using ConvertOpToLLVMPattern<vector::ScatterOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(vector::ScatterOp scatter, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto loc = scatter->getLoc();
     MemRefType memRefType = scatter.getMemRefType();
 
-    if (failed(isMemRefTypeSupported(memRefType)))
+    if (failed(isMemRefTypeSupported(memRefType, *this->getTypeConverter())))
       return failure();
 
     // Resolve alignment.
     unsigned align;
     if (failed(getMemRefAlignment(*getTypeConverter(), memRefType, align)))
       return failure();
 
     // Resolve address.
     VectorType vType = scatter.getVectorType();
     Value ptr = getStridedElementPtr(loc, memRefType, adaptor.getBase(),
                                      adaptor.getIndices(), rewriter);
-    Value ptrs =
-        getIndexedPtrs(rewriter, loc, memRefType, adaptor.getBase(), ptr,
-                       adaptor.getIndexVec(), /*vLen=*/vType.getDimSize(0));
+    Value ptrs = getIndexedPtrs(
+        rewriter, loc, *this->getTypeConverter(), memRefType, adaptor.getBase(),
+        ptr, adaptor.getIndexVec(), /*vLen=*/vType.getDimSize(0));
 
     // Replace with the scatter intrinsic.
     rewriter.replaceOpWithNewOp<LLVM::masked_scatter>(
         scatter, adaptor.getValueToStore(), ptrs, adaptor.getMask(),
         rewriter.getI32IntegerAttr(align));
     return success();
   }
 };

mlir/lib/Dialect/AMDGPU/IR/AMDGPUDialect.cpp

 //===- AMDGPUDialect.cpp - MLIR AMDGPU dialect implementation --------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the AMDGPU dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/AMDGPU/AMDGPUDialect.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"

 }
 
 //===----------------------------------------------------------------------===//
 // RawBuffer*Op
 //===----------------------------------------------------------------------===//
 template <typename T>
 static LogicalResult verifyRawBufferOp(T &op) {
   MemRefType bufferType = op.getMemref().getType().template cast<MemRefType>();
-  if (bufferType.getMemorySpaceAsInt() != 0)
+  Attribute memorySpace = bufferType.getMemorySpace();
+  bool isGlobal = false;
+  if (!memorySpace)
+    isGlobal = true;
+  else if (auto intMemorySpace = memorySpace.dyn_cast<IntegerAttr>())
+    isGlobal = intMemorySpace.getInt() == 0 || intMemorySpace.getInt() == 1;
+  else if (auto gpuMemorySpace = memorySpace.dyn_cast<gpu::AddressSpaceAttr>())
+    isGlobal = gpuMemorySpace.getValue() == gpu::AddressSpace::Global;
+
+  if (!isGlobal)
     return op.emitOpError(
         "Buffer ops must operate on a memref in global memory");
   if (!bufferType.hasRank())
     return op.emitOpError(
         "Cannot meaningfully buffer_store to an unranked memref");
   if (static_cast<int64_t>(op.getIndices().size()) != bufferType.getRank())
     return op.emitOpError("Expected " + Twine(bufferType.getRank()) +
                           " indices to memref");

mlir/lib/Dialect/AMDGPU/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRAMDGPUDialect
   AMDGPUDialect.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/AMDGPU
 
   DEPENDS
   MLIRAMDGPUEnumsGen
   MLIRAMDGPUAttributesIncGen
   MLIRAMDGPUIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
+  # Needed for GPU address space enum definition
+  MLIRGPUOps
   MLIRIR
   MLIRSideEffectInterfaces
   )

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp
   Transforms/SerializeToBlob.cpp
   Transforms/SerializeToCubin.cpp
   Transforms/SerializeToHsaco.cpp
-  Transforms/LowerMemorySpaceAttributes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/GPU
 
   LINK_COMPONENTS
   Core
   MC
   ${NVPTX_LIBS}

mlir/lib/Dialect/GPU/Transforms/LowerMemorySpaceAttributes.cpp

-//===- LowerMemorySpaceAttributes.cpp  ------------------------------------===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-///
-/// Implementation of a pass that rewrites the IR so that uses of
-/// `gpu::AddressSpaceAttr` in memref memory space annotations are replaced
-/// with caller-specified numeric values.
-///
-//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/GPU/IR/GPUDialect.h"
-#include "mlir/Dialect/GPU/Transforms/Passes.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/DialectConversion.h"
-#include "llvm/Support/Debug.h"
-
-namespace mlir {
-#define GEN_PASS_DEF_GPULOWERMEMORYSPACEATTRIBUTESPASS
-#include "mlir/Dialect/GPU/Transforms/Passes.h.inc"
-} // namespace mlir
-
-using namespace mlir;
-using namespace mlir::gpu;
-
-//===----------------------------------------------------------------------===//
-// Conversion Target
-//===----------------------------------------------------------------------===//
-
-/// Returns true if the given `type` is considered as legal during memory space
-/// attribute lowering.
-static bool isLegalType(Type type) {
-  if (auto memRefType = type.dyn_cast<BaseMemRefType>()) {
-    return !memRefType.getMemorySpace()
-                .isa_and_nonnull<gpu::AddressSpaceAttr>();
-  }
-  return true;
-}
-
-/// Returns true if the given `attr` is considered legal during memory space
-/// attribute lowering.
-static bool isLegalAttr(Attribute attr) {
-  if (auto typeAttr = attr.dyn_cast<TypeAttr>())
-    return isLegalType(typeAttr.getValue());
-  return true;
-}
-
-/// Returns true if the given `op` is legal during memory space attribute
-/// lowering.
-static bool isLegalOp(Operation *op) {
-  if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
-    return llvm::all_of(funcOp.getArgumentTypes(), isLegalType) &&
-           llvm::all_of(funcOp.getResultTypes(), isLegalType) &&
-           llvm::all_of(funcOp.getFunctionBody().getArgumentTypes(),
-                        isLegalType);
-  }
-
-  auto attrs = llvm::map_range(op->getAttrs(), [](const NamedAttribute &attr) {
-    return attr.getValue();
-  });
-
-  return llvm::all_of(op->getOperandTypes(), isLegalType) &&
-         llvm::all_of(op->getResultTypes(), isLegalType) &&
-         llvm::all_of(attrs, isLegalAttr);
-}
-
-void gpu::populateLowerMemorySpaceOpLegality(ConversionTarget &target) {
-  target.markUnknownOpDynamicallyLegal(isLegalOp);
-}
-
-//===----------------------------------------------------------------------===//
-// Type Converter
-//===----------------------------------------------------------------------===//
-
-IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
-  return IntegerAttr::get(IntegerType::get(ctx, 64), space);
-}
-
-void mlir::gpu::populateMemorySpaceAttributeTypeConversions(
-    TypeConverter &typeConverter, const MemorySpaceMapping &mapping) {
-  typeConverter.addConversion([mapping](Type type) -> std::optional<Type> {
-    auto subElementType = type.dyn_cast_or_null<SubElementTypeInterface>();
-    if (!subElementType)
-      return type;
-    Type newType = subElementType.replaceSubElements(
-        [mapping](Attribute attr) -> std::optional<Attribute> {
-          auto memorySpaceAttr = attr.dyn_cast_or_null<gpu::AddressSpaceAttr>();
-          if (!memorySpaceAttr)
-            return std::nullopt;
-          auto newValue = wrapNumericMemorySpace(
-              attr.getContext(), mapping(memorySpaceAttr.getValue()));
-          return newValue;
-        });
-    return newType;
-  });
-}
-
-namespace {
-
-/// Converts any op that has operands/results/attributes with numeric MemRef
-/// memory spaces.
-struct LowerMemRefAddressSpacePattern final : public ConversionPattern {
-  LowerMemRefAddressSpacePattern(MLIRContext *context, TypeConverter &converter)
-      : ConversionPattern(converter, MatchAnyOpTypeTag(), 1, context) {}
-
-  LogicalResult
-  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
-                  ConversionPatternRewriter &rewriter) const override {
-    SmallVector<NamedAttribute> newAttrs;
-    newAttrs.reserve(op->getAttrs().size());
-    for (auto attr : op->getAttrs()) {
-      if (auto typeAttr = attr.getValue().dyn_cast<TypeAttr>()) {
-        auto newAttr = getTypeConverter()->convertType(typeAttr.getValue());
-        newAttrs.emplace_back(attr.getName(), TypeAttr::get(newAttr));
-      } else {
-        newAttrs.push_back(attr);
-      }
-    }
-
-    SmallVector<Type> newResults;
-    (void)getTypeConverter()->convertTypes(op->getResultTypes(), newResults);
-
-    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,
-                         newResults, newAttrs, op->getSuccessors());
-
-    for (Region &region : op->getRegions()) {
-      Region *newRegion = state.addRegion();
-      rewriter.inlineRegionBefore(region, *newRegion, newRegion->begin());
-      TypeConverter::SignatureConversion result(newRegion->getNumArguments());
-      (void)getTypeConverter()->convertSignatureArgs(
-          newRegion->getArgumentTypes(), result);
-      rewriter.applySignatureConversion(newRegion, result);
-    }
-
-    Operation *newOp = rewriter.create(state);
-    rewriter.replaceOp(op, newOp->getResults());
-    return success();
-  }
-};
-} // namespace
-
-void mlir::gpu::populateMemorySpaceLoweringPatterns(
-    TypeConverter &typeConverter, RewritePatternSet &patterns) {
-  patterns.add<LowerMemRefAddressSpacePattern>(patterns.getContext(),
-                                               typeConverter);
-}
-
-namespace {
-class LowerMemorySpaceAttributesPass
-    : public mlir::impl::GPULowerMemorySpaceAttributesPassBase<
-          LowerMemorySpaceAttributesPass> {
-public:
-  using Base::Base;
-  void runOnOperation() override {
-    MLIRContext *context = &getContext();
-    Operation *op = getOperation();
-
-    ConversionTarget target(getContext());
-    populateLowerMemorySpaceOpLegality(target);
-
-    TypeConverter typeConverter;
-    typeConverter.addConversion([](Type t) { return t; });
-    populateMemorySpaceAttributeTypeConversions(
-        typeConverter, [this](AddressSpace space) -> unsigned {
-          switch (space) {
-          case AddressSpace::Global:
-            return globalAddrSpace;
-          case AddressSpace::Workgroup:
-            return workgroupAddrSpace;
-          case AddressSpace::Private:
-            return privateAddrSpace;
-          }
-          llvm_unreachable("unknown address space enum value");
-          return 0;
-        });
-    RewritePatternSet patterns(context);
-    populateMemorySpaceLoweringPatterns(typeConverter, patterns);
-    if (failed(applyFullConversion(op, target, std::move(patterns))))
-      return signalPassFailure();
-  }
-};
-} // namespace

mlir/lib/Transforms/Utils/DialectConversion.cpp

 auto TypeConverter::convertBlockSignature(Block *block)
     -> std::optional<SignatureConversion> {
   SignatureConversion conversion(block->getNumArguments());
   if (failed(convertSignatureArgs(block->getArgumentTypes(), conversion)))
     return std::nullopt;
   return conversion;
 }
 
+Attribute TypeConverter::convertMemorySpace(Attribute memorySpace) {
+  if (!memorySpace)
+    return getDefaultMemorySpaceConversion();
+  auto existingIt = cachedMemorySpaceConversions.find(memorySpace);
+  if (existingIt != cachedMemorySpaceConversions.end())
+    return existingIt->second;
+  for (MemorySpaceConversionCallbackFn &fn :
+       llvm::reverse(memorySpaceConversions)) {
+    if (std::optional<Attribute> converted = fn(memorySpace)) {
+      cachedMemorySpaceConversions.try_emplace(memorySpace, *converted);
+      return *converted;
+    }
+  }
+  return memorySpace;
+}
+
 //===----------------------------------------------------------------------===//
 // FunctionOpInterfaceSignatureConversion
 //===----------------------------------------------------------------------===//
 
 static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
                                         TypeConverter &typeConverter,
                                         ConversionPatternRewriter &rewriter) {
   FunctionType type = funcOp.getFunctionType().cast<FunctionType>();

mlir/test/Conversion/GPUCommon/lower-memory-space-attrs.mlir

+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-rocdl | FileCheck %s --check-prefixes=CHECK,ROCDL
+// RUN: mlir-opt %s -split-input-file -convert-gpu-to-nvvm | FileCheck %s --check-prefixes=CHECK,NVVM
+
+gpu.module @kernel {
+  gpu.func @private(%arg0: f32) private(%arg1: memref<4xf32, #gpu.address_space<private>>) {
+    %c0 = arith.constant 0 : index
+    memref.store %arg0, %arg1[%c0] : memref<4xf32, #gpu.address_space<private>>
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @private
+//      CHECK:  llvm.store
+// ROCDL-SAME:   : !llvm.ptr<f32, 5>
+//  NVVM-SAME:   : !llvm.ptr<f32>
+
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @workgroup(%arg0: f32) workgroup(%arg1: memref<4xf32, #gpu.address_space<workgroup>>) {
+    %c0 = arith.constant 0 : index
+    memref.store %arg0, %arg1[%c0] : memref<4xf32, #gpu.address_space<workgroup>>
+    gpu.return
+  }
+}
+
+// CHECK-LABEL:  llvm.func @workgroup
+//       CHECK:  llvm.store
+//  CHECK-SAME:   : !llvm.ptr<f32, 3>
+
+// -----
+
+gpu.module @kernel {
+  gpu.func @nested_memref(%arg0: memref<4xmemref<4xf32, #gpu.address_space<global>>, #gpu.address_space<global>>) -> f32 {
+    %c0 = arith.constant 0 : index
+    %inner = memref.load %arg0[%c0] : memref<4xmemref<4xf32, #gpu.address_space<global>>, #gpu.address_space<global>>
+    %value = memref.load %inner[%c0] : memref<4xf32, #gpu.address_space<global>>
+    gpu.return %value : f32
+  }
+}
+
+// CHECK-LABEL:  llvm.func @nested_memref
+//       CHECK:  llvm.load
+//  CHECK-SAME:   : !llvm.ptr<{{.*}}, 1>
+//       CHECK: [[value:%.+]] = llvm.load
+//  CHECK-SAME:   : !llvm.ptr<f32, 1>
+//       CHECK: llvm.return [[value]]

mlir/test/Dialect/GPU/lower-memory-space-attrs.mlir

-// RUN: mlir-opt %s -split-input-file -gpu-lower-memory-space-attributes | FileCheck %s
-// RUN: mlir-opt %s -split-input-file -gpu-lower-memory-space-attributes="private=0 global=0" | FileCheck %s --check-prefix=CUDA
-
-gpu.module @kernel {
-  gpu.func @private(%arg0: f32) private(%arg1: memref<4xf32, #gpu.address_space<private>>) {
-    %c0 = arith.constant 0 : index
-    memref.store %arg0, %arg1[%c0] : memref<4xf32, #gpu.address_space<private>>
-    gpu.return
-  }
-}
-
-//      CHECK:  gpu.func @private
-// CHECK-SAME:    private(%{{.+}}: memref<4xf32, 5>)
-//      CHECK:  memref.store
-// CHECK-SAME:   : memref<4xf32, 5>
-
-//      CUDA:  gpu.func @private
-// CUDA-SAME:    private(%{{.+}}: memref<4xf32>)
-//      CUDA:  memref.store
-// CUDA-SAME:   : memref<4xf32>
-
-// -----
-
-gpu.module @kernel {
-  gpu.func @workgroup(%arg0: f32) workgroup(%arg1: memref<4xf32, #gpu.address_space<workgroup>>) {
-    %c0 = arith.constant 0 : index
-    memref.store %arg0, %arg1[%c0] : memref<4xf32, #gpu.address_space<workgroup>>
-    gpu.return
-  }
-}
-
-//      CHECK:  gpu.func @workgroup
-// CHECK-SAME:    workgroup(%{{.+}}: memref<4xf32, 3>)
-//      CHECK:  memref.store
-// CHECK-SAME:   : memref<4xf32, 3>
-
-// -----
-
-gpu.module @kernel {
-  gpu.func @nested_memref(%arg0: memref<4xmemref<4xf32, #gpu.address_space<global>>, #gpu.address_space<global>>) {
-    %c0 = arith.constant 0 : index
-    memref.load %arg0[%c0] : memref<4xmemref<4xf32, #gpu.address_space<global>>, #gpu.address_space<global>>
-    gpu.return
-  }
-}
-
-//      CHECK:  gpu.func @nested_memref
-// CHECK-SAME:    (%{{.+}}: memref<4xmemref<4xf32, 1>, 1>)
-//      CHECK:  memref.load
-// CHECK-SAME:   : memref<4xmemref<4xf32, 1>, 1>
-
-//      CUDA:  gpu.func @nested_memref
-// CUDA-SAME:    (%{{.+}}: memref<4xmemref<4xf32>>)
-//      CUDA:  memref.load
-// CUDA-SAME:   : memref<4xmemref<4xf32>>