b813b0b4e4
This PR adds code to FoldMemRefAliasOps / --fold-memref-alias-ops to use the new IndexedMemoryAccessOpInterface and IndexedMemCopyOpInterface and implement those operations for relevant operations in the memref dialect. This is a reordering of the changes planned in #177014 and #177016 to make them more testable. There are no behavior changes expected for how memref.load and memref.store behave within the alias ops folding pass, though support for new operations, like memref.prefetch, has been added. Some error messages have been updated because certain laws of memref.load/memref.store have been moved to IndexedAccessOpInterface. Assisted-by: Claude 4.6 (helped deal with some of the boilerplate in the rewrite patterns and with extracting the patch)
924 lines
37 KiB
C++
924 lines
37 KiB
C++
//===- FoldMemRefAliasOps.cpp - Fold memref alias ops -----===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This transformation pass folds loading/storing from/to subview ops into
|
|
// loading/storing from/to the original memref.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
|
|
#include "mlir/Dialect/Arith/Utils/Utils.h"
|
|
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemRef.h"
|
|
#include "mlir/Dialect/MemRef/IR/MemoryAccessOpInterfaces.h"
|
|
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
|
|
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
|
|
#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
|
|
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "llvm/ADT/STLExtras.h"
|
|
#include "llvm/ADT/SmallBitVector.h"
|
|
#include "llvm/ADT/TypeSwitch.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include <cstdint>
|
|
|
|
#define DEBUG_TYPE "fold-memref-alias-ops"
|
|
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
|
|
|
|
namespace mlir {
|
|
namespace memref {
|
|
#define GEN_PASS_DEF_FOLDMEMREFALIASOPSPASS
|
|
#include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
|
|
} // namespace memref
|
|
} // namespace mlir
|
|
|
|
using namespace mlir;
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Utility functions
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
/// Deterimine if the last N indices of `reassocitaion` are trivial - that is,
|
|
/// check if they all contain exactly one dimension to collape/expand into.
|
|
static bool
|
|
hasTrivialReassociationSuffix(ArrayRef<ReassociationIndices> reassocs,
|
|
int64_t n) {
|
|
if (n <= 0)
|
|
return true;
|
|
return llvm::all_of(
|
|
reassocs.take_back(n),
|
|
[&](const ReassociationIndices &indices) { return indices.size() == 1; });
|
|
}
|
|
|
|
static bool hasTrailingUnitStrides(memref::SubViewOp subview, int64_t n) {
|
|
if (n <= 0)
|
|
return true;
|
|
return llvm::all_of(subview.getStaticStrides().take_back(n),
|
|
[](int64_t s) { return s == 1; });
|
|
}
|
|
|
|
/// Helpers to access the memref operand for each op.
|
|
template <typename LoadOrStoreOpTy>
|
|
static Value getMemRefOperand(LoadOrStoreOpTy op) {
|
|
return op.getMemref();
|
|
}
|
|
|
|
static Value getMemRefOperand(vector::TransferReadOp op) {
|
|
return op.getBase();
|
|
}
|
|
|
|
static Value getMemRefOperand(nvgpu::LdMatrixOp op) {
|
|
return op.getSrcMemref();
|
|
}
|
|
|
|
static Value getMemRefOperand(vector::LoadOp op) { return op.getBase(); }
|
|
|
|
static Value getMemRefOperand(vector::StoreOp op) { return op.getBase(); }
|
|
|
|
static Value getMemRefOperand(vector::MaskedLoadOp op) { return op.getBase(); }
|
|
|
|
static Value getMemRefOperand(vector::MaskedStoreOp op) { return op.getBase(); }
|
|
|
|
static Value getMemRefOperand(vector::TransferWriteOp op) {
|
|
return op.getBase();
|
|
}
|
|
|
|
static Value getMemRefOperand(gpu::SubgroupMmaLoadMatrixOp op) {
|
|
return op.getSrcMemref();
|
|
}
|
|
|
|
static Value getMemRefOperand(gpu::SubgroupMmaStoreMatrixOp op) {
|
|
return op.getDstMemref();
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Patterns
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
/// Merges subview operation with load/transferRead operation.
|
|
template <typename OpTy>
|
|
class LoadOpOfSubViewOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy loadOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges expand_shape operation with load/transferRead operation.
|
|
template <typename OpTy>
|
|
class LoadOpOfExpandShapeOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy loadOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges collapse_shape operation with load/transferRead operation.
|
|
template <typename OpTy>
|
|
class LoadOpOfCollapseShapeOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy loadOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges subview operation with store/transferWriteOp operation.
|
|
template <typename OpTy>
|
|
class StoreOpOfSubViewOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy storeOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges expand_shape operation with store/transferWriteOp operation.
|
|
template <typename OpTy>
|
|
class StoreOpOfExpandShapeOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy storeOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges collapse_shape operation with store/transferWriteOp operation.
|
|
template <typename OpTy>
|
|
class StoreOpOfCollapseShapeOpFolder final : public OpRewritePattern<OpTy> {
|
|
public:
|
|
using OpRewritePattern<OpTy>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(OpTy storeOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Folds subview(subview(x)) to a single subview(x).
|
|
class SubViewOfSubViewFolder : public OpRewritePattern<memref::SubViewOp> {
|
|
public:
|
|
using OpRewritePattern<memref::SubViewOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(memref::SubViewOp subView,
|
|
PatternRewriter &rewriter) const override {
|
|
auto srcSubView = subView.getSource().getDefiningOp<memref::SubViewOp>();
|
|
if (!srcSubView)
|
|
return failure();
|
|
|
|
// TODO: relax unit stride assumption.
|
|
if (!subView.hasUnitStride()) {
|
|
return rewriter.notifyMatchFailure(subView, "requires unit strides");
|
|
}
|
|
if (!srcSubView.hasUnitStride()) {
|
|
return rewriter.notifyMatchFailure(srcSubView, "requires unit strides");
|
|
}
|
|
|
|
// Resolve sizes according to dropped dims.
|
|
SmallVector<OpFoldResult> resolvedSizes;
|
|
llvm::SmallBitVector srcDroppedDims = srcSubView.getDroppedDims();
|
|
affine::resolveSizesIntoOpWithSizes(srcSubView.getMixedSizes(),
|
|
subView.getMixedSizes(), srcDroppedDims,
|
|
resolvedSizes);
|
|
|
|
// Resolve offsets according to source offsets and strides.
|
|
SmallVector<Value> resolvedOffsets;
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, subView.getLoc(), srcSubView.getMixedOffsets(),
|
|
srcSubView.getMixedStrides(), srcDroppedDims, subView.getMixedOffsets(),
|
|
resolvedOffsets);
|
|
|
|
// Replace original op.
|
|
rewriter.replaceOpWithNewOp<memref::SubViewOp>(
|
|
subView, subView.getType(), srcSubView.getSource(),
|
|
getAsOpFoldResult(resolvedOffsets), resolvedSizes,
|
|
srcSubView.getMixedStrides());
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
/// Folds nvgpu.device_async_copy subviews into the copy itself. This pattern
|
|
/// is folds subview on src and dst memref of the copy.
|
|
class NVGPUAsyncCopyOpSubViewOpFolder final
|
|
: public OpRewritePattern<nvgpu::DeviceAsyncCopyOp> {
|
|
public:
|
|
using OpRewritePattern<nvgpu::DeviceAsyncCopyOp>::OpRewritePattern;
|
|
|
|
LogicalResult matchAndRewrite(nvgpu::DeviceAsyncCopyOp copyOp,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges subview operations with load/store like operations unless such a
|
|
/// merger would cause the strides between dimensions accessed by that operaton
|
|
/// to change.
|
|
struct AccessOpOfSubViewOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedAccessOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedAccessOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merge a memref.expand_shape operation with an operation that accesses a
|
|
/// memref by index unless that operation accesss more than one dimension of
|
|
/// memory and any dimension other than the outermost dimension accessed this
|
|
/// way would be merged. This prevents issuses from arising with, say, a
|
|
/// vector.load of a 4x2 vector having the two trailing dimensions of the access
|
|
/// get merged.
|
|
struct AccessOpOfExpandShapeOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedAccessOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedAccessOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges an operation that accesses a memref by index with a
|
|
/// memref.collapse_shape, unless this would break apart a dimension other than
|
|
/// the outermost one that an operation accesses. This prevents, for example,
|
|
/// transforming a load of a 3x8 vector from a 6x8 memref into a load
|
|
/// from a 3x4x2 memref (as this would require special handling and could lead
|
|
/// to invalid IR if that higher-dimensional memref comes from a subview) but
|
|
/// does permit turning a load of a length-8 vector from a 3x8 memref into a
|
|
/// load from a 3x2x8 one.
|
|
struct AccessOpOfCollapseShapeOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedAccessOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedAccessOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges memref.subview operations present on the source or destination
|
|
/// operands of indexed memory copy operations (DMA operations) into those
|
|
/// operations. This is perfromed unconditionally, since folding in a subview
|
|
/// cannot change the starting position of the copy, which is what the
|
|
/// memref/index pair represent in DMA operations.
|
|
struct IndexedMemCopyOpOfSubViewOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedMemCopyOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedMemCopyOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges memref.expand_shape operations that are present on the source or
|
|
/// destination of an indexed memory copy/DMA into the memref/index arguments of
|
|
/// that DMA. As with subviews, this can be done unconditionally.
|
|
struct IndexedMemCopyOpOfExpandShapeOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedMemCopyOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedMemCopyOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
/// Merges memref.collapse_shape operations that are present on the source or
|
|
/// destination of an indexed memory copy/DMA into the memref/index arguments of
|
|
/// that DMA. As with subviews, this can be done unconditionally.
|
|
struct IndexedMemCopyOpOfCollapseShapeOpFolder final
|
|
: OpInterfaceRewritePattern<memref::IndexedMemCopyOpInterface> {
|
|
using Base::Base;
|
|
|
|
LogicalResult matchAndRewrite(memref::IndexedMemCopyOpInterface op,
|
|
PatternRewriter &rewriter) const override;
|
|
};
|
|
} // namespace
|
|
|
|
template <typename XferOp>
|
|
static LogicalResult
|
|
preconditionsFoldSubViewOpImpl(RewriterBase &rewriter, XferOp xferOp,
|
|
memref::SubViewOp subviewOp) {
|
|
static_assert(
|
|
!llvm::is_one_of<vector::TransferReadOp, vector::TransferWriteOp>::value,
|
|
"must be a vector transfer op");
|
|
if (xferOp.hasOutOfBoundsDim())
|
|
return rewriter.notifyMatchFailure(xferOp, "out of bounds transfer dim");
|
|
if (!subviewOp.hasUnitStride()) {
|
|
return rewriter.notifyMatchFailure(
|
|
xferOp, "non-1 stride subview, need to track strides in folded memref");
|
|
}
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
|
|
Operation *op,
|
|
memref::SubViewOp subviewOp) {
|
|
return success();
|
|
}
|
|
|
|
static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
|
|
vector::TransferReadOp readOp,
|
|
memref::SubViewOp subviewOp) {
|
|
return preconditionsFoldSubViewOpImpl(rewriter, readOp, subviewOp);
|
|
}
|
|
|
|
static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
|
|
vector::TransferWriteOp writeOp,
|
|
memref::SubViewOp subviewOp) {
|
|
return preconditionsFoldSubViewOpImpl(rewriter, writeOp, subviewOp);
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult LoadOpOfSubViewOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy loadOp, PatternRewriter &rewriter) const {
|
|
auto subViewOp =
|
|
getMemRefOperand(loadOp).template getDefiningOp<memref::SubViewOp>();
|
|
|
|
if (!subViewOp)
|
|
return rewriter.notifyMatchFailure(loadOp, "not a subview producer");
|
|
|
|
LogicalResult preconditionResult =
|
|
preconditionsFoldSubViewOp(rewriter, loadOp, subViewOp);
|
|
if (failed(preconditionResult))
|
|
return preconditionResult;
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, loadOp.getLoc(), subViewOp.getMixedOffsets(),
|
|
subViewOp.getMixedStrides(), subViewOp.getDroppedDims(),
|
|
loadOp.getIndices(), sourceIndices);
|
|
|
|
llvm::TypeSwitch<Operation *, void>(loadOp)
|
|
.Case([&](memref::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::LoadOp>(
|
|
loadOp, subViewOp.getSource(), sourceIndices, op.getNontemporal());
|
|
})
|
|
.Case([&](vector::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::LoadOp>(
|
|
op, op.getType(), subViewOp.getSource(), sourceIndices);
|
|
})
|
|
.Case([&](vector::MaskedLoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedLoadOp>(
|
|
op, op.getType(), subViewOp.getSource(), sourceIndices,
|
|
op.getMask(), op.getPassThru());
|
|
})
|
|
.Case([&](vector::TransferReadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
|
|
op, op.getVectorType(), subViewOp.getSource(), sourceIndices,
|
|
AffineMapAttr::get(expandDimsToRank(
|
|
op.getPermutationMap(), subViewOp.getSourceType().getRank(),
|
|
subViewOp.getDroppedDims())),
|
|
op.getPadding(), op.getMask(), op.getInBoundsAttr());
|
|
})
|
|
.Case([&](gpu::SubgroupMmaLoadMatrixOp op) {
|
|
rewriter.replaceOpWithNewOp<gpu::SubgroupMmaLoadMatrixOp>(
|
|
op, op.getType(), subViewOp.getSource(), sourceIndices,
|
|
op.getLeadDimension(), op.getTransposeAttr());
|
|
})
|
|
.Case([&](nvgpu::LdMatrixOp op) {
|
|
rewriter.replaceOpWithNewOp<nvgpu::LdMatrixOp>(
|
|
op, op.getType(), subViewOp.getSource(), sourceIndices,
|
|
op.getTranspose(), op.getNumTiles());
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
return success();
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult LoadOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy loadOp, PatternRewriter &rewriter) const {
|
|
auto expandShapeOp =
|
|
getMemRefOperand(loadOp).template getDefiningOp<memref::ExpandShapeOp>();
|
|
|
|
if (!expandShapeOp)
|
|
return failure();
|
|
|
|
// For vector::TransferReadOp, validate preconditions before creating any IR.
|
|
// resolveSourceIndicesExpandShape creates new ops, so all checks that can
|
|
// fail must happen before that call to avoid "pattern returned failure but
|
|
// IR did change" errors (caught by MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS).
|
|
SmallVector<AffineExpr> transferReadNewResults;
|
|
if (auto transferOp =
|
|
dyn_cast<vector::TransferReadOp>(loadOp.getOperation())) {
|
|
const int64_t vectorRank = transferOp.getVectorType().getRank();
|
|
const int64_t sourceRank =
|
|
cast<MemRefType>(expandShapeOp.getViewSource().getType()).getRank();
|
|
if (sourceRank < vectorRank)
|
|
return failure();
|
|
|
|
// We can only fold if the permutation map uses only the least significant
|
|
// dimension from each expanded reassociation group.
|
|
for (AffineExpr result : transferOp.getPermutationMap().getResults()) {
|
|
bool foundExpr = false;
|
|
for (auto reassocationIndices :
|
|
llvm::enumerate(expandShapeOp.getReassociationIndices())) {
|
|
auto reassociation = reassocationIndices.value();
|
|
AffineExpr dim = getAffineDimExpr(
|
|
reassociation[reassociation.size() - 1], rewriter.getContext());
|
|
if (dim == result) {
|
|
transferReadNewResults.push_back(getAffineDimExpr(
|
|
reassocationIndices.index(), rewriter.getContext()));
|
|
foundExpr = true;
|
|
break;
|
|
}
|
|
}
|
|
if (!foundExpr)
|
|
return failure();
|
|
}
|
|
}
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
// memref.load guarantees that indexes start inbounds while the vector
|
|
// operations don't. This impacts if our linearization is `disjoint`
|
|
resolveSourceIndicesExpandShape(loadOp.getLoc(), rewriter, expandShapeOp,
|
|
loadOp.getIndices(), sourceIndices,
|
|
isa<memref::LoadOp>(loadOp.getOperation()));
|
|
|
|
return llvm::TypeSwitch<Operation *, LogicalResult>(loadOp)
|
|
.Case([&](memref::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::LoadOp>(
|
|
loadOp, expandShapeOp.getViewSource(), sourceIndices,
|
|
op.getNontemporal());
|
|
return success();
|
|
})
|
|
.Case([&](vector::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::LoadOp>(
|
|
op, op.getType(), expandShapeOp.getViewSource(), sourceIndices,
|
|
op.getNontemporal());
|
|
return success();
|
|
})
|
|
.Case([&](vector::MaskedLoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedLoadOp>(
|
|
op, op.getType(), expandShapeOp.getViewSource(), sourceIndices,
|
|
op.getMask(), op.getPassThru());
|
|
return success();
|
|
})
|
|
.Case([&](vector::TransferReadOp op) {
|
|
const int64_t sourceRank = sourceIndices.size();
|
|
auto newMap = AffineMap::get(sourceRank, 0, transferReadNewResults,
|
|
op.getContext());
|
|
rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
|
|
op, op.getVectorType(), expandShapeOp.getViewSource(),
|
|
sourceIndices, newMap, op.getPadding(), op.getMask(),
|
|
op.getInBounds());
|
|
return success();
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult LoadOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy loadOp, PatternRewriter &rewriter) const {
|
|
auto collapseShapeOp = getMemRefOperand(loadOp)
|
|
.template getDefiningOp<memref::CollapseShapeOp>();
|
|
|
|
if (!collapseShapeOp)
|
|
return failure();
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
resolveSourceIndicesCollapseShape(loadOp.getLoc(), rewriter, collapseShapeOp,
|
|
loadOp.getIndices(), sourceIndices);
|
|
llvm::TypeSwitch<Operation *, void>(loadOp)
|
|
.Case([&](memref::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::LoadOp>(
|
|
loadOp, collapseShapeOp.getViewSource(), sourceIndices,
|
|
op.getNontemporal());
|
|
})
|
|
.Case([&](vector::LoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::LoadOp>(
|
|
op, op.getType(), collapseShapeOp.getViewSource(), sourceIndices,
|
|
op.getNontemporal());
|
|
})
|
|
.Case([&](vector::MaskedLoadOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedLoadOp>(
|
|
op, op.getType(), collapseShapeOp.getViewSource(), sourceIndices,
|
|
op.getMask(), op.getPassThru());
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
return success();
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult StoreOpOfSubViewOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy storeOp, PatternRewriter &rewriter) const {
|
|
auto subViewOp =
|
|
getMemRefOperand(storeOp).template getDefiningOp<memref::SubViewOp>();
|
|
|
|
if (!subViewOp)
|
|
return rewriter.notifyMatchFailure(storeOp, "not a subview producer");
|
|
|
|
LogicalResult preconditionResult =
|
|
preconditionsFoldSubViewOp(rewriter, storeOp, subViewOp);
|
|
if (failed(preconditionResult))
|
|
return preconditionResult;
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, storeOp.getLoc(), subViewOp.getMixedOffsets(),
|
|
subViewOp.getMixedStrides(), subViewOp.getDroppedDims(),
|
|
storeOp.getIndices(), sourceIndices);
|
|
|
|
llvm::TypeSwitch<Operation *, void>(storeOp)
|
|
.Case([&](memref::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::StoreOp>(
|
|
op, op.getValue(), subViewOp.getSource(), sourceIndices,
|
|
op.getNontemporal());
|
|
})
|
|
.Case([&](vector::TransferWriteOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
|
|
op, op.getValue(), subViewOp.getSource(), sourceIndices,
|
|
AffineMapAttr::get(expandDimsToRank(
|
|
op.getPermutationMap(), subViewOp.getSourceType().getRank(),
|
|
subViewOp.getDroppedDims())),
|
|
op.getMask(), op.getInBoundsAttr());
|
|
})
|
|
.Case([&](vector::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::StoreOp>(
|
|
op, op.getValueToStore(), subViewOp.getSource(), sourceIndices);
|
|
})
|
|
.Case([&](vector::MaskedStoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
|
|
op, subViewOp.getSource(), sourceIndices, op.getMask(),
|
|
op.getValueToStore());
|
|
})
|
|
.Case([&](gpu::SubgroupMmaStoreMatrixOp op) {
|
|
rewriter.replaceOpWithNewOp<gpu::SubgroupMmaStoreMatrixOp>(
|
|
op, op.getSrc(), subViewOp.getSource(), sourceIndices,
|
|
op.getLeadDimension(), op.getTransposeAttr());
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
return success();
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult StoreOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy storeOp, PatternRewriter &rewriter) const {
|
|
auto expandShapeOp =
|
|
getMemRefOperand(storeOp).template getDefiningOp<memref::ExpandShapeOp>();
|
|
|
|
if (!expandShapeOp)
|
|
return failure();
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
// memref.store guarantees that indexes start inbounds while the vector
|
|
// operations don't. This impacts if our linearization is `disjoint`
|
|
resolveSourceIndicesExpandShape(storeOp.getLoc(), rewriter, expandShapeOp,
|
|
storeOp.getIndices(), sourceIndices,
|
|
isa<memref::StoreOp>(storeOp.getOperation()));
|
|
llvm::TypeSwitch<Operation *, void>(storeOp)
|
|
.Case([&](memref::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::StoreOp>(
|
|
storeOp, op.getValueToStore(), expandShapeOp.getViewSource(),
|
|
sourceIndices, op.getNontemporal());
|
|
})
|
|
.Case([&](vector::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::StoreOp>(
|
|
op, op.getValueToStore(), expandShapeOp.getViewSource(),
|
|
sourceIndices, op.getNontemporal());
|
|
})
|
|
.Case([&](vector::MaskedStoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
|
|
op, expandShapeOp.getViewSource(), sourceIndices, op.getMask(),
|
|
op.getValueToStore());
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
return success();
|
|
}
|
|
|
|
template <typename OpTy>
|
|
LogicalResult StoreOpOfCollapseShapeOpFolder<OpTy>::matchAndRewrite(
|
|
OpTy storeOp, PatternRewriter &rewriter) const {
|
|
auto collapseShapeOp = getMemRefOperand(storeOp)
|
|
.template getDefiningOp<memref::CollapseShapeOp>();
|
|
|
|
if (!collapseShapeOp)
|
|
return failure();
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
resolveSourceIndicesCollapseShape(storeOp.getLoc(), rewriter, collapseShapeOp,
|
|
storeOp.getIndices(), sourceIndices);
|
|
llvm::TypeSwitch<Operation *, void>(storeOp)
|
|
.Case([&](memref::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<memref::StoreOp>(
|
|
storeOp, op.getValueToStore(), collapseShapeOp.getViewSource(),
|
|
sourceIndices, op.getNontemporal());
|
|
})
|
|
.Case([&](vector::StoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::StoreOp>(
|
|
op, op.getValueToStore(), collapseShapeOp.getViewSource(),
|
|
sourceIndices, op.getNontemporal());
|
|
})
|
|
.Case([&](vector::MaskedStoreOp op) {
|
|
rewriter.replaceOpWithNewOp<vector::MaskedStoreOp>(
|
|
op, collapseShapeOp.getViewSource(), sourceIndices, op.getMask(),
|
|
op.getValueToStore());
|
|
})
|
|
.DefaultUnreachable("unexpected operation");
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
AccessOpOfSubViewOpFolder::matchAndRewrite(memref::IndexedAccessOpInterface op,
|
|
PatternRewriter &rewriter) const {
|
|
auto subview = op.getAccessedMemref().getDefiningOp<memref::SubViewOp>();
|
|
if (!subview)
|
|
return rewriter.notifyMatchFailure(op, "not accessing a subview");
|
|
|
|
SmallVector<int64_t> accessedShape = op.getAccessedShape();
|
|
// Note the subtle difference between accesedShape = {1} and accessedShape =
|
|
// {} here. The former prevents us from fdolding in a subview that doesn't
|
|
// have a unit stride on the final dimension, while the latter does not (since
|
|
// it indices scalar accesss).
|
|
int64_t accessedDims = accessedShape.size();
|
|
if (!hasTrailingUnitStrides(subview, accessedDims))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "non-unit stride on accessed dimensions");
|
|
|
|
llvm::SmallBitVector droppedDims = subview.getDroppedDims();
|
|
int64_t sourceRank = subview.getSourceType().getRank();
|
|
|
|
// Ignore outermost access dimension - we only care about dropped dimensions
|
|
// between the accessed op's results, as those could break the accessing op's
|
|
// sematics.
|
|
int64_t secondAccessedDim = sourceRank - (accessedDims - 1);
|
|
if (secondAccessedDim < sourceRank) {
|
|
for (int64_t d : llvm::seq(secondAccessedDim, sourceRank)) {
|
|
if (droppedDims.test(d))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "reintroducing dropped dimension " + Twine(d) +
|
|
" would break access op semantics");
|
|
}
|
|
}
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, op.getLoc(), subview.getMixedOffsets(),
|
|
subview.getMixedStrides(), droppedDims, op.getIndices(), sourceIndices);
|
|
|
|
std::optional<SmallVector<Value>> newValues =
|
|
op.updateMemrefAndIndices(rewriter, subview.getSource(), sourceIndices);
|
|
if (newValues)
|
|
rewriter.replaceOp(op, *newValues);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult AccessOpOfExpandShapeOpFolder::matchAndRewrite(
|
|
memref::IndexedAccessOpInterface op, PatternRewriter &rewriter) const {
|
|
auto expand = op.getAccessedMemref().getDefiningOp<memref::ExpandShapeOp>();
|
|
if (!expand)
|
|
return rewriter.notifyMatchFailure(op, "not accessing an expand_shape");
|
|
|
|
SmallVector<int64_t> rawAccessedShape = op.getAccessedShape();
|
|
ArrayRef<int64_t> accessedShape = rawAccessedShape;
|
|
// Cut off the leading dimension, since we don't care about monifying its
|
|
// strides.
|
|
if (!accessedShape.empty())
|
|
accessedShape = accessedShape.drop_front();
|
|
|
|
SmallVector<ReassociationIndices, 4> reassocs =
|
|
expand.getReassociationIndices();
|
|
if (!hasTrivialReassociationSuffix(reassocs, accessedShape.size()))
|
|
return rewriter.notifyMatchFailure(
|
|
op,
|
|
"expand_shape folding would merge semanvtically important dimensions");
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
memref::resolveSourceIndicesExpandShape(op.getLoc(), rewriter, expand,
|
|
op.getIndices(), sourceIndices,
|
|
op.hasInboundsIndices());
|
|
|
|
std::optional<SmallVector<Value>> newValues = op.updateMemrefAndIndices(
|
|
rewriter, expand.getViewSource(), sourceIndices);
|
|
if (newValues)
|
|
rewriter.replaceOp(op, *newValues);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult AccessOpOfCollapseShapeOpFolder::matchAndRewrite(
|
|
memref::IndexedAccessOpInterface op, PatternRewriter &rewriter) const {
|
|
auto collapse =
|
|
op.getAccessedMemref().getDefiningOp<memref::CollapseShapeOp>();
|
|
if (!collapse)
|
|
return rewriter.notifyMatchFailure(op, "not accessing a collapse_shape");
|
|
|
|
SmallVector<int64_t> rawAccessedShape = op.getAccessedShape();
|
|
ArrayRef<int64_t> accessedShape = rawAccessedShape;
|
|
// Cut off the leading dimension, since we don't care about its strides being
|
|
// modified and we know that the dimensions within its reassociation group, if
|
|
// it's non-trivial, must be contiguous.
|
|
if (!accessedShape.empty())
|
|
accessedShape = accessedShape.drop_front();
|
|
|
|
SmallVector<ReassociationIndices, 4> reassocs =
|
|
collapse.getReassociationIndices();
|
|
if (!hasTrivialReassociationSuffix(reassocs, accessedShape.size()))
|
|
return rewriter.notifyMatchFailure(op,
|
|
"collapse_shape folding would merge "
|
|
"semanvtically important dimensions");
|
|
|
|
SmallVector<Value> sourceIndices;
|
|
memref::resolveSourceIndicesCollapseShape(op.getLoc(), rewriter, collapse,
|
|
op.getIndices(), sourceIndices);
|
|
|
|
std::optional<SmallVector<Value>> newValues = op.updateMemrefAndIndices(
|
|
rewriter, collapse.getViewSource(), sourceIndices);
|
|
if (newValues)
|
|
rewriter.replaceOp(op, *newValues);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult IndexedMemCopyOpOfSubViewOpFolder::matchAndRewrite(
|
|
memref::IndexedMemCopyOpInterface op, PatternRewriter &rewriter) const {
|
|
auto srcSubview = op.getSrc().getDefiningOp<memref::SubViewOp>();
|
|
auto dstSubview = op.getDst().getDefiningOp<memref::SubViewOp>();
|
|
if (!srcSubview && !dstSubview)
|
|
return rewriter.notifyMatchFailure(
|
|
op, "no subviews found on indexed copy inputs");
|
|
|
|
Value newSrc = op.getSrc();
|
|
SmallVector<Value> newSrcIndices = llvm::to_vector(op.getSrcIndices());
|
|
Value newDst = op.getDst();
|
|
SmallVector<Value> newDstIndices = llvm::to_vector(op.getDstIndices());
|
|
if (srcSubview) {
|
|
newSrc = srcSubview.getSource();
|
|
newSrcIndices.clear();
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, op.getLoc(), srcSubview.getMixedOffsets(),
|
|
srcSubview.getMixedStrides(), srcSubview.getDroppedDims(),
|
|
op.getSrcIndices(), newSrcIndices);
|
|
}
|
|
if (dstSubview) {
|
|
newDst = dstSubview.getSource();
|
|
newDstIndices.clear();
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, op.getLoc(), dstSubview.getMixedOffsets(),
|
|
dstSubview.getMixedStrides(), dstSubview.getDroppedDims(),
|
|
op.getDstIndices(), newDstIndices);
|
|
}
|
|
op.setMemrefsAndIndices(rewriter, newSrc, newSrcIndices, newDst,
|
|
newDstIndices);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult IndexedMemCopyOpOfExpandShapeOpFolder::matchAndRewrite(
|
|
memref::IndexedMemCopyOpInterface op, PatternRewriter &rewriter) const {
|
|
auto srcExpand = op.getSrc().getDefiningOp<memref::ExpandShapeOp>();
|
|
auto dstExpand = op.getDst().getDefiningOp<memref::ExpandShapeOp>();
|
|
if (!srcExpand && !dstExpand)
|
|
return rewriter.notifyMatchFailure(
|
|
op, "no expand_shapes found on indexed copy inputs");
|
|
|
|
Value newSrc = op.getSrc();
|
|
SmallVector<Value> newSrcIndices = llvm::to_vector(op.getSrcIndices());
|
|
Value newDst = op.getDst();
|
|
SmallVector<Value> newDstIndices = llvm::to_vector(op.getDstIndices());
|
|
if (srcExpand) {
|
|
newSrc = srcExpand.getViewSource();
|
|
newSrcIndices.clear();
|
|
memref::resolveSourceIndicesExpandShape(op.getLoc(), rewriter, srcExpand,
|
|
op.getSrcIndices(), newSrcIndices,
|
|
/*startsInbounds=*/true);
|
|
}
|
|
if (dstExpand) {
|
|
newDst = dstExpand.getViewSource();
|
|
newDstIndices.clear();
|
|
memref::resolveSourceIndicesExpandShape(op.getLoc(), rewriter, dstExpand,
|
|
op.getDstIndices(), newDstIndices,
|
|
/*startsInbounds=*/true);
|
|
}
|
|
op.setMemrefsAndIndices(rewriter, newSrc, newSrcIndices, newDst,
|
|
newDstIndices);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult IndexedMemCopyOpOfCollapseShapeOpFolder::matchAndRewrite(
|
|
memref::IndexedMemCopyOpInterface op, PatternRewriter &rewriter) const {
|
|
auto srcCollapse = op.getSrc().getDefiningOp<memref::CollapseShapeOp>();
|
|
auto dstCollapse = op.getDst().getDefiningOp<memref::CollapseShapeOp>();
|
|
if (!srcCollapse && !dstCollapse)
|
|
return rewriter.notifyMatchFailure(
|
|
op, "no collapse_shapes found on indexed copy inputs");
|
|
|
|
Value newSrc = op.getSrc();
|
|
SmallVector<Value> newSrcIndices = llvm::to_vector(op.getSrcIndices());
|
|
Value newDst = op.getDst();
|
|
SmallVector<Value> newDstIndices = llvm::to_vector(op.getDstIndices());
|
|
if (srcCollapse) {
|
|
newSrc = srcCollapse.getViewSource();
|
|
newSrcIndices.clear();
|
|
memref::resolveSourceIndicesCollapseShape(
|
|
op.getLoc(), rewriter, srcCollapse, op.getSrcIndices(), newSrcIndices);
|
|
}
|
|
if (dstCollapse) {
|
|
newDst = dstCollapse.getViewSource();
|
|
newDstIndices.clear();
|
|
memref::resolveSourceIndicesCollapseShape(
|
|
op.getLoc(), rewriter, dstCollapse, op.getDstIndices(), newDstIndices);
|
|
}
|
|
op.setMemrefsAndIndices(rewriter, newSrc, newSrcIndices, newDst,
|
|
newDstIndices);
|
|
return success();
|
|
}
|
|
|
|
LogicalResult NVGPUAsyncCopyOpSubViewOpFolder::matchAndRewrite(
|
|
nvgpu::DeviceAsyncCopyOp copyOp, PatternRewriter &rewriter) const {
|
|
|
|
LLVM_DEBUG(DBGS() << "copyOp : " << copyOp << "\n");
|
|
|
|
auto srcSubViewOp =
|
|
copyOp.getSrc().template getDefiningOp<memref::SubViewOp>();
|
|
auto dstSubViewOp =
|
|
copyOp.getDst().template getDefiningOp<memref::SubViewOp>();
|
|
|
|
if (!(srcSubViewOp || dstSubViewOp))
|
|
return rewriter.notifyMatchFailure(copyOp, "does not use subview ops for "
|
|
"source or destination");
|
|
|
|
// If the source is a subview, we need to resolve the indices.
|
|
SmallVector<Value> foldedSrcIndices(copyOp.getSrcIndices().begin(),
|
|
copyOp.getSrcIndices().end());
|
|
|
|
if (srcSubViewOp) {
|
|
LLVM_DEBUG(DBGS() << "srcSubViewOp : " << srcSubViewOp << "\n");
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, copyOp.getLoc(), srcSubViewOp.getMixedOffsets(),
|
|
srcSubViewOp.getMixedStrides(), srcSubViewOp.getDroppedDims(),
|
|
copyOp.getSrcIndices(), foldedSrcIndices);
|
|
}
|
|
|
|
// If the destination is a subview, we need to resolve the indices.
|
|
SmallVector<Value> foldedDstIndices(copyOp.getDstIndices().begin(),
|
|
copyOp.getDstIndices().end());
|
|
|
|
if (dstSubViewOp) {
|
|
LLVM_DEBUG(DBGS() << "dstSubViewOp : " << dstSubViewOp << "\n");
|
|
affine::resolveIndicesIntoOpWithOffsetsAndStrides(
|
|
rewriter, copyOp.getLoc(), dstSubViewOp.getMixedOffsets(),
|
|
dstSubViewOp.getMixedStrides(), dstSubViewOp.getDroppedDims(),
|
|
copyOp.getDstIndices(), foldedDstIndices);
|
|
}
|
|
|
|
// Replace the copy op with a new copy op that uses the source and destination
|
|
// of the subview.
|
|
rewriter.replaceOpWithNewOp<nvgpu::DeviceAsyncCopyOp>(
|
|
copyOp, nvgpu::DeviceAsyncTokenType::get(copyOp.getContext()),
|
|
(dstSubViewOp ? dstSubViewOp.getSource() : copyOp.getDst()),
|
|
foldedDstIndices,
|
|
(srcSubViewOp ? srcSubViewOp.getSource() : copyOp.getSrc()),
|
|
foldedSrcIndices, copyOp.getDstElements(), copyOp.getSrcElements(),
|
|
copyOp.getBypassL1Attr());
|
|
|
|
return success();
|
|
}
|
|
|
|
void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
|
|
patterns.add<
|
|
// Interface-based patterns to which we will be migrating.
|
|
AccessOpOfSubViewOpFolder, AccessOpOfExpandShapeOpFolder,
|
|
AccessOpOfCollapseShapeOpFolder, IndexedMemCopyOpOfSubViewOpFolder,
|
|
IndexedMemCopyOpOfExpandShapeOpFolder,
|
|
IndexedMemCopyOpOfCollapseShapeOpFolder,
|
|
// The old way of doing things. Don't add more of these.
|
|
LoadOpOfSubViewOpFolder<nvgpu::LdMatrixOp>,
|
|
LoadOpOfSubViewOpFolder<vector::LoadOp>,
|
|
LoadOpOfSubViewOpFolder<vector::MaskedLoadOp>,
|
|
LoadOpOfSubViewOpFolder<vector::TransferReadOp>,
|
|
LoadOpOfSubViewOpFolder<gpu::SubgroupMmaLoadMatrixOp>,
|
|
StoreOpOfSubViewOpFolder<vector::TransferWriteOp>,
|
|
StoreOpOfSubViewOpFolder<vector::StoreOp>,
|
|
StoreOpOfSubViewOpFolder<vector::MaskedStoreOp>,
|
|
StoreOpOfSubViewOpFolder<gpu::SubgroupMmaStoreMatrixOp>,
|
|
LoadOpOfExpandShapeOpFolder<vector::LoadOp>,
|
|
LoadOpOfExpandShapeOpFolder<vector::MaskedLoadOp>,
|
|
LoadOpOfExpandShapeOpFolder<vector::TransferReadOp>,
|
|
StoreOpOfExpandShapeOpFolder<vector::StoreOp>,
|
|
StoreOpOfExpandShapeOpFolder<vector::MaskedStoreOp>,
|
|
LoadOpOfCollapseShapeOpFolder<vector::LoadOp>,
|
|
LoadOpOfCollapseShapeOpFolder<vector::MaskedLoadOp>,
|
|
StoreOpOfCollapseShapeOpFolder<vector::StoreOp>,
|
|
StoreOpOfCollapseShapeOpFolder<vector::MaskedStoreOp>,
|
|
SubViewOfSubViewFolder, NVGPUAsyncCopyOpSubViewOpFolder>(
|
|
patterns.getContext());
|
|
}
|
|
|
|
//===----------------------------------------------------------------------===//
|
|
// Pass registration
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
namespace {
|
|
|
|
struct FoldMemRefAliasOpsPass final
|
|
: public memref::impl::FoldMemRefAliasOpsPassBase<FoldMemRefAliasOpsPass> {
|
|
void runOnOperation() override;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
void FoldMemRefAliasOpsPass::runOnOperation() {
|
|
RewritePatternSet patterns(&getContext());
|
|
memref::populateFoldMemRefAliasOpPatterns(patterns);
|
|
(void)applyPatternsGreedily(getOperation(), std::move(patterns));
|
|
}
|