Add basic lexing and parsing support for SSA operands and definitions. This

isn't actually constructing IR objects yet, it is eating the tokens and
discarding them.

PiperOrigin-RevId: 203616265

mlir/lib/Parser/Lexer.cpp

@@ -103,7 +103,10 @@ Token Lexer::lexToken() {
 
   case ';': return lexComment();
   case '@': return lexAtIdentifier(tokStart);
-  case '#': return lexAffineMapId(tokStart);
+  case '#':
+    LLVM_FALLTHROUGH;
+  case '%':
+    return lexPrefixedIdentifier(tokStart);
   case '"': return lexString(tokStart);
 
   case '0': case '1': case '2': case '3': case '4':
@@ -182,14 +185,28 @@ Token Lexer::lexAtIdentifier(const char *tokStart) {
   return formToken(Token::at_identifier, tokStart);
 }
 
-/// Lex an '#foo' identifier.
+/// Lex an identifier that starts with a prefix followed by suffix-id.
 ///
 ///   affine-map-id ::= `#` suffix-id
+///   ssa-id        ::= '%' suffix-id
 ///   suffix-id ::= digit+ | (letter|id-punct) (letter|id-punct|digit)*
 ///
-// TODO(andydavis) Consider moving suffix-id parsing to a shared function
-// so it can be re-used to parse %suffix-id.
-Token Lexer::lexAffineMapId(const char *tokStart) {
+Token Lexer::lexPrefixedIdentifier(const char *tokStart) {
+  Token::Kind kind;
+  StringRef errorKind;
+  switch (*tokStart) {
+  case '#':
+    kind = Token::hash_identifier;
+    errorKind = "invalid affine map name";
+    break;
+  case '%':
+    kind = Token::percent_identifier;
+    errorKind = "invalid SSA name";
+    break;
+  default:
+    llvm_unreachable("invalid caller");
+  }
+
   // Parse suffix-id.
   if (isdigit(*curPtr)) {
     // If suffix-id starts with a digit, the rest must be digits.
@@ -201,9 +218,10 @@ Token Lexer::lexAffineMapId(const char *tokStart) {
       ++curPtr;
     } while (isalpha(*curPtr) || isdigit(*curPtr) || isPunct(*curPtr));
   } else {
-    return emitError(curPtr-1, "invalid affine map id");
+    return emitError(curPtr - 1, errorKind);
   }
-  return formToken(Token::affine_map_identifier, tokStart);
+
+  return formToken(kind, tokStart);
 }
 
 /// Lex an integer literal.

mlir/lib/Parser/Lexer.h

@@ -60,7 +60,7 @@ private:
   Token lexComment();
   Token lexBareIdentifierOrKeyword(const char *tokStart);
   Token lexAtIdentifier(const char *tokStart);
-  Token lexAffineMapId(const char *tokStart);
+  Token lexPrefixedIdentifier(const char *tokStart);
   Token lexNumber(const char *tokStart);
   Token lexString(const char *tokStart);
 };

mlir/lib/Parser/Parser.cpp

@@ -188,17 +188,23 @@ private:
                                         const AffineMapParserState &state,
                                         AffineExpr *&result);
 
+  // SSA
+  ParseResult parseSSAUse();
+  ParseResult parseOptionalSSAUseList(Token::Kind endToken);
+  ParseResult parseSSAUseAndType();
+  ParseResult parseOptionalSSAUseAndTypeList(Token::Kind endToken);
+
   // Functions.
   ParseResult parseFunctionSignature(StringRef &name, FunctionType *&type);
   ParseResult parseExtFunc();
   ParseResult parseCFGFunc();
-  ParseResult parseMLFunc();
   ParseResult parseBasicBlock(CFGFunctionParserState &functionState);
   Statement *parseStatement(ParentType parent);
 
   OperationInst *parseCFGOperation(CFGFunctionParserState &functionState);
   TerminatorInst *parseTerminator(CFGFunctionParserState &functionState);
 
+  ParseResult parseMLFunc();
   ForStmt *parseForStmt(ParentType parent);
   IfStmt *parseIfStmt(ParentType parent);
   ParseResult parseNestedStatements(NodeStmt *parent);
@@ -658,7 +664,7 @@ ParseResult Parser::parseAttributeDict(
 ///  affine-map-def ::= affine-map-id `=` affine-map-inline
 ///
 ParseResult Parser::parseAffineMapDef() {
-  assert(curToken.is(Token::affine_map_identifier));
+  assert(curToken.is(Token::hash_identifier));
 
   StringRef affineMapId = curToken.getSpelling().drop_front();
 
@@ -667,7 +673,7 @@ ParseResult Parser::parseAffineMapDef() {
   if (entry)
     return emitError("redefinition of affine map id '" + affineMapId + "'");
 
-  consumeToken(Token::affine_map_identifier);
+  consumeToken(Token::hash_identifier);
 
   // Parse the '='
   if (!consumeIf(Token::equal))
@@ -1092,9 +1098,67 @@ AffineMap *Parser::parseAffineMapInline(StringRef mapId) {
 }
 
 //===----------------------------------------------------------------------===//
-// Functions
+// SSA
 //===----------------------------------------------------------------------===//
 
+/// Parse a SSA operand for an instruction or statement.
+///
+///   ssa-use ::= ssa-id | ssa-constant
+///
+ParseResult Parser::parseSSAUse() {
+  if (curToken.is(Token::percent_identifier)) {
+    StringRef name = curToken.getSpelling().drop_front();
+    consumeToken(Token::percent_identifier);
+    // TODO: Return this use.
+    (void)name;
+    return ParseSuccess;
+  }
+
+  // TODO: Parse SSA constants.
+
+  return emitError("expected SSA operand");
+}
+
+/// Parse a (possibly empty) list of SSA operands.
+///
+///   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(); });
+}
+
+/// Parse an SSA use with an associated type.
+///
+///   ssa-use-and-type ::= ssa-use `:` type
+ParseResult Parser::parseSSAUseAndType() {
+  if (parseSSAUse())
+    return ParseFailure;
+
+  if (!consumeIf(Token::colon))
+    return emitError("expected ':' and type for SSA operand");
+
+  if (!parseType())
+    return ParseFailure;
+
+  return ParseSuccess;
+}
+
+/// 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(); });
+}
+
+//===----------------------------------------------------------------------===//
+// Functions
+//===----------------------------------------------------------------------===//
 
 /// Parse a function signature, starting with a name and including the parameter
 /// list.
@@ -1237,7 +1301,13 @@ ParseResult Parser::parseBasicBlock(CFGFunctionParserState &functionState) {
   // Add the block to the function.
   functionState.function->push_back(block);
 
-  // TODO: parse bb argument list.
+  // If an argument list is present, parse it.
+  if (consumeIf(Token::l_paren)) {
+    if (parseOptionalSSAUseAndTypeList(Token::r_paren))
+      return ParseFailure;
+
+    // TODO: attach it.
+  }
 
   if (!consumeIf(Token::colon))
     return emitError("expected ':' after basic block name");
@@ -1280,7 +1350,13 @@ ParseResult Parser::parseBasicBlock(CFGFunctionParserState &functionState) {
 OperationInst *Parser::
 parseCFGOperation(CFGFunctionParserState &functionState) {
 
-  // TODO: parse ssa-id.
+  StringRef resultID;
+  if (curToken.is(Token::percent_identifier)) {
+    resultID = curToken.getSpelling().drop_front();
+    consumeToken();
+    if (!consumeIf(Token::equal))
+      return (emitError("expected '=' after SSA name"), nullptr);
+  }
 
   if (curToken.isNot(Token::string))
     return (emitError("expected operation name in quotes"), nullptr);
@@ -1294,9 +1370,8 @@ parseCFGOperation(CFGFunctionParserState &functionState) {
   if (!consumeIf(Token::l_paren))
     return (emitError("expected '(' to start operand list"), nullptr);
 
-  // TODO: Parse operands.
-  if (!consumeIf(Token::r_paren))
-    return (emitError("expected ')' in operand list"), nullptr);
+  // Parse the operand list.
+  parseOptionalSSAUseList(Token::r_paren);
 
   SmallVector<NamedAttribute, 4> attributes;
   if (curToken.is(Token::l_brace)) {
@@ -1304,6 +1379,7 @@ parseCFGOperation(CFGFunctionParserState &functionState) {
       return nullptr;
   }
 
+  // TODO: Don't drop result name and operand names on the floor.
   auto nameId = Identifier::get(name, context);
   return new OperationInst(nameId, attributes, context);
 }
@@ -1334,6 +1410,7 @@ TerminatorInst *Parser::parseTerminator(CFGFunctionParserState &functionState) {
       return (emitError("expected basic block name"), nullptr);
     return new BranchInst(destBB);
   }
+    // TODO: cond_br.
   }
 }
 
@@ -1501,7 +1578,7 @@ Module *Parser::parseModule() {
       if (parseCFGFunc()) return nullptr;
       break;
 
-    case Token::affine_map_identifier:
+    case Token::hash_identifier:
       if (parseAffineMapDef()) return nullptr;
       break;
 

mlir/lib/Parser/TokenKinds.def

@@ -50,9 +50,10 @@ TOK_MARKER(eof)
 TOK_MARKER(error)
 
 // Identifiers.
-TOK_IDENTIFIER(bare_identifier)         // foo
-TOK_IDENTIFIER(at_identifier)           // @foo
-TOK_IDENTIFIER(affine_map_identifier)   // #foo
+TOK_IDENTIFIER(bare_identifier)    // foo
+TOK_IDENTIFIER(at_identifier)      // @foo
+TOK_IDENTIFIER(hash_identifier)    // #foo
+TOK_IDENTIFIER(percent_identifier) // %foo
 // TODO: @@foo, etc.
 
 // Literals

mlir/test/IR/parser.mlir

@@ -34,13 +34,13 @@ extfunc @memrefs(memref<1x?x4x?x?xaffineint>, memref<i8>)
 extfunc @functions((memref<1x?x4x?x?xaffineint>, memref<i8>) -> (), ()->())
 
 
-; CHECK-LABEL: cfgfunc @simpleCFG() {
-cfgfunc @simpleCFG() {
-bb42:       ; CHECK: bb0:
-  "foo"()   ; CHECK: "foo"()
-  "bar"()   ; CHECK: "bar"()
-  return    ; CHECK: return
-}           ; CHECK: }
+; CHECK-LABEL: cfgfunc @simpleCFG(i32, f32) {
+cfgfunc @simpleCFG(i32, f32) {
+bb42(%0: i32, %f: f32):     ; CHECK: bb0:
+  %1 = "foo"(%0)            ; CHECK: "foo"()
+  "bar"(%1, %f)             ; CHECK: "bar"()
+  return                    ; CHECK: return
+}                           ; CHECK: }
 
 ; CHECK-LABEL: cfgfunc @multiblock() -> i32 {
 cfgfunc @multiblock() -> i32 {
@@ -99,7 +99,7 @@ bb42:       ; CHECK: bb0:
 cfgfunc @standard_instrs() {
 bb42:       ; CHECK: bb0:
   ; CHECK: dim xxx, 2 : sometype
-  "dim"(){index: 2}
+  %a = "dim"(%42){index: 2}
 
   ; CHECK: addf xx, yy : sometype
   "addf"()