First steps towards TF/XLA control flow lowering: functional if lowering.
- Implement support for the TensorFlow 'If' op, the first TF op definition.
- Fill in some missing basic infra, including the ability to split a basic block, the ability to create a branch with operands, etc.
- Implement basic lowering for some simple forms of If, where the condition is a zero-D bool tensor and when all the types line up. Future patches will generalize this.
There is still much to be done here. I'd like to get some example graphs coming from the converter to play with to direct this work.
PiperOrigin-RevId: 210198760
diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index bd78929..ef468de 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -167,10 +167,20 @@
// CFG function elements.
//===----------------------------------------------------------------------===//
-// Basic block.
-BasicBlock *CFGFuncBuilder::createBlock() {
+/// Add new basic block and set the insertion point to the end of it. If an
+/// 'insertBefore' basic block is passed, the block will be placed before the
+/// specified block. If not, the block will be appended to the end of the
+/// current function.
+BasicBlock *CFGFuncBuilder::createBlock(BasicBlock *insertBefore) {
BasicBlock *b = new BasicBlock();
- function->push_back(b);
+
+ // If we are supposed to insert before a specific block, do so, otherwise add
+ // the block to the end of the function.
+ if (insertBefore)
+ function->getBlocks().insert(CFGFunction::iterator(insertBefore), b);
+ else
+ function->push_back(b);
+
setInsertionPoint(b);
return b;
}