[mlir][spirv] Add CL/Kernel memory mapping

[antiagainst]

LGTM; just one thing to fix. Feel free to land afterwards!

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

95

Don't think OpenCL supports Input/Output; we should drop them.

[raikonenfnu]

Removed Input/Output from CL memory mapping.

[raikonenfnu]

In D132428#3742762, @antiagainst wrote:

LGTM; just one thing to fix. Feel free to land afterwards!

Done! thanks :)

[raikonenfnu]

In D132428#3742762, @antiagainst wrote:

LGTM; just one thing to fix. Feel free to land afterwards!

Although I don't have commit access to land yet 😅

[antiagainst]

Actually there is one more issue. I'll land once that's fixed.

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

95

Actually more to fix here: Uniform needs Shader capability; so it's not available for OpenCL. You might want to move UniformConstant here and map it to 4.

[raikonenfnu]

Replaced Uniform with UniformConstant for memoryspace 4.

[raikonenfnu]

fix comment in test.

[raikonenfnu]

Minor linting

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

 
 /// 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);
 /// 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);
+/// 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:
   explicit MemorySpaceToStorageClassConverter(
       const MemorySpaceToStorageClassMap &memorySpaceMap);
 
 private:

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

 //===----------------------------------------------------------------------===//
 
 /// 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)                                         \
+#define VULKAN_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)
 
 Optional<spirv::StorageClass>
 spirv::mapMemorySpaceToVulkanStorageClass(unsigned memorySpace) {
 #define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   case space:                                                                  \
     return storage;
 
   switch (memorySpace) {
-    STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
+    VULKAN_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     break;
   }
   return llvm::None;
 
 #undef STORAGE_SPACE_MAP_FN
 }
 
 Optional<unsigned>
 spirv::mapVulkanStorageClassToMemorySpace(spirv::StorageClass storageClass) {
 #define STORAGE_SPACE_MAP_FN(storage, space)                                   \
   case storage:                                                                \
     return space;
 
   switch (storageClass) {
-    STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
+    VULKAN_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
   default:
     break;
   }
   return llvm::None;
 
 #undef STORAGE_SPACE_MAP_FN
 }
 
-#undef STORAGE_SPACE_MAP_LIST
+#undef VULKAN_STORAGE_SPACE_MAP_LIST
+
+#define OPENCL_STORAGE_SPACE_MAP_LIST(MAP_FN)                                  \
+  MAP_FN(spirv::StorageClass::CrossWorkgroup, 0)                               \
+  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) {
+#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
+  case space:                                                                  \
+    return storage;
+
+  switch (memorySpace) {
+    OPENCL_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
+  default:
+    break;
+  }
+  return llvm::None;
+
+#undef STORAGE_SPACE_MAP_FN
+}
+
+Optional<unsigned>
+spirv::mapOpenCLStorageClassToMemorySpace(spirv::StorageClass storageClass) {
+#define STORAGE_SPACE_MAP_FN(storage, space)                                   \
+  case storage:                                                                \
+    return space;
+
+  switch (storageClass) {
+    OPENCL_STORAGE_SPACE_MAP_LIST(STORAGE_SPACE_MAP_FN)
+  default:
+    break;
+  }
+  return llvm::None;
+
+#undef STORAGE_SPACE_MAP_FN
+}
+
+#undef OPENCL_STORAGE_SPACE_MAP_LIST
 
 //===----------------------------------------------------------------------===//
 // Type Converter
 //===----------------------------------------------------------------------===//
 
 spirv::MemorySpaceToStorageClassConverter::MemorySpaceToStorageClassConverter(
     const spirv::MemorySpaceToStorageClassMap &memorySpaceMap)
     : memorySpaceMap(memorySpaceMap) {
private:
   spirv::MemorySpaceToStorageClassMap memorySpaceMap;
 };
 } // namespace
 
 LogicalResult MapMemRefStorageClassPass::initializeOptions(StringRef options) {
   if (failed(Pass::initializeOptions(options)))
     return failure();
 
-  if (clientAPI != "vulkan")
+  if (clientAPI == "opencl") {
+    memorySpaceMap = spirv::mapMemorySpaceToOpenCLStorageClass;
+  }
+
+  if (clientAPI != "vulkan" && clientAPI != "opencl")
     return failure();
 
   return success();
 }
 
 void MapMemRefStorageClassPass::runOnOperation() {
   MLIRContext *context = &getContext();
   Operation *op = getOperation();

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

 // 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
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='client-api=opencl' -verify-diagnostics %s -o - | FileCheck %s --check-prefix=OPENCL
 
 // Vulkan Mappings:
 //   0 -> StorageBuffer
 //   1 -> Generic
 //   2 -> [null]
 //   3 -> Workgroup
 //   4 -> Uniform
 
+// OpenCL Mappings:
+//   0 -> CrossWorkgroup
+//   1 -> Generic
+//   2 -> [null]
+//   3 -> Workgroup
+//   4 -> UniformConstant
+
 // VULKAN-LABEL: func @operand_result
+// OPENCL-LABEL: func @operand_result
 func.func @operand_result() {
   // VULKAN: memref<f32, #spv.storage_class<StorageBuffer>>
+  // OPENCL: memref<f32, #spv.storage_class<CrossWorkgroup>>
   %0 = "dialect.memref_producer"() : () -> (memref<f32>)
   // VULKAN: memref<4xi32, #spv.storage_class<Generic>>
+  // OPENCL: memref<4xi32, #spv.storage_class<Generic>>
   %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 1>)
   // VULKAN: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  // OPENCL: memref<?x4xf16, #spv.storage_class<Workgroup>>
   %2 = "dialect.memref_producer"() : () -> (memref<?x4xf16, 3>)
   // VULKAN: memref<*xf16, #spv.storage_class<Uniform>>
+  // OPENCL: memref<*xf16, #spv.storage_class<UniformConstant>>
   %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
 
 
   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
   // VULKAN: memref<4xi32, #spv.storage_class<Generic>>
+  // OPENCL: memref<4xi32, #spv.storage_class<Generic>>
   "dialect.memref_consumer"(%1) : (memref<4xi32, 1>) -> ()
   // VULKAN: memref<?x4xf16, #spv.storage_class<Workgroup>>
+  // OPENCL: memref<?x4xf16, #spv.storage_class<Workgroup>>
   "dialect.memref_consumer"(%2) : (memref<?x4xf16, 3>) -> ()
   // VULKAN: memref<*xf16, #spv.storage_class<Uniform>>
+  // OPENCL: memref<*xf16, #spv.storage_class<UniformConstant>>
   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
 
   return
 }
 
 // -----
 
 // VULKAN-LABEL: func @type_attribute
+// OPENCL-LABEL: func @type_attribute
 func.func @type_attribute() {
   // VULKAN: attr = memref<i32, #spv.storage_class<Generic>>
+  // OPENCL: attr = memref<i32, #spv.storage_class<Generic>>
   "dialect.memref_producer"() { attr = memref<i32, 1> } : () -> ()
   return
 }
 
 // -----
 
 // VULKAN-LABEL: func.func @function_io
+// OPENCL-LABEL: func.func @function_io
 func.func @function_io
   // VULKAN-SAME: (%{{.+}}: memref<f64, #spv.storage_class<Generic>>, %{{.+}}: memref<4xi32, #spv.storage_class<Workgroup>>)
+  // OPENCL-SAME: (%{{.+}}: memref<f64, #spv.storage_class<Generic>>, %{{.+}}: memref<4xi32, #spv.storage_class<Workgroup>>)
   (%arg0: memref<f64, 1>, %arg1: memref<4xi32, 3>)
   // VULKAN-SAME: -> (memref<f64, #spv.storage_class<Generic>>, memref<4xi32, #spv.storage_class<Workgroup>>)
+  // OPENCL-SAME: -> (memref<f64, #spv.storage_class<Generic>>, memref<4xi32, #spv.storage_class<Workgroup>>)
   -> (memref<f64, 1>, memref<4xi32, 3>) {
   return %arg0, %arg1: memref<f64, 1>, memref<4xi32, 3>
 }
 
 // -----
 
 gpu.module @kernel {
 // VULKAN-LABEL: gpu.func @function_io
+// OPENCL-LABEL: gpu.func @function_io
 // VULKAN-SAME: memref<8xi32, #spv.storage_class<StorageBuffer>>
+// OPENCL-SAME: memref<8xi32, #spv.storage_class<CrossWorkgroup>>
 gpu.func @function_io(%arg0 : memref<8xi32>) kernel { gpu.return }
 }
 
 // -----
 
 // VULKAN-LABEL: func.func @region
+// OPENCL-LABEL: func.func @region
 func.func @region(%cond: i1, %arg0: memref<f32, 1>) {
   scf.if %cond {
     //      VULKAN: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, #spv.storage_class<Workgroup>>}
+    //      OPENCL: "dialect.memref_consumer"(%{{.+}}) {attr = memref<i64, #spv.storage_class<Workgroup>>}
     // VULKAN-SAME: (memref<f32, #spv.storage_class<Generic>>) -> memref<f32, #spv.storage_class<Generic>>
+    // OPENCL-SAME: (memref<f32, #spv.storage_class<Generic>>) -> memref<f32, #spv.storage_class<Generic>>
     %0 = "dialect.memref_consumer"(%arg0) { attr = memref<i64, 3> } : (memref<f32, 1>) -> (memref<f32, 1>)
   }
   return
 }
 
 // -----
 
 // VULKAN-LABEL: func @non_memref_types
+// OPENCL-LABEL: func @non_memref_types
 func.func @non_memref_types(%arg: f32) -> f32 {
   // VULKAN: "dialect.op"(%{{.+}}) {attr = 16 : i64} : (f32) -> f32
+  // OPENCL: "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, 2>)
   return
 }