[Mlir-commits] [mlir] [mlir][spirv] Add verification for Bias operand (PR #134231)

[Igor Wodiany]

https://github.com/IgWod-IMG created https://github.com/llvm/llvm-project/pull/134231

None

>From 412bb5969e824aec631a92d9e1bf40769fb4fb94 Mon Sep 17 00:00:00 2001
From: Igor Wodiany <igor.wodiany at imgtec.com>
Date: Tue, 1 Apr 2025 14:46:42 +0100
Subject: [PATCH] [mlir][spirv] Add verification for Bias operand

---
 mlir/lib/Dialect/SPIRV/IR/ImageOps.cpp    | 50 ++++++++++++++++++-----
 mlir/test/Dialect/SPIRV/IR/image-ops.mlir | 32 ++++++++++++++-
 2 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/mlir/lib/Dialect/SPIRV/IR/ImageOps.cpp b/mlir/lib/Dialect/SPIRV/IR/ImageOps.cpp
index b198294b9bdd6..d39379de3dd83 100644
--- a/mlir/lib/Dialect/SPIRV/IR/ImageOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/ImageOps.cpp
@@ -42,6 +42,36 @@ static LogicalResult verifyImageOperands(Operation *imageOp,
   // The order we process operands is important. In case of multiple argument
   // taking operands, the arguments are ordered starting with operands having
   // smaller-numbered bits first.
+  if (spirv::bitEnumContainsAny(attr.getValue(), spirv::ImageOperands::Bias)) {
+    if (!isa<spirv::ImplicitLodOpInterface>(imageOp))
+      return imageOp->emitError(
+          "Bias is only valid with implicit-lod instructions");
+
+    if (index + 1 > operands.size())
+      return imageOp->emitError("Bias operand requires 1 argument");
+
+    if (!isa<mlir::FloatType>(operands[index].getType()))
+      return imageOp->emitError("Bias must be a floating-point type scalar");
+
+    auto samplingOp = cast<spirv::SamplingOpInterface>(imageOp);
+    auto sampledImageType =
+        cast<spirv::SampledImageType>(samplingOp.getSampledImage().getType());
+    auto imageType = cast<spirv::ImageType>(sampledImageType.getImageType());
+
+    if (!llvm::is_contained({spirv::Dim::Dim1D, spirv::Dim::Dim2D,
+                             spirv::Dim::Dim3D, spirv::Dim::Cube},
+                            imageType.getDim()))
+      return imageOp->emitError(
+          "Bias must only be used with an image type that has "
+          "a dim operand of 1D, 2D, 3D, or Cube");
+
+    if (imageType.getSamplingInfo() != spirv::ImageSamplingInfo::SingleSampled)
+      return imageOp->emitError("Bias must only be used with an image type "
+                                "that has a MS operand of 0");
+
+    ++index;
+  }
+
   if (spirv::bitEnumContainsAny(attr.getValue(), spirv::ImageOperands::Lod)) {
     if (!isa<spirv::ExplicitLodOpInterface>(imageOp) &&
         !isa<spirv::FetchOpInterface>(imageOp))
@@ -74,12 +104,13 @@ static LogicalResult verifyImageOperands(Operation *imageOp,
     if (!llvm::is_contained({spirv::Dim::Dim1D, spirv::Dim::Dim2D,
                              spirv::Dim::Dim3D, spirv::Dim::Cube},
                             imageType.getDim()))
-      return imageOp->emitError("Lod only be used with an image type that has "
-                                "a dim operand of 1D, 2D, 3D, or Cube");
+      return imageOp->emitError(
+          "Lod must only be used with an image type that has "
+          "a dim operand of 1D, 2D, 3D, or Cube");
 
     if (imageType.getSamplingInfo() != spirv::ImageSamplingInfo::SingleSampled)
-      return imageOp->emitError(
-          "Lod only be used with an image type that has a MS operand of 0");
+      return imageOp->emitError("Lod must only be used with an image type that "
+                                "has a MS operand of 0");
 
     ++index;
   }
@@ -99,8 +130,8 @@ static LogicalResult verifyImageOperands(Operation *imageOp,
     auto imageType = cast<spirv::ImageType>(sampledImageType.getImageType());
 
     if (imageType.getSamplingInfo() != spirv::ImageSamplingInfo::SingleSampled)
-      return imageOp->emitError(
-          "Grad only be used with an image type that has a MS operand of 0");
+      return imageOp->emitError("Grad must only be used with an image type "
+                                "that has a MS operand of 0");
 
     int64_t numberOfComponents = 0;
 
@@ -147,10 +178,9 @@ static LogicalResult verifyImageOperands(Operation *imageOp,
 
   // TODO: Add the validation rules for the following Image Operands.
   spirv::ImageOperands noSupportOperands =
-      spirv::ImageOperands::Bias | spirv::ImageOperands::ConstOffset |
-      spirv::ImageOperands::Offset | spirv::ImageOperands::ConstOffsets |
-      spirv::ImageOperands::Sample | spirv::ImageOperands::MinLod |
-      spirv::ImageOperands::MakeTexelAvailable |
+      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;
 
diff --git a/mlir/test/Dialect/SPIRV/IR/image-ops.mlir b/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
index 9a0b8b79e3e01..1ebdfdb41de1b 100644
--- a/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/image-ops.mlir
@@ -276,6 +276,36 @@ func.func @sample_implicit_proj_dref(%arg0 : !spirv.sampled_image<!spirv.image<f
 
 // -----
 
+//===----------------------------------------------------------------------===//
+// spirv.ImageOperands: Bias
+//===----------------------------------------------------------------------===//
+
+func.func @bias_too_many_arguments(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : f32, %arg2 : f32) -> () {
+  // expected-error @+1 {{too many image operand arguments have been provided}}
+  %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 ["Bias"], %arg2, %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, f32, f32, f32 -> vector<4xf32>
+  spirv.Return
+}
+
+// -----
+
+func.func @bias_too_many_arguments(%arg0 : !spirv.sampled_image<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : f32, %arg2 : i32) -> () {
+  // expected-error @+1 {{Bias must be a floating-point type scalar}}
+  %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 ["Bias"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Dim1D, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, f32, i32 -> vector<4xf32>
+  spirv.Return
+}
+
+// -----
+
+func.func @bias_with_rect(%arg0 : !spirv.sampled_image<!spirv.image<f32, Rect, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : f32, %arg2 : f32) -> () {
+  // expected-error @+1 {{Bias must only be used with an image type that has a dim operand of 1D, 2D, 3D, or Cube}}
+  %0 = spirv.ImageSampleImplicitLod %arg0, %arg1 ["Bias"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Rect, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, f32, f32 -> vector<4xf32>
+  spirv.Return
+}
+
+// TODO: We cannot currently test Bias with MS != 0 as all implemented implicit operations already check for that.
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // spirv.ImageOperands: Lod
 //===----------------------------------------------------------------------===//
@@ -305,7 +335,7 @@ func.func @lod_too_many_arguments(%arg0 : !spirv.sampled_image<!spirv.image<f32,
 // -----
 
 func.func @lod_with_rect(%arg0 : !spirv.sampled_image<!spirv.image<f32, Rect, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, %arg1 : vector<2xf32>, %arg2 : f32) -> () {
-  // expected-error @+1 {{Lod only be used with an image type that has a dim operand of 1D, 2D, 3D, or Cube}}
+  // expected-error @+1 {{Lod must only be used with an image type that has a dim operand of 1D, 2D, 3D, or Cube}}
   %0 = spirv.ImageSampleExplicitLod %arg0, %arg1 ["Lod"], %arg2 : !spirv.sampled_image<!spirv.image<f32, Rect, NoDepth, NonArrayed, SingleSampled, NeedSampler, Rgba8>>, vector<2xf32>, f32 -> vector<4xf32>
   spirv.Return
 }