[Mlir-commits] [mlir] Let memref.collapse_shape implement ReifyRankedShapedTypeOpInterface. (PR #107752)

[llvmlistbot at llvm.org]

================
@@ -2497,6 +2497,42 @@ MemRefType CollapseShapeOp::computeCollapsedType(
                          srcType.getMemorySpace());
 }
 
+static bool isDynamicInGroup(ReassociationIndices group,
+                             ArrayRef<int64_t> sourceShape) {
+  return llvm::any_of(group, [sourceShape](int64_t dim) {
+    return ShapedType::isDynamic(sourceShape[dim]);
+  });
+}
+
+// This method supports following cases only:
+// - There is dynamic dimension in reassociation groups with single element.
+LogicalResult CollapseShapeOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedResultShapes) {
+  SmallVector<ReassociationIndices, 4> reassociationArray =
+      getReassociationIndices();
+  Value source = getSrc();
+  auto sourceShape = cast<MemRefType>(source.getType()).getShape();
+  if (!ShapedType::isDynamicShape(sourceShape))
+    return failure();
+  for (auto group : enumerate(reassociationArray)) {
+    bool isDynInGroup = isDynamicInGroup(group.value(), sourceShape);
+    if (isDynInGroup && group.value().size() > 1)
----------------
MaheshRavishankar wrote:

I dont think we need this restriction either. You can always construct an expression for the result dimension in terms of source dimensions... So there seems to be little reason to bail

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