[mlir][spirv] Define spv.Image Operation

co-authered-by: Alan Liu <alanliu.yf@gmail.com>

Reviewed By: antiagainst

Differential Revision: https://reviews.llvm.org/D98270

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td

@@ -3030,9 +3030,11 @@ class SignlessOrUnsignedIntOfWidths<list<int> widths> :
 def SPV_IsArrayType : CPred<"$_self.isa<::mlir::spirv::ArrayType>()">;
 def SPV_IsCooperativeMatrixType :
   CPred<"$_self.isa<::mlir::spirv::CooperativeMatrixNVType>()">;
+def SPV_IsImageType : CPred<"$_self.isa<::mlir::spirv::ImageType>()">;
 def SPV_IsMatrixType : CPred<"$_self.isa<::mlir::spirv::MatrixType>()">;
 def SPV_IsPtrType : CPred<"$_self.isa<::mlir::spirv::PointerType>()">;
 def SPV_IsRTArrayType : CPred<"$_self.isa<::mlir::spirv::RuntimeArrayType>()">;
+def SPV_IsSampledImageType : CPred<"$_self.isa<::mlir::spirv::SampledImageType>()">;
 def SPV_IsStructType : CPred<"$_self.isa<::mlir::spirv::StructType>()">;
 
 
@@ -3056,12 +3058,16 @@ def SPV_AnyArray : DialectType<SPIRV_Dialect, SPV_IsArrayType,
 def SPV_AnyCooperativeMatrix : DialectType<SPIRV_Dialect,
                                SPV_IsCooperativeMatrixType,
                                "any SPIR-V cooperative matrix type">;
+def SPV_AnyImage : DialectType<SPIRV_Dialect, SPV_IsImageType,
+                                "any SPIR-V image type">;
 def SPV_AnyMatrix : DialectType<SPIRV_Dialect, SPV_IsMatrixType,
                                 "any SPIR-V matrix type">;
 def SPV_AnyRTArray : DialectType<SPIRV_Dialect, SPV_IsRTArrayType,
                                  "any SPIR-V runtime array type">;
 def SPV_AnyStruct : DialectType<SPIRV_Dialect, SPV_IsStructType,
                                 "any SPIR-V struct type">;
+def SPV_AnySampledImage : DialectType<SPIRV_Dialect, SPV_IsSampledImageType,
+                                "any SPIR-V sampled image type">;
 
 def SPV_Numerical : AnyTypeOf<[SPV_Integer, SPV_Float]>;
 def SPV_Scalar : AnyTypeOf<[SPV_Numerical, SPV_Bool]>;
@@ -3193,6 +3199,7 @@ def SPV_OC_OpCompositeConstruct        : I32EnumAttrCase<"OpCompositeConstruct",
 def SPV_OC_OpCompositeExtract          : I32EnumAttrCase<"OpCompositeExtract", 81>;
 def SPV_OC_OpCompositeInsert           : I32EnumAttrCase<"OpCompositeInsert", 82>;
 def SPV_OC_OpTranspose                 : I32EnumAttrCase<"OpTranspose", 84>;
+def SPV_OC_OpImage                     : I32EnumAttrCase<"OpImage", 100>;
 def SPV_OC_OpConvertFToU               : I32EnumAttrCase<"OpConvertFToU", 109>;
 def SPV_OC_OpConvertFToS               : I32EnumAttrCase<"OpConvertFToS", 110>;
 def SPV_OC_OpConvertSToF               : I32EnumAttrCase<"OpConvertSToF", 111>;
@@ -3331,13 +3338,13 @@ def SPV_OpcodeAttr :
       SPV_OC_OpMemberDecorate, SPV_OC_OpVectorExtractDynamic,
       SPV_OC_OpVectorInsertDynamic, SPV_OC_OpVectorShuffle,
       SPV_OC_OpCompositeConstruct, SPV_OC_OpCompositeExtract,
-      SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose, SPV_OC_OpConvertFToU,
+      SPV_OC_OpCompositeInsert, SPV_OC_OpTranspose, SPV_OC_OpImage,
-      SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF, SPV_OC_OpConvertUToF,
+      SPV_OC_OpConvertFToU, SPV_OC_OpConvertFToS, SPV_OC_OpConvertSToF,
-      SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert, SPV_OC_OpBitcast,
+      SPV_OC_OpConvertUToF, SPV_OC_OpUConvert, SPV_OC_OpSConvert, SPV_OC_OpFConvert,
-      SPV_OC_OpSNegate, SPV_OC_OpFNegate, SPV_OC_OpIAdd, SPV_OC_OpFAdd,
+      SPV_OC_OpBitcast, SPV_OC_OpSNegate, SPV_OC_OpFNegate, SPV_OC_OpIAdd,
-      SPV_OC_OpISub, SPV_OC_OpFSub, SPV_OC_OpIMul, SPV_OC_OpFMul, SPV_OC_OpUDiv,
+      SPV_OC_OpFAdd, SPV_OC_OpISub, SPV_OC_OpFSub, SPV_OC_OpIMul, SPV_OC_OpFMul,
-      SPV_OC_OpSDiv, SPV_OC_OpFDiv, SPV_OC_OpUMod, SPV_OC_OpSRem, SPV_OC_OpSMod,
+      SPV_OC_OpUDiv, SPV_OC_OpSDiv, SPV_OC_OpFDiv, SPV_OC_OpUMod, SPV_OC_OpSRem,
-      SPV_OC_OpFRem, SPV_OC_OpFMod, SPV_OC_OpMatrixTimesScalar,
+      SPV_OC_OpSMod, SPV_OC_OpFRem, SPV_OC_OpFMod, SPV_OC_OpMatrixTimesScalar,
       SPV_OC_OpMatrixTimesMatrix, SPV_OC_OpIsNan, SPV_OC_OpIsInf, SPV_OC_OpOrdered,
       SPV_OC_OpUnordered, SPV_OC_OpLogicalEqual, SPV_OC_OpLogicalNotEqual,
       SPV_OC_OpLogicalOr, SPV_OC_OpLogicalAnd, SPV_OC_OpLogicalNot, SPV_OC_OpSelect,

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td

@@ -0,0 +1,58 @@
+//===-- SPIRVImageOps.td - MLIR SPIR-V Image Ops ------*- tablegen -*------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains image ops for the SPIR-V dialect. It corresponds
+// to "3.37.10. Image Instructions" of the SPIR-V specification.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IR_IMAGE_OPS
+#define MLIR_DIALECT_SPIRV_IR_IMAGE_OPS
+
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+// -----
+
+def SPV_ImageOp : SPV_Op<"Image",
+    [NoSideEffect,
+     TypesMatchWith<"type of 'result' matches image type of 'sampledimage'",
+                    "sampledimage", "result",
+                    "$_self.cast<spirv::SampledImageType>().getImageType()">]> {
+  let summary = "Extract the image from a sampled image.";
+
+  let description = [{
+    Result Type must be OpTypeImage.
+
+    Sampled Image must have type OpTypeSampledImage whose Image Type is the
+    same as Result Type.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %0 = spv.Image %1 : !spv.sampled_image<!spv.image<f32, Cube, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>
+    ```
+  }];
+
+  let arguments = (ins
+    SPV_AnySampledImage:$sampledimage
+  );
+
+  let results = (outs
+    SPV_AnyImage:$result
+  );
+
+  let assemblyFormat = "attr-dict $sampledimage `:` type($sampledimage)";
+
+  let verifier = ?;
+
+}
+
+#endif // MLIR_DIALECT_SPIRV_IR_GLSL_OPS

mlir/include/mlir/Dialect/SPIRV/IR/SPIRVOps.td

@@ -32,6 +32,7 @@ include "mlir/Dialect/SPIRV/IR/SPIRVControlFlowOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVCooperativeMatrixOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVGLSLOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVGroupOps.td"
+include "mlir/Dialect/SPIRV/IR/SPIRVImageOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVLogicalOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVMatrixOps.td"
 include "mlir/Dialect/SPIRV/IR/SPIRVMemoryOps.td"

mlir/test/Dialect/SPIRV/IR/image-ops.mlir

@@ -0,0 +1,11 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// spv.Image
+//===----------------------------------------------------------------------===//
+
+func @image(%arg0 : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>) -> () {
+  // CHECK: spv.Image {{.*}} : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
+  %0 = spv.Image %arg0 : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
+  return
+}

mlir/test/Target/SPIRV/image-ops.mlir

@@ -0,0 +1,9 @@
+// RUN: mlir-translate -test-spirv-roundtrip %s | FileCheck %s
+
+spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
+  spv.func @image(%arg0 : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>) "None" {
+    // CHECK: {{%.*}} = spv.Image {{%.*}} : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
+    %0 = spv.Image %arg0 : !spv.sampled_image<!spv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
+    spv.Return
+  }
+}