finalizeRootUpdate(root);
}
+ /// Find uses of `from` and replace it with `to`. It also marks every modified
+ /// uses and notifies the rewriter that an in-place operation modification is
+ /// about to happen.
+ void replaceAllUsesWith(Value from, Value to);
+
/// Used to notify the rewriter that the IR failed to be rewritten because of
/// a match failure, and provide a callback to populate a diagnostic with the
/// reason why the failure occurred. This method allows for derived rewriters
sourceBlock.getOperations());
// Step 5. RAUW thread indices to thread ops.
- for (Value blockIdx : foreachThreadOp.getThreadIndices()) {
- Value val = bvm.lookup(blockIdx);
- SmallVector<OpOperand *> uses;
- for (OpOperand &use : blockIdx.getUses())
- uses.push_back(&use);
- for (OpOperand *operand : uses) {
- Operation *op = operand->getOwner();
- rewriter.updateRootInPlace(op, [&]() { operand->set(val); });
- }
+ for (Value loopIndex : foreachThreadOp.getThreadIndices()) {
+ Value blockIdx = bvm.lookup(loopIndex);
+ rewriter.replaceAllUsesWith(loopIndex, blockIdx);
}
// Step 6. Erase old op.
sourceBlock.getOperations());
// Step 6. RAUW thread indices to thread ops.
- for (Value threadIdx : foreachThreadOp.getThreadIndices()) {
- Value val = bvm.lookup(threadIdx);
- SmallVector<OpOperand *> uses;
- for (OpOperand &use : threadIdx.getUses())
- uses.push_back(&use);
- for (OpOperand *operand : uses) {
- Operation *op = operand->getOwner();
- rewriter.updateRootInPlace(op, [&]() { operand->set(val); });
- }
+ for (Value loopIndex : foreachThreadOp.getThreadIndices()) {
+ Value threadIdx = bvm.lookup(loopIndex);
+ rewriter.replaceAllUsesWith(loopIndex, threadIdx);
}
// Step 7. syncthreads.
source->erase();
}
+/// Find uses of `from` and replace it with `to`
+void RewriterBase::replaceAllUsesWith(Value from, Value to) {
+ for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
+ Operation *op = operand.getOwner();
+ updateRootInPlace(op, [&]() { operand.set(to); });
+ }
+}
+
// Merge the operations of block 'source' before the operation 'op'. Source
// block should not have existing predecessors or successors.
void RewriterBase::mergeBlockBefore(Block *source, Operation *op,