[Mlir-commits] [mlir] Fixes in 'tosa.reshape' lowering and folder (PR #85798)

[Spenser Bauman]

================
@@ -62,174 +115,113 @@ static bool findIntermediateShape(ArrayRef<int64_t> lhsShape,
     currLhsDim++;
   }
 
-  // If the iterators didn't reach the end and their leftover dimensions are not
-  // equal to 1 an intermediate shape was not found.
-  while (currLhsDim < lhsShape.size()) {
-    if (lhsShape[currLhsDim++] != 1) {
-      return false;
-    }
+  // Static shapes are guaranteed to be compatible by the op verifier, so all
+  // leftover dimensions should be 1.
+  for (; currLhsDim < lhsShape.size(); currLhsDim++) {
+    assert(lhsShape[currLhsDim] == 1);
   }
-
-  while (currRhsDim < rhsShape.size()) {
-    if (rhsShape[currRhsDim++] != 1) {
-      return false;
-    }
+  for (; currRhsDim < rhsShape.size(); currRhsDim++) {
+    assert(rhsShape[currRhsDim] == 1);
   }
-
-  return true;
+  
+  return lhsType.clone(intermediateShape);
 }
 
-static bool createReassociationMapsForCollapse(
-    PatternRewriter &rewriter, ArrayRef<int64_t> srcShape,
-    ArrayRef<int64_t> dstShape,
-    SmallVector<ReassociationExprs, 4> &reassociationMap, bool isDynamic) {
-
-  // If the shape is dynamic, create a map for collapsing into one dimension.
-  if (isDynamic) {
+SmallVector<ReassociationExprs>
+createReassociationMapForCollapse(OpBuilder &builder,
+                                  ArrayRef<int64_t> srcShape,
+                                  ArrayRef<int64_t> dstShape) {
+  if (ShapedType::isDynamicShape(srcShape) || ShapedType::isDynamicShape(dstShape)) {
+    assert(dstShape.size() == 1);
     SmallVector<AffineExpr, 2> exprs;
-    for (int i = 0, s = srcShape.size(); i < s; ++i)
-      exprs.push_back(rewriter.getAffineDimExpr(i));
-    reassociationMap = {exprs};
-    return true;
+    for (int64_t i = 0, s = srcShape.size(); i < s; ++i)
----------------
sabauma wrote:

I generally find `llvm::seq` more readable than fiddly low-level loops.

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