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
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 @@
     }
   }
 }
-
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())
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)