[mlir][spirv] Support attribute in MapMemRefStorageClassPass

[antiagainst]

Summary

MemRef memory space actually can be an attribute. Update the
map function signature to accept an attribute. The default
mappings can still only covers numeric ones, but this allows
downstream callers to extend with custom memory spaces.

[kuhar]

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

169–170

nit: can you spell out the exact type? It's not obvious to me just by looking at the RHS.

[antiagainst]

Address comments

mlir/include/mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h

 #include <memory>
 
 namespace mlir {
 class SPIRVTypeConverter;
 
 namespace spirv {
 /// Mapping from numeric MemRef memory spaces into SPIR-V symbolic ones.
 using MemorySpaceToStorageClassMap =
-    std::function<Optional<spirv::StorageClass>(unsigned)>;
+    std::function<Optional<spirv::StorageClass>(Attribute)>;
 
 /// Maps MemRef memory spaces to storage classes for Vulkan-flavored SPIR-V
 /// using the default rule. Returns None if the memory space is unknown.
-Optional<spirv::StorageClass> mapMemorySpaceToVulkanStorageClass(unsigned);
+Optional<spirv::StorageClass> mapMemorySpaceToVulkanStorageClass(Attribute);
 /// Maps storage classes for Vulkan-flavored SPIR-V to MemRef memory spaces
 /// using the default rule. Returns None if the storage class is unsupported.
 Optional<unsigned> mapVulkanStorageClassToMemorySpace(spirv::StorageClass);
 
 /// Maps MemRef memory spaces to storage classes for OpenCL-flavored SPIR-V
 /// using the default rule. Returns None if the memory space is unknown.
-Optional<spirv::StorageClass> mapMemorySpaceToOpenCLStorageClass(unsigned);
+Optional<spirv::StorageClass> mapMemorySpaceToOpenCLStorageClass(Attribute);
 /// Maps storage classes for OpenCL-flavored SPIR-V to MemRef memory spaces
 /// using the default rule. Returns None if the storage class is unsupported.
 Optional<unsigned> mapOpenCLStorageClassToMemorySpace(spirv::StorageClass);
 
 /// Type converter for converting numeric MemRef memory spaces into SPIR-V
 /// symbolic ones.
 class MemorySpaceToStorageClassConverter : public TypeConverter {
 public:

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

 
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
 
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 #define GEN_PASS_DEF_MAPMEMREFSTORAGECLASS
#define VULKAN_STORAGE_SPACE_MAP_LIST(MAP_FN)                                  \
   MAP_FN(spirv::StorageClass::Private, 5)                                      \
   MAP_FN(spirv::StorageClass::Function, 6)                                     \
   MAP_FN(spirv::StorageClass::PushConstant, 7)                                 \
   MAP_FN(spirv::StorageClass::UniformConstant, 8)                              \
   MAP_FN(spirv::StorageClass::Input, 9)                                        \
   MAP_FN(spirv::StorageClass::Output, 10)
 
 Optional<spirv::StorageClass>
-spirv::mapMemorySpaceToVulkanStorageClass(unsigned memorySpace) {
+spirv::mapMemorySpaceToVulkanStorageClass(Attribute memorySpaceAttr) {
+  // Handle null memory space attribute specially.
+  if (!memorySpaceAttr)
+    return spirv::StorageClass::StorageBuffer;
+
+  // Unknown dialect custom attributes are not supported by default.
+  // Downstream callers should plug in more specialized ones.
+  auto intAttr = memorySpaceAttr.dyn_cast<IntegerAttr>();
+  if (!intAttr)
+    return llvm::None;
+  unsigned memorySpace = intAttr.getInt();
+
 #define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   case space:                                                                  \
     return storage;
 
   switch (memorySpace) {
     VULKAN_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     break;
#define OPENCL_STORAGE_SPACE_MAP_LIST(MAP_FN)                                  \
   MAP_FN(spirv::StorageClass::Generic, 1)                                      \
   MAP_FN(spirv::StorageClass::Workgroup, 3)                                    \
   MAP_FN(spirv::StorageClass::UniformConstant, 4)                              \
   MAP_FN(spirv::StorageClass::Private, 5)                                      \
   MAP_FN(spirv::StorageClass::Function, 6)                                     \
   MAP_FN(spirv::StorageClass::Image, 7)
 
 Optional<spirv::StorageClass>
-spirv::mapMemorySpaceToOpenCLStorageClass(unsigned memorySpace) {
+spirv::mapMemorySpaceToOpenCLStorageClass(Attribute memorySpaceAttr) {
+  // Handle null memory space attribute specially.
+  if (!memorySpaceAttr)
+    return spirv::StorageClass::CrossWorkgroup;
+
+  // Unknown dialect custom attributes are not supported by default.
+  // Downstream callers should plug in more specialized ones.
+  auto intAttr = memorySpaceAttr.dyn_cast<IntegerAttr>();
+  if (!intAttr)
+    return llvm::None;
+  unsigned memorySpace = intAttr.getInt();
+
 #define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   case space:                                                                  \
     return storage;
 
   switch (memorySpace) {
     OPENCL_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     break;

 
 spirv::MemorySpaceToStorageClassConverter::MemorySpaceToStorageClassConverter(
     const spirv::MemorySpaceToStorageClassMap &memorySpaceMap)
     : memorySpaceMap(memorySpaceMap) {
   // Pass through for all other types.
   addConversion([](Type type) { return type; });
 
   addConversion([this](BaseMemRefType memRefType) -> Optional<Type> {
-    // Expect IntegerAttr memory spaces. The attribute can be missing for the
-    // case of memory space == 0.
-    Attribute spaceAttr = memRefType.getMemorySpace();
-    if (spaceAttr && !spaceAttr.isa<IntegerAttr>()) {
-      LLVM_DEBUG(llvm::dbgs() << "cannot convert " << memRefType
-                              << " due to non-IntegerAttr memory space\n");
-      return llvm::None;
-    }
-
-    unsigned space = memRefType.getMemorySpaceAsInt();
-    auto storage = this->memorySpaceMap(space);
+    Optional<spirv::StorageClass> storage =
+        this->memorySpaceMap(memRefType.getMemorySpace());
     if (!storage) {
       LLVM_DEBUG(llvm::dbgs()
                  << "cannot convert " << memRefType
                  << " due to being unable to find memory space in map\n");
       return llvm::None;
     }
 
     auto storageAttr =