[mlir][spirv] Verify SampledImageType Dim (#159397)

This patches adds check for: "It [ImageType] must not have a Dim of
SubpassData. Additionally, starting with version 1.6, it must not have a
Dim of Buffer." as defined in "3.3.6. Type-Declaration Instructions" of
SPIR-V spec.

Author: IgWod-IMG

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

@@ -251,13 +251,21 @@ static Type parseAndVerifySampledImageType(SPIRVDialect const &dialect,
   if (parser.parseType(type))
     return Type();
 
-  if (!llvm::isa<ImageType>(type)) {
+  auto imageType = dyn_cast<ImageType>(type);
+  if (!imageType) {
     parser.emitError(typeLoc,
                      "sampled image must be composed using image type, got ")
         << type;
     return Type();
   }
 
+  if (llvm::is_contained({Dim::SubpassData, Dim::Buffer}, imageType.getDim())) {
+    parser.emitError(
+        typeLoc, "sampled image Dim must not be SubpassData or Buffer, got ")
+        << stringifyDim(imageType.getDim());
+    return Type();
+  }
+
   return type;
 }
 

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

@@ -805,9 +805,16 @@ Type SampledImageType::getImageType() const { return getImpl()->imageType; }
 LogicalResult
 SampledImageType::verifyInvariants(function_ref<InFlightDiagnostic()> emitError,
                                    Type imageType) {
-  if (!llvm::isa<ImageType>(imageType))
+  auto image = dyn_cast<ImageType>(imageType);
+  if (!image)
     return emitError() << "expected image type";
 
+  // As per SPIR-V spec: "It [ImageType] must not have a Dim of SubpassData.
+  // Additionally, starting with version 1.6, it must not have a Dim of Buffer.
+  // ("3.3.6. Type-Declaration Instructions")
+  if (llvm::is_contained({Dim::SubpassData, Dim::Buffer}, image.getDim()))
+    return emitError() << "Dim must not be SubpassData or Buffer";
+
   return success();
 }
 

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

@@ -192,6 +192,9 @@ func.func @image_write_texel_type_mismatch(%arg0 : !spirv.image<f32, Dim2D, NoDe
 // spirv.ImageSampleExplicitLod
 //===----------------------------------------------------------------------===//
 
+// No need to have a negative test for Dim being Buffer as this is already handled
+// by SampledImageType logic.
+
 func.func @sample_explicit_lod(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>, %arg2 : f32) -> () {
   // CHECK: {{%.*}} = spirv.ImageSampleExplicitLod {{%.*}}, {{%.*}} ["Lod"], {{%.*}} : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32>, f32 -> vector<4xf32>
   %0 = spirv.ImageSampleExplicitLod %arg0, %arg1 ["Lod"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32>, f32 -> vector<4xf32>
@@ -200,14 +203,6 @@ func.func @sample_explicit_lod(%arg0 : !spirv.sampled_image<!spirv.image<f32, Di
 
 // -----
 
-func.func @sample_explicit_lod_buffer_dim(%arg0 : !spirv.sampled_image<!spirv.image<f32, Buffer, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>, %arg2 : f32) -> () {
-  // expected-error @+1 {{the Dim operand of the underlying image must not be Buffer}}
-  %0 = spirv.ImageSampleExplicitLod %arg0, %arg1 ["Lod"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Buffer, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32>, f32 -> vector<4xf32>
-  spirv.Return
-}
-
-// -----
-
 func.func @sample_explicit_lod_multi_sampled(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>, %arg2 : f32) -> () {
   // expected-error @+1 {{the MS operand of the underlying image type must be SingleSampled}}
   %0 = spirv.ImageSampleExplicitLod %arg0, %arg1 ["Lod"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, vector<2xf32>, f32 -> vector<4xf32>
@@ -236,6 +231,9 @@ func.func @sample_explicit_lod_no_lod(%arg0 : !spirv.sampled_image<!spirv.image<
 // spirv.ImageSampleImplicitLod
 //===----------------------------------------------------------------------===//
 
+// No need to have a negative test for Dim being Buffer as this is already handled
+// by SampledImageType logic.
+
 func.func @sample_implicit_lod(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>) -> () {
   // CHECK: {{%.*}} = spirv.ImageSampleImplicitLod {{%.*}}, {{%.*}} : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32> -> vector<4xf32>
   %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32> -> vector<4xf32>
@@ -244,14 +242,6 @@ func.func @sample_implicit_lod(%arg0 : !spirv.sampled_image<!spirv.image<f32, Di
 
 // -----
 
-func.func @sample_implicit_lod_buffer(%arg0 : !spirv.sampled_image<!spirv.image<f32, Buffer, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>) -> () {
-  // expected-error @+1 {{the Dim operand of the underlying image must not be Buffer}}
-  %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 : !spirv.sampled_image<!spirv.image<f32, Buffer, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32> -> vector<4xf32>
-  spirv.Return
-}
-
-// -----
-
 func.func @sample_implicit_lod_multi_sampled(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>) -> () {
   // expected-error @+1 {{the MS operand of the underlying image type must be SingleSampled}}
   %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 : !spirv.sampled_image<!spirv.image<f32, Dim2D, NoDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, vector<2xf32> -> vector<4xf32>
@@ -272,6 +262,9 @@ func.func @sample_implicit_lod_wrong_result(%arg0 : !spirv.sampled_image<!spirv.
 // spirv.ImageSampleProjDrefImplicitLod
 //===----------------------------------------------------------------------===//
 
+// No need to have a negative test for Dim being Buffer as this is already handled
+// by SampledImageType logic.
+
 func.func @sample_implicit_proj_dref(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
   // CHECK: {{%.*}} = spirv.ImageSampleProjDrefImplicitLod {{%.*}}, {{%.*}}, {{%.*}} : !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<4xf32>, f32 -> f32
   %0 = spirv.ImageSampleProjDrefImplicitLod %arg0, %arg1, %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<4xf32>, f32 -> f32
@@ -280,14 +273,6 @@ func.func @sample_implicit_proj_dref(%arg0 : !spirv.sampled_image<!spirv.image<f
 
 // -----
 
-func.func @sample_implicit_proj_dref_buffer_dim(%arg0 : !spirv.sampled_image<!spirv.image<f32, Buffer, IsDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
-  // expected-error @+1 {{the Dim operand of the underlying image must not be Buffer}}
-  %0 = spirv.ImageSampleProjDrefImplicitLod %arg0, %arg1, %arg2 : !spirv.sampled_image<!spirv.image<f32, Buffer, IsDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<4xf32>, f32 -> f32
-  spirv.Return
-}
-
-// -----
-
 func.func @sample_implicit_proj_dref_multi_sampled(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, %arg1 : vector<4xf32>, %arg2 : f32) -> () {
   // expected-error @+1 {{the MS operand of the underlying image type must be SingleSampled}}
   %0 = spirv.ImageSampleProjDrefImplicitLod %arg0, %arg1, %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim2D, IsDepth, NonArrayed, MultiSampled, NeedSampler, Rgba8>>, vector<4xf32>, f32 -> f32

mlir/test/Dialect/SPIRV/IR/types.mlir

@@ -238,6 +238,16 @@ func.func private @samped_image_type_invaid_type(!spirv.sampled_image<f32>) -> (
 
 // -----
 
+// expected-error @+1 {{sampled image Dim must not be SubpassData or Buffer, got Buffer}}
+func.func private @sampled_image_type(!spirv.sampled_image<!spirv.image<f32, Buffer, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>) -> ()
+
+// -----
+
+// expected-error @+1 {{sampled image Dim must not be SubpassData or Buffer, got SubpassData}}
+func.func private @sampled_image_type(!spirv.sampled_image<!spirv.image<f32, SubpassData, NoDepth, NonArrayed, SingleSampled, NoSampler, Unknown>>) -> ()
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // StructType
 //===----------------------------------------------------------------------===//