[mlir][tosa] Tosa shape propagation for tosa.cond_if

[rsuderman]

Summary

We can propagate the shape from tosa.cond_if operands into the true/false
regions then through the connected blocks. Then, using the tosa.yield ops
we can determine what all possible return types are.

[rsuderman]

Added missing test.

[jpienaar]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

19

/// ? (for doxygen/documentation comments)

28

How would you feel about us using standard lattice terminology here? E.g., meet, join, bottom (well could also use least)

50

Is context here for future?

55

Elide trivial braces ? (I know not your favorite style convention :)). You could consider a return with ternary operator too

73

Passing nullptr as context here feels weird - I would have expected both lhs & rhs to be from same context (I mean they are context adjacent but you know what I mean) and then passing either of them along.

98

[up to you] for range with llvm::zip ?

(then you could do

std::tie(lhsSize, rhsSize, resultSize) =

inside and avoid indexing directly)

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

935

This is where meet would have read easier for me :) [it only sometimes expand, and expand makes me think of something else]

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

64–69

Could this be more of a action? (don't know how to phrase it, this just feels like a noun and makes me go "propagate what? and where is this thing passed in that needs to be propagated?")

132

Is this what this pass does?

[rsuderman]

Addressed clang-tidy errors and most of jpienaar@ comments

[rsuderman]

Rest of jpienaar comments.

[rsuderman]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

19

Added in additional slash. Not entirely familiar with doxygen documentation comments so if this is off update me.

28

Changed to meet. I wasn't familiar with lattice terminology so I took a stab at expand.

50

Removed. The original npcomp implementation retained the context. Not sure why.

98

In this case, because we grab a reference to the result type I think we are stuck with direct indexing. Whenever I see an alternative it feels more confusing than adding the few extra lines.

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

64–69

Rephrased to sound more like an action.

132

Updated comment. Must have been leftover from the initial commit copypasta

[rsuderman]

Moved header guard

[jpienaar]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

28

NP, it isn't necessarily the most intuitive and so you could have made a judgement call for more readable :) (I always forget which one is which and also one can set it up in different directions - I mostly try and bring it back to And and Or [given the symbol reuse] and ask myself which one is more of an And vs more of an Or, but that is subjective)

113

I think this join is the same as Shape_Join, except for the handling where the static dims don't match. There it is an error, here you are returning ?. These are used differently (there it is an equality constraint really) so makes sense.

In ValueKnowledge there is no error state correct? So if I wanted to say shape x == shape y, then I couldn't here, it would have to be in shape function?

149

And effectively needed to differentiate between scalar or not?

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

52

Wouldn't it actually be an error if a static dim mismatches here?

[rsuderman]

Updated to change hasSizes to hasRank, add an error state, and handle
the newly added error state.

[rsuderman]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

113

I changed to include an error state. I think the premise was to just continue with whatever we can but it feels reasonable to restrict down as best as possible.

149

Less differentiate between scalar or not, with unranked and ranked. I changed to hasRank to clarify.

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

52

Yup you are right. I added a check for error state and continued in cases where it can't propagate safely.

[hanchung]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

84–87

Should swap or merge the checks? If lhs or rhs is null, we will crash in the first if-statement right?

92–94

style nit: drop the braces

113

IIRC, we prefer for (auto i : llvm::seq<unsigned>(0, result.sizes.size())) style.

139–142

ditto

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

910–916

llvm style nit: Use braces on the outer block because there are more than two levels of nesting.

for ... {
  for ...
    if ...
}

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

[rsuderman]

Updated for @hanchung comments.

[rsuderman]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

84–87

Technically its safe because ValueKnowledge is *not* part of the mlir type system. Its bool check is whether the value type is an error, not that it exists. Updated anyway.

[rsuderman]

Added missing include for llvm::seq

[jpienaar]

Looks good thanks!

[RKSimon]

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

137

@rsuderman duplicate check

if (!lhs || !rhs || lhs.dtype != rhs.dtype)

Reported on https://pvs-studio.com/en/blog/posts/cpp/1003/ (N2)

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

 
 //===----------------------------------------------------------------------===//
 // Operator: cond_if
 //===----------------------------------------------------------------------===//
 //===----------------------------------------------------------------------===//
 // Further described in docs/Rationale/RationaleTOSADialect.md .
 //===----------------------------------------------------------------------===//
 def Tosa_IfOp : Tosa_Op<"cond_if", [
+      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                                ["inferReturnTypeComponents"]>,
        SingleBlockImplicitTerminator<"YieldOp">,
        RecursiveSideEffects]> {
   let summary = "Conditional if operator";
 
   let description = [{
     Evaluates a Boolean condition and then takes one of two distinct execution
     paths. This implements the semantic If-then-else structure.
   }];

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

+//===-- ShapeUtils.h - TOSA shape support declarations ----------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Class declarations for shape utilities meant to assist shape propagation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H
+#define MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Types.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+namespace tosa {
+/// Statically known information for a particular Value.
+///
+/// This struct currently tracks only information relevant for tensor/array-like
+/// shaped types. It is fine to associate a `ValueKnowledge` with a non-shaped
+/// type as long as it is in the default "no knowledge" state returned by
+/// `getPessimisticValueState`. The important invariant is that we cannot
+/// claim to know something about a value which is false.
+///
+/// This class could also be called "dataflow facts", "lattice value", etc.
+struct ValueKnowledge {
+  ValueKnowledge() = delete;
+  ValueKnowledge(bool hasRank, llvm::ArrayRef<int64_t> newSizes, Type dtype)
+      : hasError(false), hasRank(hasRank), dtype(dtype) {
+    sizes.reserve(newSizes.size());
+    for (auto size : newSizes)
+      sizes.push_back(size);
+  }
+
+  operator bool() const { return !hasError; }
+
+  // Get the static knowledge intrinsic to `type`.
+  static ValueKnowledge getKnowledgeFromType(Type type) {
+    ValueKnowledge result = getPessimisticValueState();
+    if (auto shapedType = type.dyn_cast<ShapedType>()) {
+      if (shapedType.hasRank()) {
+        result.hasRank = true;
+        result.sizes.reserve(shapedType.getRank());
+        for (auto dim : shapedType.getShape())
+          result.sizes.push_back(dim);
+      }
+      result.dtype = shapedType.getElementType();
+    }
+    return result;
+  }
+
+  // Return a pessimistic/conservative value state without assuming any knowlege
+  // about the IR.
+  static ValueKnowledge getPessimisticValueState() {
+    return ValueKnowledge(false, {}, Type());
+  }
+
+  Type getType() const {
+    if (hasRank)
+      return RankedTensorType::get(llvm::makeArrayRef(sizes), dtype);
+    return UnrankedTensorType::get(dtype);
+  }
+
+  bool operator==(const ValueKnowledge &rhs) const {
+    return hasRank == rhs.hasRank && sizes == rhs.sizes && dtype == rhs.dtype;
+  }
+
+  // Given two pieces of static knowledge, calculate conservatively the
+  // information we can be sure about.
+  static ValueKnowledge join(const ValueKnowledge &lhs,
+                             const ValueKnowledge &rhs) {
+    // Mental model: All conditions are checking how to change from the safe "no
+    // knowledge" default-initialized state to a state with more knowledge
+    // consistent with lhs and rhs.
+    ValueKnowledge result = getPessimisticValueState();
+    result.hasError = true;
+
+    if (!lhs || !rhs || lhs.dtype != rhs.dtype)
+      return result;
+
+    result.hasError = false;
+    result.dtype = lhs.dtype;
+
+    if (!lhs.hasRank && !rhs.hasRank)
+      return result;
+
+    if (!rhs.hasRank) {
+      result.hasRank = true;
+      result.sizes = lhs.sizes;
+      return result;
+    }
+
+    if (!lhs.hasRank) {
+      result.hasRank = true;
+      result.sizes = rhs.sizes;
+      return result;
+    }
+
+    if (lhs.sizes.size() != rhs.sizes.size())
+      return result;
+
+    result.hasRank = true;
+    result.sizes.resize(lhs.sizes.size(), ShapedType::kDynamicSize);
+    for (auto i : llvm::seq<unsigned>(0, result.sizes.size())) {
+      int64_t lhsSize = lhs.sizes[i];
+      int64_t rhsSize = rhs.sizes[i];
+      int64_t &resultSize = result.sizes[i];
+      if (lhsSize == ShapedType::kDynamicSize) {
+        resultSize = rhsSize;
+      } else if (rhsSize == ShapedType::kDynamicSize) {
+        resultSize = lhsSize;
+      } else if (lhsSize == rhsSize) {
+        resultSize = lhsSize;
+      } else {
+        result.hasError = true;
+      }
+    }
+
+    return result;
+  }
+
+  // Given to types, generate a new ValueKnowledge that meets to cover both
+  // cases. E.g. if the rank of the LHS and RHS differ, the resulting tensor
+  // has unknown rank.
+  static ValueKnowledge meet(const ValueKnowledge &lhs,
+                             const ValueKnowledge &rhs) {
+    ValueKnowledge result = getPessimisticValueState();
+    result.hasError = true;
+
+    if (!rhs || !rhs || lhs.dtype != rhs.dtype)
+      return result;
+
+    result.hasError = false;
+    result.dtype = lhs.dtype;
+
+    if (!lhs.hasRank || !rhs.hasRank) {
+      result.hasRank = false;
+      return result;
+    }
+
+    if (lhs.sizes.size() != rhs.sizes.size()) {
+      result.hasRank = false;
+      return result;
+    }
+
+    result.hasRank = true;
+    result.sizes.resize(lhs.sizes.size(), ShapedType::kDynamicSize);
+    for (int i = 0, e = lhs.sizes.size(); i < e; i++) {
+      if (lhs.sizes[i] == rhs.sizes[i]) {
+        result.sizes[i] = lhs.sizes[i];
+      }
+    }
+
+    return result;
+  }
+
+  // Whether the value information has an error.
+  bool hasError;
+  // Whether the value has known rank.
+  bool hasRank;
+  // If `hasRank`, the sizes along each rank. Unknown sizes are represented as
+  // `ShapedType::kDynamicSize`.
+  llvm::SmallVector<int64_t> sizes;
+  // The dtype of a tensor.
+  // This is equal to nullptr if we don't know that it is a specific concrete
+  // type.
+  Type dtype;
+};
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 // This file implements the TOSA Specification:
 // https://developer.mlplatform.org/w/tosa/
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
+#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
 
 using namespace mlir;
 using namespace mlir::tosa;

 
 LogicalResult MaxPool2dOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   return poolingInferReturnTypes(operands, attributes, inferredReturnShapes);
 }
 
+LogicalResult IfOp::inferReturnTypeComponents(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  llvm::SmallVector<tosa::YieldOp> yieldOps;
+  for (Region *region : regions) {
+    for (auto &block : *region)
+      if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
+        yieldOps.push_back(returnOp);
+  }
+
+  if (yieldOps.empty())
+    return failure();
+
+  // Get the initial type information for the yield op.
+  llvm::SmallVector<ValueKnowledge> resultKnowledge;
+  resultKnowledge.reserve(yieldOps.front().getNumOperands());
+  for (auto operand : yieldOps.front().getOperands()) {
+    resultKnowledge.push_back(
+        ValueKnowledge::getKnowledgeFromType(operand.getType()));
+  }
+
+  for (auto yieldOp : yieldOps) {
+    if (resultKnowledge.size() != yieldOp.getNumOperands())
+      return failure();
+
+    for (auto it : llvm::enumerate(yieldOp.getOperands())) {
+      int32_t index = it.index();
+      auto meet = ValueKnowledge::meet(
+          resultKnowledge[index],
+          ValueKnowledge::getKnowledgeFromType(it.value().getType()));
+      if (!meet)
+        continue;
+      resultKnowledge[index] = meet;
+    }
+  }
+
+  for (auto result : resultKnowledge) {
+    if (result.hasRank) {
+      inferredReturnShapes.push_back(ShapedTypeComponents(result.sizes));
+    } else {
+      inferredReturnShapes.push_back(ShapedTypeComponents());
+    }
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/DataFlowAnalysis.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
-// -----------------------------------------------------------------------------
-// Analysis.
-// -----------------------------------------------------------------------------
+void propagateShapesInRegion(Region &region);
 
-static Type joinElementTypes(Type lhs, Type rhs) {
-  return lhs == rhs ? lhs : Type();
-}
+void propagateShapesToTosaIf(Operation &op) {
+  tosa::IfOp ifOp = dyn_cast<tosa::IfOp>(op);
+  if (!ifOp)
+    return;
 
-namespace {
-// Statically known information for a particular Value.
-//
-// This struct currently tracks only information relevant for tensor/array-like
-// shaped types. It is fine to associate a `ValueKnowledge` with a non-shaped
-// type as long as it is in the default "no knowledge" state returned by
-// `getPessimisticValueState`. The important invariant is that we cannot
-// claim to know something about a value which is false.
-//
-// This class could also be called "dataflow facts", "lattice value", etc.
-struct ValueKnowledge {
-  ValueKnowledge() = delete;
-  ValueKnowledge(bool hasSizes, std::vector<int64_t> sizes, Type dtype)
-      : hasSizes(hasSizes), sizes(sizes), dtype(dtype) {
-    assert(sizes.size() == 0 || hasSizes);
-  }
-
-  // Get the static knowledge intrinsic to `type`.
-  static ValueKnowledge getKnowledgeFromType(Type type) {
-    ValueKnowledge result = getPessimisticValueState(type.getContext());
-    if (auto shapedType = type.dyn_cast<ShapedType>()) {
-      if (shapedType.hasRank()) {
-        result.hasSizes = true;
-        result.sizes = shapedType.getShape();
-      }
-      result.dtype = shapedType.getElementType();
-    }
-    return result;
-  }
-
-  // Return a pessimistic/conservative value state without assuming any knowlege
-  // about the IR.
-  static ValueKnowledge getPessimisticValueState(MLIRContext *context) {
-    return ValueKnowledge(false, {}, Type());
-  }
-
-  Type getType() const {
-    if (hasSizes) {
-      return RankedTensorType::get(llvm::makeArrayRef(sizes), dtype);
-    }
-    return UnrankedTensorType::get(dtype);
-  }
-
-  bool operator==(const ValueKnowledge &rhs) const {
-    return std::make_tuple(hasSizes, sizes, dtype) ==
-           std::make_tuple(rhs.hasSizes, rhs.sizes, rhs.dtype);
-  }
-
-  // Given two pieces of static knowledge, calculate conservatively the
-  // information we can be sure about.
-  static ValueKnowledge join(const ValueKnowledge &lhs,
-                             const ValueKnowledge &rhs) {
-    // Mental model: All conditions are checking how to change from the safe "no
-    // knowledge" default-initialized state to a state with more knowledge
-    // consistent with lhs and rhs.
-    ValueKnowledge result = getPessimisticValueState(nullptr);
-
-    if (lhs.hasSizes && !rhs.hasSizes) {
-      result.hasSizes = true;
-      result.sizes = lhs.sizes;
-    } else if (!lhs.hasSizes && rhs.hasSizes) {
-      result.hasSizes = true;
-      result.sizes = rhs.sizes;
-    } else if (lhs.hasSizes && rhs.hasSizes &&
-               lhs.sizes.size() == rhs.sizes.size()) {
-      result.hasSizes = true;
-      result.sizes.resize(lhs.sizes.size(), ShapedType::kDynamicSize);
-      for (int i = 0, e = result.sizes.size(); i != e; i++) {
-        int64_t lhsSize = lhs.sizes[i];
-        int64_t rhsSize = rhs.sizes[i];
-        int64_t &resultSize = result.sizes[i];
-        if (lhsSize == ShapedType::kDynamicSize) {
-          resultSize = rhsSize;
-        } else if (rhsSize == ShapedType::kDynamicSize) {
-          resultSize = lhsSize;
-        } else if (lhsSize == rhsSize) {
-          resultSize = lhsSize;
-        }
-      }
-    }
-
-    result.dtype = joinElementTypes(lhs.dtype, rhs.dtype);
-    return result;
-  }
-
-  // Whether the Value is known to have a list of sizes.
-  bool hasSizes;
-  // If `hasSizes`, the sizes along each rank. Unknown sizes are represented as
-  // `ShapedType::kDynamicSize`.
-  std::vector<int64_t> sizes;
-  // The dtype of a tensor.
-  // This is equal to nullptr if we don't know that it is a specific concrete
-  // type.
-  Type dtype;
-};
+  for (auto &region : op.getRegions()) {
+    Block &frontBlock = region.front();
+    if (frontBlock.getNumArguments() + 1 != ifOp.getNumOperands())
+      return;
 
-} // namespace
+    for (int i = 0, e = frontBlock.getNumArguments(); i < e; i++) {
+      ValueKnowledge operandKnowledge = ValueKnowledge::getKnowledgeFromType(
+          ifOp.getOperand(i + 1).getType());
+      ValueKnowledge blockKnowledge = ValueKnowledge::getKnowledgeFromType(
+          frontBlock.getArgument(i).getType());
+      ValueKnowledge joinedKnowledge =
+          ValueKnowledge::join(operandKnowledge, blockKnowledge);
+      if (!joinedKnowledge)
+        continue;
+      frontBlock.getArgument(i).setType(joinedKnowledge.getType());
+    }
 
-/// Pass that enables broadcast by making all input arrays have the same
-/// number of dimensions. Insert RESHAPE operations to lower rank operand
-struct TosaInferShapes : public TosaInferShapesBase<TosaInferShapes> {
-public:
-  void runOnFunction() override {
-    FuncOp func = getOperation();
+    propagateShapesInRegion(region);
+  }
 
-    IRRewriter rewriter(func.getContext());
+  return;
+}
 
-    func.walk([&](Operation *op) {
-      if (op->getDialect()->getNamespace() !=
+void propagateShapesInRegion(Region &region) {
+  for (auto &block : region) {
+    for (Operation &op : block) {
+      if (op.getDialect()->getNamespace() !=
           tosa::TosaDialect::getDialectNamespace())
-        return;
+        continue;
+
+      propagateShapesToTosaIf(op);
+
       InferShapedTypeOpInterface shapeInterface =
           dyn_cast<InferShapedTypeOpInterface>(op);
       if (!shapeInterface)
-        return;
+        continue;
 
       SmallVector<ShapedTypeComponents> returnedShapes;
       if (shapeInterface
               .inferReturnTypeComponents(
-                  op->getContext(), op->getLoc(), op->getOperands(),
-                  op->getAttrDictionary(), op->getRegions(), returnedShapes)
+                  op.getContext(), op.getLoc(), op.getOperands(),
+                  op.getAttrDictionary(), op.getRegions(), returnedShapes)
               .succeeded()) {
-        for (auto it : llvm::zip(op->getResults(), returnedShapes)) {
+        for (auto it : llvm::zip(op.getResults(), returnedShapes)) {
           Value result = std::get<0>(it);
           ShapedTypeComponents predictedShape = std::get<1>(it);
 
           // Check whether this use case is replaceable. We define an op as
           // being replaceable if it is used by a ReturnOp or a TosaOp.
           bool replaceable = true;
           for (auto user : result.getUsers()) {
             if (isa<ReturnOp>(user))
               continue;
             if (user->getDialect()->getNamespace() ==
                 tosa::TosaDialect::getDialectNamespace())
               continue;
 
             replaceable = false;
           }
 
           // Determine the knowledge based on the output type.
           Type resultTy = result.getType();
           auto currentKnowledge =
               ValueKnowledge::getKnowledgeFromType(resultTy);
 
           // Compute the knowledge based on the inferred type.
-          auto inferredKnowledge =
-              ValueKnowledge::getPessimisticValueState(op->getContext());
+          auto inferredKnowledge = ValueKnowledge::getPessimisticValueState();
           inferredKnowledge.dtype =
               resultTy.cast<ShapedType>().getElementType();
-          inferredKnowledge.hasSizes = predictedShape.hasRank();
+          inferredKnowledge.hasRank = predictedShape.hasRank();
           if (predictedShape.hasRank()) {
             for (auto dim : predictedShape.getDims()) {
               inferredKnowledge.sizes.push_back(dim);
             }
           }
 
           if (!replaceable)
             continue;
 
           // Compute the new type based on the joined version.
           auto newKnowledge =
               ValueKnowledge::join(currentKnowledge, inferredKnowledge);
+          if (!newKnowledge)
+            continue;
           result.setType(newKnowledge.getType());
         }
       }
-    });
+    }
+  }
+}
+
+/// Pass that performs shape propagation across TOSA operations. This includes
+/// migrating to within the regions of if/while operations.
+struct TosaInferShapes : public TosaInferShapesBase<TosaInferShapes> {
+public:
+  void runOnFunction() override {
+    FuncOp func = getOperation();
+
+    IRRewriter rewriter(func.getContext());
+
+    propagateShapesInRegion(func.body());
 
     // Insert UnrealizedConversionCasts to guarantee ReturnOp agress with
     // the FuncOp type.
     func.walk([&](ReturnOp op) {
       FuncOp parent = dyn_cast<FuncOp>(op->getParentOp());
       if (!parent)
         return;
 

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

 func @test_pool_stride(%arg0: tensor<3x11x12x7xf32>) {
   // CHECK: -> tensor<3x4x4x7xf32>
   %0 = "tosa.avg_pool2d"(%arg0) {kernel = [4, 3], pad = [0, 0, 0, 0], stride = [2, 3]} : (tensor<3x11x12x7xf32>) -> tensor<?x?x?x?xf32>
 
   // CHECK: -> tensor<3x4x4x7xf32>
   %1 = "tosa.max_pool2d"(%arg0) {kernel = [4, 3], pad = [0, 0, 0, 0], stride = [2, 3]} : (tensor<3x11x12x7xf32>) -> tensor<?x?x?x?xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: @if_test_simple
+func @if_test_simple(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> () {
+  // CHECK: (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+  ^bb1(%arg3 : tensor<f32>, %arg4 : tensor<f32>):
+    "tosa.yield"(%arg3) : (tensor<f32>) -> ()
+  }, {
+  ^bb1(%arg5 : tensor<f32>, %arg6 : tensor<f32>):
+    "tosa.yield"(%arg6) : (tensor<f32>) -> ()
+  }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<*xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @if_test_dynamic
+func @if_test_dynamic(%arg0 : tensor<2xf32>, %arg1 : tensor<3xf32>, %arg2 : tensor<i1>) -> () {
+  // CHECK: (tensor<i1>, tensor<2xf32>, tensor<3xf32>) -> tensor<?xf32>
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+  ^bb1(%arg3 : tensor<2xf32>, %arg4 : tensor<3xf32>):
+    "tosa.yield"(%arg3) : (tensor<2xf32>) -> ()
+  }, {
+  ^bb1(%arg5 : tensor<2xf32>, %arg6 : tensor<3xf32>):
+    "tosa.yield"(%arg6) : (tensor<3xf32>) -> ()
+  }) : (tensor<i1>, tensor<2xf32>, tensor<3xf32>) -> (tensor<*xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @if_test_unranked
+func @if_test_unranked(%arg0 : tensor<f32>, %arg1 : tensor<3xf32>, %arg2 : tensor<i1>) -> () {
+  // CHECK: (tensor<i1>, tensor<f32>, tensor<3xf32>) -> tensor<*xf32>
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+  ^bb1(%arg3 : tensor<f32>, %arg4 : tensor<3xf32>):
+    "tosa.yield"(%arg3) : (tensor<f32>) -> ()
+  }, {
+  ^bb1(%arg5 : tensor<f32>, %arg6 : tensor<3xf32>):
+    "tosa.yield"(%arg6) : (tensor<3xf32>) -> ()
+  }) : (tensor<i1>, tensor<f32>, tensor<3xf32>) -> (tensor<*xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @if_test_propagate
+func @if_test_propagate(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> () {
+  // CHECK: (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+  ^bb1(%arg3 : tensor<*xf32>, %arg4 : tensor<*xf32>):
+    %1 = "tosa.add"(%arg3, %arg4) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+    "tosa.yield"(%1) : (tensor<*xf32>) -> ()
+  }, {
+  ^bb1(%arg5 : tensor<*xf32>, %arg6 : tensor<*xf32>):
+    %1 = "tosa.sub"(%arg5, %arg6) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+    "tosa.yield"(%1) : (tensor<*xf32>) -> ()
+  }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<*xf32>)
+  return
+}