[Mlir-commits] [mlir] [mlir][tosa] Add pass to assign static input shape to TOSA functions (PR #171156)

[Sayan Saha]

================
@@ -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/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) {
----------------
sahas3 wrote:

I don't have a strong use-case at the moment, I wanted to just confirm my understanding. At MathWorks we have an option where users can specify which dims they want to be static/dynamic but I don't think there's a practical usage where user specifies only some dims as static and others as dynamic.

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