[Mlir-commits] [mlir] 96d60c0 - [mlir][spirv] Verify matching of entry block arguments and function signature (#133167)

[llvmlistbot at llvm.org]

Author: AdityaK
Date: 2025-04-01T14:10:17-07:00
New Revision: 96d60c00e5ed5bddedad0eab83a089957a9cf388

URL: https://github.com/llvm/llvm-project/commit/96d60c00e5ed5bddedad0eab83a089957a9cf388
DIFF: https://github.com/llvm/llvm-project/commit/96d60c00e5ed5bddedad0eab83a089957a9cf388.diff

LOG: [mlir][spirv] Verify matching of entry block arguments and function signature (#133167)

Fixes: #132894

Added: 
    

Modified: 
    mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
    mlir/test/Dialect/SPIRV/IR/function-decorations.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
index da9855b02860d..16e91b0cb2cfc 100644
--- a/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
+++ b/mlir/lib/Dialect/SPIRV/IR/SPIRVOps.cpp
@@ -1021,6 +1021,24 @@ LogicalResult spirv::FuncOp::verifyType() {
 
 LogicalResult spirv::FuncOp::verifyBody() {
   FunctionType fnType = getFunctionType();
+  if (!isExternal()) {
+    Block &entryBlock = front();
+
+    unsigned numArguments = this->getNumArguments();
+    if (entryBlock.getNumArguments() != numArguments)
+      return emitOpError("entry block must have ")
+             << numArguments << " arguments to match function signature";
+
+    for (auto [index, fnArgType, blockArgType] :
+         llvm::enumerate(getArgumentTypes(), entryBlock.getArgumentTypes())) {
+      if (blockArgType != fnArgType) {
+        return emitOpError("type of entry block argument #")
+               << index << '(' << blockArgType
+               << ") must match the type of the corresponding argument in "
+               << "function signature(" << fnArgType << ')';
+      }
+    }
+  }
 
   auto walkResult = walk([fnType](Operation *op) -> WalkResult {
     if (auto retOp = dyn_cast<spirv::ReturnOp>(op)) {

diff  --git a/mlir/test/Dialect/SPIRV/IR/function-decorations.mlir b/mlir/test/Dialect/SPIRV/IR/function-decorations.mlir
index 07e187e6a7d68..f09767a416f6b 100644
--- a/mlir/test/Dialect/SPIRV/IR/function-decorations.mlir
+++ b/mlir/test/Dialect/SPIRV/IR/function-decorations.mlir
@@ -73,3 +73,20 @@ spirv.func @no_decoration_name_attr(%arg0 : !spirv.ptr<i32, PhysicalStorageBuffe
 spirv.func @no_decoration_name_attr(%arg0 : !spirv.ptr<i32, PhysicalStorageBuffer> { spirv.decoration = #spirv.decoration<Restrict>, random_attr = #spirv.decoration<Aliased> }) "None" {
   spirv.Return
 }
+
+// -----
+
+// expected-error @+1 {{'spirv.func' op entry block must have 1 arguments to match function signature}}
+spirv.func @f(f32) "None" {
+  %c0 = arith.constant 0 : index
+  spirv.Return
+}
+
+// -----
+
+// expected-error @+1 {{'spirv.func' op type of entry block argument #0('f64') must match the type of the corresponding argument in function signature('f32')}}
+spirv.func @f(f32) "None" {
+  ^bb0(%arg0: f64):
+  %c0 = arith.constant 0 : index
+  spirv.Return
+}