[mlir] Fix infinite loop in collapse

[guraypp]

Summary

Incrementing counter variable is inside the if statement. If the code does not enter there, the while loop will iterate infinitely. This revision moves the codes outside of if statement.

[guraypp]

update

[mravishankar]

Thanks! Could you add a test here?

mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp

getOperandReassociation(AffineMap indexingMap,
   SmallVector<ReassociationIndices> operandReassociation;
   auto origOpToCollapsedOpMapping =
       collapsingInfo.getOrigOpToCollapsedOpMapping();
   auto collapsedOpToOrigOpMapping =
       collapsingInfo.getCollapsedOpToOrigOpMapping();
   while (counter < indexingMap.getNumResults()) {
     unsigned dim =
         indexingMap.getResult(counter).cast<AffineDimExpr>().getPosition();
-    if (origOpToCollapsedOpMapping[dim].second == 0) {
-      // This is the start of a collapsed dimensions of the iteration that
-      // is gauranteed to be preserved in the indexing map. The number of folded
-      // dims is obtained from the collapsed op to original op mapping.
-      unsigned numFoldedDims =
-          collapsedOpToOrigOpMapping[origOpToCollapsedOpMapping[dim].first]
-              .size();
-      auto range = llvm::seq<unsigned>(counter, counter + numFoldedDims);
-      operandReassociation.emplace_back(range.begin(), range.end());
-      counter += numFoldedDims;
-    }
-  }
   return operandReassociation;
 }
 
 /// Get the new value to use for a given `OpOperand` in the collapsed operation.
 static Value getCollapsedOpOperand(Location loc, GenericOp genericOp,
                                    OpOperand *opOperand,
                                    const CollapsingInfo &collapsingInfo,
                                    OpBuilder &builder) {