Scaffolding for convertToCFG pass that replaces all instances of ML functions with equivalent CFG functions. Traverses module MLIR, generates CFG functions (empty for now) and removes ML functions. Adds Transforms library and tests.

PiperOrigin-RevId: 205848367

mlir/include/mlir/IR/Builders.h

@@ -123,6 +123,10 @@ public:
     insertPoint = block->end();
   }
 
+  // Add new basic block and set the insertion point to the end of it.
+  BasicBlock *createBlock();
+
+  // Instructions.
   OperationInst *createOperation(Identifier name, ArrayRef<CFGValue *> operands,
                                  ArrayRef<Type *> resultTypes,
                                  ArrayRef<NamedAttribute> attributes) {

mlir/include/mlir/IR/MLFunction.h

@@ -32,6 +32,7 @@ namespace mlir {
 class MLFunction : public Function, public StmtBlock {
 public:
   MLFunction(StringRef name, FunctionType *type);
+  ~MLFunction();
 
   // TODO: add function arguments and return values once
   // SSA values are implemented

mlir/include/mlir/IR/Statement.h

@@ -51,8 +51,8 @@ public:
   /// Returns the function that this statement is part of.
   MLFunction *getFunction() const;
 
-  /// Destroys the argument statement or one of its subclasses
-  static void destroy(Statement *stmt);
+  /// Destroys this statement and its subclass data.
+  void destroy();
 
   void print(raw_ostream &os) const;
   void dump() const;
@@ -90,9 +90,7 @@ struct ilist_traits<::mlir::Statement> {
   using Statement = ::mlir::Statement;
   using stmt_iterator = simple_ilist<Statement>::iterator;
 
-  static void deleteNode(Statement *stmt) {
-    Statement::destroy(stmt);
-  }
+  static void deleteNode(Statement *stmt) { stmt->destroy(); }
 
   void addNodeToList(Statement *stmt);
   void removeNodeFromList(Statement *stmt);

mlir/include/mlir/IR/Statements.h

@@ -58,7 +58,7 @@ public:
       : Statement(Kind::For), StmtBlock(StmtBlockKind::For),
         lowerBound(lowerBound), upperBound(upperBound), step(step) {}
 
-  //TODO: delete nested statements or assert that they are gone.
+  // Loop bounds and step are immortal objects and don't need to be deleted.
   ~ForStmt() {}
 
   // TODO: represent induction variable
@@ -95,7 +95,6 @@ public:
     return block->getStmtBlockKind() == StmtBlockKind::IfClause;
   }
 
-  //TODO: delete nested statements or assert that they are gone.
   ~IfClause() {}
 
   /// Returns the if statement that contains this clause.

mlir/include/mlir/Transforms/ConvertToCFG.h

@@ -0,0 +1,34 @@
+//===- ConvertToCFG.h - Convert ML functions to CFG ones --------*- C++ -*-===//
+//
+// 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 defines APIs to convert ML functions into CFG functions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TRANSFORMS_CONVERTTOCFG_H
+#define MLIR_TRANSFORMS_CONVERTTOCFG_H
+
+namespace mlir {
+class Module;
+
+/// Replaces all ML functions in the module with equivalent CFG functions.
+/// Function references are appropriately patched to refer only
+/// to CFG functions.
+void convertToCFG(Module *module);
+
+} // namespace mlir
+#endif // MLIR_TRANSFORMS_CONVERTTOCFG_H

mlir/lib/IR/Builders.cpp

@@ -141,7 +141,19 @@ AffineExpr *Builder::getCeilDivExpr(AffineExpr *lhs, AffineExpr *rhs) {
 }
 
 //===----------------------------------------------------------------------===//
-// Statements
+// CFG function elements.
+//===----------------------------------------------------------------------===//
+
+// Basic block.
+BasicBlock *CFGFuncBuilder::createBlock() {
+  BasicBlock *b = new BasicBlock();
+  function->push_back(b);
+  setInsertionPoint(b);
+  return b;
+}
+
+//===----------------------------------------------------------------------===//
+// Statements.
 //===----------------------------------------------------------------------===//
 
 ForStmt *MLFuncBuilder::createFor(AffineConstantExpr *lowerBound,
@@ -149,7 +161,7 @@ ForStmt *MLFuncBuilder::createFor(AffineConstantExpr *lowerBound,
                                   AffineConstantExpr *step) {
   if (!step)
     step = getConstantExpr(1);
-  auto stmt = new ForStmt(lowerBound, upperBound, step);
+  auto *stmt = new ForStmt(lowerBound, upperBound, step);
   block->getStatements().push_back(stmt);
   return stmt;
 }

mlir/lib/IR/Statement.cpp

@@ -30,16 +30,16 @@ Statement::~Statement() {
 }
 
 /// Destroy this statement or one of its subclasses.
-void Statement::destroy(Statement *stmt) {
-  switch (stmt->getKind()) {
+void Statement::destroy() {
+  switch (this->getKind()) {
   case Kind::Operation:
-    delete cast<OperationStmt>(stmt);
+    delete cast<OperationStmt>(this);
     break;
   case Kind::For:
-    delete cast<ForStmt>(stmt);
+    delete cast<ForStmt>(this);
     break;
   case Kind::If:
-    delete cast<IfStmt>(stmt);
+    delete cast<IfStmt>(this);
     break;
   }
 }
@@ -102,7 +102,6 @@ void Statement::eraseFromBlock() {
 //===----------------------------------------------------------------------===//
 
 IfStmt::~IfStmt() {
-  // TODO: correctly delete StmtBlocks under then and else clauses
   delete thenClause;
   if (elseClause != nullptr)
     delete elseClause;

mlir/lib/Transforms/ConvertToCFG.cpp

@@ -0,0 +1,166 @@
+//===- ConvertToCFG.cpp - ML function to CFG function converstion ---------===//
+//
+// 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 implements APIs to convert ML functions into CFG functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Transforms/ConvertToCFG.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/CFGFunction.h"
+#include "mlir/IR/MLFunction.h"
+#include "mlir/IR/Module.h"
+#include "llvm/ADT/DenseSet.h"
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// ML function converter
+//===----------------------------------------------------------------------===//
+
+namespace {
+// Generates CFG function equivalent to the given ML function.
+class FunctionConverter {
+public:
+  FunctionConverter(CFGFunction *cfgFunc)
+      : cfgFunc(cfgFunc), builder(cfgFunc) {}
+  CFGFunction *convert(const MLFunction *mlFunc);
+
+private:
+  CFGFunction *cfgFunc;
+  CFGFuncBuilder builder;
+};
+} // namespace
+
+CFGFunction *FunctionConverter::convert(const MLFunction *mlFunc) {
+  builder.createBlock();
+
+  // Creates return instruction with no operands.
+  // TODO: convert return operands.
+  builder.createReturnInst({});
+
+  // TODO: convert ML function body.
+
+  return cfgFunc;
+}
+
+//===----------------------------------------------------------------------===//
+// Module converter
+//===----------------------------------------------------------------------===//
+namespace {
+// ModuleConverter class does CFG conversion for the whole module.
+class ModuleConverter {
+public:
+  explicit ModuleConverter(Module *module) : module(module) {}
+  void run();
+
+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);
+  // 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();
+
+  // Map from ML functions to generated CFG functions.
+  llvm::DenseMap<const MLFunction *, CFGFunction *> generatedFuncs;
+  Module *module;
+};
+} // end anonymous namespace
+
+// 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.
+void ModuleConverter::run() {
+  convertMLFunctions();
+  replaceReferences();
+}
+
+void ModuleConverter::convertMLFunctions() {
+  for (Function *fn : module->functionList) {
+    if (auto mlFunc = dyn_cast<MLFunction>(fn))
+      generatedFuncs[mlFunc] = convert(mlFunc);
+  }
+}
+
+// Creates CFG function equivalent to the given ML function.
+CFGFunction *ModuleConverter::convert(const MLFunction *mlFunc) {
+  // TODO: ensure that CFG function name is unique.
+  CFGFunction *cfgFunc =
+      new CFGFunction(mlFunc->getName() + "_cfg", mlFunc->getType());
+  module->functionList.push_back(cfgFunc);
+
+  // Generates the body of the CFG function.
+  return FunctionConverter(cfgFunc).convert(mlFunc);
+}
+
+void ModuleConverter::replaceReferences() {
+  for (Function *fn : module->functionList) {
+    switch (fn->getKind()) {
+    case Function::Kind::CFGFunc:
+      replaceReferences(cast<CFGFunction>(fn));
+      break;
+    case Function::Kind::MLFunc:
+      replaceReferences(cast<MLFunction>(fn));
+      break;
+    case Function::Kind::ExtFunc:
+      // nothing to do for external functions
+      break;
+    }
+  }
+}
+
+void ModuleConverter::replaceReferences(CFGFunction *func) {
+  // TODO: NOP for now since function attributes are not yet implemented.
+}
+
+void ModuleConverter::replaceReferences(MLFunction *func) {
+  // TODO: NOP for now since function attributes are not yet implemented.
+}
+
+// Removes all ML functions from the module.
+void ModuleConverter::removeMLFunctions() {
+  std::vector<Function *> &fnList = module->functionList;
+
+  // Delete ML functions and its data.
+  for (auto &fn : fnList) {
+    if (auto mlFunc = dyn_cast<MLFunction>(fn)) {
+      delete mlFunc;
+      fn = nullptr;
+    }
+  }
+
+  // Remove ML functions from the function list.
+  fnList.erase(std::remove_if(fnList.begin(), fnList.end(),
+                              [](Function *fn) { return !fn; }),
+               fnList.end());
+}
+
+//===----------------------------------------------------------------------===//
+// Entry point method
+//===----------------------------------------------------------------------===//
+
+void mlir::convertToCFG(Module *module) {
+  ModuleConverter moduleConverter(module);
+  moduleConverter.run();
+  module->verify();
+}

mlir/test/Transforms/convert2cfg.mlir

@@ -0,0 +1,7 @@
+// RUN: %S/../../mlir-opt %s -o - -convert-to-cfg | FileCheck %s
+
+// CHECK-LABEL: cfgfunc @empty_cfg() {
+mlfunc @empty() {
+             // CHECK: bb0:
+  return     // CHECK:  return
+}            // CHECK: }

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -24,6 +24,7 @@
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Parser.h"
+#include "mlir/Transforms/ConvertToCFG.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/InitLLVM.h"
@@ -44,6 +45,10 @@ static cl::opt<bool>
 checkParserErrors("check-parser-errors", cl::desc("Check for parser errors"),
                   cl::init(false));
 
+static cl::opt<bool> convertToCFGOpt(
+    "convert-to-cfg",
+    cl::desc("Convert all ML functions in the module to CFG ones"));
+
 enum OptResult { OptSuccess, OptFailure };
 
 /// Open the specified output file and return it, exiting if there is any I/O or
@@ -61,7 +66,8 @@ static std::unique_ptr<ToolOutputFile> getOutputStream() {
 }
 
 /// Parses the memory buffer and, if successfully parsed, prints the parsed
-/// output.
+/// output. Optionally, convert ML functions into CFG functions.
+/// TODO: pull parsing and printing into separate functions.
 OptResult parseAndPrintMemoryBuffer(std::unique_ptr<MemoryBuffer> buffer) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
@@ -73,6 +79,10 @@ OptResult parseAndPrintMemoryBuffer(std::unique_ptr<MemoryBuffer> buffer) {
   if (!module)
     return OptFailure;
 
+  // Convert ML functions into CFG functions
+  if (convertToCFGOpt)
+    convertToCFG(module.get());
+
   // Print the output.
   auto output = getOutputStream();
   module->print(output->os());