[mlir][linalg] Support dropping unit dimensions for init tensors

[antiagainst]

Summary

init tensor operands also has indexing map and generally follow
the same constraints we expect for non-init-tensor operands.

Depends On D99114

[mravishankar]

LGTM after minor nit is addressed

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

297

Could you leave the TODO for reductions? AFAIK, this is something you are going to look into as well?

[antiagainst]

Add TODO

[nicolasvasilache]

this is related to another CL I saw pass by, I think in general dropping 1 dimensions can/should be done with rank-reducing subtensor / subview as they have clear inplace semantics, including in the dynamic case.
reshape is unfortunately much trickier and if we can avoid using it when alternatives exist that would be preferred.

I do not know the deeper implications of systematically dropping the reshapes on IREE though.

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

 namespace {
 
 /// Pattern to replace tensors operands/results that are unit extents.
 template <typename GenericOpTy>
 struct ReplaceUnitExtentTensors : public OpRewritePattern<GenericOpTy> {
   using OpRewritePattern<GenericOpTy>::OpRewritePattern;
   LogicalResult matchAndRewrite(GenericOpTy op,
                                 PatternRewriter &rewriter) const override {
-    // TODO: support init_tensors and reductions.
-    if (!op.hasTensorSemantics() || op.getNumInitTensors() != 0)
+    if (!op.hasTensorSemantics())
       return failure();
 
     MLIRContext *context = rewriter.getContext();
     Location loc = op.getLoc();
 
     SmallVector<AffineMap, 4> newIndexingMaps;
     SmallVector<ArrayAttr, 4> reassociationMaps;
     SmallVector<ShapedType, 4> newInputOutputTypes;

mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir

%2 = linalg.generic {i64, indexing_maps = [#map1, #map0],
     } -> tensor<1x2x5xf32>
   %3 = linalg.tensor_reshape %2 [#map3, #map4]
     : tensor<1x2x5xf32> into tensor<2x5xf32>
   return %3 : tensor<2x5xf32>
 }
 // CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op
 //       CHECK:   %[[RESULT:.+]] = linalg.generic
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+func @fold_unit_dim_for_init_tensor(%input: tensor<1x1000xf32>) -> tensor<1xf32> {
+  %cst = constant 0.0 : f32
+  %init = linalg.init_tensor [1] : tensor<1xf32>
+  %fill = linalg.fill(%init, %cst) : tensor<1xf32>, f32 -> tensor<1xf32>
+  %add = linalg.generic {
+      indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
+      iterator_types = ["parallel", "reduction"]}
+    ins(%input : tensor<1x1000xf32>)outs(%fill : tensor<1xf32>) {
+  ^bb0(%arg1: f32, %arg2: f32):
+    %1823 = addf %arg1, %arg2 : f32
+    linalg.yield %1823 : f32
+  } -> tensor<1xf32>
+  return %add : tensor<1xf32>
+}
+
+
+//   CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>
+//   CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>
+//   CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>
+
+//       CHECK: func @fold_unit_dim_for_init_tensor
+
+//       CHECK: %[[INPUT_RESHAPE:.+]] = linalg.tensor_reshape %{{.+}} [#[[MAP0]]] : tensor<1x1000xf32> into tensor<1000xf32>
+//       CHECK: %[[INIT_RESHAPE:.+]] = linalg.tensor_reshape %{{.+}} [] : tensor<1xf32> into tensor<f32>
+//       CHECK: %[[GENERIC:.+]] = linalg.generic
+//  CHECK-SAME:     indexing_maps = [#[[MAP1]], #[[MAP2]]]
+//  CHECK-SAME:     iterator_types = ["reduction"]
+//  CHECK-SAME:   ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>)
+//  CHECK-SAME:   outs(%[[INIT_RESHAPE]] : tensor<f32>)
+//       CHECK: %[[GENERIC_RESHAPE:.+]] = linalg.tensor_reshape %[[GENERIC]] [] : tensor<f32> into tensor<1xf32>
+//       CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32>