Implement parser/IR support for CFG functions, basic blocks and return instruction.

This is pretty much minimal scaffolding for this step.  Basic block arguments,
instructions, other terminators, a proper IR representation for
blocks/instructions, etc are all coming.

PiperOrigin-RevId: 201826439

mlir/include/mlir/IR/BasicBlock.h

@@ -0,0 +1,57 @@
+//===- BasicBlock.h - MLIR BasicBlock Class ---------------------*- 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.
+// =============================================================================
+
+#ifndef MLIR_IR_BASICBLOCK_H
+#define MLIR_IR_BASICBLOCK_H
+
+#include "mlir/IR/Instructions.h"
+
+namespace mlir {
+
+/// Each basic block in a CFG function contains a list of basic block arguments,
+/// normal instructions, and a terminator instruction.
+///
+/// Basic blocks form a graph (the CFG) which can be traversed through
+/// predecessor and successor edges.
+class BasicBlock {
+public:
+  explicit BasicBlock(CFGFunction *function);
+
+  /// Return the function that a BasicBlock is part of.
+  CFGFunction *getFunction() const {
+    return function;
+  }
+
+  // TODO: bb arguments, instruction list.
+
+  void setTerminator(TerminatorInst *inst) {
+    terminator = inst;
+  }
+  TerminatorInst *getTerminator() const { return terminator; }
+
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+private:
+  CFGFunction *const function;
+  // FIXME: wrong representation and API, leaks memory etc.
+  TerminatorInst *terminator = nullptr;
+};
+
+} // end namespace mlir
+
+#endif  // MLIR_IR_BASICBLOCK_H

mlir/include/mlir/IR/CFGFunction.h

@@ -0,0 +1,47 @@
+//===- CFGFunction.h - MLIR CFGFunction Class -------------------*- 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.
+// =============================================================================
+
+#ifndef MLIR_IR_CFGFUNCTION_H
+#define MLIR_IR_CFGFUNCTION_H
+
+#include "mlir/IR/Function.h"
+#include "mlir/IR/BasicBlock.h"
+#include <vector>
+
+namespace mlir {
+
+// This kind of function is defined in terms of a "Control Flow Graph" of basic
+// blocks, each of which includes instructions.
+class CFGFunction : public Function {
+public:
+  CFGFunction(StringRef name, FunctionType *type);
+
+  // FIXME: wrong representation and API, leaks memory etc.
+  std::vector<BasicBlock*> blockList;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(const Function *func) {
+    return func->getKind() == Kind::CFGFunc;
+  }
+
+  void print(raw_ostream &os) const;
+};
+
+
+} // end namespace mlir
+
+#endif  // MLIR_IR_CFGFUNCTION_H

mlir/include/mlir/IR/Function.h

@@ -29,16 +29,52 @@
 namespace mlir {
   class FunctionType;
 
-  /// This is the base class for all of the MLIR function types
-  class Function {
-    std::string name;
-    FunctionType *const type;
-  public:
-    explicit Function(StringRef name, FunctionType *type);
-
-    void print(raw_ostream &os);
-    void dump();
+/// This is the base class for all of the MLIR function types.
+class Function {
+public:
+  enum class Kind {
+    ExtFunc, CFGFunc
   };
+
+  Kind getKind() const { return kind; }
+
+  /// Return the name of this function, without the @.
+  const std::string &getName() const { return name; }
+
+  /// Return the type of this function.
+  FunctionType *getType() const {
+    return type;
+  }
+
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+protected:
+  Function(StringRef name, FunctionType *type, Kind kind);
+  ~Function() {}
+private:
+  Kind kind;
+  std::string name;
+  FunctionType *const type;
+
+  void operator=(const Function&) = delete;
+};
+
+/// An extfunc declaration is a declaration of a function signature that is
+/// defined in some other module.
+class ExtFunction : public Function {
+public:
+  ExtFunction(StringRef name, FunctionType *type);
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(const Function *func) {
+    return func->getKind() == Kind::ExtFunc;
+  }
+
+  void print(raw_ostream &os) const;
+};
+
+
 } // end namespace mlir
 
 #endif  // MLIR_IR_FUNCTION_H

mlir/include/mlir/IR/Instructions.h

@@ -0,0 +1,73 @@
+//===- Instructions.h - MLIR CFG Instruction Classes ------------*- 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 the classes for CFGFunction instructions.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_INSTRUCTIONS_H
+#define MLIR_IR_INSTRUCTIONS_H
+
+#include "mlir/Support/LLVM.h"
+
+namespace mlir {
+  class BasicBlock;
+  class CFGFunction;
+
+
+/// Terminator instructions are the last part of a basic block, used to
+/// represent control flow and returns.
+class TerminatorInst {
+public:
+  enum class Kind {
+    Return
+  };
+
+  Kind getKind() const { return kind; }
+
+  /// Return the BasicBlock that contains this terminator instruction.
+  BasicBlock *getBlock() const {
+    return block;
+  }
+  CFGFunction *getFunction() const;
+
+  void print(raw_ostream &os) const;
+  void dump() const;
+
+protected:
+  TerminatorInst(Kind kind, BasicBlock *block) : kind(kind), block(block) {}
+
+private:
+  Kind kind;
+  BasicBlock *block;
+};
+
+class ReturnInst : public TerminatorInst {
+public:
+  explicit ReturnInst(BasicBlock *block);
+  // TODO: Flesh this out.
+
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(const TerminatorInst *inst) {
+    return inst->getKind() == Kind::Return;
+  }
+};
+
+} // end namespace mlir
+
+#endif  // MLIR_IR_INSTRUCTIONS_H

mlir/include/mlir/IR/Module.h

@@ -34,8 +34,8 @@ public:
   std::vector<Function*> functionList;
 
 
-  void print(raw_ostream &os);
-  void dump();
+  void print(raw_ostream &os) const;
+  void dump() const;
 };
 } // end namespace mlir
 

mlir/include/mlir/IR/Types.h

@@ -43,6 +43,8 @@ enum class TypeKind {
   F32,
   F64,
 
+  /// This is a marker for the last primitive type.  The range of primitive
+  /// types is expected to be this element and earlier.
   LAST_PRIMITIVE_TYPE = F64,
 
   // Derived types.
@@ -129,8 +131,8 @@ public:
   static PrimitiveType *get(TypeKind kind, MLIRContext *context);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(const Type *T) {
-    return T->getKind() <= TypeKind::LAST_PRIMITIVE_TYPE;
+  static bool classof(const Type *type) {
+    return type->getKind() <= TypeKind::LAST_PRIMITIVE_TYPE;
   }
 private:
   PrimitiveType(TypeKind kind, MLIRContext *context);
@@ -174,11 +176,6 @@ public:
   static FunctionType *get(ArrayRef<Type*> inputs, ArrayRef<Type*> results,
                            MLIRContext *context);
 
-  /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(const Type *T) {
-    return T->getKind() == TypeKind::Function;
-  }
-
   ArrayRef<Type*> getInputs() const {
     return ArrayRef<Type*>(inputsAndResults, getSubclassData());
   }
@@ -187,6 +184,11 @@ public:
     return ArrayRef<Type*>(inputsAndResults+getSubclassData(), numResults);
   }
 
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(const Type *type) {
+    return type->getKind() == TypeKind::Function;
+  }
+
 private:
   unsigned numResults;
   Type *const *inputsAndResults;
@@ -211,8 +213,8 @@ public:
   }
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool classof(const Type *T) {
-    return T->getKind() == TypeKind::Vector;
+  static bool classof(const Type *type) {
+    return type->getKind() == TypeKind::Vector;
   }
 
 private:

mlir/include/mlir/Support/LLVM.h

@@ -48,6 +48,12 @@ namespace llvm {
   template<typename T> class TinyPtrVector;
   template<typename T> class Optional;
   template <typename PT1, typename PT2> class PointerUnion;
+  namespace detail {
+    template <typename KeyT, typename ValueT> struct DenseMapPair;
+  }
+  template<typename T> struct DenseMapInfo;
+  template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
+  class DenseMap;
 
   // Other common classes.
   class raw_ostream;
@@ -79,6 +85,7 @@ namespace mlir {
   using llvm::MutableArrayRef;
   using llvm::TinyPtrVector;
   using llvm::PointerUnion;
+  using llvm::DenseMap;
 
   // Other common classes.
   using llvm::raw_ostream;

mlir/lib/IR/AsmPrinter.cpp

@@ -0,0 +1,175 @@
+//===- AsmPrinter.cpp - MLIR Assembly Printer Implementation --------------===//
+//
+// 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 the MLIR AsmPrinter class, which is used to implement
+// the various print() methods on the core IR objects.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/CFGFunction.h"
+#include "mlir/IR/Module.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/STLExtras.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/ADT/DenseMap.h"
+using namespace mlir;
+
+
+//===----------------------------------------------------------------------===//
+// Function printing
+//===----------------------------------------------------------------------===//
+
+static void printFunctionSignature(const Function *fn, raw_ostream &os) {
+  auto type = fn->getType();
+
+  os << "@" << fn->getName() << '(';
+  interleave(type->getInputs(),
+             [&](Type *eltType) { os << *eltType; },
+             [&]() { os << ", "; });
+  os << ')';
+
+  switch (type->getResults().size()) {
+  case 0: break;
+  case 1:
+    os << " -> " << *type->getResults()[0];
+    break;
+  default:
+    os << " -> (";
+    interleave(type->getResults(),
+               [&](Type *eltType) { os << *eltType; },
+               [&]() { os << ", "; });
+    os << ')';
+    break;
+  }
+}
+
+void ExtFunction::print(raw_ostream &os) const {
+  os << "extfunc ";
+  printFunctionSignature(this, os);
+  os << "\n";
+}
+
+//===----------------------------------------------------------------------===//
+// CFG Function printing
+//===----------------------------------------------------------------------===//
+
+namespace {
+class CFGFunctionState {
+public:
+  CFGFunctionState(const CFGFunction *function, raw_ostream &os);
+
+  const CFGFunction *getFunction() const { return function; }
+
+  void print();
+  void print(const BasicBlock *block);
+  void print(const TerminatorInst *inst);
+
+  unsigned getBBID(const BasicBlock *block) {
+    auto it = basicBlockIDs.find(block);
+    assert(it != basicBlockIDs.end() && "Block not in this function?");
+    return it->second;
+  }
+
+private:
+  const CFGFunction *function;
+  raw_ostream &os;
+  DenseMap<BasicBlock*, unsigned> basicBlockIDs;
+};
+} // end anonymous namespace
+
+CFGFunctionState::CFGFunctionState(const CFGFunction *function, raw_ostream &os)
+  : function(function), os(os) {
+
+  // Each basic block gets a unique ID per function.
+  unsigned blockID = 0;
+  for (auto *block : function->blockList)
+    basicBlockIDs[block] = blockID++;
+}
+
+void CFGFunctionState::print() {
+  os << "cfgfunc ";
+  printFunctionSignature(this->getFunction(), os);
+  os << " {\n";
+
+  for (auto *block : function->blockList)
+    print(block);
+  os << "}\n\n";
+}
+
+void CFGFunctionState::print(const BasicBlock *block) {
+  os << "bb" << getBBID(block) << ":\n";
+
+  // TODO Print arguments and instructions.
+
+  print(block->getTerminator());
+}
+
+void CFGFunctionState::print(const TerminatorInst *inst) {
+  switch (inst->getKind()) {
+  case TerminatorInst::Kind::Return:
+    os << "  return\n";
+    break;
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// print and dump methods
+//===----------------------------------------------------------------------===//
+
+void TerminatorInst::print(raw_ostream &os) const {
+  CFGFunctionState state(getFunction(), os);
+  state.print(this);
+}
+
+void TerminatorInst::dump() const {
+  print(llvm::errs());
+}
+
+void BasicBlock::print(raw_ostream &os) const {
+  CFGFunctionState state(getFunction(), os);
+  state.print();
+}
+
+void BasicBlock::dump() const {
+  print(llvm::errs());
+}
+
+void Function::print(raw_ostream &os) const {
+  switch (getKind()) {
+  case Kind::ExtFunc: return cast<ExtFunction>(this)->print(os);
+  case Kind::CFGFunc: return cast<CFGFunction>(this)->print(os);
+  }
+}
+
+void Function::dump() const {
+  print(llvm::errs());
+}
+
+void CFGFunction::print(raw_ostream &os) const {
+  CFGFunctionState state(this, os);
+  state.print();
+}
+
+void Module::print(raw_ostream &os) const {
+  for (auto *fn : functionList)
+    fn->print(os);
+}
+
+void Module::dump() const {
+  print(llvm::errs());
+}
+

mlir/lib/IR/BasicBlock.cpp

@@ -0,0 +1,22 @@
+//===- BasicBlock.cpp - MLIR BasicBlock Class -----------------------------===//
+//
+// 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.
+// =============================================================================
+
+#include "mlir/IR/BasicBlock.h"
+using namespace mlir;
+
+BasicBlock::BasicBlock(CFGFunction *function) : function(function) {
+}

mlir/lib/IR/Function.cpp

@@ -15,40 +15,26 @@
 // limitations under the License.
 // =============================================================================
 
-#include "mlir/IR/Function.h"
-#include "mlir/IR/Types.h"
-#include "mlir/Support/STLExtras.h"
-#include "llvm/Support/raw_ostream.h"
+#include "mlir/IR/CFGFunction.h"
+#include "llvm/ADT/StringRef.h"
 using namespace mlir;
 
-Function::Function(StringRef name, FunctionType *type)
-  : name(name.str()), type(type) {
+Function::Function(StringRef name, FunctionType *type, Kind kind)
+  : kind(kind), name(name.str()), type(type) {
 }
 
-void Function::print(raw_ostream &os) {
-  os << "extfunc @" << name << '(';
-  interleave(type->getInputs(),
-             [&](Type *eltType) { os << *eltType; },
-             [&]() { os << ", "; });
-  os << ')';
+//===----------------------------------------------------------------------===//
+// ExtFunction implementation.
+//===----------------------------------------------------------------------===//
 
-  switch (type->getResults().size()) {
-  case 0: break;
-  case 1:
-    os << " -> " << *type->getResults()[0];
-    break;
-  default:
-    os << " -> (";
-    interleave(type->getResults(),
-               [&](Type *eltType) { os << *eltType; },
-               [&]() { os << ", "; });
-    os << ')';
-    break;
-  }
-
-  os << "\n";
+ExtFunction::ExtFunction(StringRef name, FunctionType *type)
+  : Function(name, type, Kind::ExtFunc) {
 }
 
-void Function::dump() {
-  print(llvm::errs());
+//===----------------------------------------------------------------------===//
+// CFGFunction implementation.
+//===----------------------------------------------------------------------===//
+
+CFGFunction::CFGFunction(StringRef name, FunctionType *type)
+  : Function(name, type, Kind::CFGFunc) {
 }

mlir/lib/IR/Instructions.cpp

@@ -0,0 +1,28 @@
+//===- Instructions.cpp - MLIR CFGFunction Instruction Classes ------------===//
+//
+// 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.
+// =============================================================================
+
+#include "mlir/IR/Instructions.h"
+#include "mlir/IR/BasicBlock.h"
+using namespace mlir;
+
+CFGFunction *TerminatorInst::getFunction() const {
+  return getBlock()->getFunction();
+}
+
+ReturnInst::ReturnInst(BasicBlock *block) : TerminatorInst(Kind::Return, block){
+}
+

mlir/lib/IR/Module.cpp

@@ -16,19 +16,8 @@
 // =============================================================================
 
 #include "mlir/IR/Module.h"
-#include "llvm/Support/raw_ostream.h"
 using namespace mlir;
 
 Module::Module() {
 }
 
-
-void Module::print(raw_ostream &os) {
-  for (auto *fn : functionList)
-    fn->print(os);
-}
-
-void Module::dump() {
-  print(llvm::errs());
-}
-

mlir/lib/Parser/Lexer.cpp

@@ -66,9 +66,12 @@ Token Lexer::lexToken() {
     // Ignore whitespace.
     return lexToken();
 
+  case ':': return formToken(Token::colon, tokStart);
   case ',': return formToken(Token::comma, tokStart);
   case '(': return formToken(Token::l_paren, tokStart);
   case ')': return formToken(Token::r_paren, tokStart);
+  case '{': return formToken(Token::l_brace, tokStart);
+  case '}': return formToken(Token::r_brace, tokStart);
   case '<': return formToken(Token::less, tokStart);
   case '>': return formToken(Token::greater, tokStart);
 
@@ -148,6 +151,7 @@ Token Lexer::lexBareIdentifierOrKeyword(const char *tokStart) {
     .Case("int", Token::kw_int)
     .Case("memref", Token::kw_memref)
     .Case("mlfunc", Token::kw_mlfunc)
+    .Case("return", Token::kw_return)
     .Case("tensor", Token::kw_tensor)
     .Case("vector", Token::kw_vector)
     .Default(Token::bare_identifier);

mlir/lib/Parser/Parser.cpp

@@ -22,12 +22,16 @@
 #include "mlir/Parser.h"
 #include "Lexer.h"
 #include "mlir/IR/Module.h"
+#include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/Types.h"
 #include "llvm/Support/SourceMgr.h"
 using namespace mlir;
 using llvm::SourceMgr;
+using llvm::SMLoc;
 
 namespace {
+class CFGFunctionParserState;
+
 /// Simple enum to make code read better in cases that would otherwise return a
 /// bool value.  Failure is "true" in a boolean context.
 enum ParseResult {
@@ -61,7 +65,10 @@ private:
   // Helper methods.
 
   /// Emit an error and return failure.
-  ParseResult emitError(const Twine &message);
+  ParseResult emitError(const Twine &message) {
+    return emitError(curToken.getLoc(), message);
+  }
+  ParseResult emitError(SMLoc loc, const Twine &message);
 
   /// Advance the current lexer onto the next token.
   void consumeToken() {
@@ -107,9 +114,11 @@ private:
   Type *parseType();
   ParseResult parseTypeList(SmallVectorImpl<Type*> &elements);
 
-  // Top level entity parsing.
+  // Functions.
   ParseResult parseFunctionSignature(StringRef &name, FunctionType *&type);
   ParseResult parseExtFunc();
+  ParseResult parseCFGFunc();
+  ParseResult parseBasicBlock(CFGFunctionParserState &functionState);
 };
 } // end anonymous namespace
 
@@ -117,7 +126,7 @@ private:
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
-ParseResult Parser::emitError(const Twine &message) {
+ParseResult Parser::emitError(SMLoc loc, const Twine &message) {
   // If we hit a parse error in response to a lexer error, then the lexer
   // already emitted an error.
   if (curToken.is(Token::error))
@@ -125,8 +134,7 @@ ParseResult Parser::emitError(const Twine &message) {
 
   // TODO(clattner): If/when we want to implement a -verify mode, this will need
   // to package up errors into SMDiagnostic and report them.
-  lex.getSourceMgr().PrintMessage(curToken.getLoc(), SourceMgr::DK_Error,
-                                  message);
+  lex.getSourceMgr().PrintMessage(loc, SourceMgr::DK_Error, message);
   return ParseFailure;
 }
 
@@ -442,10 +450,10 @@ ParseResult Parser::parseTypeList(SmallVectorImpl<Type*> &elements) {
   return ParseSuccess;
 }
 
+//===----------------------------------------------------------------------===//
+// Functions
+//===----------------------------------------------------------------------===//
 
-//===----------------------------------------------------------------------===//
-// Top-level entity parsing.
-//===----------------------------------------------------------------------===//
 
 /// Parse a function signature, starting with a name and including the parameter
 /// list.
@@ -493,11 +501,119 @@ ParseResult Parser::parseExtFunc() {
 
 
   // Okay, the external function definition was parsed correctly.
-  module->functionList.push_back(new Function(name, type));
+  module->functionList.push_back(new ExtFunction(name, type));
   return ParseSuccess;
 }
 
 
+namespace {
+/// This class represents the transient parser state for the internals of a
+/// function as we are parsing it, e.g. the names for basic blocks.  It handles
+/// forward references.
+class CFGFunctionParserState {
+public:
+  CFGFunctionParserState(CFGFunction *function) : function(function) {}
+
+  /// Get the basic block with the specified name, creating it if it doesn't
+  /// already exist.
+  BasicBlock *getBlockNamed(StringRef name) {
+    auto *&block = blocksByName[name];
+    if (!block) {
+      block = new BasicBlock(function);
+      // TODO: Should be automatic when we have the right function
+      // representation.
+      function->blockList.push_back(block);
+    }
+    return block;
+  }
+private:
+  CFGFunction *function;  
+  llvm::StringMap<BasicBlock*> blocksByName;
+};
+} // end anonymous namespace
+
+
+/// CFG function declarations.
+///
+///   cfg-func ::= `cfgfunc` function-signature `{` basic-block+ `}`
+///
+ParseResult Parser::parseCFGFunc() {
+  consumeToken(Token::kw_cfgfunc);
+
+  StringRef name;
+  FunctionType *type = nullptr;
+  if (parseFunctionSignature(name, type))
+    return ParseFailure;
+
+  if (!consumeIf(Token::l_brace))
+    return emitError("expected '{' in CFG function");
+
+  // Okay, the CFG function signature was parsed correctly, create the function.
+  auto function = new CFGFunction(name, type);
+
+  // Make sure we have at least one block.
+  if (curToken.is(Token::r_brace))
+    return emitError("CFG functions must have at least one basic block");
+
+  CFGFunctionParserState functionState(function);
+
+  // Parse the list of blocks.
+  while (!consumeIf(Token::r_brace))
+    if (parseBasicBlock(functionState))
+      return ParseFailure;
+
+  module->functionList.push_back(function);
+  return ParseSuccess;
+}
+
+/// Basic block declaration.
+///
+///   basic-block ::= bb-label instruction* terminator-stmt
+///   bb-label    ::= bb-id bb-arg-list? `:`
+///   bb-id       ::= bare-id
+///   bb-arg-list ::= `(` ssa-id-and-type-list? `)`
+///
+ParseResult Parser::parseBasicBlock(CFGFunctionParserState &functionState) {
+  SMLoc nameLoc = curToken.getLoc();
+  auto name = curToken.getSpelling();
+  if (!consumeIf(Token::bare_identifier))
+    return emitError("expected basic block name");
+  auto block = functionState.getBlockNamed(name);
+
+  // If this block has already been parsed, then this is a redefinition with the
+  // same block name.
+  if (block->getTerminator())
+    return emitError(nameLoc, "redefinition of block named '" +
+                     name.str() + "'");
+
+  // TODO: parse bb argument list.
+
+  if (!consumeIf(Token::colon))
+    return emitError("expected ':' after basic block name");
+
+
+  // TODO(clattner): Verify block hasn't already been parsed (this would be a
+  // redefinition of the same name) once we have a body implementation.
+
+  // TODO(clattner): Move block to the end of the list, once we have a proper
+  // block list representation in CFGFunction.
+
+  // TODO: parse instruction list.
+
+  // TODO: Generalize this once instruction list parsing is built out.
+  if (!consumeIf(Token::kw_return))
+    return emitError("expected 'return' at end of basic block");
+
+  block->setTerminator(new ReturnInst(block));
+
+  return ParseSuccess;
+}
+
+
+//===----------------------------------------------------------------------===//
+// Top-level entity parsing.
+//===----------------------------------------------------------------------===//
+
 /// This is the top-level module parser.
 Module *Parser::parseModule() {
   while (1) {
@@ -517,11 +633,14 @@ Module *Parser::parseModule() {
       return nullptr;
 
     case Token::kw_extfunc:
-      if (parseExtFunc())
-        return nullptr;
+      if (parseExtFunc()) return nullptr;
       break;
 
-    // TODO: cfgfunc, mlfunc, affine entity declarations, etc.
+    case Token::kw_cfgfunc:
+      if (parseCFGFunc()) return nullptr;
+      break;
+
+    // TODO: mlfunc, affine entity declarations, etc.
     }
   }
 }

mlir/lib/Parser/Token.h

@@ -40,10 +40,12 @@ public:
 
     // Punctuation.
     arrow,              // ->
+    colon,              // :
     comma,              // ,
     question,           // ?
     questionquestion,   // ??
     l_paren, r_paren,   // ( )
+    l_brace, r_brace,   // { }
     less, greater,      // < >
     // TODO: More punctuation.
 
@@ -62,6 +64,7 @@ public:
     kw_int,
     kw_memref,
     kw_mlfunc,
+    kw_return,
     kw_tensor,
     kw_vector,
   };

mlir/test/IR/parser.mlir

@@ -29,3 +29,18 @@ extfunc @memrefs(memref<1x?x4x?x?xint>, memref<i8>)
 
 ; CHECK: extfunc @functions((i1, i1) -> (), () -> ())
 extfunc @functions((memref<1x?x4x?x?xint>, memref<i8>) -> (), ()->())
+
+
+; CHECK-LABEL: cfgfunc @simpleCFG() {
+cfgfunc @simpleCFG() {
+bb42:       ; CHECK: bb0:
+  return    ; CHECK: return
+}           ; CHECK: }
+
+; CHECK-LABEL: cfgfunc @multiblock() -> i32 {
+cfgfunc @multiblock() -> i32 {
+bb0:         ; CHECK: bb0:
+  return     ; CHECK: return
+bb4:         ; CHECK: bb1:
+  return     ; CHECK: return
+}            ; CHECK: }