lib/Transforms/LowerIfAndFor.cpp - platform/external/tensorflow - Git at Google

 //===- LowerIfAndFor.cpp - Lower If and For instructions to CFG -----------===//
 //
 // Copyright 2019 The MLIR Authors.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
 // You may obtain a copy of the License at
 //
 //   http://www.apache.org/licenses/LICENSE-2.0
 //
 // Unless required by applicable law or agreed to in writing, software
 // distributed under the License is distributed on an "AS IS" BASIS,
 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // =============================================================================
 //
 // This file lowers If and For instructions within a function into their lower
 // level CFG equivalent blocks.
 //
 //===----------------------------------------------------------------------===//

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Pass.h"
 #include "mlir/StandardOps/StandardOps.h"
 #include "mlir/Transforms/Passes.h"
 using namespace mlir;

 namespace {
 class LowerIfAndForPass : public FunctionPass {
 public:
   LowerIfAndForPass() : FunctionPass(&passID) {}
   PassResult runOnFunction(Function *function) override;

   void lowerForInst(ForInst *forInst);
   void lowerIfInst(IfInst *ifInst);

   static char passID;
 };
 } // end anonymous namespace

 char LowerIfAndForPass::passID = 0;

 // Given a range of values, emit the code that reduces them with "min" or "max"
 // depending on the provided comparison predicate.  The predicate defines which
 // comparison to perform, "lt" for "min", "gt" for "max" and is used for the
 // `cmpi` operation followed by the `select` operation:
 //
 //   %cond   = cmpi "predicate" %v0, %v1
 //   %result = select %cond, %v0, %v1
 //
 // Multiple values are scanned in a linear sequence.  This creates a data
 // dependences that wouldn't exist in a tree reduction, but is easier to
 // recognize as a reduction by the subsequent passes.
 static Value *buildMinMaxReductionSeq(
     Location loc, CmpIPredicate predicate,
     llvm::iterator_range<OperationInst::result_iterator> values,
     FuncBuilder &builder) {
   assert(!llvm::empty(values) && "empty min/max chain");

   auto valueIt = values.begin();
   Value *value = *valueIt++;
   for (; valueIt != values.end(); ++valueIt) {
     auto cmpOp = builder.create<CmpIOp>(loc, predicate, value, *valueIt);
     value = builder.create<SelectOp>(loc, cmpOp->getResult(), value, *valueIt);
   }

   return value;
 }

 // Convert a "for" loop to a flow of blocks.
 //
 // Create an SESE region for the loop (including its body) and append it to the
 // end of the current region.  The loop region consists of the initialization
 // block that sets up the initial value of the loop induction variable (%iv) and
 // computes the loop bounds that are loop-invariant in MLFunctions; the
 // condition block that checks the exit condition of the loop; the body SESE
 // region; and the end block that post-dominates the loop.  The end block of the
 // loop becomes the new end of the current SESE region.  The body of the loop is
 // constructed recursively after starting a new region (it may be, for example,
 // a nested loop).  Induction variable modification is appended to the body SESE
 // region that always loops back to the condition block.
 //
 //      +--------------------------------+
 //      |   <code before the ForInst>    |
 //      |   <compute initial %iv value>  |
 //      |   br cond(%iv)                 |
 //      +--------------------------------+
 //             |
 //  -------|   |
 //  |      v   v
 //  |   +--------------------------------+
 //  |   | cond(%iv):                     |
 //  |   |   <compare %iv to upper bound> |
 //  |   |   cond_br %r, body, end        |
 //  |   +--------------------------------+
 //  |          |               |
 //  |          |               -------------|
 //  |          v                            |
 //  |   +--------------------------------+  |
 //  |   | body:                          |  |
 //  |   |   <body contents>              |  |
 //  |   |   %new_iv =<add step to %iv>   |  |
 //  |   |   br cond(%new_iv)             |  |
 //  |   +--------------------------------+  |
 //  |          |                            |
 //  |-----------        |--------------------
 //                      v
 //      +--------------------------------+
 //      | end:                           |
 //      |   <code after the ForInst>     |
 //      +--------------------------------+
 //
 void LowerIfAndForPass::lowerForInst(ForInst *forInst) {
   auto loc = forInst->getLoc();

   // Start by splitting the block containing the 'for' into two parts.  The part
   // before will get the init code, the part after will be the end point.
   auto *initBlock = forInst->getBlock();
   auto *endBlock = initBlock->splitBlock(forInst);

   // Create the condition block, with its argument for the loop induction
   // variable.  We set it up below.
   auto *conditionBlock = new Block();
   conditionBlock->insertBefore(endBlock);
   auto *iv = conditionBlock->addArgument(IndexType::get(forInst->getContext()));

   // Create the body block, moving the body of the forInst over to it.
   auto *bodyBlock = new Block();
   bodyBlock->insertBefore(endBlock);

   auto *oldBody = forInst->getBody();
   bodyBlock->getInstructions().splice(bodyBlock->begin(),
                                       oldBody->getInstructions(),
                                       oldBody->begin(), oldBody->end());

   // The code in the body of the forInst now uses 'iv' as its indvar.
   forInst->replaceAllUsesWith(iv);

   // Append the induction variable stepping logic and branch back to the exit
   // condition block.  Construct an affine map f : (x -> x+step) and apply this
   // map to the induction variable.
   FuncBuilder builder(bodyBlock);
   auto affStep = builder.getAffineConstantExpr(forInst->getStep());
   auto affDim = builder.getAffineDimExpr(0);
   auto affStepMap = builder.getAffineMap(1, 0, {affDim + affStep}, {});
   auto stepOp = builder.create<AffineApplyOp>(loc, affStepMap, iv);
   builder.create<BranchOp>(loc, conditionBlock, stepOp->getResult(0));

   // Now that the body block done, fill in the code to compute the bounds of the
   // induction variable in the init block.
   builder.setInsertionPointToEnd(initBlock);

   // Compute loop bounds using an affine_apply.
   SmallVector<Value *, 8> operands(forInst->getLowerBoundOperands());
   auto lbAffineApply =
       builder.create<AffineApplyOp>(loc, forInst->getLowerBoundMap(), operands);
   Value *lowerBound = buildMinMaxReductionSeq(
       loc, CmpIPredicate::SGT, lbAffineApply->getResults(), builder);

   operands.assign(forInst->getUpperBoundOperands().begin(),
                   forInst->getUpperBoundOperands().end());
   auto ubAffineApply =
       builder.create<AffineApplyOp>(loc, forInst->getUpperBoundMap(), operands);
   Value *upperBound = buildMinMaxReductionSeq(
       loc, CmpIPredicate::SLT, ubAffineApply->getResults(), builder);
   builder.create<BranchOp>(loc, conditionBlock, lowerBound);

   // With the body block done, we can fill in the condition block.
   builder.setInsertionPointToEnd(conditionBlock);
   auto comparison =
       builder.create<CmpIOp>(loc, CmpIPredicate::SLT, iv, upperBound);
   builder.create<CondBranchOp>(loc, comparison, bodyBlock, ArrayRef<Value *>(),
                                endBlock, ArrayRef<Value *>());

   // Ok, we're done!
   forInst->erase();
 }

 // Convert an "if" instruction into a flow of basic blocks.
 //
 // Create an SESE region for the if instruction (including its "then" and
 // optional "else" instruction blocks) and append it to the end of the current
 // region.  The conditional region consists of a sequence of condition-checking
 // blocks that implement the short-circuit scheme, followed by a "then" SESE
 // region and an "else" SESE region, and the continuation block that
 // post-dominates all blocks of the "if" instruction.  The flow of blocks that
 // correspond to the "then" and "else" clauses are constructed recursively,
 // enabling easy nesting of "if" instructions and if-then-else-if chains.
 //
 //      +--------------------------------+
 //      | <code before the IfInst>       |
 //      | %zero = constant 0 : index     |
 //      | %v = affine_apply #expr1(%ops) |
 //      | %c = cmpi "sge" %v, %zero      |
 //      | cond_br %c, %next, %else       |
 //      +--------------------------------+
 //             |              |
 //             |              --------------|
 //             v                            |
 //      +--------------------------------+  |
 //      | next:                          |  |
 //      |   <repeat the check for expr2> |  |
 //      |   cond_br %c, %next2, %else    |  |
 //      +--------------------------------+  |
 //             |              |             |
 //            ...             --------------|
 //             |   <Per-expression checks>  |
 //             v                            |
 //      +--------------------------------+  |
 //      | last:                          |  |
 //      |   <repeat the check for exprN> |  |
 //      |   cond_br %c, %then, %else     |  |
 //      +--------------------------------+  |
 //             |              |             |
 //             |              --------------|
 //             v                            |
 //      +--------------------------------+  |
 //      | then:                          |  |
 //      |   <then contents>              |  |
 //      |   br continue                  |  |
 //      +--------------------------------+  |
 //             |                            |
 //   |----------               |-------------
 //   |                         V
 //   |  +--------------------------------+
 //   |  | else:                          |
 //   |  |   <else contents>              |
 //   |  |   br continue                  |
 //   |  +--------------------------------+
 //   |         |
 //   ------|   |
 //         v   v
 //      +--------------------------------+
 //      | continue:                      |
 //      |   <code after the IfInst>      |
 //      +--------------------------------+
 //
 void LowerIfAndForPass::lowerIfInst(IfInst *ifInst) {
   auto loc = ifInst->getLoc();

   // Start by splitting the block containing the 'if' into two parts.  The part
   // before will contain the condition, the part after will be the continuation
   // point.
   auto *condBlock = ifInst->getBlock();
   auto *continueBlock = condBlock->splitBlock(ifInst);

   // Create a block for the 'then' code, inserting it between the cond and
   // continue blocks.  Move the instructions over from the IfInst and add a
   // branch to the continuation point.
   Block *thenBlock = new Block();
   thenBlock->insertBefore(continueBlock);

   auto *oldThen = ifInst->getThen();
   thenBlock->getInstructions().splice(thenBlock->begin(),
                                       oldThen->getInstructions(),
                                       oldThen->begin(), oldThen->end());
   FuncBuilder builder(thenBlock);
   builder.create<BranchOp>(loc, continueBlock);

   // Handle the 'else' block the same way, but we skip it if we have no else
   // code.
   Block *elseBlock = continueBlock;
   if (auto *oldElse = ifInst->getElse()) {
     elseBlock = new Block();
     elseBlock->insertBefore(continueBlock);

     elseBlock->getInstructions().splice(elseBlock->begin(),
                                         oldElse->getInstructions(),
                                         oldElse->begin(), oldElse->end());
     builder.setInsertionPointToEnd(elseBlock);
     builder.create<BranchOp>(loc, continueBlock);
    }

   // Ok, now we just have to handle the condition logic.
   auto integerSet = ifInst->getCondition().getIntegerSet();

   // Implement short-circuit logic.  For each affine expression in the 'if'
   // condition, convert it into an affine map and call `affine_apply` to obtain
   // the resulting value.  Perform the equality or the greater-than-or-equality
   // test between this value and zero depending on the equality flag of the
   // condition.  If the test fails, jump immediately to the false branch, which
   // may be the else block if it is present or the continuation block otherwise.
   // If the test succeeds, jump to the next block testing the next conjunct of
   // the condition in the similar way.  When all conjuncts have been handled,
   // jump to the 'then' block instead.
   builder.setInsertionPointToEnd(condBlock);
   Value *zeroConstant = builder.create<ConstantIndexOp>(loc, 0);

   for (auto tuple :
        llvm::zip(integerSet.getConstraints(), integerSet.getEqFlags())) {
     AffineExpr constraintExpr = std::get<0>(tuple);
     bool isEquality = std::get<1>(tuple);

     // Create the fall-through block for the next condition right before the
     // 'thenBlock'.
     auto *nextBlock = new Block();
     nextBlock->insertBefore(thenBlock);

     // Build and apply an affine map.
     auto affineMap =
         builder.getAffineMap(integerSet.getNumDims(),
                              integerSet.getNumSymbols(), constraintExpr, {});
     SmallVector<Value *, 8> operands(ifInst->getOperands());
     auto affineApplyOp =
         builder.create<AffineApplyOp>(loc, affineMap, operands);
     Value *affResult = affineApplyOp->getResult(0);

     // Compare the result of the apply and branch.
     auto comparisonOp = builder.create<CmpIOp>(
         loc, isEquality ? CmpIPredicate::EQ : CmpIPredicate::SGE, affResult,
         zeroConstant);
     builder.create<CondBranchOp>(loc, comparisonOp->getResult(), nextBlock,
                                  /*trueArgs*/ ArrayRef<Value *>(), elseBlock,
                                  /*falseArgs*/ ArrayRef<Value *>());
     builder.setInsertionPointToEnd(nextBlock);
   }

   // We will have ended up with an empty block as our continuation block (or, in
   // the degenerate case where there were zero conditions, we have the original
   // condition block).  Redirect that to the thenBlock.
   condBlock = builder.getInsertionBlock();
   if (condBlock->empty()) {
     condBlock->replaceAllUsesWith(thenBlock);
     condBlock->eraseFromFunction();
   } else {
     builder.create<BranchOp>(loc, thenBlock);
   }

   // Ok, we're done!
   ifInst->erase();
 }

 // Entry point of the function convertor.
 //
 // Conversion is performed by recursively visiting instructions of a Function.
 // It reasons in terms of single-entry single-exit (SESE) regions that are not
 // materialized in the code.  Instead, the pointer to the last block of the
 // region is maintained throughout the conversion as the insertion point of the
 // IR builder since we never change the first block after its creation.  "Block"
 // instructions such as loops and branches create new SESE regions for their
 // bodies, and surround them with additional basic blocks for the control flow.
 // Individual operations are simply appended to the end of the last basic block
 // of the current region.  The SESE invariant allows us to easily handle nested
 // structures of arbitrary complexity.
 //
 // During the conversion, we maintain a mapping between the Values present in
 // the original function and their Value images in the function under
 // construction.  When an Value is used, it gets replaced with the
 // corresponding Value that has been defined previously.  The value flow
 // starts with function arguments converted to basic block arguments.
 PassResult LowerIfAndForPass::runOnFunction(Function *function) {
   SmallVector<Instruction *, 8> instsToRewrite;

   // Collect all the If and For statements.  We do this as a prepass to avoid
   // invalidating the walker with our rewrite.
   function->walkInsts([&](Instruction *inst) {
     if (isa<IfInst>(inst) || isa<ForInst>(inst))
       instsToRewrite.push_back(inst);
   });

   // Rewrite all of the ifs and fors.  We walked the instructions in preorder,
   // so we know that we will rewrite them in the same order.
   for (auto *inst : instsToRewrite)
     if (auto *ifInst = dyn_cast<IfInst>(inst))
       lowerIfInst(ifInst);
     else
       lowerForInst(cast<ForInst>(inst));

   // Change the kind of the function to indicate it has no If's or For's.
   function->setKind(Function::Kind::CFGFunc);
   return success();
 }

 /// Lowers If and For instructions within a function into their lower level CFG
 /// equivalent blocks.
 FunctionPass *mlir::createLowerIfAndForPass() {
   return new LowerIfAndForPass();
 }

 static PassRegistration<LowerIfAndForPass>
     pass("lower-if-and-for",
          "Lower If and For instructions to CFG equivalents");
	//===- LowerIfAndFor.cpp - Lower If and For instructions to CFG -----------===//
	//
	// Copyright 2019 The MLIR Authors.
	//
	// Licensed under the Apache License, Version 2.0 (the "License");
	// you may not use this file except in compliance with the License.
	// You may obtain a copy of the License at
	//
	// http://www.apache.org/licenses/LICENSE-2.0
	//
	// Unless required by applicable law or agreed to in writing, software
	// distributed under the License is distributed on an "AS IS" BASIS,
	// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	// See the License for the specific language governing permissions and
	// limitations under the License.
	// =============================================================================
	//
	// This file lowers If and For instructions within a function into their lower
	// level CFG equivalent blocks.
	//
	//===----------------------------------------------------------------------===//

	#include "mlir/IR/Builders.h"
	#include "mlir/IR/BuiltinOps.h"
	#include "mlir/Pass.h"
	#include "mlir/StandardOps/StandardOps.h"
	#include "mlir/Transforms/Passes.h"
	using namespace mlir;

	namespace {
	class LowerIfAndForPass : public FunctionPass {
	public:
	LowerIfAndForPass() : FunctionPass(&passID) {}
	PassResult runOnFunction(Function *function) override;

	void lowerForInst(ForInst *forInst);
	void lowerIfInst(IfInst *ifInst);

	static char passID;
	};
	} // end anonymous namespace

	char LowerIfAndForPass::passID = 0;

	// Given a range of values, emit the code that reduces them with "min" or "max"
	// depending on the provided comparison predicate. The predicate defines which
	// comparison to perform, "lt" for "min", "gt" for "max" and is used for the
	// `cmpi` operation followed by the `select` operation:
	//
	// %cond = cmpi "predicate" %v0, %v1
	// %result = select %cond, %v0, %v1
	//
	// Multiple values are scanned in a linear sequence. This creates a data
	// dependences that wouldn't exist in a tree reduction, but is easier to
	// recognize as a reduction by the subsequent passes.
	static Value *buildMinMaxReductionSeq(
	Location loc, CmpIPredicate predicate,
	llvm::iterator_range<OperationInst::result_iterator> values,
	FuncBuilder &builder) {
	assert(!llvm::empty(values) && "empty min/max chain");

	auto valueIt = values.begin();
	Value value = valueIt++;
	for (; valueIt != values.end(); ++valueIt) {
	auto cmpOp = builder.create<CmpIOp>(loc, predicate, value, *valueIt);
	value = builder.create<SelectOp>(loc, cmpOp->getResult(), value, *valueIt);
	}

	return value;
	}

	// Convert a "for" loop to a flow of blocks.
	//
	// Create an SESE region for the loop (including its body) and append it to the
	// end of the current region. The loop region consists of the initialization
	// block that sets up the initial value of the loop induction variable (%iv) and
	// computes the loop bounds that are loop-invariant in MLFunctions; the
	// condition block that checks the exit condition of the loop; the body SESE
	// region; and the end block that post-dominates the loop. The end block of the
	// loop becomes the new end of the current SESE region. The body of the loop is
	// constructed recursively after starting a new region (it may be, for example,
	// a nested loop). Induction variable modification is appended to the body SESE
	// region that always loops back to the condition block.
	//
	// +--------------------------------+
	// \| <code before the ForInst> \|
	// \| <compute initial %iv value> \|
	// \| br cond(%iv) \|
	// +--------------------------------+
	// \|
	// -------\| \|
	// \| v v
	// \| +--------------------------------+
	// \| \| cond(%iv): \|
	// \| \| <compare %iv to upper bound> \|
	// \| \| cond_br %r, body, end \|
	// \| +--------------------------------+
	// \| \| \|
	// \| \| -------------\|
	// \| v \|
	// \| +--------------------------------+ \|
	// \| \| body: \| \|
	// \| \| <body contents> \| \|
	// \| \| %new_iv =<add step to %iv> \| \|
	// \| \| br cond(%new_iv) \| \|
	// \| +--------------------------------+ \|
	// \| \| \|
	// \|----------- \|--------------------
	// v
	// +--------------------------------+
	// \| end: \|
	// \| <code after the ForInst> \|
	// +--------------------------------+
	//
	void LowerIfAndForPass::lowerForInst(ForInst *forInst) {
	auto loc = forInst->getLoc();

	// Start by splitting the block containing the 'for' into two parts. The part
	// before will get the init code, the part after will be the end point.
	auto *initBlock = forInst->getBlock();
	auto *endBlock = initBlock->splitBlock(forInst);

	// Create the condition block, with its argument for the loop induction
	// variable. We set it up below.
	auto *conditionBlock = new Block();
	conditionBlock->insertBefore(endBlock);
	auto *iv = conditionBlock->addArgument(IndexType::get(forInst->getContext()));

	// Create the body block, moving the body of the forInst over to it.
	auto *bodyBlock = new Block();
	bodyBlock->insertBefore(endBlock);

	auto *oldBody = forInst->getBody();
	bodyBlock->getInstructions().splice(bodyBlock->begin(),
	oldBody->getInstructions(),
	oldBody->begin(), oldBody->end());

	// The code in the body of the forInst now uses 'iv' as its indvar.
	forInst->replaceAllUsesWith(iv);

	// Append the induction variable stepping logic and branch back to the exit
	// condition block. Construct an affine map f : (x -> x+step) and apply this
	// map to the induction variable.
	FuncBuilder builder(bodyBlock);
	auto affStep = builder.getAffineConstantExpr(forInst->getStep());
	auto affDim = builder.getAffineDimExpr(0);
	auto affStepMap = builder.getAffineMap(1, 0, {affDim + affStep}, {});
	auto stepOp = builder.create<AffineApplyOp>(loc, affStepMap, iv);
	builder.create<BranchOp>(loc, conditionBlock, stepOp->getResult(0));

	// Now that the body block done, fill in the code to compute the bounds of the
	// induction variable in the init block.
	builder.setInsertionPointToEnd(initBlock);

	// Compute loop bounds using an affine_apply.
	SmallVector<Value *, 8> operands(forInst->getLowerBoundOperands());
	auto lbAffineApply =
	builder.create<AffineApplyOp>(loc, forInst->getLowerBoundMap(), operands);
	Value *lowerBound = buildMinMaxReductionSeq(
	loc, CmpIPredicate::SGT, lbAffineApply->getResults(), builder);

	operands.assign(forInst->getUpperBoundOperands().begin(),
	forInst->getUpperBoundOperands().end());
	auto ubAffineApply =
	builder.create<AffineApplyOp>(loc, forInst->getUpperBoundMap(), operands);
	Value *upperBound = buildMinMaxReductionSeq(
	loc, CmpIPredicate::SLT, ubAffineApply->getResults(), builder);
	builder.create<BranchOp>(loc, conditionBlock, lowerBound);

	// With the body block done, we can fill in the condition block.
	builder.setInsertionPointToEnd(conditionBlock);
	auto comparison =
	builder.create<CmpIOp>(loc, CmpIPredicate::SLT, iv, upperBound);
	builder.create<CondBranchOp>(loc, comparison, bodyBlock, ArrayRef<Value *>(),
	endBlock, ArrayRef<Value *>());

	// Ok, we're done!
	forInst->erase();
	}

	// Convert an "if" instruction into a flow of basic blocks.
	//
	// Create an SESE region for the if instruction (including its "then" and
	// optional "else" instruction blocks) and append it to the end of the current
	// region. The conditional region consists of a sequence of condition-checking
	// blocks that implement the short-circuit scheme, followed by a "then" SESE
	// region and an "else" SESE region, and the continuation block that
	// post-dominates all blocks of the "if" instruction. The flow of blocks that
	// correspond to the "then" and "else" clauses are constructed recursively,
	// enabling easy nesting of "if" instructions and if-then-else-if chains.
	//
	// +--------------------------------+
	// \| <code before the IfInst> \|
	// \| %zero = constant 0 : index \|
	// \| %v = affine_apply #expr1(%ops) \|
	// \| %c = cmpi "sge" %v, %zero \|
	// \| cond_br %c, %next, %else \|
	// +--------------------------------+
	// \| \|
	// \| --------------\|
	// v \|
	// +--------------------------------+ \|
	// \| next: \| \|
	// \| <repeat the check for expr2> \| \|
	// \| cond_br %c, %next2, %else \| \|
	// +--------------------------------+ \|
	// \| \| \|
	// ... --------------\|
	// \| <Per-expression checks> \|
	// v \|
	// +--------------------------------+ \|
	// \| last: \| \|
	// \| <repeat the check for exprN> \| \|
	// \| cond_br %c, %then, %else \| \|
	// +--------------------------------+ \|
	// \| \| \|
	// \| --------------\|
	// v \|
	// +--------------------------------+ \|
	// \| then: \| \|
	// \| <then contents> \| \|
	// \| br continue \| \|
	// +--------------------------------+ \|
	// \| \|
	// \|---------- \|-------------
	// \| V
	// \| +--------------------------------+
	// \| \| else: \|
	// \| \| <else contents> \|
	// \| \| br continue \|
	// \| +--------------------------------+
	// \| \|
	// ------\| \|
	// v v
	// +--------------------------------+
	// \| continue: \|
	// \| <code after the IfInst> \|
	// +--------------------------------+
	//
	void LowerIfAndForPass::lowerIfInst(IfInst *ifInst) {
	auto loc = ifInst->getLoc();

	// Start by splitting the block containing the 'if' into two parts. The part
	// before will contain the condition, the part after will be the continuation
	// point.
	auto *condBlock = ifInst->getBlock();
	auto *continueBlock = condBlock->splitBlock(ifInst);

	// Create a block for the 'then' code, inserting it between the cond and
	// continue blocks. Move the instructions over from the IfInst and add a
	// branch to the continuation point.
	Block *thenBlock = new Block();
	thenBlock->insertBefore(continueBlock);

	auto *oldThen = ifInst->getThen();
	thenBlock->getInstructions().splice(thenBlock->begin(),
	oldThen->getInstructions(),
	oldThen->begin(), oldThen->end());
	FuncBuilder builder(thenBlock);
	builder.create<BranchOp>(loc, continueBlock);

	// Handle the 'else' block the same way, but we skip it if we have no else
	// code.
	Block *elseBlock = continueBlock;
	if (auto *oldElse = ifInst->getElse()) {
	elseBlock = new Block();
	elseBlock->insertBefore(continueBlock);

	elseBlock->getInstructions().splice(elseBlock->begin(),
	oldElse->getInstructions(),
	oldElse->begin(), oldElse->end());
	builder.setInsertionPointToEnd(elseBlock);
	builder.create<BranchOp>(loc, continueBlock);
	}

	// Ok, now we just have to handle the condition logic.
	auto integerSet = ifInst->getCondition().getIntegerSet();

	// Implement short-circuit logic. For each affine expression in the 'if'
	// condition, convert it into an affine map and call `affine_apply` to obtain
	// the resulting value. Perform the equality or the greater-than-or-equality
	// test between this value and zero depending on the equality flag of the
	// condition. If the test fails, jump immediately to the false branch, which
	// may be the else block if it is present or the continuation block otherwise.
	// If the test succeeds, jump to the next block testing the next conjunct of
	// the condition in the similar way. When all conjuncts have been handled,
	// jump to the 'then' block instead.
	builder.setInsertionPointToEnd(condBlock);
	Value *zeroConstant = builder.create<ConstantIndexOp>(loc, 0);

	for (auto tuple :
	llvm::zip(integerSet.getConstraints(), integerSet.getEqFlags())) {
	AffineExpr constraintExpr = std::get<0>(tuple);
	bool isEquality = std::get<1>(tuple);

	// Create the fall-through block for the next condition right before the
	// 'thenBlock'.
	auto *nextBlock = new Block();
	nextBlock->insertBefore(thenBlock);

	// Build and apply an affine map.
	auto affineMap =
	builder.getAffineMap(integerSet.getNumDims(),
	integerSet.getNumSymbols(), constraintExpr, {});
	SmallVector<Value *, 8> operands(ifInst->getOperands());
	auto affineApplyOp =
	builder.create<AffineApplyOp>(loc, affineMap, operands);
	Value *affResult = affineApplyOp->getResult(0);

	// Compare the result of the apply and branch.
	auto comparisonOp = builder.create<CmpIOp>(
	loc, isEquality ? CmpIPredicate::EQ : CmpIPredicate::SGE, affResult,
	zeroConstant);
	builder.create<CondBranchOp>(loc, comparisonOp->getResult(), nextBlock,
	/trueArgs/ ArrayRef<Value *>(), elseBlock,
	/falseArgs/ ArrayRef<Value *>());
	builder.setInsertionPointToEnd(nextBlock);
	}

	// We will have ended up with an empty block as our continuation block (or, in
	// the degenerate case where there were zero conditions, we have the original
	// condition block). Redirect that to the thenBlock.
	condBlock = builder.getInsertionBlock();
	if (condBlock->empty()) {
	condBlock->replaceAllUsesWith(thenBlock);
	condBlock->eraseFromFunction();
	} else {
	builder.create<BranchOp>(loc, thenBlock);
	}

	// Ok, we're done!
	ifInst->erase();
	}

	// Entry point of the function convertor.
	//
	// Conversion is performed by recursively visiting instructions of a Function.
	// It reasons in terms of single-entry single-exit (SESE) regions that are not
	// materialized in the code. Instead, the pointer to the last block of the
	// region is maintained throughout the conversion as the insertion point of the
	// IR builder since we never change the first block after its creation. "Block"
	// instructions such as loops and branches create new SESE regions for their
	// bodies, and surround them with additional basic blocks for the control flow.
	// Individual operations are simply appended to the end of the last basic block
	// of the current region. The SESE invariant allows us to easily handle nested
	// structures of arbitrary complexity.
	//
	// During the conversion, we maintain a mapping between the Values present in
	// the original function and their Value images in the function under
	// construction. When an Value is used, it gets replaced with the
	// corresponding Value that has been defined previously. The value flow
	// starts with function arguments converted to basic block arguments.
	PassResult LowerIfAndForPass::runOnFunction(Function *function) {
	SmallVector<Instruction *, 8> instsToRewrite;

	// Collect all the If and For statements. We do this as a prepass to avoid
	// invalidating the walker with our rewrite.
	function->walkInsts([&](Instruction *inst) {
	if (isa<IfInst>(inst) \|\| isa<ForInst>(inst))
	instsToRewrite.push_back(inst);
	});

	// Rewrite all of the ifs and fors. We walked the instructions in preorder,
	// so we know that we will rewrite them in the same order.
	for (auto *inst : instsToRewrite)
	if (auto *ifInst = dyn_cast<IfInst>(inst))
	lowerIfInst(ifInst);
	else
	lowerForInst(cast<ForInst>(inst));

	// Change the kind of the function to indicate it has no If's or For's.
	function->setKind(Function::Kind::CFGFunc);
	return success();
	}

	/// Lowers If and For instructions within a function into their lower level CFG
	/// equivalent blocks.
	FunctionPass *mlir::createLowerIfAndForPass() {
	return new LowerIfAndForPass();
	}

	static PassRegistration<LowerIfAndForPass>
	pass("lower-if-and-for",
	"Lower If and For instructions to CFG equivalents");