[mlir][MemRef] Add ExtractStridedMetadataOpCollapseShapeFolder #89954

dcaballe · 2024-04-24T17:16:01Z

This PR adds a new pattern to the set of patterns used to resolve the offset, sizes and stride of a memref. Similar to ExtractStridedMetadataOpSubviewFolder, the new pattern resolves strided_metadata(collapse_shape) directly, without introduce a reshape_cast op.

This PR adds a new pattern to the set of patterns used to resolve the offset, sizes and stride of a memref. Similar to `ExtractStridedMetadataOpSubviewFolder`, the new pattern resolves strided_metadata(collapse_shape) directly, without introduce a reshape_cast op.

llvmbot · 2024-04-24T17:16:35Z

@llvm/pr-subscribers-mlir-memref

Author: Diego Caballero (dcaballe)

Changes

This PR adds a new pattern to the set of patterns used to resolve the offset, sizes and stride of a memref. Similar to ExtractStridedMetadataOpSubviewFolder, the new pattern resolves strided_metadata(collapse_shape) directly, without introduce a reshape_cast op.

Full diff: https://github.com/llvm/llvm-project/pull/89954.diff

1 Files Affected:

(modified) mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp (+130-59)

diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
index 96eb7cfd2db690..b5578a58468e9c 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
@@ -550,6 +550,78 @@ getCollapsedStride(memref::CollapseShapeOp collapseShape, OpBuilder &builder,
   return {makeComposedFoldedAffineMin(builder, collapseShape.getLoc(), minMap,
                                       groupStrides)};
 }
+
+template <typename ReassociativeReshapeLikeOp,
+          SmallVector<OpFoldResult> (*getReshapedSizes)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/, unsigned /*groupId*/),
+          SmallVector<OpFoldResult> (*getReshapedStrides)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/,
+              ArrayRef<OpFoldResult> /*origStrides*/, unsigned /*groupId*/)>
+static FailureOr<StridedMetadata>
+resolveReshapeStridedMetadata(RewriterBase &rewriter,
+                              ReassociativeReshapeLikeOp reshape) {
+  // Build a plain extract_strided_metadata(memref) from
+  // extract_strided_metadata(reassociative_reshape_like(memref)).
+  Location origLoc = reshape.getLoc();
+  Value source = reshape.getSrc();
+  auto sourceType = cast<MemRefType>(source.getType());
+  unsigned sourceRank = sourceType.getRank();
+
+  auto newExtractStridedMetadata =
+      rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
+
+  // Collect statically known information.
+  auto [strides, offset] = getStridesAndOffset(sourceType);
+  MemRefType reshapeType = reshape.getResultType();
+  unsigned reshapeRank = reshapeType.getRank();
+
+  OpFoldResult offsetOfr =
+      ShapedType::isDynamic(offset)
+          ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
+          : rewriter.getIndexAttr(offset);
+
+  // Get the special case of 0-D out of the way.
+  if (sourceRank == 0) {
+    SmallVector<OpFoldResult> ones(reshapeRank, rewriter.getIndexAttr(1));
+    return StridedMetadata{newExtractStridedMetadata.getBaseBuffer(), offsetOfr,
+                           /*sizes=*/ones, /*strides=*/ones};
+  }
+
+  SmallVector<OpFoldResult> finalSizes;
+  finalSizes.reserve(reshapeRank);
+  SmallVector<OpFoldResult> finalStrides;
+  finalStrides.reserve(reshapeRank);
+
+  // Compute the reshaped strides and sizes from the base strides and sizes.
+  SmallVector<OpFoldResult> origSizes =
+      getAsOpFoldResult(newExtractStridedMetadata.getSizes());
+  SmallVector<OpFoldResult> origStrides =
+      getAsOpFoldResult(newExtractStridedMetadata.getStrides());
+  unsigned idx = 0, endIdx = reshape.getReassociationIndices().size();
+  for (; idx != endIdx; ++idx) {
+    SmallVector<OpFoldResult> reshapedSizes =
+        getReshapedSizes(reshape, rewriter, origSizes, /*groupId=*/idx);
+    SmallVector<OpFoldResult> reshapedStrides = getReshapedStrides(
+        reshape, rewriter, origSizes, origStrides, /*groupId=*/idx);
+
+    unsigned groupSize = reshapedSizes.size();
+    for (unsigned i = 0; i < groupSize; ++i) {
+      finalSizes.push_back(reshapedSizes[i]);
+      finalStrides.push_back(reshapedStrides[i]);
+    }
+  }
+  assert(((isa<memref::ExpandShapeOp>(reshape) && idx == sourceRank) ||
+          (isa<memref::CollapseShapeOp>(reshape) && idx == reshapeRank)) &&
+         "We should have visited all the input dimensions");
+  assert(finalSizes.size() == reshapeRank &&
+         "We should have populated all the values");
+
+  return StridedMetadata{newExtractStridedMetadata.getBaseBuffer(), offsetOfr,
+                         finalSizes, finalStrides};
+}
+
 /// Replace `baseBuffer, offset, sizes, strides =
 ///              extract_strided_metadata(reshapeLike(memref))`
 /// With
@@ -580,68 +652,66 @@ struct ReshapeFolder : public OpRewritePattern<ReassociativeReshapeLikeOp> {
 
   LogicalResult matchAndRewrite(ReassociativeReshapeLikeOp reshape,
                                 PatternRewriter &rewriter) const override {
-    // Build a plain extract_strided_metadata(memref) from
-    // extract_strided_metadata(reassociative_reshape_like(memref)).
-    Location origLoc = reshape.getLoc();
-    Value source = reshape.getSrc();
-    auto sourceType = cast<MemRefType>(source.getType());
-    unsigned sourceRank = sourceType.getRank();
-
-    auto newExtractStridedMetadata =
-        rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
-
-    // Collect statically known information.
-    auto [strides, offset] = getStridesAndOffset(sourceType);
-    MemRefType reshapeType = reshape.getResultType();
-    unsigned reshapeRank = reshapeType.getRank();
-
-    OpFoldResult offsetOfr =
-        ShapedType::isDynamic(offset)
-            ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
-            : rewriter.getIndexAttr(offset);
-
-    // Get the special case of 0-D out of the way.
-    if (sourceRank == 0) {
-      SmallVector<OpFoldResult> ones(reshapeRank, rewriter.getIndexAttr(1));
-      auto memrefDesc = rewriter.create<memref::ReinterpretCastOp>(
-          origLoc, reshapeType, newExtractStridedMetadata.getBaseBuffer(),
-          offsetOfr, /*sizes=*/ones, /*strides=*/ones);
-      rewriter.replaceOp(reshape, memrefDesc.getResult());
-      return success();
+    FailureOr<StridedMetadata> stridedMetadata =
+        resolveReshapeStridedMetadata<ReassociativeReshapeLikeOp,
+                                      getReshapedSizes, getReshapedStrides>(
+            rewriter, reshape);
+    if (failed(stridedMetadata)) {
+      return rewriter.notifyMatchFailure(reshape,
+                                         "failed to resolve reshape metadata");
     }
 
-    SmallVector<OpFoldResult> finalSizes;
-    finalSizes.reserve(reshapeRank);
-    SmallVector<OpFoldResult> finalStrides;
-    finalStrides.reserve(reshapeRank);
-
-    // Compute the reshaped strides and sizes from the base strides and sizes.
-    SmallVector<OpFoldResult> origSizes =
-        getAsOpFoldResult(newExtractStridedMetadata.getSizes());
-    SmallVector<OpFoldResult> origStrides =
-        getAsOpFoldResult(newExtractStridedMetadata.getStrides());
-    unsigned idx = 0, endIdx = reshape.getReassociationIndices().size();
-    for (; idx != endIdx; ++idx) {
-      SmallVector<OpFoldResult> reshapedSizes =
-          getReshapedSizes(reshape, rewriter, origSizes, /*groupId=*/idx);
-      SmallVector<OpFoldResult> reshapedStrides = getReshapedStrides(
-          reshape, rewriter, origSizes, origStrides, /*groupId=*/idx);
-
-      unsigned groupSize = reshapedSizes.size();
-      for (unsigned i = 0; i < groupSize; ++i) {
-        finalSizes.push_back(reshapedSizes[i]);
-        finalStrides.push_back(reshapedStrides[i]);
-      }
+    rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
+        reshape, reshape.getType(), stridedMetadata->basePtr,
+        stridedMetadata->offset, stridedMetadata->sizes,
+        stridedMetadata->strides);
+    return success();
+  }
+};
+
+/// Pattern to replace `extract_strided_metadata(collapse_shape)`
+/// With
+///
+/// \verbatim
+/// baseBuffer, baseOffset, baseSizes, baseStrides =
+///     extract_strided_metadata(memref)
+/// strides#i = baseStrides#i * subSizes#i
+/// offset = baseOffset + sum(subOffset#i * baseStrides#i)
+/// sizes = subSizes
+/// \verbatim
+///
+/// with `baseBuffer`, `offset`, `sizes` and `strides` being
+/// the replacements for the original `extract_strided_metadata`.
+struct ExtractStridedMetadataOpCollapseShapeFolder
+    : OpRewritePattern<memref::ExtractStridedMetadataOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
+                                PatternRewriter &rewriter) const override {
+    auto collapseShapeOp =
+        op.getSource().getDefiningOp<memref::CollapseShapeOp>();
+    if (!collapseShapeOp)
+      return failure();
+
+    FailureOr<StridedMetadata> stridedMetadata =
+        resolveReshapeStridedMetadata<memref::CollapseShapeOp, getCollapsedSize,
+                                      getCollapsedStride>(rewriter,
+                                                          collapseShapeOp);
+    if (failed(stridedMetadata)) {
+      return rewriter.notifyMatchFailure(
+          op, "failed to resolve metadata in terms of source collapse_shape op");
     }
-    assert(((isa<memref::ExpandShapeOp>(reshape) && idx == sourceRank) ||
-            (isa<memref::CollapseShapeOp>(reshape) && idx == reshapeRank)) &&
-           "We should have visited all the input dimensions");
-    assert(finalSizes.size() == reshapeRank &&
-           "We should have populated all the values");
-    auto memrefDesc = rewriter.create<memref::ReinterpretCastOp>(
-        origLoc, reshapeType, newExtractStridedMetadata.getBaseBuffer(),
-        offsetOfr, finalSizes, finalStrides);
-    rewriter.replaceOp(reshape, memrefDesc.getResult());
+
+    Location loc = collapseShapeOp.getLoc();
+    SmallVector<Value> results;
+    results.push_back(stridedMetadata->basePtr);
+    results.push_back(getValueOrCreateConstantIndexOp(rewriter, loc,
+                                                      stridedMetadata->offset));
+    results.append(
+        getValueOrCreateConstantIndexOp(rewriter, loc, stridedMetadata->sizes));
+    results.append(getValueOrCreateConstantIndexOp(rewriter, loc,
+                                                   stridedMetadata->strides));
+    rewriter.replaceOp(op, results);
     return success();
   }
 };
@@ -1030,6 +1100,7 @@ void memref::populateResolveExtractStridedMetadataPatterns(
     RewritePatternSet &patterns) {
   patterns.add<ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
                ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
+               ExtractStridedMetadataOpCollapseShapeFolder,
                ExtractStridedMetadataOpGetGlobalFolder,
                ExtractStridedMetadataOpSubviewFolder,
                RewriteExtractAlignedPointerAsIndexOfViewLikeOp,

llvmbot · 2024-04-24T17:16:35Z

@llvm/pr-subscribers-mlir

Author: Diego Caballero (dcaballe)

Changes

This PR adds a new pattern to the set of patterns used to resolve the offset, sizes and stride of a memref. Similar to ExtractStridedMetadataOpSubviewFolder, the new pattern resolves strided_metadata(collapse_shape) directly, without introduce a reshape_cast op.

Full diff: https://github.com/llvm/llvm-project/pull/89954.diff

1 Files Affected:

(modified) mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp (+130-59)

diff --git a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
index 96eb7cfd2db690..b5578a58468e9c 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp
@@ -550,6 +550,78 @@ getCollapsedStride(memref::CollapseShapeOp collapseShape, OpBuilder &builder,
   return {makeComposedFoldedAffineMin(builder, collapseShape.getLoc(), minMap,
                                       groupStrides)};
 }
+
+template <typename ReassociativeReshapeLikeOp,
+          SmallVector<OpFoldResult> (*getReshapedSizes)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/, unsigned /*groupId*/),
+          SmallVector<OpFoldResult> (*getReshapedStrides)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/,
+              ArrayRef<OpFoldResult> /*origStrides*/, unsigned /*groupId*/)>
+static FailureOr<StridedMetadata>
+resolveReshapeStridedMetadata(RewriterBase &rewriter,
+                              ReassociativeReshapeLikeOp reshape) {
+  // Build a plain extract_strided_metadata(memref) from
+  // extract_strided_metadata(reassociative_reshape_like(memref)).
+  Location origLoc = reshape.getLoc();
+  Value source = reshape.getSrc();
+  auto sourceType = cast<MemRefType>(source.getType());
+  unsigned sourceRank = sourceType.getRank();
+
+  auto newExtractStridedMetadata =
+      rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
+
+  // Collect statically known information.
+  auto [strides, offset] = getStridesAndOffset(sourceType);
+  MemRefType reshapeType = reshape.getResultType();
+  unsigned reshapeRank = reshapeType.getRank();
+
+  OpFoldResult offsetOfr =
+      ShapedType::isDynamic(offset)
+          ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
+          : rewriter.getIndexAttr(offset);
+
+  // Get the special case of 0-D out of the way.
+  if (sourceRank == 0) {
+    SmallVector<OpFoldResult> ones(reshapeRank, rewriter.getIndexAttr(1));
+    return StridedMetadata{newExtractStridedMetadata.getBaseBuffer(), offsetOfr,
+                           /*sizes=*/ones, /*strides=*/ones};
+  }
+
+  SmallVector<OpFoldResult> finalSizes;
+  finalSizes.reserve(reshapeRank);
+  SmallVector<OpFoldResult> finalStrides;
+  finalStrides.reserve(reshapeRank);
+
+  // Compute the reshaped strides and sizes from the base strides and sizes.
+  SmallVector<OpFoldResult> origSizes =
+      getAsOpFoldResult(newExtractStridedMetadata.getSizes());
+  SmallVector<OpFoldResult> origStrides =
+      getAsOpFoldResult(newExtractStridedMetadata.getStrides());
+  unsigned idx = 0, endIdx = reshape.getReassociationIndices().size();
+  for (; idx != endIdx; ++idx) {
+    SmallVector<OpFoldResult> reshapedSizes =
+        getReshapedSizes(reshape, rewriter, origSizes, /*groupId=*/idx);
+    SmallVector<OpFoldResult> reshapedStrides = getReshapedStrides(
+        reshape, rewriter, origSizes, origStrides, /*groupId=*/idx);
+
+    unsigned groupSize = reshapedSizes.size();
+    for (unsigned i = 0; i < groupSize; ++i) {
+      finalSizes.push_back(reshapedSizes[i]);
+      finalStrides.push_back(reshapedStrides[i]);
+    }
+  }
+  assert(((isa<memref::ExpandShapeOp>(reshape) && idx == sourceRank) ||
+          (isa<memref::CollapseShapeOp>(reshape) && idx == reshapeRank)) &&
+         "We should have visited all the input dimensions");
+  assert(finalSizes.size() == reshapeRank &&
+         "We should have populated all the values");
+
+  return StridedMetadata{newExtractStridedMetadata.getBaseBuffer(), offsetOfr,
+                         finalSizes, finalStrides};
+}
+
 /// Replace `baseBuffer, offset, sizes, strides =
 ///              extract_strided_metadata(reshapeLike(memref))`
 /// With
@@ -580,68 +652,66 @@ struct ReshapeFolder : public OpRewritePattern<ReassociativeReshapeLikeOp> {
 
   LogicalResult matchAndRewrite(ReassociativeReshapeLikeOp reshape,
                                 PatternRewriter &rewriter) const override {
-    // Build a plain extract_strided_metadata(memref) from
-    // extract_strided_metadata(reassociative_reshape_like(memref)).
-    Location origLoc = reshape.getLoc();
-    Value source = reshape.getSrc();
-    auto sourceType = cast<MemRefType>(source.getType());
-    unsigned sourceRank = sourceType.getRank();
-
-    auto newExtractStridedMetadata =
-        rewriter.create<memref::ExtractStridedMetadataOp>(origLoc, source);
-
-    // Collect statically known information.
-    auto [strides, offset] = getStridesAndOffset(sourceType);
-    MemRefType reshapeType = reshape.getResultType();
-    unsigned reshapeRank = reshapeType.getRank();
-
-    OpFoldResult offsetOfr =
-        ShapedType::isDynamic(offset)
-            ? getAsOpFoldResult(newExtractStridedMetadata.getOffset())
-            : rewriter.getIndexAttr(offset);
-
-    // Get the special case of 0-D out of the way.
-    if (sourceRank == 0) {
-      SmallVector<OpFoldResult> ones(reshapeRank, rewriter.getIndexAttr(1));
-      auto memrefDesc = rewriter.create<memref::ReinterpretCastOp>(
-          origLoc, reshapeType, newExtractStridedMetadata.getBaseBuffer(),
-          offsetOfr, /*sizes=*/ones, /*strides=*/ones);
-      rewriter.replaceOp(reshape, memrefDesc.getResult());
-      return success();
+    FailureOr<StridedMetadata> stridedMetadata =
+        resolveReshapeStridedMetadata<ReassociativeReshapeLikeOp,
+                                      getReshapedSizes, getReshapedStrides>(
+            rewriter, reshape);
+    if (failed(stridedMetadata)) {
+      return rewriter.notifyMatchFailure(reshape,
+                                         "failed to resolve reshape metadata");
     }
 
-    SmallVector<OpFoldResult> finalSizes;
-    finalSizes.reserve(reshapeRank);
-    SmallVector<OpFoldResult> finalStrides;
-    finalStrides.reserve(reshapeRank);
-
-    // Compute the reshaped strides and sizes from the base strides and sizes.
-    SmallVector<OpFoldResult> origSizes =
-        getAsOpFoldResult(newExtractStridedMetadata.getSizes());
-    SmallVector<OpFoldResult> origStrides =
-        getAsOpFoldResult(newExtractStridedMetadata.getStrides());
-    unsigned idx = 0, endIdx = reshape.getReassociationIndices().size();
-    for (; idx != endIdx; ++idx) {
-      SmallVector<OpFoldResult> reshapedSizes =
-          getReshapedSizes(reshape, rewriter, origSizes, /*groupId=*/idx);
-      SmallVector<OpFoldResult> reshapedStrides = getReshapedStrides(
-          reshape, rewriter, origSizes, origStrides, /*groupId=*/idx);
-
-      unsigned groupSize = reshapedSizes.size();
-      for (unsigned i = 0; i < groupSize; ++i) {
-        finalSizes.push_back(reshapedSizes[i]);
-        finalStrides.push_back(reshapedStrides[i]);
-      }
+    rewriter.replaceOpWithNewOp<memref::ReinterpretCastOp>(
+        reshape, reshape.getType(), stridedMetadata->basePtr,
+        stridedMetadata->offset, stridedMetadata->sizes,
+        stridedMetadata->strides);
+    return success();
+  }
+};
+
+/// Pattern to replace `extract_strided_metadata(collapse_shape)`
+/// With
+///
+/// \verbatim
+/// baseBuffer, baseOffset, baseSizes, baseStrides =
+///     extract_strided_metadata(memref)
+/// strides#i = baseStrides#i * subSizes#i
+/// offset = baseOffset + sum(subOffset#i * baseStrides#i)
+/// sizes = subSizes
+/// \verbatim
+///
+/// with `baseBuffer`, `offset`, `sizes` and `strides` being
+/// the replacements for the original `extract_strided_metadata`.
+struct ExtractStridedMetadataOpCollapseShapeFolder
+    : OpRewritePattern<memref::ExtractStridedMetadataOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::ExtractStridedMetadataOp op,
+                                PatternRewriter &rewriter) const override {
+    auto collapseShapeOp =
+        op.getSource().getDefiningOp<memref::CollapseShapeOp>();
+    if (!collapseShapeOp)
+      return failure();
+
+    FailureOr<StridedMetadata> stridedMetadata =
+        resolveReshapeStridedMetadata<memref::CollapseShapeOp, getCollapsedSize,
+                                      getCollapsedStride>(rewriter,
+                                                          collapseShapeOp);
+    if (failed(stridedMetadata)) {
+      return rewriter.notifyMatchFailure(
+          op, "failed to resolve metadata in terms of source collapse_shape op");
     }
-    assert(((isa<memref::ExpandShapeOp>(reshape) && idx == sourceRank) ||
-            (isa<memref::CollapseShapeOp>(reshape) && idx == reshapeRank)) &&
-           "We should have visited all the input dimensions");
-    assert(finalSizes.size() == reshapeRank &&
-           "We should have populated all the values");
-    auto memrefDesc = rewriter.create<memref::ReinterpretCastOp>(
-        origLoc, reshapeType, newExtractStridedMetadata.getBaseBuffer(),
-        offsetOfr, finalSizes, finalStrides);
-    rewriter.replaceOp(reshape, memrefDesc.getResult());
+
+    Location loc = collapseShapeOp.getLoc();
+    SmallVector<Value> results;
+    results.push_back(stridedMetadata->basePtr);
+    results.push_back(getValueOrCreateConstantIndexOp(rewriter, loc,
+                                                      stridedMetadata->offset));
+    results.append(
+        getValueOrCreateConstantIndexOp(rewriter, loc, stridedMetadata->sizes));
+    results.append(getValueOrCreateConstantIndexOp(rewriter, loc,
+                                                   stridedMetadata->strides));
+    rewriter.replaceOp(op, results);
     return success();
   }
 };
@@ -1030,6 +1100,7 @@ void memref::populateResolveExtractStridedMetadataPatterns(
     RewritePatternSet &patterns) {
   patterns.add<ExtractStridedMetadataOpAllocFolder<memref::AllocOp>,
                ExtractStridedMetadataOpAllocFolder<memref::AllocaOp>,
+               ExtractStridedMetadataOpCollapseShapeFolder,
                ExtractStridedMetadataOpGetGlobalFolder,
                ExtractStridedMetadataOpSubviewFolder,
                RewriteExtractAlignedPointerAsIndexOfViewLikeOp,

dcaballe · 2024-04-24T17:16:59Z

Interesting enough, I can't find where memref::populateExpandStridedMetadataPatterns is being tested???

github-actions · 2024-04-24T17:18:59Z

✅ With the latest revision this PR passed the C/C++ code formatter.

qcolombet · 2024-04-24T19:19:36Z

Interesting enough, I can't find where memref::populateExpandStridedMetadataPatterns is being tested???

mlir/test/Dialect/MemRef/expand-strided-metadata.mlir, no?

qcolombet · 2024-04-24T19:22:12Z

mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp

@@ -550,6 +550,78 @@ getCollapsedStride(memref::CollapseShapeOp collapseShape, OpBuilder &builder,
  return {makeComposedFoldedAffineMin(builder, collapseShape.getLoc(), minMap,
                                      groupStrides)};
 }
+
+template <typename ReassociativeReshapeLikeOp,


Could you move the comment explaining what this does?

qcolombet · 2024-04-24T19:37:10Z

Interesting enough, I can't find where memref::populateExpandStridedMetadataPatterns is being tested???

mlir/test/Dialect/MemRef/expand-strided-metadata.mlir, no?

I guess you meant populateResolveExtractStridedMetadataPatterns since that's the one you're modifying.
I thought that populateResolveExtractStridedMetadataPatterns was a subset of populateExpandStridedMetadataPatterns but looks like it changed with https://reviews.llvm.org/D147393 and I didn't notice.

I see two paths forward

use populateResolveExtractStridedMetadataPatterns inside populateExpandStridedMetadataPatterns so that we can be sure Resolve is a subset of Expand. I think this is true today (modulo what https://reviews.llvm.org/D147393 added, but re-converging remains good IMHO) and we can revisit when they need to diverge.
integrate a test in the narrowing type test pass, since this is the only place where Resolve is used.

qcolombet

LGTM modulo:

add a test
move some comments on the refactored helper function

ftynse · 2024-04-25T10:36:41Z

mlir/lib/Dialect/MemRef/Transforms/ExpandStridedMetadata.cpp

+          SmallVector<OpFoldResult> (*getReshapedSizes)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/, unsigned /*groupId*/),
+          SmallVector<OpFoldResult> (*getReshapedStrides)(
+              ReassociativeReshapeLikeOp, OpBuilder &,
+              ArrayRef<OpFoldResult> /*origSizes*/,
+              ArrayRef<OpFoldResult> /*origStrides*/, unsigned /*groupId*/)>


A fly-by nit: is there an observable benefit to using these template arguments as opposed to passing in function_ref callbacks?

Probably not. I mostly followed the style of the caller

dcaballe · 2024-04-25T20:57:28Z

I see two paths forward

use populateResolveExtractStridedMetadataPatterns inside populateExpandStridedMetadataPatterns so that we can be sure Resolve is a subset of Expand. I think this is true today (modulo what reviews.llvm.org/D147393 added, but re-converging remains good IMHO) and we can revisit when they need to diverge.

integrate a test in the narrowing type test pass, since this is the only place where Resolve is used.

Thanks for the suggesting. Ok, let me try 1.

- Add test - Add doc - Use function_ref

…89954) This PR adds a new pattern to the set of patterns used to resolve the offset, sizes and stride of a memref. Similar to `ExtractStridedMetadataOpSubviewFolder`, the new pattern resolves strided_metadata(collapse_shape) directly, without introduce a reshape_cast op.

dcaballe requested review from tpopp and qcolombet April 24, 2024 17:16

llvmbot added mlir mlir:memref labels Apr 24, 2024

qcolombet reviewed Apr 24, 2024

View reviewed changes

qcolombet approved these changes Apr 24, 2024

View reviewed changes

ftynse reviewed Apr 25, 2024

View reviewed changes

dcaballe added 4 commits April 26, 2024 09:56

Review feedback

19c46b5

- Add test - Add doc - Use function_ref

Merge remote-tracking branch 'github/main' into expand-metadata-collapse

ca6cd4e

Fix doc

3500630

Final fixups!

e3bd384

dcaballe merged commit 450ac01 into llvm:main Apr 26, 2024
3 of 4 checks passed

dcaballe deleted the expand-metadata-collapse branch April 26, 2024 14:20

This was referenced Apr 29, 2024

main #90439

Closed

[AArch64] Add support for Cortex-R82AE and improve Cortex-R82 #90440

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][MemRef] Add ExtractStridedMetadataOpCollapseShapeFolder #89954

[mlir][MemRef] Add ExtractStridedMetadataOpCollapseShapeFolder #89954

dcaballe commented Apr 24, 2024

llvmbot commented Apr 24, 2024

llvmbot commented Apr 24, 2024

dcaballe commented Apr 24, 2024

github-actions bot commented Apr 24, 2024 •

edited

Loading

qcolombet commented Apr 24, 2024

qcolombet Apr 24, 2024

qcolombet commented Apr 24, 2024

qcolombet left a comment

ftynse Apr 25, 2024

dcaballe Apr 25, 2024

dcaballe commented Apr 25, 2024

[mlir][MemRef] Add ExtractStridedMetadataOpCollapseShapeFolder #89954

[mlir][MemRef] Add ExtractStridedMetadataOpCollapseShapeFolder #89954

Conversation

dcaballe commented Apr 24, 2024

llvmbot commented Apr 24, 2024

llvmbot commented Apr 24, 2024

dcaballe commented Apr 24, 2024

github-actions bot commented Apr 24, 2024 • edited Loading

qcolombet commented Apr 24, 2024

qcolombet Apr 24, 2024

Choose a reason for hiding this comment

qcolombet commented Apr 24, 2024

qcolombet left a comment

Choose a reason for hiding this comment

ftynse Apr 25, 2024

Choose a reason for hiding this comment

dcaballe Apr 25, 2024

Choose a reason for hiding this comment

dcaballe commented Apr 25, 2024

github-actions bot commented Apr 24, 2024 •

edited

Loading