[SPIRV][ImageOperand] Add Image Operands for Image Instructions (NFC)

[Weiwei-2021]

Summary

This patch is to add Image Operands in SPIR-V Dialect. And let ImageDrefGather to use Image Operands as an example.

Image Operands are used in many image instructions. The tricky part is that "Image Operands encodes what oprands follow, as per Image Operands".
And ususally, they are optional to image instructions.

I don't want to change the existing assemblyFormat too much. So this is the format I create for now:

%0 = spv.ImageDrefGather %1, ... %3 : f32 ["Bias|Lod"](%4, %5 : f32, f32) -> ...

Any comment about the Image Operands is welcome since it will impact almost all image instructions' format. Thank you in advance.

Anoter thing is that Image Operands have a lot of validation rules (See Section 3.14 in SPIR-V Spec). It would be a hugh amuont of work.
Therefore I leave an TODO in verifyImageOperands function.

co-authored: Alan Liu <[e-mail]>

[antiagainst]

I guess you are sending this out for early comments. ODS and assembly LGTM. Need to have some minimal verification and (de)serialization to land though. :) Some more comments following.

Image Operands are used in many image instructions. The tricky part is that "Image Operands encodes what oprands follow, as per Image Operands".
And ususally, they are optional to image instructions.

Indeed. Using

OptionalAttr<SPV_ImageOperandAttr>:$imageoperands,
Variadic<SPV_Type>:$values

Is a good approach to me.

I don't want to change the existing assemblyFormat too much. So this is the format I create for now:
%0 = spv.ImageDrefGather %1, ... %3 : f32 ["Bias|Lod"](%4, %5 : f32, f32) -> ...
Any comment about the Image Operands is welcome since it will impact almost all image instructions' format. Thank you in advance.

This LGTM!

Another thing is that Image Operands have a lot of validation rules (See Section 3.14 in SPIR-V Spec). It would be a huge amount of work.

Therefore I leave an TODO in verifyImageOperands function.

It's fine to implement them piece by piece. Like we can support those image operands one by one; erroring out for the rest.

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

72

Nit: values is too generic. What about calling it operand_arguments?

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

337

This isn't necessary?

342

If attr is null but values isn't empty, that should be a validation error?

345

Nit: In MLIR we prefer early return when it simplifies code. So here if (!attr) return success().

348

FYI: we have bitEnumContains utility functions generated for probing whether a bit is set.

351

I assume this part needs to be rewritten? Not sure I follow what's verified here..

[antiagainst]

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

348

And also operator| and operator&.

[Weiwei-2021]

Updated the verify part.

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

351

Sorry about the confusion. Initially, I want to check the "Bias" and "Lod" will have one following operand for each. And after reading the detail, I realized that they cannot show up at the same time. :(
Many image operands has explicit requirements on image instruction using it. And we don't even support those instructions. Hence I add an llvm_unreachable for all non-supported image operands in the updated patch.

[Weiwei-2021]

fix some problems.

[Weiwei-2021]

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

348

Thanks for pointing it out. :)

[Weiwei-2021]

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

351

Ping @antiagainst to check whether this versino is ok.

[antiagainst]

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

345

This else isn't necesary.

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

SPV_I32EnumAttr<"ImageFormat", "valid SPIR-V ImageFormat", [
       SPV_IF_R16, SPV_IF_R8, SPV_IF_Rgba16Snorm, SPV_IF_Rg16Snorm, SPV_IF_Rg8Snorm,
       SPV_IF_R16Snorm, SPV_IF_R8Snorm, SPV_IF_Rgba32i, SPV_IF_Rgba16i, SPV_IF_Rgba8i,
       SPV_IF_R32i, SPV_IF_Rg32i, SPV_IF_Rg16i, SPV_IF_Rg8i, SPV_IF_R16i, SPV_IF_R8i,
       SPV_IF_Rgba32ui, SPV_IF_Rgba16ui, SPV_IF_Rgba8ui, SPV_IF_R32ui,
       SPV_IF_Rgb10a2ui, SPV_IF_Rg32ui, SPV_IF_Rg16ui, SPV_IF_Rg8ui, SPV_IF_R16ui,
       SPV_IF_R8ui
     ]>;
 
+def SPV_IO_None               : BitEnumAttrCase<"None", 0x0000>;
+def SPV_IO_Bias               : BitEnumAttrCase<"Bias", 0x0001> {
+  list<Availability> availability = [
+    Capability<[SPV_C_Shader]>
+  ];
+}
+def SPV_IO_Lod                : BitEnumAttrCase<"Lod", 0x0002>;
+def SPV_IO_Grad               : BitEnumAttrCase<"Grad", 0x0004>;
+def SPV_IO_ConstOffset        : BitEnumAttrCase<"ConstOffset", 0x0008>;
+def SPV_IO_Offset             : BitEnumAttrCase<"Offset", 0x0010> {
+  list<Availability> availability = [
+    Capability<[SPV_C_ImageGatherExtended]>
+  ];
+}
+def SPV_IO_ConstOffsets       : BitEnumAttrCase<"ConstOffsets", 0x0020> {
+  list<Availability> availability = [
+    Capability<[SPV_C_ImageGatherExtended]>
+  ];
+}
+def SPV_IO_Sample             : BitEnumAttrCase<"Sample", 0x0040>;
+def SPV_IO_MinLod             : BitEnumAttrCase<"MinLod", 0x0080> {
+  list<Availability> availability = [
+    Capability<[SPV_C_MinLod]>
+  ];
+}
+def SPV_IO_MakeTexelAvailable : BitEnumAttrCase<"MakeTexelAvailable", 0x0100> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Capability<[SPV_C_VulkanMemoryModel]>
+  ];
+}
+def SPV_IO_MakeTexelVisible   : BitEnumAttrCase<"MakeTexelVisible", 0x0200> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Capability<[SPV_C_VulkanMemoryModel]>
+  ];
+}
+def SPV_IO_NonPrivateTexel    : BitEnumAttrCase<"NonPrivateTexel", 0x0400> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Capability<[SPV_C_VulkanMemoryModel]>
+  ];
+}
+def SPV_IO_VolatileTexel      : BitEnumAttrCase<"VolatileTexel", 0x0800> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_5>,
+    Capability<[SPV_C_VulkanMemoryModel]>
+  ];
+}
+def SPV_IO_SignExtend         : BitEnumAttrCase<"SignExtend", 0x1000> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_4>,
+  ];
+}
+def SPV_IO_ZeroExtend         : BitEnumAttrCase<"ZeroExtend", 0x2000> {
+  list<Availability> availability = [
+    MinVersion<SPV_V_1_4>,
+  ];
+}
+
+def SPV_ImageOperandAttr :
+    SPV_BitEnumAttr<"ImageOperands", "valid SPIR-V ImageOperands", [
+      SPV_IO_None, SPV_IO_Bias, SPV_IO_Lod, SPV_IO_Grad, SPV_IO_ConstOffset,
+      SPV_IO_Offset, SPV_IO_ConstOffsets, SPV_IO_Sample, SPV_IO_MinLod,
+      SPV_IO_MakeTexelAvailable, SPV_IO_MakeTexelVisible, SPV_IO_NonPrivateTexel,
+      SPV_IO_VolatileTexel, SPV_IO_SignExtend, SPV_IO_ZeroExtend
+    ]>;
+
 def SPV_LT_Export : I32EnumAttrCase<"Export", 0> {
   list<Availability> availability = [
     Capability<[SPV_C_Linkage]>
   ];
 }
 def SPV_LT_Import : I32EnumAttrCase<"Import", 1> {
   list<Availability> availability = [
     Capability<[SPV_C_Linkage]>

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

let description = [{
     Sampled Image.
 
     Dref is the depth-comparison reference value. It must be a 32-bit
     floating-point type scalar.
 
     Image Operands encodes what operands follow, as per Image Operands.
 
     <!-- End of AutoGen section -->
-
+    ```
+    image-operands ::= `"None"` | `"Bias"` | `"Lod"` | `"Grad"`
+                      | `"ConstOffset"` | `"Offser"` | `"ConstOffsets"`
+                      | `"Sample"` | `"MinLod"` | `"MakeTexelAvailable"`
+                      | `"MakeTexelVisible"` | `"NonPrivateTexel"`
+                      | `"VolatileTexel"` | `"SignExtend"` | `"ZeroExtend"`
     #### Example:
+    ```
 
     ```mlir
     %0 = spv.ImageDrefGather %1 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %2 : vector<4xf32>, %3 : f32 -> vector<4xi32>
+    %0 = spv.ImageDrefGather %1 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %2 : vector<4xf32>, %3 : f32 ["Bias|Lod"] %4, %5 : f32, f32 -> vector<4xi32>
     ```
   }];
 
   let availability = [
     MinVersion<SPV_V_1_0>,
     MaxVersion<SPV_V_1_5>,
     Extension<[]>,
     Capability<[SPV_C_Shader]>
   ];
 
   let arguments = (ins
     SPV_AnySampledImage:$sampledimage,
     SPV_ScalarOrVectorOf<SPV_Float>:$coordinate,
-    SPV_Float:$dref
+    SPV_Float:$dref,
+    OptionalAttr<SPV_ImageOperandAttr>:$imageoperands,
+    Variadic<SPV_Type>:$operand_arguments
   );
 
   let results = (outs
     SPV_Vector:$result
   );
 
-  let assemblyFormat = "attr-dict $sampledimage `:` type($sampledimage) `,` $coordinate `:` type($coordinate) `,` $dref `:` type($dref) `->` type($result)";
+  let assemblyFormat = [{$sampledimage `:` type($sampledimage) `,`
+                         $coordinate `:` type($coordinate) `,` $dref `:` type($dref)
+                         custom<ImageOperands>($imageoperands)
+                         ( `(` $operand_arguments^ `:` type($operand_arguments) `)`)?
+                         attr-dict
+                         `->` type($result)}];
 
   let verifier = [{ return ::verify(*this); }];
 }
 
 // -----
 
 def SPV_ImageQuerySizeOp : SPV_Op<"ImageQuerySize", [NoSideEffect]> {
   let summary = "Query the dimensions of Image, with no level of detail.";

mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp

if (spirv::bitEnumContains(*memAccess, spirv::MemoryAccess::Aligned)) {
         printer << ", " << alignment;
       }
     }
     printer << "]";
   }
   elidedAttrs.push_back(spirv::attributeName<spirv::StorageClass>());
 }
 
+static ParseResult parseImageOperands(OpAsmParser &parser,
+                                      spirv::ImageOperandsAttr &attr) {
+  // Expect image operands
+  if (parser.parseOptionalLSquare())
+    return success();
+
+  spirv::ImageOperands imageOperands;
+  if (parseEnumStrAttr(imageOperands, parser))
+    return failure();
+
+  attr = spirv::ImageOperandsAttr::get(parser.getBuilder().getContext(),
+                                       imageOperands);
+
+  return parser.parseRSquare();
+}
+
+static void printImageOperands(OpAsmPrinter &printer, Operation *imageOp,
+                               spirv::ImageOperandsAttr attr) {
+  if (attr) {
+    auto strImageOperands = stringifyImageOperands(attr.getValue());
+    printer << "[\"" << strImageOperands << "\"]";
+  }
+
+  return;
+}
+
+template <typename Op>
+static LogicalResult verifyImageOperands(Op imageOp,
+                                         spirv::ImageOperandsAttr attr,
+                                         Operation::operand_range operands) {
+  if (!attr) {
+    if (operands.empty())
+      return success();
+    else
+      return imageOp.emitError("the Image Operands should encode what operands "
+                               "follow, as per Image Operands");
+  }
+
+  // TODO: Add the validation rules for the following Image Operands.
+  spirv::ImageOperands noSupportOperands =
+      spirv::ImageOperands::Bias | spirv::ImageOperands::Lod |
+      spirv::ImageOperands::Grad | spirv::ImageOperands::ConstOffset |
+      spirv::ImageOperands::Offset | spirv::ImageOperands::ConstOffsets |
+      spirv::ImageOperands::Sample | spirv::ImageOperands::MinLod |
+      spirv::ImageOperands::MakeTexelAvailable |
+      spirv::ImageOperands::MakeTexelVisible |
+      spirv::ImageOperands::SignExtend | spirv::ImageOperands::ZeroExtend;
+
+  if (spirv::bitEnumContains(attr.getValue(), noSupportOperands))
+    llvm_unreachable("unimplemented operands of Image Operands");
+
+  return success();
+}
+
 static LogicalResult verifyCastOp(Operation *op,
                                   bool requireSameBitWidth = true,
                                   bool skipBitWidthCheck = false) {
   // Some CastOps have no limit on bit widths for result and operand type.
   if (skipBitWidthCheck)
     return success();
 
   Type operandType = op->getOperand(0).getType();
static LogicalResult verify(spirv::GLSLLdexpOp ldexpOp) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // spv.ImageDrefGather
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(spirv::ImageDrefGatherOp imageDrefGatherOp) {
-  // TODO: Support optional operands.
   VectorType resultType =
       imageDrefGatherOp.result().getType().cast<VectorType>();
   auto sampledImageType = imageDrefGatherOp.sampledimage()
                               .getType()
                               .cast<spirv::SampledImageType>();
   auto imageType = sampledImageType.getImageType().cast<spirv::ImageType>();
 
   if (resultType.getNumElements() != 4)
if (imageDim != spirv::Dim::Dim2D && imageDim != spirv::Dim::Cube &&
     return imageDrefGatherOp.emitOpError(
         "the Dim operand of the underlying image type must be 2D, Cube, or "
         "Rect");
 
   if (imageMS != spirv::ImageSamplingInfo::SingleSampled)
     return imageDrefGatherOp.emitOpError(
         "the MS operand of the underlying image type must be 0");
 
-  return success();
+  spirv::ImageOperandsAttr attr = imageDrefGatherOp.imageoperandsAttr();
+  auto operandArguments = imageDrefGatherOp.operand_arguments();
+
+  return verifyImageOperands(imageDrefGatherOp, attr, operandArguments);
 }
 
 //===----------------------------------------------------------------------===//
 // spv.ImageQuerySize
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(spirv::ImageQuerySizeOp imageQuerySizeOp) {
   spirv::ImageType imageType =

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

 // RUN: mlir-opt -split-input-file -verify-diagnostics %s | FileCheck %s
 
-
 //===----------------------------------------------------------------------===//
 // spv.ImageDrefGather
 //===----------------------------------------------------------------------===//
 
 func @image_dref_gather(%arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
   // CHECK: spv.ImageDrefGather {{.*}} : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, {{.*}} : vector<4xf32>, {{.*}} : f32 -> vector<4xi32>
   %0 = spv.ImageDrefGather %arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32 -> vector<4xi32>
   spv.Return
 }
 
 // -----
 
+func @image_dref_gather_with_single_imageoperands(%arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
+  // CHECK: spv.ImageDrefGather {{.*}} ["NonPrivateTexel"] -> vector<4xi32>
+  %0 = spv.ImageDrefGather %arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32 ["NonPrivateTexel"] -> vector<4xi32>
+  spv.Return
+}
+
+// -----
+
+func @image_dref_gather_with_mismatch_imageoperands(%arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
+  // expected-error @+1 {{the Image Operands should encode what operands follow, as per Image Operands}}
+  %0 = spv.ImageDrefGather %arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32 (%arg2, %arg2 : f32, f32) -> vector<4xi32>
+  spv.Return
+}
+
+// -----
+
 func @image_dref_gather_error_result_type(%arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
   // expected-error @+1 {{result type must be a vector of four components}}
   %0 = spv.ImageDrefGather %arg0 : !spv.sampled_image<!spv.image<i32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32 -> vector<3xi32>
   spv.Return
 }
 
 // -----