[Mlir-commits] [mlir] [mlir][tosa] Handle function declarations in tosa input shape pass (PR #205359)

[Jiaqi He]

https://github.com/heturing created https://github.com/llvm/llvm-project/pull/205359

Fixes https://github.com/llvm/llvm-project/issues/205063.

The `tosa-experimental-input-shape` pass currently does not handle function declarations correctly. The pass may run on declarations, but the current implementation assumes that every function has a body and unconditionally accesses the entry block and the last block when updating argument and result types.

This patch checks whether the function has a body before accessing body blocks. For declarations, the pass updates the function signature input types and preserves the original result types, since there is no return operation from which result types can be inferred.

A regression test is added for the declaration case.

>From 68a0ccad4e885a7864f9b4ddcf91bdea139bdd97 Mon Sep 17 00:00:00 2001
From: Jiaqi He <heturing at gmail.com>
Date: Tue, 23 Jun 2026 22:18:57 +0800
Subject: [PATCH] [mlir][Tosa] Handle function declarations in tosa input shape
 pass

---
 .../Tosa/Transforms/TosaInputShape.cpp        | 41 ++++++++++---------
 mlir/test/Dialect/Tosa/tosa-input-shape.mlir  | 13 ++++++
 2 files changed, 35 insertions(+), 19 deletions(-)

diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
index e733fb8d378f5..801542116358f 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
@@ -132,16 +132,18 @@ struct TosaInputShape : public tosa::impl::TosaInputShapeBase<TosaInputShape> {
       return tensorType.cloneWith(requestedShape, tensorType.getElementType());
     };
 
-    // Update argument shapes in the entry block
-    Block &entryBlock = func.getBody().front();
-    const SmallVector<Type> argTypes(entryBlock.getArgumentTypes());
-    for (const auto &[argIdx, shape] : argsParsed) {
-      FailureOr<Type> newTensorType =
-          getUpdatedTensorType(argIdx, argTypes, shape);
-      if (failed(newTensorType))
-        return signalPassFailure();
-
-      entryBlock.getArgument(argIdx).setType(newTensorType.value());
+    // Update argument shapes in the entry block if the function has body.
+    if (!func.getBody().empty()) {
+      Block &entryBlock = func.getBody().front();
+      const SmallVector<Type> argTypes(entryBlock.getArgumentTypes());
+      for (const auto &[argIdx, shape] : argsParsed) {
+        FailureOr<Type> newTensorType =
+            getUpdatedTensorType(argIdx, argTypes, shape);
+        if (failed(newTensorType))
+          return signalPassFailure();
+
+        entryBlock.getArgument(argIdx).setType(newTensorType.value());
+      }
     }
 
     // Get new func argument types
@@ -158,16 +160,17 @@ struct TosaInputShape : public tosa::impl::TosaInputShapeBase<TosaInputShape> {
     }
 
     // Update function signature
-    Block &lastBlock = func.getBody().back();
-    const Operation *terminator = lastBlock.getTerminator();
-    SmallVector<Type> newResults;
-    if (auto returnOp = dyn_cast_or_null<func::ReturnOp>(terminator)) {
-      const auto types = returnOp.getOperandTypes();
-      newResults.assign(types.begin(), types.end());
-    } else {
-      const auto types = oldFunctionType.getResults();
-      newResults.assign(types.begin(), types.end());
+    const auto oldResultTypes = oldFunctionType.getResults();
+    SmallVector<Type> newResults(oldResultTypes.begin(), oldResultTypes.end());
+    if (!func.getBody().empty()) {
+      Block &lastBlock = func.getBody().back();
+      const Operation *terminator = lastBlock.getTerminator();
+      if (auto returnOp = dyn_cast_or_null<func::ReturnOp>(terminator)) {
+        const auto returnTypes = returnOp.getOperandTypes();
+        newResults.assign(returnTypes.begin(), returnTypes.end());
+      }
     }
+
     const FunctionType newFunctionType =
         oldFunctionType.clone(newInputs, newResults);
     func.setFunctionType(newFunctionType);
diff --git a/mlir/test/Dialect/Tosa/tosa-input-shape.mlir b/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
index ee8e7aac609d5..8b390a1826240 100644
--- a/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
@@ -70,3 +70,16 @@ func.func @test_incompatible_input_shape(%arg0: tensor<1x?xf32>, %arg1: tensor<1
     // expected-error at -1 {{arg0 has incompatible shape with requested input shape (2, 16), got 'tensor<1x?xf32>'}}
     return %arg0 : tensor<1x?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: test_func_decl
+func.func private @test_func_decl(
+    // CHECK: tensor<2x16xi32>
+    %arg0: tensor<2x?xi32>,
+    // CHECK: f32
+    %arg1: f32,
+    // CHECK: tensor<64x9xi32>
+    %arg2: tensor<?x9xi32>) ->
+    // CHECK: tensor<?x256xf32>
+    tensor<?x256xf32>