[mlir][spirv] Add default Vulkan memory space to storage class mapping

[kuhar]

LGTM provided that the base patch gets accepted in the current form

mlir/include/mlir/Conversion/Passes.td

547

nit: 'to use to populate storage class mappings' or `to use for populating mappings'?

[ThomasRaoux]

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

41

Can we remove the matching list:
https://github.com/llvm/llvm-project/blob/954de25a92d08d83149c38e17f628d424b587bb5/mlir/lib/Dialect/SPIRV/Transforms/SPIRVConversion.cpp#L129

[antiagainst]

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

41

Yup that one will be removed in the next patch, which will integrate and enforce this pass.

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 = DenseMap<unsigned, spirv::StorageClass>;
+/// Returns the default map for targeting Vulkan-flavored SPIR-V.
+MemorySpaceToStorageClassMap getDefaultVulkanStorageClassMap();
 
 /// Type converter for converting numeric MemRef memory spaces into SPIR-V
 /// symbolic ones.
 class MemorySpaceToStorageClassConverter : public TypeConverter {
 public:
   explicit MemorySpaceToStorageClassConverter(
       const MemorySpaceToStorageClassMap &memorySpaceMap);
 

mlir/include/mlir/Conversion/Passes.td

 // MemRefToSPIRV
 //===----------------------------------------------------------------------===//
 
 def MapMemRefStorageClass : Pass<"map-memref-spirv-storage-class", "ModuleOp"> {
   let summary = "Map numeric MemRef memory spaces to SPIR-V storage classes";
   let constructor = "mlir::createMapMemRefStorageClassPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [
-    Option<"mappings", "mappings", "std::string", /*default=*/"",
-           "A comma-separated list of memory space to storage class mappings; "
-           "for example, '0=StorageClass,1=Uniform,2=Workgroup'">
+    Option<"clientAPI", "client-api", "std::string", /*default=*/"\"vulkan\"",
+           "The client API to use for populating mappings">
   ];
 }
 
 def ConvertMemRefToSPIRV : Pass<"convert-memref-to-spirv", "ModuleOp"> {
   let summary = "Convert MemRef dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertMemRefToSPIRVPass()";
   let dependentDialects = ["spirv::SPIRVDialect"];
   let options = [

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRV.h"
 #include "mlir/Conversion/MemRefToSPIRV/MemRefToSPIRVPass.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "mlir-map-memref-storage-class"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
-// Utility Functions
+// Mappings
 //===----------------------------------------------------------------------===//
 
-/// Parses the  memory space mapping string `memorySpaceMapStr `and writes the
-/// mappings encoded inside to `memorySpaceMap`.
-static bool
-parseMappingStr(StringRef memorySpaceMapStr,
-                spirv::MemorySpaceToStorageClassMap &memorySpaceMap) {
-  memorySpaceMap.clear();
-  SmallVector<StringRef> mappings;
-  llvm::SplitString(memorySpaceMapStr, mappings, ",");
-  for (StringRef mapping : mappings) {
-    StringRef key, value;
-    std::tie(key, value) = mapping.split('=');
-    unsigned space;
-    if (!llvm::to_integer(key, space)) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "failed to parse mapping string key: " << key << "\n");
-      memorySpaceMap.clear();
-      return false;
-    }
-    Optional<spirv::StorageClass> storage = spirv::symbolizeStorageClass(value);
-    if (!storage) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "failed to parse mapping string value: " << value << "\n");
-      memorySpaceMap.clear();
-      return false;
-    }
-    memorySpaceMap[space] = *storage;
-  }
-  return true;
+spirv::MemorySpaceToStorageClassMap spirv::getDefaultVulkanStorageClassMap() {
+/// Mapping between SPIR-V storage classes to memref memory spaces.
+///
+/// Note: memref does not have a defined semantics for each memory space; it
+/// depends on the context where it is used. There are no particular reasons
+/// behind the number assignments; we try to follow NVVM conventions and largely
+/// give common storage classes a smaller number.
+#define STORAGE_SPACE_MAP_LIST(MAP_FN)                                         \
+  MAP_FN(spirv::StorageClass::StorageBuffer, 0)                                \
+  MAP_FN(spirv::StorageClass::Generic, 1)                                      \
+  MAP_FN(spirv::StorageClass::Workgroup, 3)                                    \
+  MAP_FN(spirv::StorageClass::Uniform, 4)                                      \
+  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)                                      \
+  MAP_FN(spirv::StorageClass::CrossWorkgroup, 11)                              \
+  MAP_FN(spirv::StorageClass::AtomicCounter, 12)                               \
+  MAP_FN(spirv::StorageClass::Image, 13)                                       \
+  MAP_FN(spirv::StorageClass::CallableDataKHR, 14)                             \
+  MAP_FN(spirv::StorageClass::IncomingCallableDataKHR, 15)                     \
+  MAP_FN(spirv::StorageClass::RayPayloadKHR, 16)                               \
+  MAP_FN(spirv::StorageClass::HitAttributeKHR, 17)                             \
+  MAP_FN(spirv::StorageClass::IncomingRayPayloadKHR, 18)                       \
+  MAP_FN(spirv::StorageClass::ShaderRecordBufferKHR, 19)                       \
+  MAP_FN(spirv::StorageClass::PhysicalStorageBuffer, 20)                       \
+  MAP_FN(spirv::StorageClass::CodeSectionINTEL, 21)                            \
+  MAP_FN(spirv::StorageClass::DeviceOnlyINTEL, 22)                             \
+  MAP_FN(spirv::StorageClass::HostOnlyINTEL, 23)
+
+#define STORAGE_SPACE_MAP_FN(storage, space) {space, storage},
+
+  return {STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)};
+
+#undef STORAGE_SPACE_MAP_FN
+#undef STORAGE_SPACE_MAP_LIST
 }
 
 //===----------------------------------------------------------------------===//
 // Type Converter
 //===----------------------------------------------------------------------===//
 
 spirv::MemorySpaceToStorageClassConverter::MemorySpaceToStorageClassConverter(
     const spirv::MemorySpaceToStorageClassMap &memorySpaceMap)
private:
   spirv::MemorySpaceToStorageClassMap memorySpaceMap;
 };
 } // namespace
 
 LogicalResult MapMemRefStorageClassPass::initializeOptions(StringRef options) {
   if (failed(Pass::initializeOptions(options)))
     return failure();
 
-  if (!parseMappingStr(mappings, memorySpaceMap))
+  if (clientAPI != "vulkan")
     return failure();
 
+  memorySpaceMap = spirv::getDefaultVulkanStorageClassMap();
+
   LLVM_DEBUG({
     llvm::dbgs() << "memory space to storage class mapping:\n";
     if (memorySpaceMap.empty())
       llvm::dbgs() << "  [empty]\n";
     for (auto kv : memorySpaceMap)
       llvm::dbgs() << "  " << kv.first << " -> "
                    << spirv::stringifyStorageClass(kv.second) << "\n";
   });

mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir

-// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='mappings=0=StorageBuffer,1=Uniform,2=Workgroup,3=PushConstant' -verify-diagnostics %s -o - | FileCheck %s
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='client-api=vulkan' -verify-diagnostics %s -o - | FileCheck %s --check-prefix=VULKAN
 
-// Mappings:
+// Vulkan Mappings:
 //   0 -> StorageBuffer (12)
-//   2 -> Workgroup (4)
-//   1 -> Uniform (2)
-//   3 -> PushConstant (9)
+//   1 -> Generic (8)
+//   3 -> Workgroup (4)
+//   4 -> Uniform (2)
 // TODO: create a StorageClass wrapper class so we can print the symbolc
 // storage class (instead of the backing IntegerAttr) and be able to
 // round trip the IR.
 
-// CHECK-LABEL: func @operand_result
+// VULKAN-LABEL: func @operand_result
 func.func @operand_result() {
-  // CHECK: memref<f32, 12 : i32>
+  // VULKAN: memref<f32, 12 : i32>
   %0 = "dialect.memref_producer"() : () -> (memref<f32>)
-  // CHECK: memref<4xi32, 2 : i32>
+  // VULKAN: memref<4xi32, 8 : i32>
   %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
-  // CHECK: memref<?x4xf16, 4 : i32>
-  %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 2>)
-  // CHECK: memref<*xf16, 9 : i32>
-  %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 3>)
+  // VULKAN: memref<?x4xf16, 4 : i32>
+  %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
+  // VULKAN: memref<*xf16, 2 : i32>
+  %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
 
 
   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
-  // CHECK: memref<4xi32, 2 : i32>
+  // VULKAN: memref<4xi32, 8 : i32>
   "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
-  // CHECK: memref<?x4xf16, 4 : i32>
-  "dialect.memref_consumer"(%2) : (memref<?x4xf16, 2>) -> ()
-  // CHECK: memref<*xf16, 9 : i32>
-  "dialect.memref_consumer"(%3) : (memref<*xf16, 3>) -> ()
+  // VULKAN: memref<?x4xf16, 4 : i32>
+  "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
+  // VULKAN: memref<*xf16, 2 : i32>
+  "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
 
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @type_attribute
+// VULKAN-LABEL: func @type_attribute
 func.func @type_attribute() {
-  // CHECK: attr = memref<i32, 2 : i32>
+  // VULKAN: attr = memref<i32, 8 : i32>
   "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> ()
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @function_io
+// VULKAN-LABEL: func @function_io
 func.func @function_io
-  // CHECK-SAME: (%{{.+}}: memref<f64, 4 : i32>, %{{.+}}: memref<4xi32, 9 : i32>)
-  (%arg0: memref<f64, 2>, %arg1: memref<4xi32, 3>)
-  // CHECK-SAME: -> (memref<f64, 4 : i32>, memref<4xi32, 9 : i32>)
-  -> (memref<f64, 2>, memref<4xi32, 3>) {
-  return %arg0, %arg1: memref<f64, 2>, memref<4xi32, 3>
+  // VULKAN-SAME: (%{{.+}}: memref<f64, 8 : i32>, %{{.+}}: memref<4xi32, 4 : i32>)
+  (%arg0: memref<f64, 1>, %arg1: memref<4xi32, 3>)
+  // VULKAN-SAME: -> (memref<f64, 8 : i32>, memref<4xi32, 4 : i32>)
+  -> (memref<f64, 1>, memref<4xi32, 3>) {
+  return %arg0, %arg1: memref<f64, 1>, memref<4xi32, 3>
 }
 
 // -----
 
-// CHECK: func @region
+// VULKAN: func @region
 func.func @region(%cond: i1, %arg0: memref<f32, 1>) {
   scf.if %cond {
-    //      CHECK: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, 4 : i32>}
-    // CHECK-SAME: (memref<f32, 2 : i32>) -> memref<f32, 2 : i32>
-    %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 2> } : (memref<f32, 1>) -> (memref<f32, 1>)
+    //      VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, 4 : i32>}
+    // VULKAN-SAME: (memref<f32, 8 : i32>) -> memref<f32, 8 : i32>
+    %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 3> } : (memref<f32, 1>) -> (memref<f32, 1>)
   }
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @non_memref_types
+// VULKAN-LABEL: func @non_memref_types
 func.func @non_memref_types(%arg: f32) -> f32 {
-  // CHECK: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
+  // VULKAN: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
   %0 = "dialect.op"(%arg) { attr = 16 } : (f32) -> (f32)
   return %0 : f32
 }
 
 // -----
 
 func.func @missing_mapping() {
   // expected-error @+1 {{failed to legalize}}
-  %0 = "dialect.memref_producer"() : () -> (memref<f32, 128>)
+  %0 = "dialect.memref_producer"() : () -> (memref<f32, 2>)
   return
 }