[Mlir-commits] [mlir] [mlir] Re-land Loosen restrictions on folding dynamic reshapes (PR #142827)

[Kevin Gleason]

GleasonK wrote:

Hello! Integrating this change into StableHLO now and am hitting some issues on tensor scalars (`tensor<f32> -> tensor<1xf32>` and vice versa).

For a given `stablehlo.reshape %0 : tensor<f32> -> tensor<1xf32>`, it seems I need to nullify the reassociation map in order to map to tensor dialect:

```c++
std::optional<SmallVector<ReassociationIndices>> reassociationMap =
            getReassociationIndicesForReshape(operandType, resultType);
// reassociationMap = {{0}}  -- used to be {{}}

// Generate expand operation.
if (operandType.getRank() == 0) {
  // Seems to be needed now?? Else hit:
  // error: 'tensor.expand_shape' op expected collapsed rank (0) to equal the number of reassociation maps (1).
  reassociationMap = SmallVector<ReassociationIndices>();
}
rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(
    reshapeOp, resultType, operand, *reassociationMap);
```

Is this intentional? These APIs used to return `{{}}` ([ref](https://github.com/llvm/llvm-project/blob/517800e37e8d3a4ee84214bef65e227612c2a98b/mlir/lib/Dialect/Utils/ReshapeOpsUtils.cpp#L89)). Looking at other uses of `getReassociationIndicesForReshape` it looks like the intent is to be forwareded to tensor expand/collapse ops, which it seems it no longer can be for scalars. (Apologies if my understanding is entirely wrong, I've inherited this code :))

Code pointer in case its helpful: [StablehloLegalizeToLinalg.cpp#L1129](https://github.com/openxla/stablehlo/blob/9d0b3cc0a79ab058b78f8932bfc707c36fd315a2/stablehlo/conversions/linalg/transforms/StablehloLegalizeToLinalg.cpp#L1129).

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