[Mlir-commits] [mlir] 6bb2043 - [mlir][tosa] Add pass to assign static input shape to TOSA functions (#171156)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Jan 15 09:31:40 PST 2026
Author: Luke Hutton
Date: 2026-01-15T17:31:34Z
New Revision: 6bb20438598bfced09b1afbb76e551fbb955060d
URL: https://github.com/llvm/llvm-project/commit/6bb20438598bfced09b1afbb76e551fbb955060d
DIFF: https://github.com/llvm/llvm-project/commit/6bb20438598bfced09b1afbb76e551fbb955060d.diff
LOG: [mlir][tosa] Add pass to assign static input shape to TOSA functions (#171156)
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 at arm.com>
Added:
mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
mlir/test/Dialect/Tosa/tosa-input-shape.mlir
Modified:
mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h
mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
Removed:
################################################################################
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 4a5f283bc66c8..03be41d684f3f 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -223,4 +223,23 @@ def TosaNarrowF64ToF32Pass : Pass<"tosa-narrow-f64-to-f32", "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 = [
+ "tensor::TensorDialect"
+ ];
+ 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 ">
+ ];
+}
+
#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 0ff68b1bb54f4..e8a76fa3a1d21 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
TosaNarrowTypes.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..e733fb8d378f5
--- /dev/null
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaInputShape.cpp
@@ -0,0 +1,182 @@
+//===- 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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.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();
+ 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
+ const FunctionType oldFunctionType = func.getFunctionType();
+ const ArrayRef<Type> oldInputTypes = oldFunctionType.getInputs();
+ SmallVector<Type> newInputs(oldInputTypes.begin(), oldInputTypes.end());
+ 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();
+ 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 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..ee8e7aac609d5
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-input-shape.mlir
@@ -0,0 +1,72 @@
+// 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>) {
+ // CHECK: ^bb0(%arg4: tensor<2x?xi32>, %arg5: tensor<?x256xf32>, %arg6: 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 {
+ // CHECK: ^bb0(%arg4: tensor<2x?xi32>, %arg5: tensor<?x256xf32>, %arg6: 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>
+ }
+ // 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>
+}
More information about the Mlir-commits
mailing list