[Mlir-commits] [mlir] [mlir][tosa] Add pass to assign static input shape to TOSA functions (PR #171156)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Dec 8 08:58:53 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Luke Hutton (lhutton1)
<details>
<summary>Changes</summary>
This commit introduces the `--tosa-experimental-input-shape` pass, which allows a user to convert dynamically shaped input arguments of TOSA functions to a user defined static shape. Here is a simple example:
```bash
func.func @<!-- -->test(%arg0: tensor<2x?xi32>, %arg1: tensor<?x256xf32>, %arg2: tensor<?x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
%0 = tosa.add %arg0, %arg0 : (tensor<2x?xi32>, tensor<2x?xi32>) -> tensor<2x?xi32>
%1 = tosa.reciprocal %arg1 : (tensor<?x256xf32>) -> tensor<?x256xf32>
%2 = tosa.sub %arg2, %arg2 : (tensor<?x9xf32>, tensor<?x9xf32>) -> tensor<?x9xf32>
return %0, %1, %2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
}
$ mlir-opt --tosa-experimental-input-shape="args=arg0:2x16,arg2:64x9" test.mlir
func.func @<!-- -->test(%arg0: tensor<2x16xi32>, %arg1: tensor<?x256xf32>, %arg2: tensor<64x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
%0 = tosa.add %arg0, %arg0 : (tensor<2x16xi32>, tensor<2x16xi32>) -> tensor<2x?xi32>
%1 = tosa.reciprocal %arg1 : (tensor<?x256xf32>) -> tensor<?x256xf32>
%2 = tosa.sub %arg2, %arg2 : (tensor<64x9xf32>, tensor<64x9xf32>) -> tensor<?x9xf32>
return %0, %1, %2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
}
```
When used in conjunction with the `--tosa-infer-shapes` pass, it can be used to resolve simple TOSA functions (those that don't include TOSA shape operations) to propagate static shape information. Continuing from the example above:
```bash
$ mlir-opt --tosa-infer-shapes test2.mlir
func.func @<!-- -->test(%arg0: tensor<2x16xi32>, %arg1: tensor<?x256xf32>, %arg2: tensor<64x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
%0 = tosa.add %arg0, %arg0 : (tensor<2x16xi32>, tensor<2x16xi32>) -> tensor<2x16xi32>
%cast = tensor.cast %0 : tensor<2x16xi32> to tensor<2x?xi32>
%1 = tosa.reciprocal %arg1 : (tensor<?x256xf32>) -> tensor<?x256xf32>
%2 = tosa.sub %arg2, %arg2 : (tensor<64x9xf32>, tensor<64x9xf32>) -> tensor<64x9xf32>
%cast_0 = tensor.cast %2 : tensor<64x9xf32> to tensor<?x9xf32>
return %cast, %1, %cast_0 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
}
```
Note: tosa-infer-shapes currently doesn't have an option to update the function signature.
Co-authored-by: Kaushik Varadharajan <kaushik.varadharajan@<!-- -->arm.com>
---
Full diff: https://github.com/llvm/llvm-project/pull/171156.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h (+2)
- (modified) mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td (+22)
- (modified) mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt (+1)
- (added) mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp (+184)
- (added) mlir/test/Dialect/Tosa/tosa-input-shape.mlir (+70)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h
index ba99d2f1d2727..e33d7c698856c 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h
@@ -41,6 +41,8 @@ void populateTosaConstantReduction(MLIRContext *ctx,
void populateTosaTypeConversion(TypeConverter &converter);
std::unique_ptr<Pass> createTosaTestQuantUtilAPIPass();
+std::unique_ptr<Pass>
+createTosaInputShapePass(std::vector<std::string> args = {});
#define GEN_PASS_REGISTRATION
#include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
index 12f520297b702..58ea47db731a6 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -198,4 +198,26 @@ def TosaNarrowI64ToI32Pass : Pass<"tosa-narrow-i64-to-i32", "func::FuncOp"> {
];
}
+def TosaInputShape : Pass<"tosa-experimental-input-shape", "func::FuncOp"> {
+ let summary = "Override dynamic function arguments to specified static shapes.";
+ let description = [{
+ Pass that overrides the dynamic input shapes of function arguments to
+ specified static shapes. If a specified static shape conflicts with the
+ static dimensions in an original input shape, an error is reported.
+ }];
+
+ let constructor = "tosa::createTosaInputShapePass()";
+ let dependentDialects = [
+ "func::FuncDialect",
+ "tensor::TensorDialect",
+ "tosa::TosaDialect",
+ ];
+ let options = [
+ ListOption<"args", "args", "std::string",
+ "Comma-separated list of shape descriptions. Each description contains the "
+ "argument name, a colon, and a shape with dimensions separated by x "
+ "(e.g. arg0:5x5,arg3:2x64)">,
+ ];
+}
+
#endif // MLIR_DIALECT_TOSA_TRANSFORMS_PASSES
diff --git a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
index 091b481d6394b..cf1e6ab55872c 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
@@ -14,6 +14,7 @@ add_mlir_dialect_library(MLIRTosaTransforms
TosaProfileCompliance.cpp
TosaValidation.cpp
TosaNarrowI64ToI32.cpp
+ TosaInputShape.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa/Transforms
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
new file mode 100644
index 0000000000000..ae7701892817a
--- /dev/null
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
@@ -0,0 +1,184 @@
+//===- TosaInputShape.cpp -------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Pass that overrides the dynamic input shapes of function arguments to
+// specified static shapes. If a specified static shape conflicts with the
+// static dimensions in an original input shape, an error is reported.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tosa {
+#define GEN_PASS_DEF_TOSAINPUTSHAPE
+#include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"
+} // namespace tosa
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+namespace {
+
+typedef std::pair<size_t, SmallVector<int64_t>> IdxAndShape;
+
+FailureOr<IdxAndShape> parseInputShape(Location loc, StringRef input) {
+ if (!input.consume_front("arg")) {
+ emitError(loc) << "expected prefix 'arg' at the start of " << input;
+ return failure();
+ }
+
+ const size_t colonPos = input.find(':');
+ if (colonPos == StringRef::npos) {
+ emitError(loc) << "expected ':' after argument index in '" << input << "'";
+ return failure();
+ }
+
+ const StringRef indexStr = input.substr(0, colonPos);
+ input = input.substr(colonPos + 1);
+
+ size_t index;
+ if (indexStr.getAsInteger(10, index) || index < 0) {
+ emitError(loc) << "invalid argument index, got " << indexStr;
+ return failure();
+ }
+
+ SmallVector<int64_t> shape;
+ while (!input.empty()) {
+ const size_t xPos = input.find("x");
+ StringRef dimStr;
+ if (xPos == StringRef::npos) {
+ dimStr = input;
+ input = "";
+ } else {
+ dimStr = input.substr(0, xPos);
+ input = input.substr(xPos + 1);
+ }
+
+ int64_t dimVal;
+ if (dimStr.getAsInteger(10, dimVal) || dimVal <= 0) {
+ return failure();
+ }
+ shape.push_back(dimVal);
+ }
+
+ const auto idxAndShape = std::make_pair(index, shape);
+ return {idxAndShape};
+}
+
+// Parse input shape arguments from command line input. Returns parsed
+// static shapes and an optional error message.
+// For example:
+// "args=arg0:5x10,arg8:3x9" => {{{0, {5, 10}}, {8, {3, 9}}}, ""}
+// "args=arg0:" => {{}, "error message"}
+FailureOr<SmallVector<IdxAndShape>> parseInputShapes(Location loc, const std::vector<std::string>& args) {
+ SmallVector<IdxAndShape> inputShapes;
+ for (const std::string& arg : args) {
+ const auto maybeInputShape = parseInputShape(loc, arg);
+ if (failed(maybeInputShape))
+ return failure();
+ inputShapes.push_back(maybeInputShape.value());
+ }
+ return inputShapes;
+}
+
+struct TosaInputShape : public tosa::impl::TosaInputShapeBase<TosaInputShape> {
+public:
+ TosaInputShape() = default;
+
+ explicit TosaInputShape(std::vector<std::string> args) : TosaInputShape() {
+ this->args = args;
+ }
+
+ void runOnOperation() override {
+ MLIRContext *context = &getContext();
+ const Location unknownLoc = UnknownLoc::get(context);
+ const auto maybeArgsParsed = parseInputShapes(unknownLoc, args);
+ if (failed(maybeArgsParsed))
+ return;
+ const SmallVector<IdxAndShape> argsParsed = maybeArgsParsed.value();
+ func::FuncOp func = getOperation();
+
+ const auto getUpdatedTensorType =
+ [&](size_t argIdx,
+ ArrayRef<Type> argTypes, ArrayRef<int64_t> requestedShape) -> FailureOr<Type> {
+ const size_t numInputs = argTypes.size();
+ if (argIdx >= numInputs)
+ return func.emitError()
+ << "provided arg index " << argIdx
+ << " is larger than number of inputs "
+ << numInputs << ".";
+
+ auto tensorType = dyn_cast<TensorType>(argTypes[argIdx]);
+ if (!tensorType)
+ return func.emitError()
+ << "expected tensor type, got " << argTypes[argIdx];
+
+ const ArrayRef<int64_t> originalShape = tensorType.getShape();
+ if (failed(verifyCompatibleShape(originalShape, requestedShape)))
+ return func.emitError()
+ << "arg" << argIdx << " has incompatible shape with requested input shape ("
+ << requestedShape << "), got " << tensorType;
+ return tensorType.cloneWith(requestedShape, tensorType.getElementType());
+ };
+
+ // Update argument shapes in the entry block
+ Block &entryBlock = func.getBody().front();
+ for (const auto &[argIdx, shape] : argsParsed) {
+ SmallVector<Type> argTypes(entryBlock.getArgumentTypes());
+ FailureOr<Type> newTensorType =
+ getUpdatedTensorType(argIdx, argTypes, shape);
+ if (failed(newTensorType))
+ return signalPassFailure();
+
+ entryBlock.getArgument(argIdx).setType(newTensorType.value());
+ }
+
+ // Get new func argument types
+ FunctionType oldFunctionType = func.getFunctionType();
+ ArrayRef<Type> oldInputTypes = oldFunctionType.getInputs();
+ const size_t numInputs = oldInputTypes.size();
+ SmallVector<Type> newInputs(oldInputTypes.begin(), oldInputTypes.end());
+ newInputs.reserve(numInputs);
+ for (const auto &[argIdx, shape] : argsParsed) {
+ FailureOr<Type> newTensorType =
+ getUpdatedTensorType(argIdx, oldInputTypes, shape);
+ if (failed(newTensorType))
+ return signalPassFailure();
+
+ newInputs[argIdx] = newTensorType.value();
+ }
+
+ // Update function signature
+ Block &lastBlock = func.getBody().back();
+ Operation *terminator = lastBlock.getTerminator();
+ SmallVector<Type> newResults;
+ if (auto returnOp = dyn_cast_or_null<func::ReturnOp>(terminator)) {
+ auto types = returnOp.getOperandTypes();
+ newResults.assign(types.begin(), types.end());
+ } else {
+ auto types = oldFunctionType.getResults();
+ newResults.assign(types.begin(), types.end());
+ }
+ const FunctionType newFunctionType = oldFunctionType.clone(
+ newInputs,
+ newResults);
+ func.setFunctionType(newFunctionType);
+ }
+};
+
+} // namespace
+
+std::unique_ptr<Pass>
+mlir::tosa::createTosaInputShapePass(std::vector<std::string> args) {
+ return std::make_unique<TosaInputShape>(args);
+}
diff --git a/mlir/test/Dialect/Tosa/tosa-input-shape.mlir b/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
new file mode 100644
index 0000000000000..520a4575d3d60
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
@@ -0,0 +1,70 @@
+// RUN: mlir-opt -split-input-file -verify-diagnostics -tosa-experimental-input-shape="args=arg0:2x16,arg2:64x9" %s | FileCheck %s
+
+// CHECK-LABEL: test_empty_func
+func.func @test_empty_func(
+ // CHECK: %arg0: tensor<2x16xi32>
+ %arg0: tensor<2x?xi32>,
+ // CHECK: %arg1: tensor<?x256xf32>
+ %arg1: tensor<?x256xf32>,
+ // CHECK: %arg2: tensor<64x9xf32>
+ %arg2: tensor<?x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
+ // CHECK: %arg0, %arg1, %arg2 : tensor<2x16xi32>, tensor<?x256xf32>, tensor<64x9xf32>
+ return %arg0, %arg1, %arg2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_func_with_ops
+func.func @test_func_with_ops(
+ // CHECK: %arg0: tensor<2x16xi32>
+ %arg0: tensor<2x?xi32>,
+ // CHECK: %arg1: tensor<?x256xf32>
+ %arg1: tensor<?x256xf32>,
+ // CHECK: %arg2: tensor<64x9xf32>
+ %arg2: tensor<?x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
+ // CHECK: %[[ADD:.*]] = tosa.add %arg0, %arg0 : (tensor<2x16xi32>, tensor<2x16xi32>)
+ %0 = tosa.add %arg0, %arg0 : (tensor<2x?xi32>, tensor<2x?xi32>) -> tensor<2x?xi32>
+ // CHECK: %[[RECIP:.*]] = tosa.reciprocal %arg1 : (tensor<?x256xf32>)
+ %1 = tosa.reciprocal %arg1 : (tensor<?x256xf32>) -> tensor<?x256xf32>
+ // CHECK: %[[SUB:.*]] = tosa.sub %arg2, %arg2 : (tensor<64x9xf32>, tensor<64x9xf32>)
+ %2 = tosa.sub %arg2, %arg2 : (tensor<?x9xf32>, tensor<?x9xf32>) -> tensor<?x9xf32>
+ return %0, %1, %2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+}
+
+// -----
+
+// CHECK-LABEL: test_controlflow
+func.func @test_controlflow(
+ // CHECK: %arg0: tensor<2x16xi32>
+ %arg0: tensor<2x?xi32>,
+ // CHECK: %arg1: tensor<?x256xf32>
+ %arg1: tensor<?x256xf32>,
+ // CHECK: %arg2: tensor<64x9xf32>
+ %arg2: tensor<?x9xf32>,
+ // CHECK: %arg3: tensor<i1>
+ %arg3: tensor<i1>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
+ // CHECK: %[[IF:.*]]:3 = tosa.cond_if %arg3 (%arg4 = %arg0, %arg5 = %arg1, %arg6 = %arg2) : tensor<i1> (tensor<2x16xi32>, tensor<?x256xf32>, tensor<64x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
+ %0:3 = tosa.cond_if %arg3 (%arg4 = %arg0, %arg5 = %arg1, %arg6 = %arg2) : tensor<i1> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) -> (tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>) {
+ ^bb0(%arg4: tensor<2x?xi32>, %arg5: tensor<?x256xf32>, %arg6: tensor<?x9xf32>):
+ tosa.yield %arg4, %arg5, %arg6 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+ } else {
+ ^bb0(%arg4: tensor<2x?xi32>, %arg5: tensor<?x256xf32>, %arg6: tensor<?x9xf32>):
+ tosa.yield %arg4, %arg5, %arg6 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+ }
+ // CHECK: return %[[IF]]#0, %[[IF]]#1, %[[IF]]#2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+ return %0#0, %0#1, %0#2 : tensor<2x?xi32>, tensor<?x256xf32>, tensor<?x9xf32>
+}
+
+// -----
+
+func.func @test_wrong_number_input_args(%arg0: tensor<2x?xf32>) -> tensor<2x?xf32> {
+ // expected-error at -1 {{provided arg index 2 is larger than number of inputs 1}}
+ return %arg0 : tensor<2x?xf32>
+}
+
+// -----
+
+func.func @test_incompatible_input_shape(%arg0: tensor<1x?xf32>, %arg1: tensor<1x?xf32>, %arg2: tensor<1x?xf32>) -> tensor<1x?xf32> {
+ // expected-error at -1 {{arg0 has incompatible shape with requested input shape (2, 16), got 'tensor<1x?xf32>'}}
+ return %arg0 : tensor<1x?xf32>
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/171156
More information about the Mlir-commits
mailing list