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.

Sign up for GitHub

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
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVBase.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4362,6 +4362,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>;
Expand Down Expand Up @@ -4558,10 +4559,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,
Expand Down
52 changes: 52 additions & 0 deletions mlir/include/mlir/Dialect/SPIRV/IR/SPIRVImageOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>
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.

```
}];

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>:$image_operands,
Variadic<SPIRV_Type>:$operand_arguments
);

let results = (outs);

let assemblyFormat = [{$image `:` type($image) `,`
$coordinate `:` type($coordinate) `,`
$texel `:` type($texel)
custom<ImageOperands>($image_operands)
( `(` $operand_arguments^ `:` type($operand_arguments) `)`)?
attr-dict}];
}

// -----

def SPIRV_ImageOp : SPIRV_Op<"Image",
[Pure,
TypesMatchWith<"type of 'result' matches image type of 'sampledimage'",
Expand Down
39 changes: 39 additions & 0 deletions mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2042,6 +2042,45 @@ LogicalResult spirv::ImageDrefGatherOp::verify() {
return verifyImageOperands(*this, attr, operandArguments);
}

//===----------------------------------------------------------------------===//
// spirv.ImageWriteOp
//===----------------------------------------------------------------------===//

LogicalResult spirv::ImageWriteOp::verify() {
ImageType imageType = cast<ImageType>(getImage().getType());
Type sampledType = imageType.getElementType();
ImageSamplerUseInfo samplerInfo = imageType.getSamplerUseInfo();

if (!llvm::is_contained({spirv::ImageSamplerUseInfo::SamplerUnknown,
spirv::ImageSamplerUseInfo::NoSampler},
samplerInfo)) {
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 = getElementTypeOrSelf(getTexel());
if (!isa<NoneType>(sampledType) && texelType != sampledType) {
return emitOpError(
"the texel component type must match the image sampled type");
}

// TODO: Ideally it should be somewhere verified that "The Image Format must
// not be Unknown, unless the StorageImageWriteWithoutFormat Capability was
// declared." This function however may not be the suitable place for such
// verification.

return verifyImageOperands(*this, getImageOperandsAttr(),
getOperandArguments());
}

//===----------------------------------------------------------------------===//
// spirv.ShiftLeftLogicalOp
//===----------------------------------------------------------------------===//
Expand Down
42 changes: 42 additions & 0 deletions mlir/test/Dialect/SPIRV/IR/image-ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
27 changes: 25 additions & 2 deletions mlir/test/Target/SPIRV/image-ops.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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>>
Expand All @@ -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
}