From 5a0d3d02046a03e64115ee4f1bf8a5c8fa50225b Mon Sep 17 00:00:00 2001
From: Alex Zinenko <zinenko@google.com>
Date: Mon, 12 Nov 2018 14:58:36 -0800
Subject: [PATCH] Basic conversion of MLFunctions to CFGFunctions.

Implement a pass converting a subset of MLFunctions to CFGFunctions.  Currently
supports arbitrarily complex imperfect loop nests with statically constant
(i.e., not affine map) bounds filled with operations.  Does NOT support
branches and non-constant loop bounds.

Conversion is performed per-function and the function names are preserved to
avoid breaking any external references to the current module.  In-memory IR is
updated to point to the right functions in direct calls and constant loads.
This behavior is tested via a really hidden flag that enables function
renaming.

Inside each function, the control flow conversion is based on single-entry
single-exit regions, i.e. subgraphs of the CFG that have exactly one incoming
and exactly one outgoing edge.  Since an MLFunction must have a single "return"
statement as per MLIR spec, it constitutes an SESE region.  Individual
operations are appended to this region.  Control flow statements are
recursively converted into such regions that are concatenated with the current
region.  Bodies of the compound statement also form SESE regions, which allows
to nest control flow statements easily.  Note that SESE regions are not
materialized in the code.  It is sufficent to keep track of the end of the
region as the current instruction insertion point as long as all recursive
calls update the insertion point in the end.

The converter maintains a mapping between SSA values in ML functions and their
CFG counterparts.  The mapping is used to find the operands for each operation
and is updated to contain the results of each operation as the conversion
continues.

PiperOrigin-RevId: 221162602
---
 mlir/include/mlir/IR/Builders.h         |  10 +
 mlir/include/mlir/IR/OperationSupport.h |   8 +-
 mlir/lib/Transforms/ConvertToCFG.cpp    | 358 +++++++++++++++++++++---
 mlir/test/Transforms/convert2cfg.mlir   | 224 ++++++++++++++-
 4 files changed, 561 insertions(+), 39 deletions(-)

diff --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index cc6de25c8288..2a1095ae14f0 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -188,6 +188,14 @@ public:
     insertPoint = BasicBlock::iterator();
   }
 
+  /// Return the block the current insertion point belongs to.  Note that the
+  /// the insertion point is not necessarily the end of the block.
+  BasicBlock *getInsertionBlock() const { return block; }
+
+  /// Return the insert position as the BasicBlock iterator.  The block itself
+  /// can be obtained by calling getInsertionBlock.
+  BasicBlock::iterator getInsertionPoint() const { return insertPoint; }
+
   /// Set the insertion point to the specified location.
   void setInsertionPoint(BasicBlock *block, BasicBlock::iterator insertPoint) {
     assert(block->getFunction() == function &&
@@ -279,6 +287,8 @@ public:
 
 private:
   template <typename T> T *insertTerminator(T *term) {
+    // FIXME: b/118738403
+    assert(!block->getTerminator() && "cannot insert the second terminator");
     block->setTerminator(term);
     return term;
   }
diff --git a/mlir/include/mlir/IR/OperationSupport.h b/mlir/include/mlir/IR/OperationSupport.h
index d692e6112f17..118ac8a2afff 100644
--- a/mlir/include/mlir/IR/OperationSupport.h
+++ b/mlir/include/mlir/IR/OperationSupport.h
@@ -218,8 +218,14 @@ public:
     types.append(newTypes.begin(), newTypes.end());
   }
 
+  /// Add an attribute with the specified name.
   void addAttribute(StringRef name, Attribute attr) {
-    attributes.push_back({Identifier::get(name, context), attr});
+    addAttribute(Identifier::get(name, context), attr);
+  }
+
+  /// Add an attribute with the specified name.
+  void addAttribute(Identifier name, Attribute attr) {
+    attributes.push_back({name, attr});
   }
 };
 
diff --git a/mlir/lib/Transforms/ConvertToCFG.cpp b/mlir/lib/Transforms/ConvertToCFG.cpp
index 8ffa23ca1025..f6a9afd5c9f5 100644
--- a/mlir/lib/Transforms/ConvertToCFG.cpp
+++ b/mlir/lib/Transforms/ConvertToCFG.cpp
@@ -1,4 +1,4 @@
-//===- ConvertToCFG.cpp - ML function to CFG function converstion ---------===//
+//===- ConvertToCFG.cpp - ML function to CFG function conversion ----------===//
 //
 // Copyright 2019 The MLIR Authors.
 //
@@ -20,12 +20,18 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/MLFunction.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
+#include "mlir/IR/StmtVisitor.h"
 #include "mlir/Pass.h"
+#include "mlir/StandardOps/StandardOps.h"
+#include "mlir/Support/Functional.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/Support/CommandLine.h"
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
@@ -34,26 +40,264 @@ using namespace mlir;
 
 namespace {
 // Generates CFG function equivalent to the given ML function.
-class FunctionConverter {
+class FunctionConverter : public StmtVisitor<FunctionConverter> {
 public:
   FunctionConverter(CFGFunction *cfgFunc)
       : cfgFunc(cfgFunc), builder(cfgFunc) {}
-  CFGFunction *convert(const MLFunction *mlFunc);
+  CFGFunction *convert(MLFunction *mlFunc);
+
+  void visitForStmt(ForStmt *forStmt);
+  void visitIfStmt(IfStmt *ifStmt);
+  void visitOperationStmt(OperationStmt *opStmt);
 
 private:
+  CFGValue *getConstantIndexValue(int64_t value);
+
   CFGFunction *cfgFunc;
   CFGFuncBuilder builder;
+
+  // Mapping between original MLValues and lowered CFGValues.
+  llvm::DenseMap<const MLValue *, CFGValue *> valueRemapping;
 };
 } // end anonymous namespace
 
-CFGFunction *FunctionConverter::convert(const MLFunction *mlFunc) {
-  builder.createBlock();
+// Return a vector of OperationStmt's arguments as the CFGValues or SSAValues
+// depending on the template argument.  For each statement operands, represented
+// as MLValue, lookup its CFGValue conterpart in the valueRemapping table.
+// The return type parameterization is necessary because some instructions
+// accept vectors of SSAValues while others accept vectors of CFGValues.
+template <typename SSAValueTy>
+static llvm::SmallVector<SSAValueTy *, 4>
+operandsAs(OperationStmt *opStmt,
+           const llvm::DenseMap<const MLValue *, CFGValue *> &valueRemapping) {
+  static_assert(std::is_same<SSAValueTy, SSAValue>::value ||
+                    std::is_same<SSAValueTy, CFGValue>::value,
+                "can only cast statement operands to CFGValue or SSAValue");
+  llvm::SmallVector<SSAValueTy *, 4> operands;
+  for (const MLValue *operand : opStmt->getOperands()) {
+    assert(valueRemapping.count(operand) != 0 && "operand is not defined");
+    operands.push_back(valueRemapping.lookup(operand));
+  }
+  return operands;
+}
+
+// Convert an operation statement into an operation instruction.
+//
+// The operation description (name, number and types of operands or results)
+// remains the same but the values must be updated to be CFGValues.  Update the
+// mapping MLValue->CFGValue as the conversion is performed.  The operation
+// instruction is appended to current block (end of SESE region).
+void FunctionConverter::visitOperationStmt(OperationStmt *opStmt) {
+  // Handle returns separately, they are transformed into a specially-typed
+  // return instruction.
+  // TODO(zinenko): after terminators and operations are merged, remove this
+  // special case and de-template operandsAs.
+  if (opStmt->getName().getStringRef() == ReturnOp::getOperationName()) {
+    builder.createReturn(opStmt->getLoc(),
+                         operandsAs<CFGValue>(opStmt, valueRemapping));
+    return;
+  }
+
+  // Set up basic operation state (context, name, operands).
+  OperationState state(cfgFunc->getContext(), opStmt->getLoc(),
+                       opStmt->getName());
+  state.addOperands(operandsAs<SSAValue>(opStmt, valueRemapping));
+
+  // Set up operation return types.  The corresponding SSAValues will become
+  // available after the operation is created.
+  state.addTypes(
+      functional::map([](SSAValue *result) { return result->getType(); },
+                      opStmt->getResults()));
+
+  // Copy attributes.
+  for (auto attr : opStmt->getAttrs()) {
+    state.addAttribute(attr.first.strref(), attr.second);
+  }
+
+  auto opInst = builder.createOperation(state);
+
+  // Make results of the operation accessible to the following operations
+  // through remapping.
+  assert(opInst->getNumResults() == opStmt->getNumResults());
+  for (unsigned i = 0, n = opInst->getNumResults(); i < n; ++i) {
+    valueRemapping.insert(
+        std::make_pair(opStmt->getResult(i), opInst->getResult(i)));
+  }
+}
+
+// Create a CFGValue for the given integer constant of index type.
+CFGValue *FunctionConverter::getConstantIndexValue(int64_t value) {
+  auto op = builder.create<ConstantIndexOp>(builder.getUnknownLoc(), value);
+  return cast<CFGValue>(op->getResult());
+}
+
+// Convert a "for" loop to a flow of basic 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.
+//
+//      +--------------------------------+
+//      | <end of current SESE region>   |
+//      | <current insertion point>      |
+//      | br init                        |
+//      +--------------------------------+
+//             |
+//             v
+//      +--------------------------------+
+//      | init:                          |
+//      |   <start of loop SESE region>  |
+//      |   <compute initial %iv value>  |
+//      |   br cond(%iv)                 |
+//      +--------------------------------+
+//             |
+//  -------|   |
+//  |      v   v
+//  |   +--------------------------------+
+//  |   | cond(%iv):                     |
+//  |   |   <compare %iv to upper bound> |
+//  |   |   cond_br %r, body, end        |
+//  |   +--------------------------------+
+//  |          |               |
+//  |          |               -------------|
+//  |          v                            |
+//  |   +--------------------------------+  |
+//  |   | body:                          |  |
+//  |   |   <body SESE region start>     |  |
+//  |   |   <...>                        |  |
+//  |   +--------------------------------+  |
+//  |          |                            |
+//  |         ... <SESE region of the body> |
+//  |          |                            |
+//  |          v                            |
+//  |   +--------------------------------+  |
+//  |   | body-end:                      |  |
+//  |   |   <body SESE region end>       |  |
+//  |   |   %new_iv =<add step to %iv>   |  |
+//  |   |   br cond(%new_iv)             |  |
+//  |   +--------------------------------+  |
+//  |          |                            |
+//  |-----------        |--------------------
+//                      v
+//      +--------------------------------+
+//      | end:                           |
+//      |   <end of loop SESE region>    |
+//      |   <new insertion point>        |
+//      +--------------------------------+
+//
+void FunctionConverter::visitForStmt(ForStmt *forStmt) {
+  // First, store the loop insertion location so that we can go back to it after
+  // creating the new blocks (block creation updates the insertion point).
+  BasicBlock *loopInsertionPoint = builder.getInsertionBlock();
+
+  // Create blocks so that they appear in more human-readable order in the
+  // output.
+  BasicBlock *loopInitBlock = builder.createBlock();
+  BasicBlock *loopConditionBlock = builder.createBlock();
+  BasicBlock *loopBodyFirstBlock = builder.createBlock();
 
-  // Creates return instruction with no operands.
-  // TODO: convert return operands.
-  builder.createReturn(mlFunc->getReturnStmt()->getLoc(), {});
+  // At the loop insertion location, branch immediately to the loop init block.
+  builder.setInsertionPoint(loopInsertionPoint);
+  builder.createBranch(builder.getUnknownLoc(), loopInitBlock);
 
-  // TODO: convert ML function body.
+  // The loop condition block has an argument for loop induction variable.
+  // Create it upfront and make the loop induction variable -> basic block
+  // argument remapping available to the following instructions.  ForStatement
+  // is-a MLValue corresponding to the loop induction variable.
+  builder.setInsertionPoint(loopConditionBlock);
+  CFGValue *iv = loopConditionBlock->addArgument(builder.getIndexType());
+  valueRemapping.insert(std::make_pair(forStmt, iv));
+
+  // Recursively construct loop body region.
+  // Walking manually because we need custom logic before and after traversing
+  // the list of children.
+  builder.setInsertionPoint(loopBodyFirstBlock);
+  for (auto &stmt : *forStmt) {
+    this->visit(&stmt);
+  }
+
+  // Builder point is currently at the last block of the loop body.  Append the
+  // induction variable stepping to this block and branch back to the exit
+  // condition block.
+  CFGValue *step = getConstantIndexValue(forStmt->getStep());
+  auto stepOp = builder.create<AddIOp>(forStmt->getLoc(), iv, step);
+  CFGValue *nextIvValue = cast<CFGValue>(stepOp->getResult());
+  builder.createBranch(builder.getUnknownLoc(), loopConditionBlock,
+                       {nextIvValue});
+
+  // Create post-loop block here so that it appears after all loop body blocks.
+  BasicBlock *postLoopBlock = builder.createBlock();
+
+  builder.setInsertionPoint(loopInitBlock);
+  // TODO(zinenko): support non-constant loop bounds
+  assert(forStmt->hasConstantBounds() && "NYI: non-constant for loop bounds");
+  CFGValue *lowerBound =
+      getConstantIndexValue(forStmt->getConstantLowerBound());
+  CFGValue *upperBound =
+      getConstantIndexValue(forStmt->getConstantUpperBound());
+  builder.createBranch(builder.getUnknownLoc(), loopConditionBlock,
+                       {lowerBound});
+
+  builder.setInsertionPoint(loopConditionBlock);
+  auto comparisonOp = builder.create<CmpIOp>(
+      forStmt->getLoc(), CmpIPredicate::SLT, iv, upperBound);
+  auto comparisonResult = cast<CFGValue>(comparisonOp->getResult());
+  builder.createCondBranch(builder.getUnknownLoc(), comparisonResult,
+                           loopBodyFirstBlock, postLoopBlock);
+
+  // Finally, make sure building can continue by setting the post-loop block
+  // (end of loop SESE region) as the insertion point.
+  builder.setInsertionPoint(postLoopBlock);
+}
+
+void FunctionConverter::visitIfStmt(IfStmt *ifStmt) {
+  // TODO(zinenko): implement
+  assert(false && "NYI: CFG lowering of if statements");
+}
+
+// Entry point of the function convertor.
+//
+// Conversion is performed by recursively visiting statements of an MLFunction.
+// 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"
+// statements 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 MLValues present in
+// the original function and their CFGValue images in the function under
+// construction.  When an MLValue is used, it gets replaced with the
+// corresponding CFGValue that has been defined previously.  The value flow
+// starts with function arguments converted to basic block arguments.
+CFGFunction *FunctionConverter::convert(MLFunction *mlFunc) {
+  auto outerBlock = builder.createBlock();
+
+  // CFGFunctions do not have explicit arguments but use the arguments to the
+  // first basic block instead.  Create those from the MLFunction arguments and
+  // set up the value remapping.
+  outerBlock->addArguments(mlFunc->getType().getInputs());
+  assert(mlFunc->getNumArguments() == outerBlock->getNumArguments());
+  for (unsigned i = 0, n = mlFunc->getNumArguments(); i < n; ++i) {
+    const MLValue *mlArgument = mlFunc->getArgument(i);
+    CFGValue *cfgArgument = outerBlock->getArgument(i);
+    valueRemapping.insert(std::make_pair(mlArgument, cfgArgument));
+  }
+
+  // Convert statements in order.
+  for (auto &stmt : *mlFunc) {
+    visit(&stmt);
+  }
 
   return cfgFunc;
 }
@@ -76,18 +320,17 @@ private:
   // Generates CFG functions for all ML functions in the module.
   void convertMLFunctions();
   // Generates CFG function for the given ML function.
-  CFGFunction *convert(const MLFunction *mlFunc);
+  CFGFunction *convert(MLFunction *mlFunc);
   // Replaces all ML function references in the module
   // with references to the generated CFG functions.
   void replaceReferences();
   // Replaces function references in the given function.
   void replaceReferences(CFGFunction *cfgFunc);
-  void replaceReferences(MLFunction *mlFunc);
-  // Removes all ML funtions from the module.
-  void removeMLFunctions();
+  // Replaces MLFunctions with their CFG counterparts in the module.
+  void replaceFunctions();
 
   // Map from ML functions to generated CFG functions.
-  llvm::DenseMap<const MLFunction *, CFGFunction *> generatedFuncs;
+  llvm::DenseMap<MLFunction *, CFGFunction *> generatedFuncs;
   Module *module = nullptr;
 };
 } // end anonymous namespace
@@ -96,11 +339,19 @@ char ModuleConverter::passID = 0;
 
 // Iterates over all functions in the module generating CFG functions
 // equivalent to ML functions and replacing references to ML functions
-// with references to the generated ML functions.
+// with references to the generated ML functions.  The names of the converted
+// functions match those of the original functions to avoid breaking any
+// external references to the current module.  Therefore, converted functions
+// are added to the module at the end of the pass, after removing the original
+// functions to avoid name clashes.  Conversion procedure has access to the
+// module as member of ModuleConverter and must not rely on the converted
+// function to belong to the module.
 PassResult ModuleConverter::runOnModule(Module *m) {
   module = m;
   convertMLFunctions();
   replaceReferences();
+  replaceFunctions();
+
   return success();
 }
 
@@ -112,17 +363,22 @@ void ModuleConverter::convertMLFunctions() {
 }
 
 // Creates CFG function equivalent to the given ML function.
-CFGFunction *ModuleConverter::convert(const MLFunction *mlFunc) {
-  // TODO: ensure that CFG function name is unique.
-  auto *cfgFunc =
-      new CFGFunction(mlFunc->getLoc(), mlFunc->getName().str() + "_cfg",
-                      mlFunc->getType(), mlFunc->getAttrs());
-  module->getFunctions().push_back(cfgFunc);
+CFGFunction *ModuleConverter::convert(MLFunction *mlFunc) {
+  // Use the same name as for ML function; do not add the converted function to
+  // the module yet to avoid collision.
+  auto name = mlFunc->getName().str();
+  auto *cfgFunc = new CFGFunction(mlFunc->getLoc(), name, mlFunc->getType(),
+                                  mlFunc->getAttrs());
 
   // Generates the body of the CFG function.
   return FunctionConverter(cfgFunc).convert(mlFunc);
 }
 
+// Replace references to MLFunctions with the references to the converted
+// CFGFunctions.  Since this all MLFunctions are converted at this point, it is
+// unnecessary to replace references in the MLFunctions that are going to be
+// removed anyway.  However, it is necessary to replace the references in the
+// converted CFGFunctions that have not been added to the module yet.
 void ModuleConverter::replaceReferences() {
   for (Function &fn : *module) {
     switch (fn.getKind()) {
@@ -130,32 +386,61 @@ void ModuleConverter::replaceReferences() {
       replaceReferences(&cast<CFGFunction>(fn));
       break;
     case Function::Kind::MLFunc:
-      replaceReferences(&cast<MLFunction>(fn));
+      // ML functions must have been converted already and will be removed.
       break;
     case Function::Kind::ExtFunc:
       // nothing to do for external functions
       break;
     }
   }
+  for (auto pair : generatedFuncs) {
+    replaceReferences(pair.second);
+  }
 }
 
-void ModuleConverter::replaceReferences(CFGFunction *func) {
-  // TODO: NOP for now since function attributes are not yet implemented.
+// Replace the value of a function attribute named "name" attached to the
+// operation "op" and containing an MLFunction-typed value with the result of
+// converting "func" to a CFGFunction.
+static inline void replaceMLFunctionAttr(
+    Operation &op, Identifier name, const Function *func,
+    const llvm::DenseMap<MLFunction *, CFGFunction *> &generatedFuncs) {
+  const auto *mlFunc = dyn_cast<MLFunction>(func);
+  if (!mlFunc)
+    return;
+
+  Builder b(op.getContext());
+  auto cfgFunc = generatedFuncs.lookup(mlFunc);
+  op.setAttr(name, b.getFunctionAttr(cfgFunc));
 }
 
-void ModuleConverter::replaceReferences(MLFunction *func) {
-  // TODO: NOP for now since function attributes are not yet implemented.
+// Replace references to MLFunctions with the references to the converted
+// CFGFunctions.  References to MLFunctions can potentially appear in any
+// function attribute (in particular, they are known to appear in the  "callee"
+// attribute of a direct call and the "value" attribute of a constant).  Replace
+// the values of these attributes to point to the converted functions.
+void ModuleConverter::replaceReferences(CFGFunction *func) {
+  for (auto &bb : *func) {
+    for (auto &inst : bb) {
+      for (auto &attr : inst.getAttrs()) {
+        // TODO(zinenko): handle nested attributes, e.g. array attributes
+        // containing functions.
+        auto functionAttr = attr.second.dyn_cast<FunctionAttr>();
+        if (!functionAttr)
+          continue;
+        replaceMLFunctionAttr(inst, attr.first, functionAttr.getValue(),
+                              generatedFuncs);
+      }
+    }
+  }
 }
 
-// Removes all ML functions from the module.
-void ModuleConverter::removeMLFunctions() {
-  // Delete ML functions from the module.
-  for (auto it = module->begin(), e = module->end(); it != e;) {
-    // Manipulate iterator carefully to avoid deleting a function we're pointing
-    // at.
-    Function &fn = *it++;
-    if (auto mlFunc = dyn_cast<MLFunction>(&fn))
-      mlFunc->eraseFromModule();
+// The CFG and ML functions have the same name.  First, erase the MLFunction.
+// Then insert the CFGFunction at the same place.
+void ModuleConverter::replaceFunctions() {
+  for (auto pair : generatedFuncs) {
+    auto &functions = module->getFunctions();
+    auto it = functions.erase(pair.first);
+    functions.insert(it, pair.second);
   }
 }
 
@@ -165,7 +450,8 @@ void ModuleConverter::removeMLFunctions() {
 
 /// Replaces all ML functions in the module with equivalent CFG functions.
 /// Function references are appropriately patched to refer to the newly
-/// generated CFG functions.
+/// generated CFG functions.  Converted functions have the same names as the
+/// original functions to preserve module linking.
 ModulePass *mlir::createConvertToCFGPass() { return new ModuleConverter(); }
 
 static PassRegistration<ModuleConverter>
diff --git a/mlir/test/Transforms/convert2cfg.mlir b/mlir/test/Transforms/convert2cfg.mlir
index 67ad7886e57c..4ed92dd9abd9 100644
--- a/mlir/test/Transforms/convert2cfg.mlir
+++ b/mlir/test/Transforms/convert2cfg.mlir
@@ -1,7 +1,227 @@
-// RUN: mlir-opt %s -convert-to-cfg | FileCheck %s
+// RUN: mlir-opt -convert-to-cfg %s | FileCheck %s
 
-// CHECK-LABEL: cfgfunc @empty_cfg() {
+// CHECK-LABEL: cfgfunc @empty() {
 mlfunc @empty() {
              // CHECK: bb0:
   return     // CHECK:  return
 }            // CHECK: }
+
+extfunc @body(index) -> ()
+
+// Simple loops are properly converted.
+// CHECK-LABEL: cfgfunc @simple_loop() {
+// CHECK-NEXT: bb0:
+// CHECK-NEXT:   br bb1
+// CHECK-NEXT: bb1:	// pred: bb0
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %c42 = constant 42 : index
+// CHECK-NEXT:   br bb2(%c1) : index
+// CHECK-NEXT: bb2(%0: index):	// 2 preds: bb1, bb3
+// CHECK-NEXT:   %1 = cmpi "slt", %0, %c42 : index
+// CHECK-NEXT:   cond_br %1, bb3, bb4
+// CHECK-NEXT: bb3:	// pred: bb2
+// CHECK-NEXT:   call @body(%0) : (index) -> ()
+// CHECK-NEXT:   %c1_0 = constant 1 : index
+// CHECK-NEXT:   %2 = addi %0, %c1_0 : index
+// CHECK-NEXT:   br bb2(%2) : index
+// CHECK-NEXT: bb4:	// pred: bb2
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+mlfunc @simple_loop() {
+  for %i = 1 to 42 {
+    call @body(%i) : (index) -> ()
+  }
+  return
+}
+
+// Direct calls get renamed if asked (IR data structures properly updated) and
+// keep the same name otherwise.
+cfgfunc @simple_caller() {
+bb0:
+// CHECK: call @simple_loop() : () -> ()
+  call @simple_loop() : () -> ()
+  return
+}
+
+// Constant loads get renamed if asked (IR data structure properly updated) and
+// keep the same name otherwise.
+cfgfunc @simple_indirect_caller() {
+bb0:
+// CHECK: %f = constant @simple_loop : () -> ()
+  %f = constant @simple_loop : () -> ()
+  call_indirect %f() : () -> ()
+  return
+}
+
+// CHECK-LABEL: cfgfunc @ml_caller() {
+mlfunc @ml_caller() {
+// Direct calls inside ML functions are renamed if asked (given that the
+// function itself is also converted).
+// CHECK: call @simple_loop() : () -> ()
+  call @simple_loop() : () -> ()
+// Direct calls to not yet declared ML functions are also renamed.
+// CHECK: call @more_imperfectly_nested_loops() : () -> ()
+  call @more_imperfectly_nested_loops() : () -> ()
+  return
+}
+
+/////////////////////////////////////////////////////////////////////
+
+extfunc @body_args(index) -> (index)
+extfunc @other(index, i32) -> (i32)
+
+// Arguments and return values of the functions are converted.
+// CHECK-LABEL: cfgfunc @mlfunc_args(i32, i32) -> (i32, i32) {
+// CHECK-NEXT: bb0(%arg0: i32, %arg1: i32):
+// CHECK-NEXT:   %c0_i32 = constant 0 : i32
+// CHECK-NEXT:   br bb1
+// CHECK-NEXT: bb1:	// pred: bb0
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c42 = constant 42 : index
+// CHECK-NEXT:   br bb2(%c0) : index
+// CHECK-NEXT: bb2(%0: index):	// 2 preds: bb1, bb3
+// CHECK-NEXT:   %1 = cmpi "slt", %0, %c42 : index
+// CHECK-NEXT:   cond_br %1, bb3, bb4
+// CHECK-NEXT: bb3:	// pred: bb2
+// CHECK-NEXT:   %2 = call @body_args(%0) : (index) -> index
+// CHECK-NEXT:   %3 = call @other(%2, %arg0) : (index, i32) -> i32
+// CHECK-NEXT:   %4 = call @other(%2, %3) : (index, i32) -> i32
+// CHECK-NEXT:   %5 = call @other(%2, %arg1) : (index, i32) -> i32
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %6 = addi %0, %c1 : index
+// CHECK-NEXT:   br bb2(%6) : index
+// CHECK-NEXT: bb4:	// pred: bb2
+// CHECK-NEXT:   %c0_0 = constant 0 : index
+// CHECK-NEXT:   %7 = call @other(%c0_0, %c0_i32) : (index, i32) -> i32
+// CHECK-NEXT:   return %c0_i32, %7 : i32, i32
+// CHECK-NEXT: }
+mlfunc @mlfunc_args(%a : i32, %b : i32) -> (i32, i32) {
+  %r1 = constant 0 : i32
+  for %i = 0 to 42 {
+    %1 = call @body_args(%i) : (index) -> (index)
+    %2 = call @other(%1, %a) : (index, i32) -> (i32)
+    %3 = call @other(%1, %2) : (index, i32) -> (i32)
+    %4 = call @other(%1, %b) : (index, i32) -> (i32)
+  }
+  %ri = constant 0 : index
+  %r2 = call @other(%ri, %r1) : (index, i32) -> (i32)
+  return %r1, %r2 : i32, i32
+}
+
+/////////////////////////////////////////////////////////////////////
+
+extfunc @pre(index) -> ()
+extfunc @body2(index, index) -> ()
+extfunc @post(index) -> ()
+
+// CHECK-LABEL: cfgfunc @imperfectly_nested_loops() {
+// CHECK-NEXT: bb0:
+// CHECK-NEXT:   br bb1
+// CHECK-NEXT: bb1:	// pred: bb0
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c42 = constant 42 : index
+// CHECK-NEXT:   br bb2(%c0) : index
+// CHECK-NEXT: bb2(%0: index):	// 2 preds: bb1, bb7
+// CHECK-NEXT:   %1 = cmpi "slt", %0, %c42 : index
+// CHECK-NEXT:   cond_br %1, bb3, bb8
+// CHECK-NEXT: bb3:	// pred: bb2
+// CHECK-NEXT:   call @pre(%0) : (index) -> ()
+// CHECK-NEXT:   br bb4
+// CHECK-NEXT: bb4:	// pred: bb3
+// CHECK-NEXT:   %c7 = constant 7 : index
+// CHECK-NEXT:   %c56 = constant 56 : index
+// CHECK-NEXT:   br bb5(%c7) : index
+// CHECK-NEXT: bb5(%2: index):	// 2 preds: bb4, bb6
+// CHECK-NEXT:   %3 = cmpi "slt", %2, %c56 : index
+// CHECK-NEXT:   cond_br %3, bb6, bb7
+// CHECK-NEXT: bb6:	// pred: bb5
+// CHECK-NEXT:   call @body2(%0, %2) : (index, index) -> ()
+// CHECK-NEXT:   %c2 = constant 2 : index
+// CHECK-NEXT:   %4 = addi %2, %c2 : index
+// CHECK-NEXT:   br bb5(%4) : index
+// CHECK-NEXT: bb7:	// pred: bb5
+// CHECK-NEXT:   call @post(%0) : (index) -> ()
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %5 = addi %0, %c1 : index
+// CHECK-NEXT:   br bb2(%5) : index
+// CHECK-NEXT: bb8:	// pred: bb2
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+mlfunc @imperfectly_nested_loops() {
+  for %i = 0 to 42 {
+    call @pre(%i) : (index) -> ()
+    for %j = 7 to 56 step 2 {
+      call @body2(%i, %j) : (index, index) -> ()
+    }
+    call @post(%i) : (index) -> ()
+  }
+  return
+}
+
+/////////////////////////////////////////////////////////////////////
+
+extfunc @mid(index) -> ()
+extfunc @body3(index, index) -> ()
+
+// CHECK-LABEL: cfgfunc @more_imperfectly_nested_loops() {
+// CHECK-NEXT: bb0:
+// CHECK-NEXT:   br bb1
+// CHECK-NEXT: bb1:	// pred: bb0
+// CHECK-NEXT:   %c0 = constant 0 : index
+// CHECK-NEXT:   %c42 = constant 42 : index
+// CHECK-NEXT:   br bb2(%c0) : index
+// CHECK-NEXT: bb2(%0: index):	// 2 preds: bb1, bb11
+// CHECK-NEXT:   %1 = cmpi "slt", %0, %c42 : index
+// CHECK-NEXT:   cond_br %1, bb3, bb12
+// CHECK-NEXT: bb3:	// pred: bb2
+// CHECK-NEXT:   call @pre(%0) : (index) -> ()
+// CHECK-NEXT:   br bb4
+// CHECK-NEXT: bb4:	// pred: bb3
+// CHECK-NEXT:   %c7 = constant 7 : index
+// CHECK-NEXT:   %c56 = constant 56 : index
+// CHECK-NEXT:   br bb5(%c7) : index
+// CHECK-NEXT: bb5(%2: index):	// 2 preds: bb4, bb6
+// CHECK-NEXT:   %3 = cmpi "slt", %2, %c56 : index
+// CHECK-NEXT:   cond_br %3, bb6, bb7
+// CHECK-NEXT: bb6:	// pred: bb5
+// CHECK-NEXT:   call @body2(%0, %2) : (index, index) -> ()
+// CHECK-NEXT:   %c2 = constant 2 : index
+// CHECK-NEXT:   %4 = addi %2, %c2 : index
+// CHECK-NEXT:   br bb5(%4) : index
+// CHECK-NEXT: bb7:	// pred: bb5
+// CHECK-NEXT:   call @mid(%0) : (index) -> ()
+// CHECK-NEXT:   br bb8
+// CHECK-NEXT: bb8:	// pred: bb7
+// CHECK-NEXT:   %c18 = constant 18 : index
+// CHECK-NEXT:   %c37 = constant 37 : index
+// CHECK-NEXT:   br bb9(%c18) : index
+// CHECK-NEXT: bb9(%5: index):	// 2 preds: bb8, bb10
+// CHECK-NEXT:   %6 = cmpi "slt", %5, %c37 : index
+// CHECK-NEXT:   cond_br %6, bb10, bb11
+// CHECK-NEXT: bb10:	// pred: bb9
+// CHECK-NEXT:   call @body3(%0, %5) : (index, index) -> ()
+// CHECK-NEXT:   %c3 = constant 3 : index
+// CHECK-NEXT:   %7 = addi %5, %c3 : index
+// CHECK-NEXT:   br bb9(%7) : index
+// CHECK-NEXT: bb11:	// pred: bb9
+// CHECK-NEXT:   call @post(%0) : (index) -> ()
+// CHECK-NEXT:   %c1 = constant 1 : index
+// CHECK-NEXT:   %8 = addi %0, %c1 : index
+// CHECK-NEXT:   br bb2(%8) : index
+// CHECK-NEXT: bb12:	// pred: bb2
+// CHECK-NEXT:   return
+// CHECK-NEXT: }
+mlfunc @more_imperfectly_nested_loops() {
+  for %i = 0 to 42 {
+    call @pre(%i) : (index) -> ()
+    for %j = 7 to 56 step 2 {
+      call @body2(%i, %j) : (index, index) -> ()
+    }
+    call @mid(%i) : (index) -> ()
+    for %k = 18 to 37 step 3 {
+      call @body3(%i, %k) : (index, index) -> ()
+    }
+    call @post(%i) : (index) -> ()
+  }
+  return
+}
-- 
2.34.1