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

[llvmlistbot at llvm.org]

Author: Jiaqi He
Date: 2026-06-24T15:57:31+01:00
New Revision: 229e42155579d191ce114d1e067a2a05aba564e3

URL: https://github.com/llvm/llvm-project/commit/229e42155579d191ce114d1e067a2a05aba564e3
DIFF: https://github.com/llvm/llvm-project/commit/229e42155579d191ce114d1e067a2a05aba564e3.diff

LOG: [mlir][tosa] Handle function declarations in tosa input shape pass (#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.

Added: 
    

Modified: 
    mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
    mlir/test/Dialect/Tosa/tosa-input-shape.mlir

Removed: 
    


################################################################################
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>