From 8bca285a4cc9160f39328aa637c7e63993d8661b Mon Sep 17 00:00:00 2001
From: Igor Wodiany <igor.wodiany@imgtec.com>
Date: Tue, 10 Jun 2025 19:53:32 +0100
Subject: [PATCH] [mlir][spirv] Include `SPIRV_AnyImage` in `SPIRV_Type`

This change is trigger by encountering the following error:

```
<unknown>:0: error: 'spirv.Load' op result #0 must be void
or bool or 8/16/32/64-bit integer or 16/32/64-bit float or
vector of bool or 8/16/32/64-bit integer or 16/32/64-bit
float values of length 2/3/4/8/16 or any SPIR-V pointer type
or any SPIR-V array type or any SPIR-V run time array type
or any SPIR-V struct type or any SPIR-V cooperative matrix
type or any SPIR-V matrix type or any SPIR-V sampled image
type, but got '!spirv.image<f32, Dim2D, NoDepth, NonArrayed,
SingleSampled, NoSampler, Rgba8>'<unknown>:0: note: see current
operation:
%126 = "spirv.Load"(%125) {relaxed_precision} : (!spirv.ptr<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>, UniformConstant>) -> !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>
```
---
 mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td |  3 ++-
 mlir/test/Dialect/SPIRV/IR/memory-ops.mlir      | 10 ++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index 8fd533db83d9a..b143cf9a5f509 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4196,7 +4196,8 @@ def SPIRV_Composite :
 def SPIRV_Type : AnyTypeOf<[
     SPIRV_Void, SPIRV_Bool, SPIRV_Integer, SPIRV_Float, SPIRV_Vector,
     SPIRV_AnyPtr, SPIRV_AnyArray, SPIRV_AnyRTArray, SPIRV_AnyStruct,
-    SPIRV_AnyCooperativeMatrix, SPIRV_AnyMatrix, SPIRV_AnySampledImage
+    SPIRV_AnyCooperativeMatrix, SPIRV_AnyMatrix, SPIRV_AnySampledImage,
+    SPIRV_AnyImage
   ]>;
 
 def SPIRV_SignedInt : SignedIntOfWidths<[8, 16, 32, 64]>;
diff --git a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
index 57ff94762ff68..a3b96c698a344 100644
--- a/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/memory-ops.mlir
@@ -356,6 +356,16 @@ spirv.module Logical GLSL450 {
 
 // -----
 
+// CHECK-LABEL: @image_load
+func.func @image_load() -> () {
+  %0 = spirv.Variable : !spirv.ptr<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>, Function>
+  // CHECK: spirv.Load "Function" %{{.*}} : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>
+  %1 = spirv.Load "Function" %0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>
+  return
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spirv.StoreOp
 //===----------------------------------------------------------------------===//