[mlir][linalg] Use top down traversal for padding.

[gysit]

Summary

Pad the operation using a top down traversal. The top down traversal unlocks folding opportunities and dim op canonicalizations due to the introduced extract slice operation after the padded operation.

Depends On D114585

[mehdi_amini]

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

142

I have concerns about this: the top down / bottom up is supposed, I believe, to be only a crude heuristic and not something used to control the behavior like it is intended here.
The Greedy rewriter is still intended to be a greedy strategy and the comment here has me concerned.

Please revert, unless I misunderstood this change.

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

void runOnFunction() override {
     if (!anchorOpName.empty()) {
       tilingAndFusionPattern.add<LinalgTileAndFuseTensorOpsPattern>(
           anchorOpName, funcOp.getContext(), options, filter);
     } else {
       tilingAndFusionPattern.add<LinalgTileAndFuseTensorOpsPattern>(
           funcOp.getContext(), options, filter);
     }
     // Search the root operation using bottom up traversal.
-    GreedyRewriteConfig grc;
-    grc.useTopDownTraversal = false;
-    (void)applyPatternsAndFoldGreedily(funcOp,
-                                       std::move(tilingAndFusionPattern), grc);
+    GreedyRewriteConfig config;
+    config.useTopDownTraversal = false;
+    (void)applyPatternsAndFoldGreedily(
+        funcOp, std::move(tilingAndFusionPattern), config);
   }
 
   LinalgTilingAndFusionOptions options;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to apply pattern-based linalg tiling.
 struct LinalgStrategyTilePass
void runOnFunction() override {
     RewritePatternSet paddingPattern(funcOp.getContext());
     if (!anchorOpName.empty()) {
       paddingPattern.add<LinalgPaddingPattern>(
           anchorOpName, funcOp.getContext(), options, filter);
     } else {
       paddingPattern.add<LinalgPaddingPattern>(funcOp.getContext(), options,
                                                filter);
     }
-    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(paddingPattern))))
+    // Traverse the operations top down to pad producers before consumers. The
+    // extract slice operation introduced after every padded operation enables
+    // padding its consumers. Padding an operation whose producers have not been
+    // padded before fails due to the missing extract slice operations. In this
+    // case, the padding pattern increments the transformation marker without
+    // padding the operation. The top down traversal is thus not only a
+    // performance optimization but also needed to pad all operations along the
+    // use-def chains.
+    GreedyRewriteConfig config;
+    config.useTopDownTraversal = true;
+    if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(paddingPattern),
+                                            config)))
       signalPassFailure();
   }
 
   LinalgPaddingOptions options;
   LinalgTransformationFilter filter;
 };
 
 /// Configurable pass to apply pattern-based linalg generalization.

mlir/test/Dialect/Linalg/pad.mlir

func @dynamic_sizes(%arg0: tensor<?x?xf32>,
   //      MATMUL:   return %[[T8]]
   return %13 : tensor<?x?xf32>
 }
 
 // -----
 
 #map0 = affine_map<()[s0] -> (64, s0)>
 
+//      FILL:  pad_multiple
+// FILL-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<64x64xf32>
+func @pad_multiple(%arg0: tensor<64x64xf32>,
+                      %iv0 : index) -> tensor<?x?xf32> {
+  %cst = arith.constant 0.0 : f32
+  %size = affine.min #map0()[%iv0]
+  %0 = tensor.extract_slice %arg0[0, 0] [%size, %size] [1, 1] : tensor<64x64xf32> to tensor<?x?xf32>
+
+  // Check both fill operations are padded by the same pad tensor operation.
+  //      FILL:  %[[T0:.*]] = linalg.pad_tensor
+  //      FILL:  %[[T1:.*]] = linalg.fill(%{{.*}}, %[[T0]])
+  //      FILL:  %[[T2:.*]] = linalg.fill(%{{.*}}, %[[T1]])
+  //      FILL:  = tensor.extract_slice %[[T2]]
+  %1 = linalg.fill(%cst, %0) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+  %2 = linalg.fill(%cst, %1) : f32, tensor<?x?xf32> -> tensor<?x?xf32>
+  return %2 : tensor<?x?xf32>
+}
+
+// -----
+
+#map0 = affine_map<()[s0] -> (64, s0)>
+
 //      MATMUL:  compose_padding
 // MATMUL-SAME:    %[[ARG0:[0-9a-zA-Z]*]]: tensor<64x64xf32>
 func @compose_padding(%arg0: tensor<64x64xf32>,
                       %iv0 : index) -> tensor<?x?xf32> {
   %cst = arith.constant 0.0 : f32
 
   //      MATMUL:  %[[SIZE:.*]] = affine.min
   %size = affine.min #map0()[%iv0]