commit c0d69303c1f0d73ec954ff51c9ce1d7529c78a0f
author Jacques Pienaar <jpienaar@google.com> Fri Jul 27 11:07:12 2018 -0700
committer jpienaar <jpienaar@google.com> Fri Mar 29 12:46:30 2019 -0700
tree 996a478c9bdef1c40db7ddbea34a1b9e4055d436
parent 081d9e7472b23fd2944ec6f899f584d61cebd4c3

Add tf_control type and allow $ in bare-id.

* Add tf_control as primitive type;
* Allow $ in bare-id to allow attributes with $ (to make it trivially to mangle a TF attribute);

PiperOrigin-RevId: 206342642

lib/IR/AsmPrinter.cpp

diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index e53962e..26cb1fc 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -287,6 +287,9 @@
   case Type::Kind::F64:
     os << "f64";
     return;
+  case Type::Kind::TFControl:
+    os << "tf_control";
+    return;
 
   case Type::Kind::Integer: {
     auto *integer = cast<IntegerType>(type);

lib/IR/Builders.cpp

diff --git a/lib/IR/Builders.cpp b/lib/IR/Builders.cpp
index d5b85cc..a3ad431 100644
--- a/lib/IR/Builders.cpp
+++ b/lib/IR/Builders.cpp
@@ -47,6 +47,10 @@
 
 PrimitiveType *Builder::getF64Type() { return Type::getF64(context); }
 
+PrimitiveType *Builder::getTFControlType() {
+  return Type::getTFControl(context);
+}
+
 IntegerType *Builder::getIntegerType(unsigned width) {
   return Type::getInteger(width, context);
 }

lib/Parser/Lexer.cpp

diff --git a/lib/Parser/Lexer.cpp b/lib/Parser/Lexer.cpp
index 011dfcb..ce99b19 100644
--- a/lib/Parser/Lexer.cpp
+++ b/lib/Parser/Lexer.cpp
@@ -151,12 +151,13 @@
 
 /// Lex a bare identifier or keyword that starts with a letter.
 ///
-///   bare-id ::= letter (letter|digit|[_])*
+///   bare-id ::= letter (letter|digit|[_$])*
 ///   integer-type ::= `i[1-9][0-9]*`
 ///
 Token Lexer::lexBareIdentifierOrKeyword(const char *tokStart) {
-  // Match the rest of the identifier regex: [0-9a-zA-Z_]*
-  while (isalpha(*curPtr) || isdigit(*curPtr) || *curPtr == '_')
+  // Match the rest of the identifier regex: [0-9a-zA-Z_$]*
+  while (isalpha(*curPtr) || isdigit(*curPtr) || *curPtr == '_' ||
+         *curPtr == '$')
     ++curPtr;
 
   // Check to see if this identifier is a keyword.
@@ -288,4 +289,3 @@
     }
   }
 }
-

lib/Parser/Parser.cpp

diff --git a/lib/Parser/Parser.cpp b/lib/Parser/Parser.cpp
index 11850ea..71737ce 100644
--- a/lib/Parser/Parser.cpp
+++ b/lib/Parser/Parser.cpp
@@ -281,6 +281,9 @@
   case Token::kw_affineint:
     consumeToken(Token::kw_affineint);
     return builder.getAffineIntType();
+  case Token::kw_tf_control:
+    consumeToken(Token::kw_tf_control);
+    return builder.getTFControlType();
   case Token::inttype: {
     auto width = getToken().getIntTypeBitwidth();
     if (!width.hasValue())

lib/Parser/TokenKinds.def

diff --git a/lib/Parser/TokenKinds.def b/lib/Parser/TokenKinds.def
index 44a40ee..7817d7b 100644
--- a/lib/Parser/TokenKinds.def
+++ b/lib/Parser/TokenKinds.def
@@ -107,6 +107,7 @@
 TOK_KEYWORD(size)
 TOK_KEYWORD(step)
 TOK_KEYWORD(tensor)
+TOK_KEYWORD(tf_control)
 TOK_KEYWORD(to)
 TOK_KEYWORD(true)
 TOK_KEYWORD(vector)