[mlir][linalg] Add support for folding pack(fill) into fill.

[hanchung]

clang-format

[qedawkins]

Mostly LGTM. Also I have a question about how best to handle patterns for named linalg ops going forward.

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

491–492

Can we forward the new pack op destination to the fill op if the result is a tensor.empty (same as we're doing for generic elementwise in D149250)? This should be fine because the padding value and fill value are the same.

492

Once padding is supported (with some toggle option) in propagation through generics, would it be better to let the linalg.fill building happen via a callback specified in BubbleUpPackOpThroughFillOpPattern? Then we can reuse the pattern for generic ops for any linalg op which could help consolidate code at the expense of extra analysis in certain cases like this.

[chelini]

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

481

nit, I would write: SmallVector<Value> operands = {fillOp.getDpsInitOperand(0)->get(), packOpDest};

486

nit: operands.append(packOp.getInnerTiles().begin(), packOp.getInnerTiles().end());

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

843

Do we plan to fold the pack into the fill op in another PR? I would avoid the packing and make the fill operation work directly on the packed layout. Do you have some use cases where you want to preserve the packing? If so, we could do this folding in populateSimplifyTensorPack.

[hanchung]

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

491–492

I think we can drop the pack op directly after I read chelini's comment. I can't think a case that we'd like to preserve the new pack op because it's only used by the new fill op. What do you think?

492

Propagation with padding value is tricky, esp. for generic ops. The fill op case is fairly simple, so I'd like to scope it to the pattern at this moment.

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

843

Sounds really good to me, maybe we can do it in this PR..

[hanchung]

address commets

[qedawkins]

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

491–492

That's probably fine too, although I think they work out to be essentially the same thing whenever the fill is on an empty tensor (with an extra floating pack that can get DCE'd). Correct me if I'm wrong about that though.

My only concern was if the fill op destination was a non-empty tensor for whatever reason and we would then want to maintain the use-def chain similar to bubbling up through pure elementwise (as noted in D149250, given that fill is just a specific kind of pure elementwise). The difference is that filling a tensor is only sensible on an empty tensor then? Basically I don't see why a named op gets to be propagated differently than an equivalent generic (without padding). At a minimum I would expect us to be using the LinalgFillOpInterface for this then.

[hanchung]

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

491–492

My only concern was if the fill op destination was a non-empty tensor for whatever reason...

It looks fine to me because all the things in tensors world are destructive. Filling on whatever tensors (maybe except bufferization.alloc_tensor op) should create a new tensor with filled value. The pass is intended to be applied at tensor graph level, which is at a higher level than bufferization. So I think it is fine.

The generic op version of the fill op is very rare to me.. I agree that it could be a case, but we don't have to expose it at this moment. I have an idea for supporting padding cases in limit cases, and it could cover the case.

At a minimum I would expect us to be using the LinalgFillOpInterface for this then.

I was not aware of the interface. I studied it, and found that we can't use it here. The property is having a scalar value operand and one output operand. The computation body can be different, e.g., we could define FillExpOp using the interface which fills the arith.exp(scalar) into the output operand. In the context, we'll need different check between padding value and filled value. This interface is only used by linalg.fill op now, I think we can scope it to linalg.fill at this moment. If we really have other needs, it should be easy to extend.

[hanchung]

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

491–492

I'm open to keep the pack chain; fold them away only if the dest tensor is a tensor.empty op. There could be other use cases that I haven't seen before.

[qedawkins]

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

491–492

I see, this is fine then and I didn't intend to block. I just figured named linalg ops should try to be as close as possible to the generic patterns (the only difference being the pack chain in this case), but agreed that it would be strange usage of a fill op that I haven't seen either. I can go either way on keeping the pack chain because the pattern looks cleaner without it. Thanks for the discussion!

[hanchung]

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

491–492

I'm happy that we have the nice discussion! (I planned to wait for a day just in case if others would like to review it. No worry at all!)

[chelini]

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

491–492

That's probably fine too, although I think they work out to be essentially the same thing whenever the fill is on an empty tensor (with an extra floating pack that can get DCE'd). Correct me if I'm wrong about that though.

My only concern was if the fill op destination was a non-empty tensor for whatever reason and we would then want to maintain the use-def chain similar to bubbling up through pure elementwise (as noted in D149250, given that fill is just a specific kind of pure elementwise). The difference is that filling a tensor is only sensible on an empty tensor then? Basically I don't see why a named op gets to be propagated differently than an equivalent generic (without padding). At a minimum I would expect us to be using the LinalgFillOpInterface for this then.

Thanks for all these discussions! In D149250, I attempted to avoid packing an 'init' operand if the 'init' is already a tensor.empty. In this case, empty carries only the shape; we can update it with the packed version. I start wondering if the into keyword in tensor.pack is confusing. Internally there was some discussion about it because people think we are writing into the destination tensor. Let's discuss this offline in a call!

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

843

Why do we want to pass %dest as arg for the fill, seems a bit wired a tensor level. I would create a tensor.empty inside the func.

859

same here.

[qedawkins]

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

491–492

I start wondering if the into keyword in tensor.pack is confusing. Internally there was some discussion about it because people think we are writing into the destination tensor. Let's discuss this offline in a call!

Agreed, that has confused me as well; also happy to discuss in a call!

[hanchung]

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

491–492

agree, into is a confusing term. We can maybe use ins and outs because it's intended to capture the shape information. happy to discuss it!

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

843

It is not a big concern to me because sometimes a function can be inlined into a program. My intention was to create a minimal test for the pass. We can do this for static cases, but not dynamic shapes. We'll have to tie the dim sizes to something. Passing it as an argument is simplest way I've found. Let's do it for static shapes, and leave the dynamic case as what it is..

[hanchung]

update tests

[hanchung]

update tests

[nicolasvasilache]

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

453

I had the impression that @qcolombet had written something like this ~1 month ago ?

Signaling in case there is a dedup that needs to happen here.

[qcolombet]

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

453

We definitely talked about it, but no I didn't land anything on that.

[qcolombet]

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

453

I think what you have in mind @nicolasvasilache is pad + fill: https://github.com/openxla/iree/blob/ef2bb52317eed3b90618a6c16616ae14e33b7803/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp#L254

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

bubbleUpPackOpThroughGenericOp(RewriterBase &rewriter, tensor::PackOp packOp,
   Value dest = (genericOp.getRegionOutputArgs()[0].use_empty())
                    ? packOpDest
                    : packedOutOperand;
 
   return packElementWiseOp(rewriter, genericOp, dest, packedOutIndexingMap,
                            *packInfo);
 }
 
+/// Folds pack(fill) into a single fill op if
+///   1. The pack op does not have padding value, or
+///   2. The filled value and padding value are the same.
+static FailureOr<FillOp>
+foldFillPackIntoFillOp(RewriterBase &rewriter, tensor::PackOp packOp,
+                       ControlPropagationFn controlFn) {
+  auto fillOp = packOp.getSource().getDefiningOp<FillOp>();
+  if (!fillOp)
+    return failure();
+
+  // User controlled propagation function.
+  if (!controlFn(fillOp))
+    return failure();
+
+  if (auto paddingValue = packOp.getPaddingValue())
+    if (!isEqualConstantIntOrValue(paddingValue, fillOp.value()))
+      return failure();
+
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(fillOp);
+
+  Value packOpDest = packOp.getDest();
+  if (!packOpDest.hasOneUse())
+    return failure();
+  if (auto emptyOp = packOpDest.getDefiningOp<tensor::EmptyOp>()) {
+    packOpDest = tensor::PackOp::createDestinationTensor(
+        rewriter, fillOp.getLoc(), fillOp.getDpsInitOperand(0)->get(),
+        packOp.getMixedTiles(), packOp.getInnerDimsPos(),
+        packOp.getOuterDimsPerm());
+  } else {
+    DominanceInfo dom(fillOp);
+    if (!dom.properlyDominates(packOpDest, fillOp))
+      return failure();
+  }
+
+  Value fillDest = packOpDest;
+  return clone(rewriter, fillOp, packOpDest.getType(),
+               {fillOp.value(), fillDest});
+}
+
 /// Wrapper pattern that applies bubbleUpPackOpThroughGenericOp method.
 struct BubbleUpPackOpThroughGenericOpPattern
     : public OpRewritePattern<tensor::PackOp> {
 public:
   BubbleUpPackOpThroughGenericOpPattern(MLIRContext *context,
                                         ControlPropagationFn fun)
       : OpRewritePattern<tensor::PackOp>(context), controlFn(std::move(fun)) {}
 
   LogicalResult matchAndRewrite(tensor::PackOp packOp,
                                 PatternRewriter &rewriter) const override {
     auto genericOp =
         bubbleUpPackOpThroughGenericOp(rewriter, packOp, controlFn);
     if (failed(genericOp))
       return failure();
     rewriter.replaceOp(packOp, genericOp->getResults());
     return success();
   }
 
 private:
   ControlPropagationFn controlFn;
 };
 
+/// Wrapper pattern that applies foldFillPackIntoFillOp method.
+struct FoldFillPackIntoFillOpPattern : public OpRewritePattern<tensor::PackOp> {
+public:
+  FoldFillPackIntoFillOpPattern(MLIRContext *context, ControlPropagationFn fun)
+      : OpRewritePattern<tensor::PackOp>(context), controlFn(std::move(fun)) {}
+
+  LogicalResult matchAndRewrite(tensor::PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    auto fillOp = foldFillPackIntoFillOp(rewriter, packOp, controlFn);
+    if (failed(fillOp))
+      return failure();
+    rewriter.replaceOp(packOp, fillOp.value().result());
+    return success();
+  }
+
+private:
+  ControlPropagationFn controlFn;
+};
+
 // TODO: Relax this restriction. We should unpack an elementwise also
 // in the presence of multiple unpack ops as producers.
 /// Return the unpacked operand, if present, for the current generic op.
 static FailureOr<OpOperand *> getUnPackedOperand(GenericOp genericOp) {
   OpOperand *unPackedOperand = nullptr;
   for (OpOperand &operand : genericOp->getOpOperands()) {
     auto unPackOp = operand.get().getDefiningOp<tensor::UnPackOp>();
     if (!unPackOp)

 };
 
 } // namespace
 
 void mlir::linalg::populateDataLayoutPropagationPatterns(
     RewritePatternSet &patterns,
     const ControlPropagationFn &controlPackUnPackPropagation) {
   patterns.insert<BubbleUpPackOpThroughGenericOpPattern,
+                  FoldFillPackIntoFillOpPattern,
                   PushDownUnPackOpThroughGenericOp, PushDownUnPackThroughPadOp>(
       patterns.getContext(), controlPackUnPackPropagation);
 }

mlir/test/Dialect/Linalg/data-layout-propagation.mlir

 // CHECK-SAME:     iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "parallel"]
 // CHECK-SAME:     ins(%[[PACK_ARG0]], %[[FILTER]]
 // CHECK-SAME:     outs(%[[INIT]]
 // CHECK:        %[[UNPACK_NEW_DEST:.+]] = tensor.empty() : tensor<16x540x960xi32>
 // CHECK:        %[[UNPACK:.+]] = tensor.unpack %[[POOL]]
 // CHECK-SAME:     inner_dims_pos = [0] inner_tiles = [16]
 // CHECK-SAME:     into %[[UNPACK_NEW_DEST]]
 // CHECK:        return %[[UNPACK]] : tensor<16x540x960xi32>
+
+// -----
+
+func.func @fill_pack() -> tensor<24x32x16x16xf32> {
+  %dest = tensor.empty() : tensor<384x512xf32>
+  %cst = arith.constant 0.000000e+00 : f32
+  %0 = tensor.empty() : tensor<24x32x16x16xf32>
+  %1 = linalg.fill ins(%cst : f32) outs(%dest : tensor<384x512xf32>) -> tensor<384x512xf32>
+  %pack = tensor.pack %1 inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %0 : tensor<384x512xf32> -> tensor<24x32x16x16xf32>
+  return %pack : tensor<24x32x16x16xf32>
+}
+// CHECK-LABEL: func.func @fill_pack
+// CHECK:         %[[PACKED_EMPTY:.+]] = tensor.empty() : tensor<24x32x16x16xf32>
+// CHECK:         %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
+// CHECK:         return %[[FILL]]
+
+// -----
+
+#map = affine_map<()[s0] -> (s0 ceildiv 16)>
+func.func @dynamic_fill_pack(%arg0: tensor<?x?xf32>) -> tensor<?x?x16x16xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = linalg.fill ins(%cst : f32) outs(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %dim = tensor.dim %0, %c0 : tensor<?x?xf32>
+  %dim_0 = tensor.dim %0, %c1 : tensor<?x?xf32>
+  %1 = affine.apply #map()[%dim]
+  %2 = affine.apply #map()[%dim_0]
+  %3 = tensor.empty(%1, %2) : tensor<?x?x16x16xf32>
+  %pack = tensor.pack %0 padding_value(%cst : f32) inner_dims_pos = [0, 1] inner_tiles = [16, 16] into %3 : tensor<?x?xf32> -> tensor<?x?x16x16xf32>
+  return %pack : tensor<?x?x16x16xf32>
+}
+// CHECK-DAG:   #[[MAP:.+]] = affine_map<()[s0] -> (s0 ceildiv 16)>
+// CHECK:       func.func @dynamic_fill_pack
+// CHECK-SAME:    %[[DEST:[a-zA-Z0-9]+]]
+// CHECK-DAG:     %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG:     %[[C1:.+]] = arith.constant 1 : index
+// CHECK:         %[[D0:.+]] = tensor.dim %[[DEST]], %[[C0]]
+// CHECK:         %[[D1:.+]] = tensor.dim %[[DEST]], %[[C1]]
+// CHECK:         %[[PACKED_D0:.+]] = affine.apply #[[MAP]]()[%[[D0]]]
+// CHECK:         %[[PACKED_D1:.+]] = affine.apply #[[MAP]]()[%[[D1]]]
+// CHECK:         %[[PACKED_EMPTY:.+]] = tensor.empty(%[[PACKED_D0]], %[[PACKED_D1]]) : tensor<?x?x16x16xf32>
+// CHECK:         %[[FILL:.+]] = linalg.fill ins(%{{.+}}) outs(%[[PACKED_EMPTY]]
+// CHECK:         return %[[FILL]]