[mlir][Linalg] Generalize the logic to compute reassociation maps

[mlir][Linalg] Generalize the logic to compute reassociation maps
while folding tensor_reshape op.

While folding reshapes that introduce unit extent dims, the logic to
compute the reassociation maps can be generalized to handle some
corner cases, for example, when the folded shape still has unit-extent
dims but corresponds to folded unit extent dims of the expanded shape.

Differential Revision: https://reviews.llvm.org/D88521