Add basic parser support for operands:

- This introduces a new FunctionParser base class to handle logic common
   between the kinds of functions we have, e.g. ssa operand/def parsing.
 - This introduces a basic symbol table (without support for forward
   references!) and links defs and uses.
 - CFG functions now parse and build operand lists for operations.  The printer
   isn't set up for them yet tho.

PiperOrigin-RevId: 205246110

mlir/include/mlir/IR/CFGValue.h

@@ -42,15 +42,16 @@ using InstOperand = SSAOperandImpl<CFGValue, Instruction>;
 
 /// CFGValue is the base class for CFG value types.
 class CFGValue : public SSAValueImpl<InstOperand, CFGValueKind> {
-protected:
-  CFGValue(CFGValueKind kind, Type *type) : SSAValueImpl(kind, type) {}
-
+public:
   static bool classof(const SSAValue *value) {
     switch (value->getKind()) {
     case SSAValueKind::InstResult:
       return true;
     }
   }
+
+protected:
+  CFGValue(CFGValueKind kind, Type *type) : SSAValueImpl(kind, type) {}
 };
 
 /// Instruction results are CFG Values.

mlir/lib/IR/AsmPrinter.cpp

@@ -395,6 +395,8 @@ void FunctionState::printOperation(const Operation *op) {
   // TODO: escape name if necessary.
   os << "  \"" << op->getName().str() << "\"()";
 
+  // FIXME: Print operand references.
+
   auto attrs = op->getAttrs();
   if (!attrs.empty()) {
     os << '{';

mlir/lib/Parser/Parser.cpp

@@ -84,8 +84,8 @@ private:
 
 namespace {
 
-typedef std::function<Operation *(Identifier, ArrayRef<Type *>,
-                                  ArrayRef<NamedAttribute>)>
+typedef std::function<Operation *(Identifier, ArrayRef<SSAValue *>,
+                                  ArrayRef<Type *>, ArrayRef<NamedAttribute>)>
     CreateOperationFunction;
 
 /// This class implement support for parsing global entities like types and
@@ -165,15 +165,6 @@ public:
   AffineMap *parseAffineMapInline();
   AffineMap *parseAffineMapReference();
 
-  // SSA
-  ParseResult parseSSAUse();
-  ParseResult parseOptionalSSAUseList(Token::Kind endToken);
-  ParseResult parseSSAUseAndType();
-  ParseResult parseOptionalSSAUseAndTypeList(Token::Kind endToken);
-
-  // Operations
-  ParseResult parseOperation(const CreateOperationFunction &createOpFunc);
-
 private:
   // The Parser is subclassed and reinstantiated.  Do not add additional
   // non-trivial state here, add it to the ParserState class.
@@ -1171,25 +1162,96 @@ AffineMap *Parser::parseAffineMapReference() {
 }
 
 //===----------------------------------------------------------------------===//
-// SSA
+// FunctionParser
 //===----------------------------------------------------------------------===//
 
+namespace {
+/// This class contains parser state that is common across CFG and ML functions,
+/// notably for dealing with operations and SSA values.
+class FunctionParser : public Parser {
+public:
+  FunctionParser(ParserState &state) : Parser(state) {}
+
+  /// This represents a use of an SSA value in the program.  This tracks
+  /// location information in case this ends up being a use of an undefined
+  /// value.
+  typedef std::pair<StringRef, SMLoc> SSAUseInfo;
+
+  /// Given a reference to an SSA value and its type, return a reference.  This
+  /// returns null on failure.
+  SSAValue *resolveSSAUse(SSAUseInfo useInfo, Type *type);
+
+  /// Register a definition of a value with the symbol table.
+  ParseResult addDefinition(SSAUseInfo useInfo, SSAValue *value);
+
+  // SSA parsing productions.
+  ParseResult parseSSAUse(SSAUseInfo &result);
+  ParseResult parseOptionalSSAUseList(Token::Kind endToken,
+                                      SmallVectorImpl<SSAUseInfo> &results);
+  SSAValue *parseSSAUseAndType();
+  ParseResult
+  parseOptionalSSAUseAndTypeList(Token::Kind endToken,
+                                 SmallVectorImpl<SSAValue *> &results);
+
+  // Operations
+  ParseResult parseOperation(const CreateOperationFunction &createOpFunc);
+
+private:
+  /// This keeps track of all of the SSA values we are tracking, indexed by
+  /// their name (either an identifier or a number).
+  llvm::StringMap<std::pair<SSAValue *, SMLoc>> values;
+};
+} // end anonymous namespace
+
+/// Given an unbound reference to an SSA value and its type, return a the value
+/// it specifies.  This returns null on failure.
+SSAValue *FunctionParser::resolveSSAUse(SSAUseInfo useInfo, Type *type) {
+  // If we have already seen a value of this name, return it.
+  auto it = values.find(useInfo.first);
+  if (it != values.end()) {
+    // Check that the type matches the other uses.
+    auto result = it->second.first;
+    if (result->getType() == type)
+      return result;
+
+    emitError(useInfo.second, "use of value '" + useInfo.first.str() +
+                                  "' expects different type than prior uses");
+    emitError(it->second.second, "prior use here");
+    return nullptr;
+  }
+
+  // Otherwise we have a forward reference.
+  // TODO: Handle forward references.
+  emitError(useInfo.second, "undeclared or forward reference");
+  return nullptr;
+}
+
+/// Register a definition of a value with the symbol table.
+ParseResult FunctionParser::addDefinition(SSAUseInfo useInfo, SSAValue *value) {
+
+  // If this is the first definition of this thing, then we are trivially done.
+  auto insertInfo = values.insert({useInfo.first, {value, useInfo.second}});
+  if (insertInfo.second)
+    return ParseSuccess;
+
+  // If we already had a value, replace it with the new one and remove the
+  // placeholder, only if it was a forward ref.
+  // TODO: Handle forward references.
+  emitError(useInfo.second, "redefinition of SSA value " + useInfo.first.str());
+  return ParseFailure;
+}
+
 /// Parse a SSA operand for an instruction or statement.
 ///
 ///   ssa-use ::= ssa-id | ssa-constant
+/// TODO: SSA Constants.
 ///
-ParseResult Parser::parseSSAUse() {
-  if (getToken().is(Token::percent_identifier)) {
-    StringRef name = getTokenSpelling().drop_front();
-    consumeToken(Token::percent_identifier);
-    // TODO: Return this use.
-    (void)name;
-    return ParseSuccess;
-  }
-
-  // TODO: Parse SSA constants.
-
-  return emitError("expected SSA operand");
+ParseResult FunctionParser::parseSSAUse(SSAUseInfo &result) {
+  result.first = getTokenSpelling();
+  result.second = getToken().getLoc();
+  if (!consumeIf(Token::percent_identifier))
+    return emitError("expected SSA operand");
+  return ParseSuccess;
 }
 
 /// Parse a (possibly empty) list of SSA operands.
@@ -1197,42 +1259,51 @@ ParseResult Parser::parseSSAUse() {
 ///   ssa-use-list ::= ssa-use (`,` ssa-use)*
 ///   ssa-use-list-opt ::= ssa-use-list?
 ///
-ParseResult Parser::parseOptionalSSAUseList(Token::Kind endToken) {
-  // TODO: Build and return this.
-  return parseCommaSeparatedList(
-      endToken, [&]() -> ParseResult { return parseSSAUse(); });
+ParseResult
+FunctionParser::parseOptionalSSAUseList(Token::Kind endToken,
+                                        SmallVectorImpl<SSAUseInfo> &results) {
+  return parseCommaSeparatedList(endToken, [&]() -> ParseResult {
+    SSAUseInfo result;
+    if (parseSSAUse(result))
+      return ParseFailure;
+    results.push_back(result);
+    return ParseSuccess;
+  });
 }
 
 /// Parse an SSA use with an associated type.
 ///
 ///   ssa-use-and-type ::= ssa-use `:` type
-ParseResult Parser::parseSSAUseAndType() {
-  if (parseSSAUse())
-    return ParseFailure;
+SSAValue *FunctionParser::parseSSAUseAndType() {
+  SSAUseInfo useInfo;
+  if (parseSSAUse(useInfo))
+    return nullptr;
 
   if (!consumeIf(Token::colon))
-    return emitError("expected ':' and type for SSA operand");
+    return (emitError("expected ':' and type for SSA operand"), nullptr);
 
-  if (!parseType())
-    return ParseFailure;
+  auto *type = parseType();
+  if (!type)
+    return nullptr;
 
-  return ParseSuccess;
+  return resolveSSAUse(useInfo, type);
 }
 
 /// Parse a (possibly empty) list of SSA operands with types.
 ///
 ///   ssa-use-and-type-list ::= ssa-use-and-type (`,` ssa-use-and-type)*
 ///
-ParseResult Parser::parseOptionalSSAUseAndTypeList(Token::Kind endToken) {
-  // TODO: Build and return this.
-  return parseCommaSeparatedList(
-      endToken, [&]() -> ParseResult { return parseSSAUseAndType(); });
+ParseResult FunctionParser::parseOptionalSSAUseAndTypeList(
+    Token::Kind endToken, SmallVectorImpl<SSAValue *> &results) {
+  return parseCommaSeparatedList(endToken, [&]() -> ParseResult {
+    if (auto *value = parseSSAUseAndType()) {
+      results.push_back(value);
+      return ParseSuccess;
+    }
+    return ParseFailure;
+  });
 }
 
-//===----------------------------------------------------------------------===//
-// Operations
-//===----------------------------------------------------------------------===//
-
 /// Parse the CFG or MLFunc operation.
 ///
 /// TODO(clattner): This is a change from the MLIR spec as written, it is an
@@ -1243,12 +1314,12 @@ ParseResult Parser::parseOptionalSSAUseAndTypeList(Token::Kind endToken) {
 ///    `:` function-type
 ///
 ParseResult
-Parser::parseOperation(const CreateOperationFunction &createOpFunc) {
+FunctionParser::parseOperation(const CreateOperationFunction &createOpFunc) {
   auto loc = getToken().getLoc();
 
   StringRef resultID;
   if (getToken().is(Token::percent_identifier)) {
-    resultID = getTokenSpelling().drop_front();
+    resultID = getTokenSpelling();
     consumeToken(Token::percent_identifier);
     if (!consumeIf(Token::equal))
       return emitError("expected '=' after SSA name");
@@ -1267,7 +1338,8 @@ Parser::parseOperation(const CreateOperationFunction &createOpFunc) {
     return emitError("expected '(' to start operand list");
 
   // Parse the operand list.
-  parseOptionalSSAUseList(Token::r_paren);
+  SmallVector<SSAUseInfo, 8> operandInfos;
+  parseOptionalSSAUseList(Token::r_paren, operandInfos);
 
   SmallVector<NamedAttribute, 4> attributes;
   if (getToken().is(Token::l_brace)) {
@@ -1286,22 +1358,50 @@ Parser::parseOperation(const CreateOperationFunction &createOpFunc) {
   if (!fnType)
     return emitError(typeLoc, "expected function type");
 
-  // TODO: Don't drop result name and operand names on the floor.
+  // Check that we have the right number of types for the operands.
+  auto operandTypes = fnType->getInputs();
+  if (operandTypes.size() != operandInfos.size()) {
+    auto plural = "s"[operandInfos.size() == 1];
+    return emitError(typeLoc, "expected " + llvm::utostr(operandInfos.size()) +
+                                  " type" + plural +
+                                  " in operand list but had " +
+                                  llvm::utostr(operandTypes.size()));
+  }
+
+  // Resolve all of the operands.
+  SmallVector<SSAValue *, 8> operands;
+  for (unsigned i = 0, e = operandInfos.size(); i != e; ++i) {
+    operands.push_back(resolveSSAUse(operandInfos[i], operandTypes[i]));
+    if (!operands.back())
+      return ParseFailure;
+  }
+
   auto nameId = builder.getIdentifier(name);
-
-  auto oper = createOpFunc(nameId, fnType->getResults(), attributes);
-
-  if (!oper)
+  auto op = createOpFunc(nameId, operands, fnType->getResults(), attributes);
+  if (!op)
     return ParseFailure;
 
   // We just parsed an operation.  If it is a recognized one, verify that it
   // is structurally as we expect.  If not, produce an error with a reasonable
   // source location.
-  if (auto *opInfo = oper->getAbstractOperation(builder.getContext())) {
-    if (auto error = opInfo->verifyInvariants(oper))
+  if (auto *opInfo = op->getAbstractOperation(builder.getContext())) {
+    if (auto error = opInfo->verifyInvariants(op))
       return emitError(loc, error);
   }
 
+  // If the instruction had a name, register it.
+  if (!resultID.empty()) {
+    // FIXME: Add result infra to handle Stmt results as well to make this
+    // generic.
+    if (auto *inst = dyn_cast<OperationInst>(op)) {
+      if (inst->getResults().empty())
+        return emitError(loc, "cannot name an operation with no results");
+
+      // TODO: This should be getResult(0)
+      addDefinition({resultID, loc}, &inst->getResults()[0]);
+    }
+  }
+
   return ParseSuccess;
 }
 
@@ -1312,10 +1412,10 @@ Parser::parseOperation(const CreateOperationFunction &createOpFunc) {
 namespace {
 /// This is a specialized parser for CFGFunction's, maintaining the state
 /// transient to their bodies.
-class CFGFunctionParser : public Parser {
+class CFGFunctionParser : public FunctionParser {
 public:
   CFGFunctionParser(ParserState &state, CFGFunction *function)
-      : Parser(state), function(function), builder(function) {}
+      : FunctionParser(state), function(function), builder(function) {}
 
   ParseResult parseFunctionBody();
 
@@ -1397,7 +1497,8 @@ ParseResult CFGFunctionParser::parseBasicBlock() {
 
   // If an argument list is present, parse it.
   if (consumeIf(Token::l_paren)) {
-    if (parseOptionalSSAUseAndTypeList(Token::r_paren))
+    SmallVector<SSAValue *, 8> bbArgs;
+    if (parseOptionalSSAUseAndTypeList(Token::r_paren, bbArgs))
       return ParseFailure;
 
     // TODO: attach it.
@@ -1409,9 +1510,14 @@ ParseResult CFGFunctionParser::parseBasicBlock() {
   // Set the insertion point to the block we want to insert new operations into.
   builder.setInsertionPoint(block);
 
-  auto createOpFunc = [this](Identifier name, ArrayRef<Type *> resultTypes,
-                             ArrayRef<NamedAttribute> attrs) -> Operation * {
-    return builder.createOperation(name, {}, resultTypes, attrs);
+  auto createOpFunc = [&](Identifier name, ArrayRef<SSAValue *> operands,
+                          ArrayRef<Type *> resultTypes,
+                          ArrayRef<NamedAttribute> attrs) -> Operation * {
+    SmallVector<CFGValue *, 8> cfgOperands;
+    cfgOperands.reserve(operands.size());
+    for (auto *op : operands)
+      cfgOperands.push_back(cast<CFGValue>(op));
+    return builder.createOperation(name, cfgOperands, resultTypes, attrs);
   };
 
   // Parse the list of operations that make up the body of the block.
@@ -1460,10 +1566,10 @@ TerminatorInst *CFGFunctionParser::parseTerminator() {
 
 namespace {
 /// Refined parser for MLFunction bodies.
-class MLFunctionParser : public Parser {
+class MLFunctionParser : public FunctionParser {
 public:
   MLFunctionParser(ParserState &state, MLFunction *function)
-      : Parser(state), function(function), builder(function) {}
+      : FunctionParser(state), function(function), builder(function) {}
 
   ParseResult parseFunctionBody();
 
@@ -1568,8 +1674,9 @@ ParseResult MLFunctionParser::parseElseClause(IfClause *elseClause) {
 /// Parse a list of statements ending with `return` or `}`
 ///
 ParseResult MLFunctionParser::parseStatements(StmtBlock *block) {
-  auto createOpFunc = [this](Identifier name, ArrayRef<Type *> resultTypes,
-                             ArrayRef<NamedAttribute> attrs) -> Operation * {
+  auto createOpFunc = [&](Identifier name, ArrayRef<SSAValue *> operands,
+                          ArrayRef<Type *> resultTypes,
+                          ArrayRef<NamedAttribute> attrs) -> Operation * {
     return builder.createOperation(name, attrs);
   };
 

mlir/test/IR/parser-errors.mlir

@@ -153,3 +153,19 @@ bb42:
 #map = (d0) -> (%  // expected-error {{invalid SSA name}}
 
 // -----
+
+cfgfunc @test() {
+bb40:
+  %1 = "foo"() : (i32)->i64 // expected-error {{expected 0 types in operand list but had 1}}
+  return
+}
+
+
+// -----
+
+cfgfunc @redef() {
+bb42:
+  %x = "dim"(){index: 0} : ()->i32
+  %x = "dim"(){index: 0} : ()->i32 // expected-error {{redefinition of SSA value %x}}
+  return
+}

mlir/test/IR/parser.mlir

@@ -69,11 +69,11 @@ extfunc @functions((memref<1x?x4x?x?xaffineint, #map0, 0>, memref<i8, #map1, 0>)
 // CHECK-LABEL: cfgfunc @simpleCFG(i32, f32) {
 cfgfunc @simpleCFG(i32, f32) {
 // CHECK: bb0:
-bb42(%0: i32, %f: f32):
+bb42: // (%0: i32, %f: f32):    TODO(clattner): implement bbargs.
   // CHECK: "foo"() : () -> i64
-  %1 = "foo"(%0) : (i32)->i64
-  // CHECK: "bar"() : () -> (i1, i1, i1)
-  "bar"(%1, %f) : (i64, f32) -> (i1,i1,i1)
+  %1 = "foo"() : ()->i64
+  // CHECK: "bar"() : (i64) -> (i1, i1, i1)
+  "bar"(%1) : (i64) -> (i1,i1,i1)
   // CHECK: return
   return
 // CHECK: }
@@ -96,10 +96,12 @@ mlfunc @emptyMLF() {
   return     // CHECK:  return
 }            // CHECK: }
 
-// CHECK-LABEL: mlfunc @mlfunc_with_ops() {
-mlfunc @mlfunc_with_ops() {
+// CHECK-LABEL: cfgfunc @cfgfunc_with_ops() {
+cfgfunc @cfgfunc_with_ops() {
+bb0:
+  %t = "getTensor"() : () -> tensor<4x4x?xf32>
   // CHECK: dim xxx, 2 : sometype
-  %a = "dim"(%42){index: 2}  : () -> affineint
+  %a = "dim"(%t){index: 2} : (tensor<4x4x?xf32>) -> affineint
 
   // CHECK: addf xx, yy : sometype
   "addf"() : () -> ()
@@ -154,10 +156,13 @@ bb42:       // CHECK: bb0:
 // CHECK-LABEL: cfgfunc @standard_instrs() {
 cfgfunc @standard_instrs() {
 bb42:       // CHECK: bb0:
-  // CHECK: dim xxx, 2 : sometype
-  %a = "dim"(%42){index: 2} : () -> affineint
+  %42 = "getTensor"() : () -> tensor<4x4x?xf32>
 
+  // CHECK: dim xxx, 2 : sometype
+  %a = "dim"(%42){index: 2} : (tensor<4x4x?xf32>) -> affineint
+
+  %f = "Const"(){value: 1} : () -> f32
   // CHECK: addf xx, yy : sometype
-  "addf"() : (f32,f32) -> f32
+  "addf"(%f, %f) : (f32,f32) -> f32
   return
 }