[AutoDiff] Support differentiation of non-active try_apply. (#33483)

Add differentiation support for non-active `try_apply` SIL instructions.

Notable pullback generation changes:
* Original basic blocks are now visited in a different order:
* starting from the original basic block, all its predecessors
* are visited in a breadth-first search order. This ensures that
* all successors of any block are visited before the block itself.

Resolves TF-433.
This commit is contained in:
Alex Efremov
2020-08-15 19:22:28 +02:00
committed by GitHub
parent fbc0463e62
commit 40c448514a
10 changed files with 183 additions and 116 deletions

View File

@@ -271,59 +271,6 @@ inline void createEntryArguments(SILFunction *f) {
}
}
/// Helper class for visiting basic blocks in post-order post-dominance order,
/// based on a worklist algorithm.
class PostOrderPostDominanceOrder {
SmallVector<DominanceInfoNode *, 16> buffer;
PostOrderFunctionInfo *postOrderInfo;
size_t srcIdx = 0;
public:
/// Constructor.
/// \p root The root of the post-dominator tree.
/// \p postOrderInfo The post-order info of the function.
/// \p capacity Should be the number of basic blocks in the dominator tree to
/// reduce memory allocation.
PostOrderPostDominanceOrder(DominanceInfoNode *root,
PostOrderFunctionInfo *postOrderInfo,
int capacity = 0)
: postOrderInfo(postOrderInfo) {
buffer.reserve(capacity);
buffer.push_back(root);
}
/// Get the next block from the worklist.
DominanceInfoNode *getNext() {
if (srcIdx == buffer.size())
return nullptr;
return buffer[srcIdx++];
}
/// Pushes the dominator children of a block onto the worklist in post-order.
void pushChildren(DominanceInfoNode *node) {
pushChildrenIf(node, [](SILBasicBlock *) { return true; });
}
/// Conditionally pushes the dominator children of a block onto the worklist
/// in post-order.
template <typename Pred>
void pushChildrenIf(DominanceInfoNode *node, Pred pred) {
SmallVector<DominanceInfoNode *, 4> children;
for (auto *child : *node)
children.push_back(child);
llvm::sort(children.begin(), children.end(),
[&](DominanceInfoNode *n1, DominanceInfoNode *n2) {
return postOrderInfo->getPONumber(n1->getBlock()) <
postOrderInfo->getPONumber(n2->getBlock());
});
for (auto *child : children) {
SILBasicBlock *childBB = child->getBlock();
if (pred(childBB))
buffer.push_back(child);
}
}
};
/// Cloner that remaps types using the target function's generic environment.
class BasicTypeSubstCloner final
: public TypeSubstCloner<BasicTypeSubstCloner, SILOptFunctionBuilder> {