[Mlir-commits] [mlir] [mlir][spirv]: Add Image to Vulkan Storage Class Map (PR #144899)

[Jack Frankland]

https://github.com/FranklandJack updated https://github.com/llvm/llvm-project/pull/144899

>From 3bf9d0e7d6043888f0c9a7bd78083e23ff716c8c Mon Sep 17 00:00:00 2001
From: Jack Frankland <jack.frankland at arm.com>
Date: Thu, 19 Jun 2025 13:58:42 +0100
Subject: [PATCH] [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 <jack.frankland at arm.com>
---
 .../MapMemRefStorageClassPass.cpp             |  3 +-
 .../MemRefToSPIRV/map-storage-class.mlir      | 32 ++++++++++++++++++-
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp b/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp
index 4cbc3dfdae223..1fbc5a03987e8 100644
--- a/mlir/lib/Conversion/MemRefToSPIRV/MapMemRefStorageClassPass.cpp
+++ b/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) {
diff --git a/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir
index f0956b62760a2..fdc69b8119994 100644
--- a/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir
+++ b/mlir/test/Conversion/MemRefToSPIRV/map-storage-class.mlir
@@ -1,5 +1,6 @@
 // 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
@@ -7,6 +8,14 @@
 //   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
@@ -14,6 +23,9 @@
 //   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
 }
-}
\ No newline at end of file
+}