[Mlir-commits] [mlir] [mlir][tosa] Enhance TosaInferShapes pass for simple shape inference (PR #178418)

[Tai Ly]

================
@@ -333,13 +161,252 @@ void validateSameOperandsAndResultRankTrait(Region &region) {
 struct TosaInferShapes
     : public tosa::impl::TosaInferShapesPassBase<TosaInferShapes> {
 public:
+  explicit TosaInferShapes() = default;
+  explicit TosaInferShapes(const TosaInferShapesPassOptions &options)
+      : TosaInferShapes() {
+    this->foldShapeExpressions = options.foldShapeExpressions;
+    this->convertFunctionBoundaries = options.convertFunctionBoundaries;
+  }
+
   void runOnOperation() override {
     func::FuncOp func = getOperation();
     TypeModificationState state;
     propagateShapesInRegion(func.getBody(), state);
     state.commit();
 
     validateSameOperandsAndResultRankTrait(func.getBody());
+
+    if (convertFunctionBoundaries)
+      convertFunctionReturnTypes(func);
----------------
Tai78641 wrote:

you might want to double check, but folding may leave dead-code ops around?
if so, you may want to insert here a walk that gets rid of dead code ops.
something like:

  // We can remove all dead code by going in reverse.
  // This is because we would remove usages before we
  // see the users.
  func.walk<WalkOrder::PostOrder, ReverseIterator>(
      [&](Operation *op) {
        if (isOpTriviallyDead(op))
          rewriter.eraseOp(op);
      });

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