Skip to content

[mlir][spirv] Add definition for ImageWriteOp #124124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 6, 2025

Conversation

IgWod-IMG
Copy link
Contributor

This Pull Request adds OpImageWrite as defined in section 3.52.10. (Image Instructions). The tests in mlir/test/Target/SPIRV/image-ops.mlir are also updated (and extended with the new op), so they now pass validation with spirv-val after serialization into SPIR-V. The test was missing ImageQuery capability and entry points. For entry points dummy main functions were added.

@llvmbot
Copy link
Member

llvmbot commented Jan 23, 2025

@llvm/pr-subscribers-mlir-spirv

@llvm/pr-subscribers-mlir

Author: Igor Wodiany (IgWod-IMG)

Changes

This Pull Request adds OpImageWrite as defined in section 3.52.10. (Image Instructions). The tests in mlir/test/Target/SPIRV/image-ops.mlir are also updated (and extended with the new op), so they now pass validation with spirv-val after serialization into SPIR-V. The test was missing ImageQuery capability and entry points. For entry points dummy main functions were added.


Full diff: https://github.com/llvm/llvm-project/pull/124124.diff

5 Files Affected:

  • (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td (+5-4)
  • (modified) mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td (+52)
  • (modified) mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp (+59)
  • (modified) mlir/test/Dialect/SPIRV/IR/image-ops.mlir (+42)
  • (modified) mlir/test/Target/SPIRV/image-ops.mlir (+25-2)
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
index c84677d26a8b69..ebe92b4964baf3 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
@@ -4354,6 +4354,7 @@ def SPIRV_OC_OpCompositeExtract             : I32EnumAttrCase<"OpCompositeExtrac
 def SPIRV_OC_OpCompositeInsert              : I32EnumAttrCase<"OpCompositeInsert", 82>;
 def SPIRV_OC_OpTranspose                    : I32EnumAttrCase<"OpTranspose", 84>;
 def SPIRV_OC_OpImageDrefGather              : I32EnumAttrCase<"OpImageDrefGather", 97>;
+def SPIRV_OC_OpImageWrite                   : I32EnumAttrCase<"OpImageWrite", 99>;
 def SPIRV_OC_OpImage                        : I32EnumAttrCase<"OpImage", 100>;
 def SPIRV_OC_OpImageQuerySize               : I32EnumAttrCase<"OpImageQuerySize", 104>;
 def SPIRV_OC_OpConvertFToU                  : I32EnumAttrCase<"OpConvertFToU", 109>;
@@ -4549,10 +4550,10 @@ def SPIRV_OpcodeAttr :
       SPIRV_OC_OpVectorInsertDynamic, SPIRV_OC_OpVectorShuffle,
       SPIRV_OC_OpCompositeConstruct, SPIRV_OC_OpCompositeExtract,
       SPIRV_OC_OpCompositeInsert, SPIRV_OC_OpTranspose, SPIRV_OC_OpImageDrefGather,
-      SPIRV_OC_OpImage, SPIRV_OC_OpImageQuerySize, SPIRV_OC_OpConvertFToU,
-      SPIRV_OC_OpConvertFToS, SPIRV_OC_OpConvertSToF, SPIRV_OC_OpConvertUToF,
-      SPIRV_OC_OpUConvert, SPIRV_OC_OpSConvert, SPIRV_OC_OpFConvert,
-      SPIRV_OC_OpConvertPtrToU, SPIRV_OC_OpConvertUToPtr,
+      SPIRV_OC_OpImageWrite, SPIRV_OC_OpImage, SPIRV_OC_OpImageQuerySize,
+      SPIRV_OC_OpConvertFToU, SPIRV_OC_OpConvertFToS, SPIRV_OC_OpConvertSToF,
+      SPIRV_OC_OpConvertUToF, SPIRV_OC_OpUConvert, SPIRV_OC_OpSConvert,
+      SPIRV_OC_OpFConvert, SPIRV_OC_OpConvertPtrToU, SPIRV_OC_OpConvertUToPtr,
       SPIRV_OC_OpPtrCastToGeneric, SPIRV_OC_OpGenericCastToPtr,
       SPIRV_OC_OpGenericCastToPtrExplicit, SPIRV_OC_OpBitcast, SPIRV_OC_OpSNegate,
       SPIRV_OC_OpFNegate, SPIRV_OC_OpIAdd, SPIRV_OC_OpFAdd, SPIRV_OC_OpISub,
diff --git a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td
index 755d26de9d4726..e07eb8f714e1bc 100644
--- a/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td
+++ b/mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td
@@ -135,6 +135,58 @@ def SPIRV_ImageQuerySizeOp : SPIRV_Op<"ImageQuerySize", [Pure]> {
 
 // -----
 
+def SPIRV_ImageWriteOp : SPIRV_Op<"ImageWrite", []> {
+  let summary = "Write a texel to an image without a sampler.";
+
+  let description = [{
+    Image must be an object whose type is OpTypeImage with a Sampled operand
+    of 0 or 2. If the Arrayed operand is 1, then additional capabilities may
+    be required; e.g., ImageCubeArray, or ImageMSArray. Its Dim operand
+    must not be SubpassData.
+
+    Coordinate must be a scalar or vector of floating-point type or integer
+    type. It contains non-normalized texel coordinates (u[, v] ... [, array
+    layer]) as needed by the definition of Image. See the client API
+    specification for handling of coordinates outside the image.
+
+    Texel is the data to write. It must be a scalar or vector with component
+    type the same as Sampled Type of the OpTypeImage (unless that Sampled
+    Type is OpTypeVoid).
+
+    The Image Format must not be Unknown, unless the
+    StorageImageWriteWithoutFormat Capability was declared.
+
+    Image Operands encodes what operands follow, as per Image Operands.
+
+    <!-- End of AutoGen section -->
+
+    #### Example:
+
+    ```mlir
+    spirv.ImageWrite %0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %1 : vector<2xsi32>, %2 : vector<4xf32>
+    ```
+  }];
+
+  let arguments = (ins
+    SPIRV_AnyImage:$image,
+    AnyTypeOf<[SPIRV_ScalarOrVectorOf<SPIRV_Float>, SPIRV_ScalarOrVectorOf<SPIRV_Integer>]>:$coordinate,
+    AnyTypeOf<[SPIRV_ScalarOrVectorOf<SPIRV_Float>, SPIRV_ScalarOrVectorOf<SPIRV_Integer>]>:$texel,
+    OptionalAttr<SPIRV_ImageOperandsAttr>:$imageoperands,
+    Variadic<SPIRV_Type>:$operand_arguments
+  );
+
+  let results = (outs);
+
+  let assemblyFormat = [{$image `:` type($image) `,`
+                         $coordinate `:` type($coordinate) `,`
+                         $texel `:` type($texel)
+                         custom<ImageOperands>($imageoperands)
+                         ( `(` $operand_arguments^ `:` type($operand_arguments) `)`)?
+                         attr-dict}];
+}
+
+// -----
+
 def SPIRV_ImageOp : SPIRV_Op<"Image",
     [Pure,
      TypesMatchWith<"type of 'result' matches image type of 'sampledimage'",
diff --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index 040bf6a34cea78..01630df2b15b7e 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
@@ -2020,6 +2020,65 @@ LogicalResult spirv::ImageDrefGatherOp::verify() {
   return verifyImageOperands(*this, attr, operandArguments);
 }
 
+//===----------------------------------------------------------------------===//
+// spirv.ImageWriteOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult spirv::ImageWriteOp::verify() {
+  ImageType imageType = llvm::cast<ImageType>(getImage().getType());
+  Type sampledType = imageType.getElementType();
+  ImageSamplerUseInfo samplerInfo = imageType.getSamplerUseInfo();
+
+  if (samplerInfo != spirv::ImageSamplerUseInfo::SamplerUnknown &&
+      samplerInfo != spirv::ImageSamplerUseInfo::NoSampler) {
+    return emitOpError(
+        "the sampled operand of the underlying image must be 0 or 2");
+  }
+
+  // TODO: Do we need check for: "If the Arrayed operand is 1, then additional
+  // capabilities may be required; e.g., ImageCubeArray, or ImageMSArray.".
+
+  if (imageType.getDim() == spirv::Dim::SubpassData) {
+    return emitOpError(
+        "the Dim operand of the underlying image must not be SubpassData");
+  }
+
+  Type texelType = getTexel().getType();
+  if (llvm::isa<VectorType>(texelType)) {
+    texelType = llvm::cast<VectorType>(texelType).getElementType();
+  }
+
+  if (!llvm::isa<NoneType>(sampledType) && texelType != sampledType) {
+    return emitOpError(
+        "the texel component type must match the image sampled type");
+  }
+
+  if (imageType.getImageFormat() == spirv::ImageFormat::Unknown) {
+    auto module = this->getOperation()->getParentOfType<spirv::ModuleOp>();
+    if (module) {
+      if (std::optional<spirv::VerCapExtAttr> triple = module.getVceTriple()) {
+        auto capabilities = (*triple).getCapabilities();
+        auto requiredCapability =
+            std::find(capabilities.begin(), capabilities.end(),
+                      spirv::Capability::StorageImageWriteWithoutFormat);
+        // It is allowed for the format to be unknown if
+        // StorageImageWriteWithoutFormat capability is present.
+        if (requiredCapability == capabilities.end())
+          return emitOpError("the image format operand of the underlying image "
+                             "must not be unknown");
+      } else {
+        return emitOpError("the image format operand of the underlying image "
+                           "must not be unknown");
+      }
+    }
+  }
+
+  spirv::ImageOperandsAttr attr = getImageoperandsAttr();
+  auto operandArguments = getOperandArguments();
+
+  return verifyImageOperands(*this, attr, operandArguments);
+}
+
 //===----------------------------------------------------------------------===//
 // spirv.ShiftLeftLogicalOp
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SPIRV/IR/image-ops.mlir b/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
index ab674ee0809ae6..1161f85563ae61 100644
--- a/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
@@ -115,3 +115,45 @@ func.func @image_query_size_error_result2(%arg0 : !spirv.image<f32, Buffer, NoDe
 }
 
 // -----
+
+//===----------------------------------------------------------------------===//
+// spirv.ImageWrite
+//===----------------------------------------------------------------------===//
+
+func.func @image_write(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>) -> () {
+  // CHECK:  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+  spirv.Return
+}
+
+// -----
+
+func.func @image_write_scalar_texel(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : f32) -> () {
+  // CHECK:  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : f32
+  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : f32
+  spirv.Return
+}
+
+// -----
+
+func.func @image_write_need_sampler(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>) -> () {
+  // expected-error @+1 {{the sampled operand of the underlying image must be 0 or 2}}
+  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+  spirv.Return
+}
+
+// -----
+
+func.func @image_write_subpass_data(%arg0 : !spirv.image<f32, SubpassData, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>) -> () {
+  // expected-error @+1 {{the Dim operand of the underlying image must not be SubpassData}}
+  spirv.ImageWrite %arg0 : !spirv.image<f32, SubpassData, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+  spirv.Return
+}
+
+// -----
+
+func.func @image_write_texel_type_mismatch(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xi32>) -> () {
+  // expected-error @+1 {{the texel component type must match the image sampled type}}
+  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %arg1 : vector<2xsi32>, %arg2 : vector<4xi32>
+  spirv.Return
+}
diff --git a/mlir/test/Target/SPIRV/image-ops.mlir b/mlir/test/Target/SPIRV/image-ops.mlir
index 92429fc8023d2a..6b52a84ba82f7e 100644
--- a/mlir/test/Target/SPIRV/image-ops.mlir
+++ b/mlir/test/Target/SPIRV/image-ops.mlir
@@ -1,6 +1,6 @@
-// RUN: mlir-translate -no-implicit-module -test-spirv-roundtrip %s | FileCheck %s
+// RUN: mlir-translate --no-implicit-module --split-input-file --test-spirv-roundtrip %s | FileCheck %s
 
-spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, ImageQuery], []> {
   spirv.func @image(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>, %arg1 : vector<4xf32>, %arg2 : f32) "None" {
     // CHECK: {{%.*}} = spirv.Image {{%.*}} : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
     %0 = spirv.Image %arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Unknown>>
@@ -13,4 +13,27 @@ spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader], []> {
     %0 = spirv.ImageQuerySize %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown> -> vector<2xi32>
     spirv.Return
   }
+  spirv.func @image_write(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>) "None" {
+    // CHECK:  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+    spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba8>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+    spirv.Return
+  }
+  spirv.func @main() "None" {
+    spirv.Return
+  }
+  spirv.EntryPoint "GLCompute" @main
+}
+
+// -----
+
+spirv.module Logical GLSL450 requires #spirv.vce<v1.0, [Shader, StorageImageWriteWithoutFormat], []> {
+  spirv.func @image_write(%arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>) "None" {
+    // CHECK:  spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+    spirv.ImageWrite %arg0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>, %arg1 : vector<2xsi32>, %arg2 : vector<4xf32>
+    spirv.Return
+  }
+  spirv.func @main() "None" {
+    spirv.Return
+  }
+  spirv.EntryPoint "GLCompute" @main
 }

@IgWod-IMG
Copy link
Contributor Author

IgWod-IMG commented Jan 23, 2025

Just to add to the description.

For the verification of the op, I am not sure whether I should check the module capabilities, as I currently do, or skip that as it's too broad for the instruction level verification. Also, I'm not sure whether I query the module capabilities in the correct way - it works with the target test.

As always, I am happy to address any feedback. And just to add I do not have committer rights, so the change has to be committed on my behalf - I will look into requesting the right.

@IgWod-IMG
Copy link
Contributor Author

Ping

@IgWod-IMG
Copy link
Contributor Author

I have pushed the updated patch, let me know if there is anything else to fix, and if not please merge it - I am waiting for commit access to be granted.

Copy link
Member

@kuhar kuhar left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good overall, just a few nits

#### Example:

```mlir
spirv.ImageWrite %0 : !spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NoSampler, Rgba16>, %1 : vector<2xsi32>, %2 : vector<4xf32>
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This syntax is a bit unusual; we typically do spirv.Name %0, %1, %2 : type0, type1, type2
I can see that other image ops use the style you adopted as well, so let's keep it as-is in this PR, but it would be nice to clean this up separately.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree, I noticed the inconsistency, but I decided to match existing image ops. There are few more image ops that are on my todo list and I actually started cleaning things up. For example, I am moving all image ops verification from mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp into a separate file. So, I'll clean up the syntax as well as a part of that.

@IgWod-IMG
Copy link
Contributor Author

I have just pushed an updated patch. Please merge if happy. Thanks!

@kuhar kuhar merged commit 8609e27 into llvm:main Feb 6, 2025
6 of 7 checks passed
@IgWod-IMG IgWod-IMG deleted the img_image_write branch February 6, 2025 14:27
Icohedron pushed a commit to Icohedron/llvm-project that referenced this pull request Feb 11, 2025
This Pull Request adds OpImageWrite as defined in section 3.52.10.
(Image Instructions). The tests in
`mlir/test/Target/SPIRV/image-ops.mlir` are also updated (and extended
with the new op), so they now pass validation with `spirv-val` after
serialization into SPIR-V. The test was missing `ImageQuery` capability
and entry points. For entry points dummy `main` functions were added.
kuhar pushed a commit that referenced this pull request Feb 25, 2025
This patch makes multiple changes to images ops:

1) The assembly format is unified with the rest of the dialect to use
`%0 = spirv.op %1, %2, %3 : f32, f32, f32` rather than having each type
directly attached to each argument.
2) The verification is moved from `SPIRVOps.cpp` to a new file so the
ops can be easier maintained.
3) Majority of C++ verification is removed and moved into ODS.
Verification of `ImageQuerySizeOp` is left in C++ due to the complexity
of rules.
4) `spirv::bitEnumContainsAll` is replaced by
`spirv::bitEnumContainsAny` in `verifyImageOperands`. In this context
`...Any` seems to be the correct function, as we want to check whether
unsupported operand is being used - in opposite to checking if all
unsupported operands are being used.
5) Simplify target tests by removing entry points and adding `Linkage`
capability to the modules.

This change is made in preparation for adding more Image ops. Change to
the assembly format was previously mentioned in #124124.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants