This is an archive of the discontinued LLVM Phabricator instance.

Differential D130465

[mlir][tensor][bufferize][NFC] Remove duplicate code
ClosedPublic

Authored by springerm on Jul 25 2022, 1:45 AM.

Details

Reviewers
ftynse
ingomueller-net
Commits
rG1defec873065: [mlir][tensor][bufferize][NFC] Remove duplicate code
Summary

InsertSliceOp and ParallelInsertSliceOp are very similar and can share some of the bufferization analysis code.

Diff Detail

Event Timeline

springerm created this revision.Jul 25 2022, 1:45 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 25 2022, 1:45 AM
Herald added subscribers: bzcheeseman, sdasgup3, wenzhicui and 18 others. · View Herald Transcript
springerm requested review of this revision.Jul 25 2022, 1:45 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 25 2022, 1:45 AM
Herald added subscribers: stephenneuendorffer, nicolasvasilache. · View Herald Transcript
Harbormaster completed remote builds in B177298: Diff 447227.Jul 25 2022, 2:01 AM
ftynse accepted this revision.Jul 25 2022, 2:27 AM

Looks mechanical.

This revision is now accepted and ready to land.Jul 25 2022, 2:27 AM
Closed by commit rG1defec873065: [mlir][tensor][bufferize][NFC] Remove duplicate code (authored by springerm). · Explain WhyJul 25 2022, 3:34 AM
This revision was automatically updated to reflect the committed changes.
springerm added a commit: rG1defec873065: [mlir][tensor][bufferize][NFC] Remove duplicate code.

Revision Contents

PathSize
mlir/
lib/
Dialect/
Tensor/
Transforms/
281 lines

Diff 447267

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

Show First 20 Lines • Show All 546 Lines • ▼ Show 20 Linesstruct InsertOpInterface
BufferRelation bufferRelation(Operation *op, OpResult opResult, BufferRelation bufferRelation(Operation *op, OpResult opResult,
const AnalysisState &state) const { const AnalysisState &state) const {
return BufferRelation::Equivalent; return BufferRelation::Equivalent;
} }
}; };
/// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e. /// Return true if the (ExtractSliceOp, InsertSliceOp) pair match (i.e.
/// equivalent operand / result and same offset/sizes/strides specification). /// equivalent operand / result and same offset/sizes/strides specification).
/// template <typename OpTy>
/// This is one particular type of relationship between ops on tensors that
/// reduce to an equivalence on buffers. This should be generalized and
/// exposed as interfaces on the proper types.
static bool areEquivalentExtractSliceOps(const AnalysisState &state, static bool areEquivalentExtractSliceOps(const AnalysisState &state,
ExtractSliceOp st, InsertSliceOp sti) { ExtractSliceOp extractSliceOp,
if (!st || !sti) OpTy insertSliceOp) {
if (!extractSliceOp || !insertSliceOp)
return false;
if (extractSliceOp != insertSliceOp &&
!state.areEquivalentBufferizedValues(extractSliceOp.getSource(),
insertSliceOp.getDest()))
return false; return false;
if (sti != sti && if (!sameOffsetsSizesAndStrides(extractSliceOp, insertSliceOp,
!state.areEquivalentBufferizedValues(st.getSource(), sti.getDest())) isEqualConstantIntOrValue))
return false;
if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
return false; return false;
return true; return true;
} }
/// Return true if `value` is originating from an ExtractSliceOp that matches /// Return true if `value` is originating from an ExtractSliceOp that matches
/// the given InsertSliceOp. /// the given InsertSliceOp.
template <typename OpTy>
static bool hasMatchingExtractSliceOp(const AnalysisState &state, Value value, static bool hasMatchingExtractSliceOp(const AnalysisState &state, Value value,
InsertSliceOp insertOp) { OpTy insertSliceOp) {
auto condition = [&](Value val) { auto condition = [&](Value val) {
if (auto extractOp = val.getDefiningOp<ExtractSliceOp>()) if (auto extractSliceOp = val.getDefiningOp<ExtractSliceOp>())
if (areEquivalentExtractSliceOps(state, extractOp, insertOp)) if (areEquivalentExtractSliceOps(state, extractSliceOp, insertSliceOp))
return true; return true;
return false; return false;
}; };
return llvm::all_of(state.findValueInReverseUseDefChain(value, condition), return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
condition); condition);
} }
/// Bufferization of tensor.insert_slice. Replace with a memory copy. Under template <typename OpTy>
/// certain circumstances, this op can also be a no-op. static bool isNotConflictingInsertSliceLikeOp(Operation *op, OpOperand *uRead,
struct InsertSliceOpInterface
: public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
tensor::InsertSliceOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return &opOperand == &op->getOpOperand(1) /*dest*/;
}
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
if (&opOperand == &op->getOpOperand(1) /*dest*/)
return {op->getResult(0)};
return {};
}
BufferRelation bufferRelation(Operation *op, OpResult opResult,
const AnalysisState &state) const {
return BufferRelation::Equivalent;
}
bool isNotConflicting(Operation *op, OpOperand *uRead,
OpOperand *uConflictingWrite, OpOperand *uConflictingWrite,
const AnalysisState &state) const { const AnalysisState &state) {
Operation *readingOp = uRead->getOwner(); Operation *readingOp = uRead->getOwner();
Operation *conflictingWritingOp = uConflictingWrite->getOwner(); Operation *conflictingWritingOp = uConflictingWrite->getOwner();
// Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If // Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
// uRead is an InsertSliceOp... // uRead is an InsertSliceOp...
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(readingOp)) { if (auto insertSliceOp = dyn_cast<OpTy>(readingOp)) {
// As an example, consider the following IR. // As an example, consider the following IR.
// //
// %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] } // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
// %1 = linalg.fill %cst, %0 {inplace= [true] } // %1 = linalg.fill %cst, %0 {inplace= [true] }
// %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1] // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
// {inplace= [true] } // {inplace= [true] }
// TODO: Use insertSliceOp.getDestOpOperand etc. when available. // TODO: Use insertSliceOp.getDestOpOperand etc. when available.
if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ && if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
hasMatchingExtractSliceOp(state, uConflictingWrite->get(), hasMatchingExtractSliceOp(state, uConflictingWrite->get(),
insertSliceOp)) insertSliceOp))
// Case 1: The main insight is that InsertSliceOp reads only part of // Case 1: The main insight is that InsertSliceOp reads only part of
// the destination tensor. The overwritten area is not read. If // the destination tensor. The overwritten area is not read. If
// uConflictingWrite writes into exactly the memory location that is // uConflictingWrite writes into exactly the memory location that is
// being read by uRead, this is not a conflict. // being read by uRead, this is not a conflict.
// //
// In the above example: // In the above example:
// uRead = OpOperand 1 (%t) of tensor.insert_slice // uRead = OpOperand 1 (%t) of tensor.insert_slice
// uConflictingWrite = OpOperand 1 (%0) of linalg.fill // uConflictingWrite = OpOperand 1 (%0) of linalg.fill
// //
// The read of %t does not conflict with the write of the FillOp // The read of %t does not conflict with the write of the FillOp
// (same aliases!) because the area that the FillOp operates on is // (same aliases!) because the area that the FillOp operates on is
// exactly the one that is *not* read via %t. // exactly the one that is *not* read via %t.
return true; return true;
if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ && if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ && uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp)) hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp))
// Case 2: The read of the source tensor and the write to the dest // Case 2: The read of the source tensor and the write to the dest
// tensor via an InsertSliceOp is not a conflict if the read is // tensor via an InsertSliceOp is not a conflict if the read is
// reading exactly that part of an equivalent tensor that the // reading exactly that part of an equivalent tensor that the
// InsertSliceOp is writing. // InsertSliceOp is writing.
// //
// In the above example: // In the above example:
// uRead = OpOperand 0 (%1) of tensor.insert_slice // uRead = OpOperand 0 (%1) of tensor.insert_slice
// uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
return true; return true;
} }
// If uConflictingWrite is an InsertSliceOp... // If uConflictingWrite is an InsertSliceOp...
if (auto insertSliceOp = dyn_cast<InsertSliceOp>(conflictingWritingOp)) if (auto insertSliceOp = dyn_cast<OpTy>(conflictingWritingOp))
// As an example, consider the following IR. // As an example, consider the following IR.
// //
// %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] } // %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
// %1 = linalg.fill %cst, %0 {inplace= [true] } // %1 = linalg.fill %cst, %0 {inplace= [true] }
// %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1] // %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
// {inplace= [true] } // {inplace= [true] }
// %3 = vector.transfer_read %1, %cst // %3 = vector.transfer_read %1, %cst
// //
// In the above example: // In the above example:
// uRead = OpOperand 0 (%1) of vector.transfer_read // uRead = OpOperand 0 (%1) of vector.transfer_read
// uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice // uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
// lastWrite = %1 // lastWrite = %1
// //
// This is not a conflict because the InsertSliceOp overwrites the // This is not a conflict because the InsertSliceOp overwrites the
// memory segment of %1 with the exact same data. (Effectively, there // memory segment of %1 with the exact same data. (Effectively, there
// is no memory write here.) // is no memory write here.)
if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ && if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
state.areEquivalentBufferizedValues(uRead->get(), state.areEquivalentBufferizedValues(uRead->get(),
insertSliceOp.getSource()) && insertSliceOp.getSource()) &&
hasMatchingExtractSliceOp(state, insertSliceOp.getSource(), hasMatchingExtractSliceOp(state, insertSliceOp.getSource(),
insertSliceOp)) insertSliceOp))
return true; return true;
return false; return false;
} }
/// Bufferization of tensor.insert_slice. Replace with a memory copy. Under
/// certain circumstances, this op can also be a no-op.
struct InsertSliceOpInterface
: public BufferizableOpInterface::ExternalModel<InsertSliceOpInterface,
tensor::InsertSliceOp> {
bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return true;
}
bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
return &opOperand == &op->getOpOperand(1) /*dest*/;
}
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const {
if (&opOperand == &op->getOpOperand(1) /*dest*/)
return {op->getResult(0)};
return {};
}
BufferRelation bufferRelation(Operation *op, OpResult opResult,
const AnalysisState &state) const {
return BufferRelation::Equivalent;
}
bool isNotConflicting(Operation *op, OpOperand *uRead,
OpOperand *uConflictingWrite,
const AnalysisState &state) const {
return isNotConflictingInsertSliceLikeOp<tensor::InsertSliceOp>(
op, uRead, uConflictingWrite, state);
}
LogicalResult bufferize(Operation *op, RewriterBase &rewriter, LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
const BufferizationOptions &options) const { const BufferizationOptions &options) const {
// insert_slice ops arise from tiling and bufferizing them out-of-place is // insert_slice ops arise from tiling and bufferizing them out-of-place is
// generally a deal breaker. When used with loops, this ends up cloning the // generally a deal breaker. When used with loops, this ends up cloning the
// whole tensor on every single iteration and is a symptom of a // whole tensor on every single iteration and is a symptom of a
// catastrophically bad scheduling decision. // catastrophically bad scheduling decision.
// TODO: be very loud about it or even consider failing the pass. // TODO: be very loud about it or even consider failing the pass.
auto insertSliceOp = cast<tensor::InsertSliceOp>(op); auto insertSliceOp = cast<tensor::InsertSliceOp>(op);
▲ Show 20 Lines • Show All 103 Lines • ▼ Show 20 Lines auto resultMemRefType = getMemRefType(
reshapeOp.getResult(), options, /*layout=*/{}, reshapeOp.getResult(), options, /*layout=*/{},
srcBuffer->getType().cast<BaseMemRefType>().getMemorySpaceAsInt()); srcBuffer->getType().cast<BaseMemRefType>().getMemorySpaceAsInt());
replaceOpWithNewBufferizedOp<memref::ReshapeOp>( replaceOpWithNewBufferizedOp<memref::ReshapeOp>(
rewriter, op, resultMemRefType, *srcBuffer, *shapeBuffer); rewriter, op, resultMemRefType, *srcBuffer, *shapeBuffer);
return success(); return success();
} }
}; };
/// Return true if the (ExtractSliceOp, ParallelInsertSliceOp) pair match (i.e.
/// equivalent operand / result and same offset/sizes/strides specification).
static bool areEquivalentExtractSliceOps(const AnalysisState &state,
ExtractSliceOp st,
ParallelInsertSliceOp sti) {
if (!st || !sti)
return false;
if (st != sti &&
!state.areEquivalentBufferizedValues(st.getSource(), sti.getDest()))
return false;
if (!sameOffsetsSizesAndStrides(st, sti, isEqualConstantIntOrValue))
return false;
return true;
}
/// Return true if `value` is originating from an ExtractSliceOp that matches
/// the given InsertSliceOp.
static bool hasMatchingExtractSliceOp(const AnalysisState &state, Value value,
ParallelInsertSliceOp insertOp) {
auto condition = [&](Value val) {
if (auto extractOp = val.getDefiningOp<ExtractSliceOp>())
if (areEquivalentExtractSliceOps(state, extractOp, insertOp))
return true;
return false;
};
return llvm::all_of(state.findValueInReverseUseDefChain(value, condition),
condition);
}
/// Analysis of ParallelInsertSliceOp. /// Analysis of ParallelInsertSliceOp.
struct ParallelInsertSliceOpInterface struct ParallelInsertSliceOpInterface
: public BufferizableOpInterface::ExternalModel< : public BufferizableOpInterface::ExternalModel<
ParallelInsertSliceOpInterface, ParallelInsertSliceOp> { ParallelInsertSliceOpInterface, ParallelInsertSliceOp> {
SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand, SmallVector<OpResult> getAliasingOpResult(Operation *op, OpOperand &opOperand,
const AnalysisState &state) const { const AnalysisState &state) const {
if (&opOperand != &op->getOpOperand(1) /*dest*/) if (&opOperand != &op->getOpOperand(1) /*dest*/)
return {}; return {};
▲ Show 20 Lines • Show All 127 Lines • ▼ Show 20 Lines LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
for (OpOperand *use : resultUses) { for (OpOperand *use : resultUses) {
rewriter.updateRootInPlace(use->getOwner(), rewriter.updateRootInPlace(use->getOwner(),
[&]() { use->set(toTensorOp); }); [&]() { use->set(toTensorOp); });
} }
rewriter.eraseOp(op); rewriter.eraseOp(op);
return success(); return success();
} }
// TODO: This is copied from TensorInterfaceImpl.cpp. Find a way to share
// the code.
bool isNotConflicting(Operation *op, OpOperand *uRead, bool isNotConflicting(Operation *op, OpOperand *uRead,
OpOperand *uConflictingWrite, OpOperand *uConflictingWrite,
const AnalysisState &state) const { const AnalysisState &state) const {
Operation *readingOp = uRead->getOwner(); return isNotConflictingInsertSliceLikeOp<tensor::ParallelInsertSliceOp>(
Operation *conflictingWritingOp = uConflictingWrite->getOwner(); op, uRead, uConflictingWrite, state);
// Special rules for matching ExtractSliceOp/InsertSliceOp pairs. If
// uRead is an InsertSliceOp...
if (auto insertSliceOp = dyn_cast<ParallelInsertSliceOp>(readingOp)) {
// As an example, consider the following IR.
//
// %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
// %1 = linalg.fill %cst, %0 {inplace= [true] }
// %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
// {inplace= [true] }
// TODO: Use insertSliceOp.getDestOpOperand etc. when available.
if (uRead == &insertSliceOp->getOpOperand(1) /*dest*/ &&
hasMatchingExtractSliceOp(state, uConflictingWrite->get(),
insertSliceOp))
// Case 1: The main insight is that InsertSliceOp reads only part of
// the destination tensor. The overwritten area is not read. If
// uConflictingWrite writes into exactly the memory location that is
// being read by uRead, this is not a conflict.
//
// In the above example:
// uRead = OpOperand 1 (%t) of tensor.insert_slice
// uConflictingWrite = OpOperand 1 (%0) of linalg.fill
//
// The read of %t does not conflict with the write of the FillOp
// (same aliases!) because the area that the FillOp operates on is
// exactly the one that is *not* read via %t.
return true;
if (uRead == &insertSliceOp->getOpOperand(0) /*source*/ &&
uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
hasMatchingExtractSliceOp(state, uRead->get(), insertSliceOp))
// Case 2: The read of the source tensor and the write to the dest
// tensor via an InsertSliceOp is not a conflict if the read is
// reading exactly that part of an equivalent tensor that the
// InsertSliceOp is writing.
//
// In the above example:
// uRead = OpOperand 0 (%1) of tensor.insert_slice
// uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
return true;
}
// If uConflictingWrite is an InsertSliceOp...
if (auto insertSliceOp =
dyn_cast<ParallelInsertSliceOp>(conflictingWritingOp))
// As an example, consider the following IR.
//
// %0 = tensor.extract_slice %t[%a, %b][%c, %d][1, 1] {inplace = [true] }
// %1 = linalg.fill %cst, %0 {inplace= [true] }
// %2 = tensor.insert_slice %1 into %t[%a, %b][%c, %d][1, 1]
// {inplace= [true] }
// %3 = vector.transfer_read %1, %cst
//
// In the above example:
// uRead = OpOperand 0 (%1) of vector.transfer_read
// uConflictingWrite = OpOperand 1 (%t) of tensor.insert_slice
// lastWrite = %1
//
// This is not a conflict because the InsertSliceOp overwrites the
// memory segment of %1 with the exact same data. (Effectively, there
// is no memory write here.)
if (uConflictingWrite == &insertSliceOp->getOpOperand(1) /*dest*/ &&
state.areEquivalentBufferizedValues(uRead->get(),
insertSliceOp.getSource()) &&
hasMatchingExtractSliceOp(state, insertSliceOp.getSource(),
insertSliceOp))
return true;
return false;
} }
}; };
} // namespace } // namespace
} // namespace tensor } // namespace tensor
} // namespace mlir } // namespace mlir
void mlir::tensor::registerBufferizableOpInterfaceExternalModels( void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
Show All 21 Lines