Page MenuHomePhabricator

[mlir][tensor] Support more cases in MergeConsecutiveExtractSlice
ClosedPublic

Authored by antiagainst on Sep 20 2022, 11:08 AM.

Details

Summary

This commit adds utility functions to perform general merging of
OffsetSizeAndStrideOpInterface by supporting producer rank
reducing and non-unit strides.

With it we can extend MergeConsecutiveExtractSlice to support
more cases.

Diff Detail

Event Timeline

antiagainst created this revision.Sep 20 2022, 11:08 AM
antiagainst requested review of this revision.Sep 20 2022, 11:08 AM
Herald added a project: Restricted Project. · View Herald TranscriptSep 20 2022, 11:08 AM
This revision is now accepted and ready to land.Sep 20 2022, 3:26 PM
This revision was landed with ongoing or failed builds.Sep 20 2022, 5:16 PM
This revision was automatically updated to reflect the committed changes.

Seems like the utility function mergeOffsetsSizesAndStrides should be added to ViewLikeInterface or Dialect/Utils because it has greater scope. The primary use that comes to mind is that the memref "fold alias ops" (used to be "fold subviews") essential does the same thing (SubView and ExtractSlice are both implement the OffsetSizeAndStrideOpInterface). Therefore, it seems like an opportunity for code de-duplication.

Edit: I tried doing this in https://reviews.llvm.org/D134393 and run into the issue with Interface/ViewLikeInterface depending on Dialect/Utils and use of Affine ops. So the only common place it could go to be used by Memref and Tensor dialects is in affine dialect utils.

mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp
23

All three of these helper functions seem superfluous because the the makeComposedFoldedAffineApply does the optimization in getAffineExpr and the expressions aren't complicated enough to require the add and mul helpers. I tested the below code:

LogicalResult mlir::mergeOffsetsSizesAndStrides(
    OpBuilder &builder, Location loc, ArrayRef<OpFoldResult> producerOffsets,
    ArrayRef<OpFoldResult> producerSizes,
    ArrayRef<OpFoldResult> producerStrides,
    const llvm::SmallBitVector &droppedProducerDims,
    ArrayRef<OpFoldResult> consumerOffsets,
    ArrayRef<OpFoldResult> consumerSizes,
    ArrayRef<OpFoldResult> consumerStrides,
    SmallVector<OpFoldResult> &combinedOffsets,
    SmallVector<OpFoldResult> &combinedSizes,
    SmallVector<OpFoldResult> &combinedStrides) {
  combinedOffsets.resize(producerOffsets.size());
  combinedSizes.resize(producerOffsets.size());
  combinedStrides.resize(producerOffsets.size());

  AffineExpr s0, s1, d0;
  bindDims(builder.getContext(), d0);
  bindSymbols(builder.getContext(), s0, s1);

  unsigned consumerPos = 0;
  for (auto i : llvm::seq<unsigned>(0, producerOffsets.size())) {
    if (droppedProducerDims.test(i)) {
      // For dropped dims, get the values from the producer.
      combinedOffsets[i] = producerOffsets[i];
      combinedSizes[i] = producerSizes[i];
      combinedStrides[i] = producerStrides[i];
      continue;
    }
    SmallVector<OpFoldResult> offsetSymbols, strideSymbols;
    // The combined offset is computed as
    //    producer_offset + consumer_offset * producer_strides.
    combinedOffsets[i] = makeComposedFoldedAffineApply(
        builder, loc, d0 * s0 + s1,
        {consumerOffsets[consumerPos], producerStrides[i], producerOffsets[i]});
    combinedSizes[i] = consumerSizes[consumerPos];
    // The combined stride is computed as
    //    consumer_stride * producer_stride.
    combinedStrides[i] = makeComposedFoldedAffineApply(
        builder, loc, d0 * s0,
        {consumerStrides[consumerPos], producerStrides[i]});

    consumerPos++;
  }
  return success();
}
mravishankar added inline comments.Sep 21 2022, 4:01 PM
mlir/lib/Dialect/Tensor/Transforms/MergeConsecutiveInsertExtractSlicePatterns.cpp
23

Thats great! This was developed in IREE first, and that time we didnt have makeComposedFoldedAffineApply. I did suspect that those would make these methods redundant.

Thanks @mehdi_amini for the revert! I'll address it tomorrow.

Thanks @christopherbate for the suggestion! Yeah I didn't put in the ViewLikeInterface impl cpp because of the affine op dependency. Didn't know about the makeComposedFoldedAffineApply; that's pretty nice! Let me update accordingly tomorrow.

Relanded https://reviews.llvm.org/rGbd81524e7f52c38ad0bc689934343a476e545265 and uploaded https://reviews.llvm.org/D134487 to move to a more common place to be shared with memref too.