Add generic type attribute mapping infrastructure, use it in GpuToX

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.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D142159

Author: krzysz00

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

@@ -147,6 +147,11 @@ class LLVMTypeConverter : public TypeConverter {
   unsigned getUnrankedMemRefDescriptorSize(UnrankedMemRefType type,
                                            const DataLayout &layout);
 
+  /// 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);
 

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

@@ -25,7 +25,8 @@ def AMDGPU_Dialect : Dialect {
 
 
   let dependentDialects = [
-    "arith::ArithDialect"
+    "arith::ArithDialect",
+    "gpu::GPUDialect"
   ];
   let useDefaultAttributePrinterParser = 1;
 }

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

@@ -61,23 +61,6 @@ inline void populateGpuRewritePatterns(RewritePatternSet &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();
 

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

@@ -37,23 +37,4 @@ def GpuMapParallelLoopsPass
   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

@@ -21,6 +21,7 @@
 namespace mlir {
 
 // Forward declarations.
+class Attribute;
 class Block;
 class ConversionPatternRewriter;
 class MLIRContext;
@@ -87,6 +88,34 @@ class TypeConverter {
     SmallVector<Type, 4> argTypes;
   };
 
+  /// The general result of a type attribute conversion callback, allowing
+  /// for early termination. The default constructor creates the na case.
+  class AttributeConversionResult {
+  public:
+    constexpr AttributeConversionResult() : impl() {}
+    AttributeConversionResult(Attribute attr) : impl(attr, resultTag) {}
+
+    static AttributeConversionResult result(Attribute attr);
+    static AttributeConversionResult na();
+    static AttributeConversionResult abort();
+
+    bool hasResult() const;
+    bool isNa() const;
+    bool isAbort() const;
+
+    Attribute getResult() const;
+
+  private:
+    AttributeConversionResult(Attribute attr, unsigned tag) : impl(attr, tag) {}
+
+    llvm::PointerIntPair<Attribute, 2> impl;
+    // Note that na is 0 so that we can use PointerIntPair's default
+    // constructor.
+    static constexpr unsigned naTag = 0;
+    static constexpr unsigned resultTag = 1;
+    static constexpr unsigned abortTag = 2;
+  };
+
   /// Register a conversion function. A conversion function must be convertible
   /// to any of the following forms(where `T` is a class derived from `Type`:
   ///   * std::optional<Type>(T)
@@ -156,6 +185,34 @@ class TypeConverter {
         wrapMaterialization<T>(std::forward<FnT>(callback)));
   }
 
+  /// Register a conversion function for attributes within types. Type
+  /// converters may call this function in order to allow hoking into the
+  /// translation of attributes that exist within types. For example, a type
+  /// converter for the `memref` type could use these conversions to convert
+  /// memory spaces or layouts in an extensible way.
+  ///
+  /// The conversion functions take a non-null Type or subclass of Type and a
+  /// non-null Attribute (or subclass of Attribute), and returns a
+  /// `AttributeConversionResult`. This result can either contan an `Attribute`,
+  /// which may be `nullptr`, representing the conversion's success,
+  /// `AttributeConversionResult::na()` (the default empty value), indicating
+  /// that the conversion function did not apply and that further conversion
+  /// functions should be checked, or `AttributeConversionResult::abort()`
+  /// indicating that the conversion process should be aborted.
+  ///
+  /// 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>,
+      typename A =
+          typename llvm::function_traits<std::decay_t<FnT>>::template arg_t<1>>
+  void addTypeAttributeConversion(FnT &&callback) {
+    registerTypeAttributeConversion(
+        wrapTypeAttributeConversion<T, A>(std::forward<FnT>(callback)));
+  }
+
   /// 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
@@ -226,6 +283,12 @@ class TypeConverter {
                                  resultType, inputs);
   }
 
+  /// Convert an attribute present `attr` from within the type `type` using
+  /// the registered conversion functions. If no applicable conversion has been
+  /// registered, return std::nullopt. Note that the empty attribute/`nullptr`
+  /// is a valid return value for this function.
+  std::optional<Attribute> convertTypeAttribute(Type type, Attribute attr);
+
 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
@@ -237,6 +300,10 @@ class TypeConverter {
   using MaterializationCallbackFn = std::function<std::optional<Value>(
       OpBuilder &, Type, ValueRange, Location)>;
 
+  /// The signature of the callback used to convert a type attribute.
+  using TypeAttributeConversionCallbackFn =
+      std::function<AttributeConversionResult(Type, Attribute)>;
+
   /// Attempt to materialize a conversion using one of the provided
   /// materialization functions.
   Value materializeConversion(
@@ -311,6 +378,32 @@ class TypeConverter {
     };
   }
 
+  /// 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 A, typename FnT>
+  TypeAttributeConversionCallbackFn
+  wrapTypeAttributeConversion(FnT &&callback) {
+    return [callback = std::forward<FnT>(callback)](
+               Type type, Attribute attr) -> AttributeConversionResult {
+      if (T derivedType = type.dyn_cast<T>()) {
+        if (A derivedAttr = attr.dyn_cast_or_null<A>())
+          return callback(derivedType, derivedAttr);
+      }
+      return AttributeConversionResult::na();
+    };
+  }
+
+  /// Register a memory space conversion, clearing caches.
+  void
+  registerTypeAttributeConversion(TypeAttributeConversionCallbackFn callback) {
+    typeAttributeConversions.emplace_back(std::move(callback));
+    // 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;
 
@@ -319,6 +412,9 @@ class TypeConverter {
   SmallVector<MaterializationCallbackFn, 2> sourceMaterializations;
   SmallVector<MaterializationCallbackFn, 2> targetMaterializations;
 
+  /// The list of registered type attribute conversion functions.
+  SmallVector<TypeAttributeConversionCallbackFn, 2> typeAttributeConversions;
+
   /// 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.

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.cpp

@@ -8,6 +8,7 @@
 
 #include "GPUOpsLowering.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "llvm/ADT/STLExtras.h"
@@ -474,3 +475,18 @@ LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
   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.addTypeAttributeConversion(
+      [mapping](BaseMemRefType type, gpu::AddressSpaceAttr memorySpaceAttr) {
+        gpu::AddressSpace memorySpace = memorySpaceAttr.getValue();
+        unsigned addressSpace = mapping(memorySpace);
+        return wrapNumericMemorySpace(memorySpaceAttr.getContext(),
+                                      addressSpace);
+      });
+}

mlir/lib/Conversion/GPUCommon/GPUOpsLowering.h

@@ -112,6 +112,14 @@ struct ScalarizeVectorOpLowering : public ConvertOpToLLVMPattern<SourceOp> {
   }
 };
 
+/// 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/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp

@@ -241,38 +241,26 @@ struct LowerGpuOpsToNVVMOpsPass
         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);
+    // 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.
+    populateGpuMemorySpaceAttributeConversions(
+        converter, [](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;
+        });
     // Lowering for MMAMatrixType.
     converter.addConversion([&](gpu::MMAMatrixType type) -> Type {
       return convertMMAToLLVMType(type);

mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -132,33 +132,21 @@ struct LowerGpuOpsToROCDLOpsPass
       (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) {
+          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;
+        });
+
     RewritePatternSet llvmPatterns(ctx);
 
     mlir::arith::populateArithToLLVMConversionPatterns(converter, llvmPatterns);

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -112,8 +112,10 @@ bool ConvertToLLVMPattern::isConvertibleAndHasIdentityMaps(
 Type ConvertToLLVMPattern::getElementPtrType(MemRefType type) const {
   auto elementType = type.getElementType();
   auto structElementType = typeConverter->convertType(elementType);
-  return getTypeConverter()->getPointerType(structElementType,
-                                            type.getMemorySpaceAsInt());
+  auto addressSpace = getTypeConverter()->getMemRefAddressSpace(type);
+  if (failed(addressSpace))
+    return {};
+  return getTypeConverter()->getPointerType(structElementType, *addressSpace);
 }
 
 void ConvertToLLVMPattern::getMemRefDescriptorSizes(