This is an archive of the discontinued LLVM Phabricator instance.

Show First 20 Lines • Show All 3,712 Lines • ▼ Show 20 Lines	static bool areAllRankReducedLeadingDim(tensor::ExtractSliceOp extractOp,
RankedTensorType inferredType = extractOp.inferResultType(		RankedTensorType inferredType = extractOp.inferResultType(
extractOp.getSourceType(), extractOp.getMixedOffsets(),		extractOp.getSourceType(), extractOp.getMixedOffsets(),
extractOp.getMixedSizes(), extractOp.getMixedStrides());		extractOp.getMixedSizes(), extractOp.getMixedStrides());
return extractOp.getType().getShape().take_back(trailingRank) ==		return extractOp.getType().getShape().take_back(trailingRank) ==
inferredType.getShape().take_back(trailingRank);		inferredType.getShape().take_back(trailingRank);
}		}

namespace {		namespace {
/// Fold transfer_reads of a tensor.extract_slice op. E.g.:		/// Store to load forwarding for transfer operations with permutation maps.
///
/// ```
/// %0 = tensor.extract_slice %t[%a, %b] [%c, %d] [1, 1]
/// : tensor<?x?xf32> to tensor<?x?xf32>
/// %1 = vector.transfer_read %0[%e, %f], %cst {in_bounds = [true, true]}
/// : tensor<?x?xf32>, vector<4x5xf32>
/// ```
/// is rewritten to:
/// ```
/// %p0 = arith.addi %a, %e : index
/// %p1 = arith.addi %b, %f : index
/// %1 = vector.transfer_read %t[%p0, %p1], %cst {in_bounds = [true, true]}
/// : tensor<?x?xf32>, vector<4x5xf32>
/// ```
// TODO: this is brittle and should be deprecated in favor of a more general
// pattern that applies on-demand.
struct FoldExtractSliceIntoTransferRead
: public OpRewritePattern<TransferReadOp> {
public:
using OpRewritePattern::OpRewritePattern;

LogicalResult matchAndRewrite(TransferReadOp xferOp,
PatternRewriter &rewriter) const override {
// TODO: support 0-d corner case.
if (xferOp.getTransferRank() == 0)
return failure();
if (xferOp.hasOutOfBoundsDim())
return failure();
if (!xferOp.getPermutationMap().isMinorIdentity())
return failure();
if (xferOp.getMask())
return failure();
auto extractOp = xferOp.getSource().getDefiningOp<tensor::ExtractSliceOp>();
if (!extractOp)
return failure();
if (!extractOp.hasUnitStride())
return failure();

// Bail on illegal rank-reduction: we need to check that the rank-reduced
// dims are exactly the leading dims. I.e. the following is illegal:
// ```
// %0 = tensor.extract_slice %t[0,0,0][2,1,4][1,1,1] :
// tensor<2x1x4xf32> to tensor<2x4xf32>
// %1 = vector.transfer_read %0[0,0], %cst :
// tensor<2x4xf32>, vector<2x4xf32>
// ```
//
// Cannot fold into:
// ```
// %0 = vector.transfer_read %t[0,0,0], %cst :
// tensor<2x1x4xf32>, vector<2x4xf32>
// ```
// For this, check the trailing `vectorRank` dims of the extract_slice
// result tensor match the trailing dims of the inferred result tensor.
if (!areAllRankReducedLeadingDim(extractOp, extractOp.getType().getRank()))
return failure();

int64_t rankReduced =
extractOp.getSourceType().getRank() - extractOp.getType().getRank();

SmallVector<Value> newIndices;
// In case this is a rank-reducing ExtractSliceOp, copy rank-reduced
// indices first.
for (int64_t i = 0; i < rankReduced; ++i) {
OpFoldResult offset = extractOp.getMixedOffsets()[i];
newIndices.push_back(getValueOrCreateConstantIndexOp(
rewriter, extractOp.getLoc(), offset));
}
for (const auto &it : llvm::enumerate(xferOp.getIndices())) {
OpFoldResult offset =
extractOp.getMixedOffsets()[it.index() + rankReduced];
newIndices.push_back(rewriter.create<arith::AddIOp>(
xferOp->getLoc(), it.value(),
getValueOrCreateConstantIndexOp(rewriter, extractOp.getLoc(),
offset)));
}
SmallVector<bool> inBounds(xferOp.getTransferRank(), true);
rewriter.replaceOpWithNewOp<TransferReadOp>(
xferOp, xferOp.getVectorType(), extractOp.getSource(), newIndices,
xferOp.getPadding(), ArrayRef<bool>{inBounds});

return success();
}
};

/// Store to load forwarding for transfer operations with permuation maps.
/// Even if the permutation maps are different we can still propagate the store		/// Even if the permutation maps are different we can still propagate the store
/// into the load if the size of the dimensions read and written match. Then we		/// into the load if the size of the dimensions read and written match. Then we
/// can replace the transfer_read + transfer_write by vector.broadcast and		/// can replace the transfer_read + transfer_write by vector.broadcast and
/// vector.transpose.		/// vector.transpose.
/// Example:		/// Example:
/// ```		/// ```
/// %w0 = vector.transfer_write %v0, %arg0[%c0, %c0, %c0]		/// %w0 = vector.transfer_write %v0, %arg0[%c0, %c0, %c0]
/// {in_bounds = [true, true],		/// {in_bounds = [true, true],
▲ Show 20 Lines • Show All 64 Lines • ▼ Show 20 Lines	rewriter.replaceOpWithNewOp<vector::TransposeOp>(readOp, vec,
transposePerm);		transposePerm);
return success();		return success();
}		}
};		};
} // namespace		} // namespace

void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,		void TransferReadOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {		MLIRContext *context) {
// clang-format off		results.add<TransferReadAfterWriteToBroadcast>(context);
results.add <
// TODO: this is brittle and should be deprecated in favor of a
// more general pattern that applies on-demand.
FoldExtractSliceIntoTransferRead,
TransferReadAfterWriteToBroadcast>(context);
// clang-format on
}		}

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// TransferWriteOp		// TransferWriteOp
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

/// 1. Builder with type inference.		/// 1. Builder with type inference.
void TransferWriteOp::build(OpBuilder &builder, OperationState &result,		void TransferWriteOp::build(OpBuilder &builder, OperationState &result,
▲ Show 20 Lines • Show All 319 Lines • ▼ Show 20 Lines	while (defWrite) {
break;		break;
writeToModify = defWrite;		writeToModify = defWrite;
defWrite = defWrite.getSource().getDefiningOp<vector::TransferWriteOp>();		defWrite = defWrite.getSource().getDefiningOp<vector::TransferWriteOp>();
}		}
return failure();		return failure();
}		}
};		};

/// Fold tensor.insert_slice into vector.transfer_write if the transfer_write
/// could directly write to the insert_slice's destination. E.g.:
///
/// ```
/// %0 = vector.transfer_write %v, %t1[%c0, %c0] {in_bounds = [true, true]}
/// : vector<4x5xf32>, tensor<4x5xf32>
/// %1 = tensor.insert_slice %0 into %t2[%a, %b] [4, 5] [1, 1]
/// : tensor<4x5xf32> into tensor<?x?xf32>
/// ```
/// is rewritten to:
/// ```
/// %1 = vector.transfer_write %v, %t2[%a, %b] {in_bounds = [true, true]}
/// : vector<4x5xf32>, tensor<?x?xf32>
/// ```
// TODO: this is brittle and should be deprecated in favor of a more general
// pattern that applies on-demand.
struct FoldInsertSliceIntoTransferWrite
: public OpRewritePattern<tensor::InsertSliceOp> {
public:
using OpRewritePattern::OpRewritePattern;

LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
PatternRewriter &rewriter) const override {
if (!insertOp.hasUnitStride())
return failure();

auto xferOp = insertOp.getSource().getDefiningOp<TransferWriteOp>();
if (!xferOp)
return failure();
// TODO: support 0-d corner case.
if (xferOp.getTransferRank() == 0)
return failure();

if (xferOp.hasOutOfBoundsDim())
return failure();
if (xferOp.getVectorType().getRank() != xferOp.getShapedType().getRank())
return failure();
if (xferOp.getMask())
return failure();
// Fold only if the TransferWriteOp completely overwrites the `source` with
// a vector. I.e., the result of the TransferWriteOp is a new tensor whose
// content is the data of the vector.
if (!llvm::equal(xferOp.getVectorType().getShape(),
xferOp.getShapedType().getShape()))
return failure();
if (!xferOp.getPermutationMap().isIdentity())
return failure();

// Bail on illegal rank-reduction: we need to check that the rank-reduced
// dims are exactly the leading dims. I.e. the following is illegal:
// ```
// %0 = vector.transfer_write %v, %t[0,0], %cst :
// vector<2x4xf32>, tensor<2x4xf32>
// %1 = tensor.insert_slice %0 into %tt[0,0,0][2,1,4][1,1,1] :
// tensor<2x4xf32> into tensor<2x1x4xf32>
// ```
//
// Cannot fold into:
// ```
// %0 = vector.transfer_write %v, %t[0,0,0], %cst :
// vector<2x4xf32>, tensor<2x1x4xf32>
// ```
// For this, check the trailing `vectorRank` dims of the insert_slice result
// tensor match the trailing dims of the inferred result tensor.
int64_t rankReduced =
insertOp.getType().getRank() - insertOp.getSourceType().getRank();
int64_t vectorRank = xferOp.getVectorType().getRank();
RankedTensorType inferredSourceTensorType =
tensor::ExtractSliceOp::inferResultType(
insertOp.getType(), insertOp.getMixedOffsets(),
insertOp.getMixedSizes(), insertOp.getMixedStrides());
auto actualSourceTensorShape = insertOp.getSourceType().getShape();
if (rankReduced > 0 &&
actualSourceTensorShape.take_back(vectorRank) !=
inferredSourceTensorType.getShape().take_back(vectorRank))
return failure();

SmallVector<Value> indices = getValueOrCreateConstantIndexOp(
rewriter, insertOp.getLoc(), insertOp.getMixedOffsets());
SmallVector<bool> inBounds(xferOp.getTransferRank(), true);
rewriter.replaceOpWithNewOp<TransferWriteOp>(insertOp, xferOp.getVector(),
insertOp.getDest(), indices,
ArrayRef<bool>{inBounds});
return success();
}
};

/// Rewrite tensor::ExtractSliceOp(vector::TransferWriteOp) to		/// Rewrite tensor::ExtractSliceOp(vector::TransferWriteOp) to
/// vector::TransferWriteOp(tensor::ExtractSliceOp) if the full slice is		/// vector::TransferWriteOp(tensor::ExtractSliceOp) if the full slice is
/// overwritten and inserted into another tensor. After this rewrite, the		/// overwritten and inserted into another tensor. After this rewrite, the
/// operations bufferize in-place since all of them work on the same slice.		/// operations bufferize in-place since all of them work on the same slice.
///		///
/// For example:		/// For example:
/// ```mlir		/// ```mlir
/// %0 = vector.transfer_write %vec, %init_tensor[%c0, %c0]		/// %0 = vector.transfer_write %vec, %init_tensor[%c0, %c0]
▲ Show 20 Lines • Show All 95 Lines • ▼ Show 20 Lines	LogicalResult matchAndRewrite(tensor::InsertSliceOp insertOp,
return success();		return success();
}		}
};		};

} // namespace		} // namespace

void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,		void TransferWriteOp::getCanonicalizationPatterns(RewritePatternSet &results,
MLIRContext *context) {		MLIRContext *context) {
// clang-format off		results.add<FoldWaw, SwapExtractSliceOfTransferWrite>(context);
results.add<FoldWaw,
// TODO: this is brittle and should be deprecated in favor of a
// more general pattern that applies on-demand.
FoldInsertSliceIntoTransferWrite,
SwapExtractSliceOfTransferWrite>(context);
// clang-format on
}		}

//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//
// LoadOp		// LoadOp
//===----------------------------------------------------------------------===//		//===----------------------------------------------------------------------===//

static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,		static LogicalResult verifyLoadStoreMemRefLayout(Operation *op,
MemRefType memRefTy) {		MemRefType memRefTy) {
▲ Show 20 Lines • Show All 1,604 Lines • Show Last 20 Lines

This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Vector] Retire brittle canonicalizations of insert/extract slice with vector.transferChanges PlannedPublic

Details

Diff Detail

Event Timeline

Revision Contents

Diff 507640

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

[mlir][Vector] Retire brittle canonicalizations of insert/extract slice with vector.transfer
Changes PlannedPublic