[mlir][spirv]: Add Image to Vulkan Storage Class Map

Extend the "storage class" <-> "memory space" map for the Vulkan SPIR-V
environment to include the Image class. 12 is chosen as the next
available value in the MemRef memory space list.

Extend the pass testing to include missing memory scopes and add a new
test file for the memory address indices which only support a mapping in
the Vuklan environment. It appears that previously there was a missing
CHECK line for the default pass behavior so that has been added.

Signed-off-by: Jack Frankland <[email protected]>

Author: FranklandJack

mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp

@@ -59,7 +59,8 @@ using namespace mlir;
   MAP_FN(spirv::StorageClass::UniformConstant, 8)                              \
   MAP_FN(spirv::StorageClass::Input, 9)                                        \
   MAP_FN(spirv::StorageClass::Output, 10)                                      \
-  MAP_FN(spirv::StorageClass::PhysicalStorageBuffer, 11)
+  MAP_FN(spirv::StorageClass::PhysicalStorageBuffer, 11)                       \
+  MAP_FN(spirv::StorageClass::Image, 12)
 
 std::optional<spirv::StorageClass>
 spirv::mapMemorySpaceToVulkanStorageClass(Attribute memorySpaceAttr) {

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

@@ -0,0 +1,27 @@
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class='client-api=vulkan' %s -o - | FileCheck %s
+
+// Vulkan Specific Mappings:
+//   8 -> UniformConstant
+//   9 -> Input
+//   10 -> Output
+//   11 -> PhysicalStorageBuffer
+//   12 -> Image
+
+/// Check that Vulkan specific memory space indices get converted into the correct
+/// SPIR-V storage class. If mappings to OpenCL address spaces are added for these
+/// indices then those test case should be moved into the common test file.
+
+// CHECK-LABEL: func @test_vk_specific_memory_spaces
+func.func @test_vk_specific_memory_spaces() {
+  // CHECK: memref<4xi32, #spirv.storage_class<UniformConstant>>
+  %1 = "dialect.memref_producer"() : () -> (memref<4xi32, 8>)
+  // CHECK: memref<4xi32, #spirv.storage_class<Input>>
+  %2 = "dialect.memref_producer"() : () -> (memref<4xi32, 9>)
+  // CHECK: memref<4xi32, #spirv.storage_class<Output>>
+  %3 = "dialect.memref_producer"() : () -> (memref<4xi32, 10>)
+  // CHECK: memref<4xi32, #spirv.storage_class<PhysicalStorageBuffer>>
+  %4 = "dialect.memref_producer"() : () -> (memref<4xi32, 11>)
+  // CHECK: memref<4xi32, #spirv.storage_class<Image>>
+  %5 = "dialect.memref_producer"() : () -> (memref<4xi32, 12>)
+  return
+}

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

@@ -1,19 +1,31 @@
 // 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
+// RUN: mlir-opt -split-input-file -allow-unregistered-dialect -map-memref-spirv-storage-class -verify-diagnostics %s -o - | FileCheck %s
 
 // Vulkan Mappings:
 //   0 -> StorageBuffer
 //   1 -> Generic
 //   2 -> [null]
 //   3 -> Workgroup
 //   4 -> Uniform
+//   5 -> Private
+//   6 -> Function
+//   7 -> PushConstant
+//   8 -> UniformConstant
+//   9 -> Input
+//   10 -> Output
+//   11 -> PhysicalStorageBuffer
+//   12 -> Image
 
 // OpenCL Mappings:
 //   0 -> CrossWorkgroup
 //   1 -> Generic
 //   2 -> [null]
 //   3 -> Workgroup
 //   4 -> UniformConstant
+//   5 -> Private
+//   6 -> Function
+//   7 -> Image
 
 // VULKAN-LABEL: func @operand_result
 // OPENCL-LABEL: func @operand_result
@@ -30,6 +42,15 @@ func.func @operand_result() {
   // VULKAN: memref<*xf16, #spirv.storage_class<Uniform>>
   // OPENCL: memref<*xf16, #spirv.storage_class<UniformConstant>>
   %3 = "dialect.memref_producer"() : () -> (memref<*xf16, 4>)
+  // VULKAN: memref<*xf16, #spirv.storage_class<Private>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Private>>
+  %4 = "dialect.memref_producer"() : () -> (memref<*xf16, 5>)
+  // VULKAN: memref<*xf16, #spirv.storage_class<Function>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Function>>
+  %5 = "dialect.memref_producer"() : () -> (memref<*xf16, 6>)
+  // VULKAN: memref<*xf16, #spirv.storage_class<PushConstant>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Image>>
+  %6 = "dialect.memref_producer"() : () -> (memref<*xf16, 7>)
 
 
   "dialect.memref_consumer"(%0) : (memref<f32>) -> ()
@@ -42,6 +63,15 @@ func.func @operand_result() {
   // VULKAN: memref<*xf16, #spirv.storage_class<Uniform>>
   // OPENCL: memref<*xf16, #spirv.storage_class<UniformConstant>>
   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
+  // VULKAN: memref<*xf16, #spirv.storage_class<Private>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Private>>
+  "dialect.memref_consumer"(%4) : (memref<*xf16, 5>) -> ()
+  // VULKAN: memref<*xf16, #spirv.storage_class<Function>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Function>>
+  "dialect.memref_consumer"(%5) : (memref<*xf16, 6>) -> ()
+  // VULKAN: memref<*xf16, #spirv.storage_class<PushConstant>>
+  // OPENCL: memref<*xf16, #spirv.storage_class<Image>>
+  "dialect.memref_consumer"(%6) : (memref<*xf16, 7>) -> ()
 
   return
 }
@@ -166,4 +196,4 @@ func.func @operand_result() {
   "dialect.memref_consumer"(%3) : (memref<*xf16, 4>) -> ()
   return
 }
-}
+}