[mlir][spirv] Add definition for selected sample operations (#129558)

This commit adds following three operations: ImageSampleImplicitLodOp,
ImageSampleExplicitLodOp and ImageSampleProjDrefImplicitLodOp

Author: IgWod-IMG

mlir/include/mlir/Dialect/SPIRV/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(IR)
+add_subdirectory(Interfaces)
 add_subdirectory(Transforms)

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

@@ -15,6 +15,7 @@
 #define MLIR_DIALECT_SPIRV_IR_IMAGE_OPS
 
 include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+include "mlir/Dialect/SPIRV/Interfaces/SPIRVImageInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // -----
@@ -48,14 +49,19 @@ class SPIRV_DimIsNot<string operand, list<string> values, string transform="$_se
 
 class SPIRV_NoneOrElementMatchImage<string operand, string image, string transform="$_self"> : PredOpTrait<
   "the " # operand # " component type must match the image sampled type",
-  CPred<"::llvm::isa<NoneType>(cast<ImageType>(" # !subst("$_self", "$" # image # ".getType()", transform) # ").getElementType()) ||"
+  CPred<"::llvm::isa<NoneType>(cast<::mlir::spirv::ImageType>(" # !subst("$_self", "$" # image # ".getType()", transform) # ").getElementType()) ||"
         "(getElementTypeOrSelf($" # operand # ")"
           "=="
-        "cast<ImageType>(" # !subst("$_self", "$" # image # ".getType()", transform) # ").getElementType())"
+        "cast<::mlir::spirv::ImageType>(" # !subst("$_self", "$" # image # ".getType()", transform) # ").getElementType())"
   >
 >;
 
-def SPIRV_SampledImageTransform : StrFunc<"llvm::cast<spirv::SampledImageType>($_self).getImageType()">;
+class SPIRV_ImageOperandIsPresent<string operand, list<string> values> : PredOpTrait<
+  "either " # !interleave(values, " or ") # " image operands must be present",
+  CPred<"::mlir::spirv::bitEnumContainsAny($" # operand # ", " # "::mlir::spirv::ImageOperands::" # !interleave(values, " | ::mlir::spirv::ImageOperands::") # ")">
+>;
+
+def SPIRV_SampledImageTransform : StrFunc<"llvm::cast<::mlir::spirv::SampledImageType>($_self).getImageType()">;
 
 // -----
 
@@ -89,7 +95,7 @@ def SPIRV_ImageDrefGatherOp : SPIRV_Op<"ImageDrefGather",
 
     ```mlir
     %0 = spirv.ImageDrefGather %1, %2, %3 : !spirv.sampled_image<!spirv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, vector<4xf32>, f32 -> vector<4xi32>
-    %0 = spirv.ImageDrefGather %1, %2, %3 : !spirv.sampled_image<!spirv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, vector<4xf32>, f32 ["NonPrivateTexel"] -> vector<4xi32>
+    %0 = spirv.ImageDrefGather %1, %2, %3 ["NonPrivateTexel"] : !spirv.sampled_image<!spirv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, vector<4xf32>, f32 -> vector<4xi32>
     ```
   }];
 
@@ -268,4 +274,214 @@ def SPIRV_ImageOp : SPIRV_Op<"Image",
   let hasVerifier = 0;
 }
 
-#endif // MLIR_DIALECT_SPIRV_IR_GL_OPS
+// -----
+
+def SPIRV_ImageSampleExplicitLodOp : SPIRV_Op<"ImageSampleExplicitLod",
+    [SPIRV_DimIsNot<"sampled_image", ["Buffer"], SPIRV_SampledImageTransform.result>,
+     SPIRV_MSOperandIs<"sampled_image", ["SingleSampled"], SPIRV_SampledImageTransform.result>,
+     SPIRV_NoneOrElementMatchImage<"result", "sampled_image", SPIRV_SampledImageTransform.result>,
+     SPIRV_ImageOperandIsPresent<"image_operands", ["Lod", "Grad"]>,
+     DeclareOpInterfaceMethods<SPIRV_ExplicitLodOpInterface>]> {
+  let summary = "Sample an image using an explicit level of detail.";
+
+  let description = [{
+    Result Type must be a vector of four components of floating-point type
+    or integer type. Its components must be the same as Sampled Type of the
+    underlying OpTypeImage (unless that underlying Sampled Type is
+    OpTypeVoid).
+
+    Sampled Image must be an object whose type is OpTypeSampledImage. Its
+    OpTypeImage must not have a Dim of Buffer. The MS operand of the
+    underlying OpTypeImage must be 0.
+
+    Coordinate must be a scalar or vector of floating-point type or integer
+    type. It contains (u[, v] ... [, array layer]) as needed by the
+    definition of Sampled Image. Unless the Kernel capability is declared,
+    it must be floating point. It may be a vector larger than needed, but
+    all unused components appear after all used components.
+
+    Image Operands encodes what operands follow, as per Image Operands.
+    Either Lod or Grad image operands must be present.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %result = spirv.ImageSampleExplicitLod %image, %coord ["Lod"](%lod) :
+      !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>,
+      vector<2xf32> (f32) -> vector<4xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_AnySampledImage:$sampled_image,
+    AnyTypeOf<[SPIRV_ScalarOrVectorOf<SPIRV_Float>, SPIRV_ScalarOrVectorOf<SPIRV_Integer>]>:$coordinate,
+    SPIRV_ImageOperandsAttr:$image_operands,
+    Variadic<SPIRV_Type>:$operand_arguments
+  );
+
+  let results = (outs
+    AnyTypeOf<[SPIRV_Vec4<SPIRV_Integer>, SPIRV_Vec4<SPIRV_Float>]>:$result
+  );
+
+  let assemblyFormat = [{
+    $sampled_image `,` $coordinate custom<ImageOperands>($image_operands) ( `,` $operand_arguments^ )? attr-dict
+      `:` type($sampled_image) `,` type($coordinate) ( `,` type($operand_arguments)^ )?
+      `->` type($result)
+  }];
+}
+
+// -----
+
+def SPIRV_ImageSampleImplicitLodOp : SPIRV_Op<"ImageSampleImplicitLod",
+    [SPIRV_DimIsNot<"sampled_image", ["Buffer"], SPIRV_SampledImageTransform.result>,
+     SPIRV_MSOperandIs<"sampled_image", ["SingleSampled"], SPIRV_SampledImageTransform.result>,
+     SPIRV_NoneOrElementMatchImage<"result", "sampled_image", SPIRV_SampledImageTransform.result>,
+     DeclareOpInterfaceMethods<SPIRV_ImplicitLodOpInterface>]> {
+  let summary = "Sample an image with an implicit level of detail.";
+
+  let description = [{
+    An invocation will not execute a dynamic instance of this instruction
+    (X') until all invocations in its derivative group have executed all
+    dynamic instances that are program-ordered before X'.
+
+    Result Type must be a vector of four components of floating-point type
+    or integer type. Its components must be the same as Sampled Type of the
+    underlying OpTypeImage (unless that underlying Sampled Type is
+    OpTypeVoid).
+
+    Sampled Image must be an object whose type is OpTypeSampledImage. Its
+    OpTypeImage must not have a Dim of Buffer. The MS operand of the
+    underlying OpTypeImage must be 0.
+
+    Coordinate must be a scalar or vector of floating-point type. It
+    contains (u[, v] ... [, array layer]) as needed by the definition of
+    Sampled Image. It may be a vector larger than needed, but all unused
+    components appear after all used components.
+
+    Image Operands encodes what operands follow, as per Image Operands.
+
+    This instruction is only valid in the Fragment Execution Model. In
+    addition, it consumes an implicit derivative that can be affected by
+    code motion.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %result = spirv.ImageSampleImplicitLod %image, %coord :
+      !spirv.sampled_image<!spirv.image<f32, Cube, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>,
+      vector<3xf32> -> vector<4xf32>
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[]>,
+    Capability<[SPIRV_C_Shader]>
+  ];
+
+  let arguments = (ins
+    SPIRV_AnySampledImage:$sampled_image,
+    SPIRV_ScalarOrVectorOf<SPIRV_Float>:$coordinate,
+    OptionalAttr<SPIRV_ImageOperandsAttr>:$image_operands,
+    Variadic<SPIRV_Type>:$operand_arguments
+  );
+
+  let results = (outs
+    AnyTypeOf<[SPIRV_Vec4<SPIRV_Integer>, SPIRV_Vec4<SPIRV_Float>]>:$result
+  );
+
+  let assemblyFormat = [{
+    $sampled_image `,` $coordinate custom<ImageOperands>($image_operands) ( `,` $operand_arguments^ )? attr-dict
+      `:` type($sampled_image) `,` type($coordinate) ( `,` type($operand_arguments)^ )?
+      `->` type($result)
+  }];
+}
+
+// -----
+
+def SPIRV_ImageSampleProjDrefImplicitLodOp : SPIRV_Op<"ImageSampleProjDrefImplicitLod",
+    [SPIRV_DimIsNot<"sampled_image", ["Buffer"], SPIRV_SampledImageTransform.result>,
+     SPIRV_MSOperandIs<"sampled_image", ["SingleSampled"], SPIRV_SampledImageTransform.result>,
+     TypesMatchWith<"type of 'result' matches image type of 'sampled_image'",
+                    "sampled_image", "result",
+                    "::llvm::cast<::mlir::spirv::ImageType>(::llvm::cast<spirv::SampledImageType>($_self).getImageType()).getElementType()">,
+     DeclareOpInterfaceMethods<SPIRV_ImplicitLodOpInterface>]> {
+
+  let summary = [{
+    Sample an image with a project coordinate, doing depth-comparison, with
+    an implicit level of detail.
+  }];
+
+  let description = [{
+    An invocation will not execute a dynamic instance of this instruction
+    (X') until all invocations in its derivative group have executed all
+    dynamic instances that are program-ordered before X'.
+
+    Result Type must be a scalar of integer type or floating-point type. It
+    must be the same as Sampled Type of the underlying OpTypeImage.
+
+    Sampled Image must be an object whose type is OpTypeSampledImage. The
+    Dim operand of the underlying OpTypeImage must be 1D, 2D, 3D, or Rect,
+    and the Arrayed and MS operands must be 0.
+
+    Coordinate must be a vector of floating-point type. It contains (u[,
+    v] [, w], q), as needed by the definition of Sampled Image, with the q
+    component consumed for the projective division. That is, the actual
+    sample coordinate is (u/q [, v/q] [, w/q]), as needed by the definition
+    of Sampled Image. It may be a vector larger than needed, but all unused
+    components appear after all used components.
+
+    Dref/q is the depth-comparison reference value. Dref must be a 32-bit
+    floating-point type scalar.
+
+    Image Operands encodes what operands follow, as per Image Operands.
+
+    This instruction is only valid in the Fragment Execution Model. In
+    addition, it consumes an implicit derivative that can be affected by
+    code motion.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    %result = spirv.ImageSampleProjDrefImplicitLod %image, %coord, %dref :
+      !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>,
+      vector<4xf16>, f32 -> f32
+    ```
+  }];
+
+  let availability = [
+    MinVersion<SPIRV_V_1_0>,
+    MaxVersion<SPIRV_V_1_6>,
+    Extension<[]>,
+    Capability<[SPIRV_C_Shader]>
+  ];
+
+  let arguments = (ins
+    SPIRV_AnySampledImage:$sampled_image,
+    AnyTypeOf<[SPIRV_ScalarOrVectorOf<SPIRV_Float>, SPIRV_ScalarOrVectorOf<SPIRV_Integer>]>:$coordinate,
+    SPIRV_Float32:$dref,
+    OptionalAttr<SPIRV_ImageOperandsAttr>:$image_operands,
+    Variadic<SPIRV_Type>:$operand_arguments
+  );
+
+  let results = (outs
+    AnyTypeOf<[SPIRV_Integer, SPIRV_Float]>:$result
+  );
+
+  let assemblyFormat = [{
+    $sampled_image `,` $coordinate `,` $dref custom<ImageOperands>($image_operands) ( `,` $operand_arguments^ )? attr-dict
+      `:` type($sampled_image) `,` type($coordinate) `,` type($dref) ( `,` type($operand_arguments)^ )?
+      `->` type($result)
+  }];
+}
+
+// -----
+
+#endif // MLIR_DIALECT_SPIRV_IR_IMAGE_OPS

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

@@ -17,6 +17,7 @@
 #include "mlir/Dialect/SPIRV/IR/SPIRVAttributes.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOpTraits.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"
+#include "mlir/Dialect/SPIRV/Interfaces/SPIRVImageInterfaces.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/Interfaces/CallInterfaces.h"

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

@@ -0,0 +1 @@
+add_mlir_interface(SPIRVImageInterfaces)

mlir/include/mlir/Dialect/SPIRV/Interfaces/CMakeLists.txt

@@ -0,0 +1,14 @@
+//===- SPIRVImageInterfaces.h -----------------------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IMAGE_INTERFACES_H_
+#define MLIR_DIALECT_SPIRV_IMAGE_INTERFACES_H_
+
+#include "mlir/Dialect/SPIRV/Interfaces/SPIRVImageInterfaces.h.inc"
+
+#endif // MLIR_DIALECT_SPIRV_IMAGE_INTERFACES_H_

mlir/include/mlir/Dialect/SPIRV/Interfaces/SPIRVImageInterfaces.h

@@ -0,0 +1,89 @@
+//===-- SPIRVImageInterfaces.td - MLIR SPIR-V Image Interfaces ------*- 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 interfaces used by image operations. 
+//
+//===------------------------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPIRV_IMAGE_INTERFACES
+#define MLIR_DIALECT_SPIRV_IMAGE_INTERFACES
+
+include "mlir/Dialect/SPIRV/IR/SPIRVBase.td"
+
+// -----
+
+def SPIRV_SampleOpInterface : OpInterface<"SamplingOpInterface"> {
+  let description = [{
+    The `SampleOpInterface` defines sampling image operations (`spirv.Image*Sample*`)
+    and provides interface methods to query operands common across all sampling
+    instructions.
+
+    Currently only getters for sampled image and coordinate are provided. The
+    default implementation will work as long as the instruction the interface is
+    attached to uses standard names for the operands: `$sampled_image` and `$coordinate`.
+  }];
+  let cppNamespace = "::mlir::spirv";
+  let methods = [
+    InterfaceMethod<
+      "Get SampledImage of the operation.",
+      "::mlir::TypedValue<::mlir::Type>", "getSampledImage", (ins), [{
+        return $_op.getSampledImage();
+    }]>,
+    InterfaceMethod<
+      "Get Coordinate of the operation.",
+      "::mlir::TypedValue<::mlir::Type>", "getCoordinate", (ins), [{
+        return $_op.getCoordinate();
+    }]>
+  ];
+}
+
+// -----
+
+def SPIRV_ExplicitLodOpInterface : OpInterface<"ExplicitLodOpInterface", [SPIRV_SampleOpInterface]> {
+  let description = [{
+    The `ExplicitLodOpInterface` defines explicit sampling lod operations (`spirv.*ExplicitLod`). Currently
+    the interface is only used to check whether the instruction is an explicit lod.
+  }];
+  let cppNamespace = "::mlir::spirv";
+}
+
+// -----
+
+def SPIRV_ImplicitLodOpInterface : OpInterface<"ImplicitLodOpInterface", [SPIRV_SampleOpInterface]> {
+  let description = [{
+    The `ImplicitLodOpInterface` defines implicit sampling lod operations (`spirv.*ImplicitLod`). Currently
+    the interface is only used to check whether the instruction is an implicit lod.
+  }];
+  let cppNamespace = "::mlir::spirv";
+}
+
+// -----
+
+def SPIRV_FetchOpInterface : OpInterface<"FetchOpInterface"> {
+  let description = [{
+    The `FetchOpInterface` defines fetch image operations (`spirv.ImageFetch` and
+    `spirv.ImageSpareFetch`) and provides interface methods to query operands common
+    across all fetch instructions.
+
+    Currently only a getter for image is provided. The default implementation
+    will work as long as the instruction the interface is attached to uses
+    standard names for the operands: `$image`.
+  }];
+  let cppNamespace = "::mlir::spirv";
+  let methods = [
+    InterfaceMethod<
+      "Get Image of the operation.",
+      "::mlir::TypedValue<::mlir::Type>", "getImage", (ins), [{
+        return $_op.getImage();
+    }]>
+  ];
+}
+
+// -----
+
+#endif // MLIR_DIALECT_SPIRV_IMAGE_INTERFACES

mlir/include/mlir/Dialect/SPIRV/Interfaces/SPIRVImageInterfaces.td

@@ -1,4 +1,5 @@
 add_subdirectory(IR)
+add_subdirectory(Interfaces)
 add_subdirectory(Linking)
 add_subdirectory(Transforms)
 add_subdirectory(Utils)

mlir/lib/Dialect/SPIRV/CMakeLists.txt

@@ -35,6 +35,7 @@ add_mlir_dialect_library(MLIRSPIRVDialect
   MLIRSPIRVCanonicalizationIncGen
   MLIRSPIRVEnumAvailabilityIncGen
   MLIRSPIRVEnumsIncGen
+  MLIRSPIRVImageInterfacesIncGen
   MLIRSPIRVOpsIncGen
 
   LINK_LIBS PUBLIC
@@ -43,6 +44,7 @@ add_mlir_dialect_library(MLIRSPIRVDialect
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
+  MLIRSPIRVImageInterfaces
   MLIRSupport
   MLIRTransforms
   MLIRUBDialect