[mlir] Make sure linearizeCollapsedDims doesn't drop input map dims

[ThomasRaoux]

Summary

The new affine map generated by linearizeCollapsedDims should not drop
dimensions. We need to make sure we create a map with at least as many
dimensions as the source map. This prevents
FoldProducerReshapeOpByLinearization from generating invalid IR.

This solves regression in IREE due to https://github.com/llvm/llvm-project/commit/e4e4da86aff5606ef792d987a3ec85639219228c

[mravishankar]

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

395

I am not entirely sure about the number of symbols. The number of symbols should be the same as the number of symbols in {resultExprs}, so this might degenerate to that in most cases. In any case not blocking this on that.

[ThomasRaoux]

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

395

I don't fully understand but this was the reason for the previous change in https://github.com/llvm/llvm-project/commit/e4e4da86aff5606ef792d987a3ec85639219228c. Unfortunately there were no tests along with the commit.

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

for (auto en : llvm::zip(sourceShape.slice(indices[0], indices.size()),
         continue;
       sizes.push_back(std::get<0>(en));
       dimExprs.push_back(std::get<1>(en));
     }
     AffineExpr linearizedExpr =
         makeCanonicalStridedLayoutExpr(sizes, dimExprs, context);
     resultExprs.push_back(linearizedExpr);
   }
-  return AffineMap::inferFromExprList({resultExprs}).front();
+  // The new affine map cannot drop unused dimension but some new symbols may
+  // have been added. Create a map with at least as many dimensions/symbols as
+  // the original affine map.
+  int64_t maxDim = -1;
+  int64_t maxSym = -1;
+  getMaxDimAndSymbol<SmallVector<AffineExpr>>({resultExprs}, maxDim, maxSym);
+  unsigned numDims = std::max(unsigned(maxDim + 1), sourceMap.getNumDims());
+  unsigned numSyms = std::max(unsigned(maxSym + 1), sourceMap.getNumSymbols());
+  return AffineMap::get(numDims, numSyms, resultExprs, context);
 }
 
 // TensorExpandShapeOp is fusable with its consumer (i.e. reshape as a
 // producer). Fusing when operand has higher rank will require use of mods and
 // divs in the indexing maps of the fused op which would make it non-invertible.
 static bool isTensorReshapeOpFoldableByLinearization(
     TensorExpandShapeOp expandOp, AffineMap useIndexMap, bool asProducer) {
   if (!asProducer)

mlir/test/Dialect/Linalg/reshape_linearization_fusion.mlir

 }
 // CHECK-LABEL: func @generic_op_reshape_consumer_nofusion
 //  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?x5xf32>
 //  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?x5xf32>
 //       CHECK:   %[[NOFUSE:.+]] = linalg.generic
 //  CHECK-SAME:     ins(%[[ARG0]], %[[ARG1]]
 //       CHECK:   %[[RESULT:.+]] = linalg.tensor_collapse_shape %[[NOFUSE]]
 //       CHECK:   return %[[RESULT]]
+
+
+// -----
+
+func @generic_op_permultation_reshape_consumer_fusion_unused_dim(%arg0 : tensor<6x1xf32>) -> tensor<6xi32> {
+  %0 = linalg.init_tensor [6, 1] : tensor<6x1xi32>
+  %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                        affine_map<(d0, d1) -> (d0, d1)>],
+   iterator_types = ["parallel", "parallel"]}
+   ins(%arg0 : tensor<6x1xf32>) outs(%0 : tensor<6x1xi32>) {
+    ^bb0(%arg3: f32, %arg4: i32):  // no predecessors
+      %5 = arith.fptosi %arg3 : f32 to i32
+      linalg.yield %5 : i32
+    } -> tensor<6x1xi32>
+    %6 = linalg.tensor_collapse_shape %1 [[0, 1]] : tensor<6x1xi32> into tensor<6xi32>
+  return %6 : tensor<6xi32>
+}
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0)>
+//       CHECK: func @generic_op_permultation_reshape_consumer_fusion_unused_dim
+//  CHECK-SAME:   %[[ARG0:.+]]: tensor<6x1xf32>
+//       CHECK:   %[[T0:.+]] = linalg.init_tensor [6, 1]
+//       CHECK:   %[[T1:.+]] = linalg.tensor_collapse_shape %[[T0]]
+//  CHECK-SAME:   [0, 1]
+//       CHECK:   linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP0]], #[[MAP1]]]
+//  CHECK-SAME:     ins(%[[ARG0]] : tensor<6x1xf32>)
+//  CHECK-SAME:     outs(%[[T1]] : tensor<6xi32>)