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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-tosa

Author: Jiaqi He (heturing)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/205359.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp (+22-19) 
- (modified) mlir/test/Dialect/Tosa/tosa-input-shape.mlir (+13) 


``````````diff
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>

``````````

</details>


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