Skip to content

Commit

Permalink
[Stream] fix ordering issue in EmplaceAllocations (iree-org#18321)
Browse files Browse the repository at this point in the history 
Fix a dominance issue when the updates are from last to first results
for a dispatch with multiple results. This is done by adding two phases
to the pass. First phase collects the update ops that can be potentially
placed in dispatch results. Then we sort and apply them in block order.
Fixes : iree-org#18308
  • Loading branch information
@nirvedhmeshram
nirvedhmeshram authored Aug 22, 2024
1 parent e7a1898 commit 292f2d4
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 54 deletions.
113 changes: 59 additions & 54 deletions compiler/src/iree/compiler/Dialect/Stream/Transforms/EmplaceAllocations.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@ replaceUsesAndTransfer(Value oldValue, Value newValue,
// first we can't then bail and leave them out-of-place.
static bool tryEmplaceDispatchOp(IREE::Stream::AsyncDispatchOp dispatchOp) {
bool didChange = false;
// Collect the Update ops and their corresponding resultIndex.
SmallVector<std::tuple<IREE::Stream::AsyncUpdateOp, int>> updateOpDataVec;
for (auto [resultIndex, result] : llvm::enumerate(dispatchOp.getResults())) {
// Ignore results with multiple users. We could potentially place these but
// that makes tracking much more complicated.
Expand All @@ -77,66 +79,68 @@ static bool tryEmplaceDispatchOp(IREE::Stream::AsyncDispatchOp dispatchOp) {
}

// Find potential update to place the dispatch result into.
Value targetResource;
Value targetResourceSize;
Value targetOffset;
Value targetEnd;
Value targetLength;
Value targetResult;
Value targetResultSize;
Attribute targetAffinityAttr;
Operation *userOp = *result.user_begin();
if (auto updateOp = dyn_cast<IREE::Stream::AsyncUpdateOp>(userOp)) {
if (updateOp.getUpdate() != result) {
// TODO(#14566): continue if sparse emplacement on multiple results.
break;
}

// Currently only allow exactly matching affinities.
// TODO(multi-device): memory compatibility - if compatible then allow.
if (updateOp.getAffinityAttr() != dispatchOp.getAffinityAttr()) {
continue;
}

// Try to move all SSA values required into the appropriate place.
// TODO(benvanik): undo this if there's a failure (or record/roll-back).
if (!IREE::Util::tryMoveProducerBefore(updateOp.getUpdateSize(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetSize(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetOffset(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetEnd(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTarget(),
dispatchOp)) {
// Failed to move while keeping valid SSA dominance.
// TODO(#14566): continue if sparse emplacement on multiple results.
break;
}

targetResource = updateOp.getTarget();
if (targetResource.getDefiningOp() == dispatchOp) {
// NOTE: we may have already replaced the update target with one of our
// results - if so we need to find the operand to capture tied to that
// new result instead of our own new result (which would make a cycle).
targetResource = dispatchOp.getTiedResultOperand(targetResource);
}
targetResourceSize = updateOp.getTargetSize();
targetOffset = updateOp.getTargetOffset();
targetEnd = updateOp.getTargetEnd();
targetLength = updateOp.getUpdateSize();
targetResult = updateOp.getResult();
targetResultSize = updateOp.getTargetSize();
targetAffinityAttr = updateOp.getAffinityAttr();
auto updateOp = dyn_cast_or_null<IREE::Stream::AsyncUpdateOp>(userOp);
if (!updateOp)
break;
if (updateOp.getUpdate() != result) {
// TODO(#14566): continue if sparse emplacement on multiple results.
break;
}

// Currently only allow exactly matching affinities.
// TODO(multi-device): memory compatibility - if compatible then allow.
if (updateOp.getAffinityAttr() != dispatchOp.getAffinityAttr()) {
continue;
}
updateOpDataVec.emplace_back(updateOp, resultIndex);
}
// Sort the update ops in block order so that we dont accidentally move them
// above the dispatch op in the next section which will cause a dominance
// issue.
llvm::sort(updateOpDataVec, [&](const auto &a, const auto &b) {
return std::get<0>(a)->isBeforeInBlock(std::get<0>(b));
});
for (auto updateOpData : updateOpDataVec) {
int resultIndex;
IREE::Stream::AsyncUpdateOp updateOp;
std::tie(updateOp, resultIndex) = updateOpData;
Value targetResource = updateOp.getTarget();
if (targetResource.getDefiningOp() == dispatchOp) {
// NOTE: we may have already replaced the update target with one of our
// results - if so we need to find the operand to capture tied to that
// new result instead of our own new result (which would make a cycle).
targetResource = dispatchOp.getTiedResultOperand(targetResource);
}
if (!targetResource) {
// TODO(#14566): continue if sparse emplacement on multiple results.
break;
}
Value targetResourceSize = updateOp.getTargetSize();
Value targetOffset = updateOp.getTargetOffset();
Value targetEnd = updateOp.getTargetEnd();
Value targetLength = updateOp.getUpdateSize();
Value targetResult = updateOp.getResult();
Value targetResultSize = updateOp.getTargetSize();

// Try to move all SSA values required into the appropriate place.
// TODO(benvanik): undo this if there's a failure (or record/roll-back).
if (!IREE::Util::tryMoveProducerBefore(updateOp.getUpdateSize(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetSize(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetOffset(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTargetEnd(),
dispatchOp) ||
!IREE::Util::tryMoveProducerBefore(updateOp.getTarget(), dispatchOp)) {
// Failed to move while keeping valid SSA dominance.
// TODO(#14566): continue if sparse emplacement on multiple results.
break;
}

// Add operand and tie the result.
operandIndex = dispatchOp.getResourceOperands().size();
auto operandIndex = dispatchOp.getResourceOperands().size();
dispatchOp.getResourceOperandsMutable().append(targetResource);
dispatchOp.getResourceOperandSizesMutable().append(targetResourceSize);
dispatchOp.getResourceOperandOffsetsMutable().append(targetOffset);
Expand All @@ -150,8 +154,9 @@ static bool tryEmplaceDispatchOp(IREE::Stream::AsyncDispatchOp dispatchOp) {
dispatchOp.getResultSizesMutable().assign(resultSizes);

// Replace users with the result of the dispatch op.
replaceUsesAndTransfer(targetResult, result, dispatchOp.getAffinityAttr());
userOp->erase();
replaceUsesAndTransfer(targetResult, updateOp.getUpdate(),
dispatchOp.getAffinityAttr());
updateOp->erase();

didChange = true;
}
Expand Down
29 changes: 29 additions & 0 deletions compiler/src/iree/compiler/Dialect/Stream/Transforms/test/emplace_allocations.mlir
View file
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,35 @@ util.func public @emplaceMultiResultDispatchInto(

// -----

// Test that multiple results with the last result updated first get placed correctly

// CHECK-LABEL: @emplaceMultiResultDispatchIntoSwapped
util.func public @emplaceMultiResultDispatchIntoSwapped(
// CHECK-SAME: %[[INPUT:arg[0-9]+]]: !stream.resource<*>, %[[INPUT_SIZE:arg[0-9]+]]: index,
%input: !stream.resource<*>, %input_size: index,
// CHECK-SAME: %[[UPDATE_SIZE:arg[0-9]+]]: index, %[[TARGET_SIZE:arg[0-9]+]]: index
%update_size: index, %target_size: index) -> !stream.resource<*> {
%c0 = arith.constant 0 : index
%c32 = arith.constant 32 : index
%c64 = arith.constant 64 : index
// CHECK: %[[TARGET:.+]] = stream.async.alloca
// CHECK: %[[DISPATCH:.+]]:2 = stream.async.dispatch @ex::@dispatch0
// CHECK-SAME: ({{.+}}, %[[TARGET]][%c0 to %c32 for %[[UPDATE_SIZE]]], %[[TARGET]][%c32 to %c64 for %[[UPDATE_SIZE]]]) :
// CHECK-SAME: ({{.+}}) -> (%[[TARGET]]{%[[TARGET_SIZE]]}, %[[TARGET]]{%[[TARGET_SIZE]]})
%update:2 = stream.async.dispatch @ex::@dispatch0(%input[%c0 to %input_size for %input_size]) :
(!stream.resource<*>{%input_size}) -> (!stream.resource<*>{%update_size}, !stream.resource<*>{%update_size})
// CHECK-NOT: stream.async.alloca
%target = stream.async.alloca : !stream.resource<*>{%target_size}
// CHECK-NOT: stream.async.update
%target0 = stream.async.update %update#1, %target[%c0 to %c32] : !stream.resource<*>{%update_size} -> %target as !stream.resource<*>{%target_size}
// CHECK-NOT: stream.async.update
%target1 = stream.async.update %update#0, %target0[%c32 to %c64] : !stream.resource<*>{%update_size} -> %target0 as !stream.resource<*>{%target_size}
// CHECK: util.return %[[DISPATCH]]#0
util.return %target1 : !stream.resource<*>
}

// -----

// TODO(#14566): multiple results with sparse ties don't work due to implicit
// operand/result ordering on the dispatch ops. Flow and stream dispatch ops and
// the executable entry points need to be reworked to remove the implicit
Expand Down

0 comments on commit 292f2d4

Please sign in to comment.