[mlir][IR] Add builtin TokenType (#195640)

Introduces a new parameterless, opaque, builtin SSA value type,
`!token`. A token cannot appear in a value-forwarding position such as
`cf.br`, `arith.select`, `scf.for` iter-args or function call/return.
Walking back from any token use reaches the producing operation without
crossing such a boundary. Tokens carry no runtime data and cannot
constant-fold.

This contract is enforced by changing the default ODS `AnyType`
predicate to exclude tokens. In addition, token-producing and
token-consuming ops must opt in by implementing the `ProducesToken` and
`ConsumesToken` traits. The IR verifier checks for the presence of these
traits.

Note: CSE, DCE, hoisting, and `OperationEquivalence` remain unchanged. A
stronger contract (e.g., uniqueness, arity, paired lifetime), if
desirable based on the semantics of token-producing and token-consuming
ops, can be expressed at the op level via existing mechanisms (side
effects, block arguments, traits, attributes), keeping the type
orthogonal to operation-level transformations. This mirrors LLVM's
`token` type.

Other changes:
* Builtin bytecode: new entry for `TokenType`.
* AsmParser / AsmPrinter: `token` keyword for the textual format.
* Remove `LLVMTokenType` from the LLVM dialect and use the builtin
`TokenType` instead.
* Async dialect: references to `async::TokenType` are qualified to
disambiguate from the new builtin `TokenType`.
* Documentation: new `mlir/docs/Tokens.md` covering the structural
contract, design rationale, and ODS integration.

No changes to `Operation`, the generic op syntax, the bytecode op
encoding, or core C++ APIs around `Operation`.

This commit is in preparation of adding support for breaking exit from
regions. (E.g., early exit from loops.)

RFC:
https://discourse.llvm.org/t/rfc-add-a-builtin-token-type-to-mlir/90706

Assisted-by: claude-opus-4.7-thinking-high

---------

Co-authored-by: Mehdi Amini <joker.eph@gmail.com>

mlir/docs/Dialects/LLVM.md

@@ -224,6 +224,7 @@ dialect-compatible types_. The following types are compatible:
     (`FloatType`).
 -   1D vectors of signless integers or floating point types - `vector<NxT>`
     (`VectorType`).
+-   Tokens (`TokenType`) - non-inspectable values associated with an operation.
 
 Note that only a subset of types that can be represented by a given class is
 compatible. For example, signed and unsigned integers are not compatible. LLVM
@@ -242,8 +243,6 @@ LLVM dialect:
 
 -   `!llvm.ppc_fp128` (`LLVMPPCFP128Type`) - 128-bit floating-point value (two
     64 bits).
--   `!llvm.token` (`LLVMTokenType`) - a non-inspectable value associated with an
-    operation.
 -   `!llvm.metadata` (`LLVMMetadataType`) - LLVM IR metadata, to be used only if
     the metadata cannot be represented as structured MLIR attributes.
 -   `!llvm.void` (`LLVMVoidType`) - does not represent any value; can only

mlir/docs/LangRef.md

@@ -740,6 +740,34 @@ The [builtin dialect](Dialects/Builtin.md) defines a set of types that are
 directly usable by any other dialect in MLIR. These types cover a range from
 primitive integer and floating-point types, function types, and more.
 
+### Token Type
+
+A *token* is an SSA value of the builtin parameterless, opaque `token` type.
+It carries no runtime data. Given a use of a token SSA value, its definition
+is guaranteed to be the semantic producer of the token. Generic transformations
+must preserve this invariant: they may not introduce a forwarding step between
+a use and its producer, nor retarget a use to a producer with different
+semantics. New uses of a token can be introduced safely. As a consequence:
+
+1. A token must not appear as a forwarded value. E.g., it cannot be used as a
+   successor operand of a `BranchOpInterface` op.
+2. A token cannot constant-fold. No constant of token type exists.
+3. The presence of tokens has no effect on standard transformations such as
+   CSE, DCE or hoisting.
+4. Use of a token is side-effect free: a token user follows the usual
+   `isTriviallyDead()` rules.
+
+These properties mirror what LLVM IR already documents for its own
+[`token` type](https://llvm.org/docs/LangRef.html#token-type).
+
+Operations must opt in to producing or consuming tokens with
+`TokenProducerTrait` and `TokenConsumerTrait`.
+
+Note: Because tokens are SSA values, they cannot cross `IsolatedFromAbove`
+region boundaries.
+
+See [Tokens](Tokens.md) for details on ODS integration and examples.
+
 ## Properties
 
 Properties are extra data members stored directly on an Operation class. They

mlir/docs/Tokens.md

@@ -0,0 +1,107 @@
+# Tokens
+
+[TOC]
+
+## Overview
+
+Intuitively, a *token* value is a pointer to an operation (via an OpResult)
+or a pointer to a region (via an entry block argument). A token cannot be
+forwarded: a token def-use chain cannot be obscured by ops with forwarding
+semantics such as `arith.select` or `cf.br`. This allows you to always walk
+back from a use and say "this token came from *that* specific op". The exact
+structural contract is specified in the
+[LangRef section on tokens](LangRef.md#token-type).
+
+A token is an SSA value that has the builtin token type. The token type is 
+parameterless, opaque and carries no runtime data. Apart from the structural
+contract specified in the LangRef, tokens are like any other SSA values.
+
+## Design Rationale
+
+The token type allows operations to refer to another operation without a new
+parallel def-use system for operations. It reuses the existing def-use
+machinery for SSA. It introduces no changes to the generic op syntax, the
+bytecode infrastructure or core C++ APIs around `Operation`.
+
+As with regular def-use chains, a token def-use chain is unidirectional. A
+token use points to the token's definition and not the other way around.
+Transformations can remove the use of a token without having to touch or
+inspect the definition of the token.
+
+Because tokens are SSA values, they cannot cross `IsolatedFromAbove` region
+boundaries. This is intentional: it allows passes to process isolated regions
+concurrently without racing on def-use chains. When a token-like dependency
+must cross such a boundary, another mechanism must be used (e.g. a symbolic
+reference using an attribute).
+
+## ODS Integration
+
+Tokens are excluded from the default `AnyType` predicate, so an op that has
+not opted in cannot accept a token as an arbitrary operand or result. This
+restriction prevents tokens from being accidentally passed as operands with
+forwarding semantics.
+
+Two predicates are provided in `CommonTypeConstraints.td`:
+
+| Predicate          | Accepts                              | Use when …                                                            |
+| ------------------ | ------------------------------------ | ----------------------------------------------------------------------|
+| `AnyType`          | any non-token type                   | the default; matches the historical meaning of "any type" pre-tokens. |
+| `Token`            | only the builtin `TokenType`         | the op specifically takes a token operand/result.                     |
+
+Example:
+
+```tablegen
+def MyProduceOp : MyDialect_Op<"produce"> {
+  let results = (outs Token:$token);
+}
+
+def MyConsumeOp : MyDialect_Op<"consume"> {
+  let arguments = (ins Token:$scope, AnyType:$value);
+}
+```
+
+ODS automatically adds `TokenProducerTrait` when an op declares a `Token`-typed
+result, and `TokenConsumerTrait` when it declares a `Token`-typed operand. The
+traits must be listed manually when tokens appear as entry block arguments and
+for ops that are not defined in ODS.
+
+Region entry block arguments of `token` type are also token producers and
+require the parent operation to define `TokenProducerTrait`. Token block
+arguments in non-entry blocks are rejected.
+
+## Examples
+
+### Non-forwarding Semantics
+
+The [LangRef](LangRef.md#token-type) requires that a token never appears as a
+forwarded value. For example, you cannot use a token like this:
+
+* a forwarded result or operand of a `CallOpInterface` op;
+* an argument or result type of a `FunctionOpInterface` op;
+* a successor operand of a `BranchOpInterface` op;
+* a block argument of a non-entry block;
+* a forwarded operand to or from any region of a `RegionBranchOpInterface`
+  op (iter-args, region results, or yielded values); or
+* the result of any op that selects or merges values it does not understand
+  (e.g. `arith.select`).
+
+### ODS-based Verification: Tokens Rejected in `AnyType` Positions
+
+`scf.yield` operands have forwarding semantics. A token cannot be yielded from
+a branch or a loop.
+
+```mlir
+// error: 'scf.if' op result #0 must be variadic of any non-token type,
+//        but got 'token'
+%t = scf.if %cond -> token {
+  %a = my.token.produce : token
+  scf.yield %a : token
+} else {
+  %b = my.token.produce : token
+  scf.yield %b : token
+}
+```
+
+`scf.if`'s results are declared with `Variadic<AnyType>` and `scf.yield`'s
+operands likewise use `AnyType`. Because `AnyType` excludes tokens, both
+`scf.if` and `scf.yield` fail verification.

mlir/docs/Traits/_index.md

@@ -349,3 +349,17 @@ This trait removes the requirement on regions held by an operation to have
 [terminator operations](../LangRef.md/#control-flow-and-ssacfg-regions) at the end of a block.
 This requires that these regions have a single block. An example of operation
 using this trait is the top-level `ModuleOp`.
+
+### TokenProducerTrait
+
+*   `OpTrait::TokenProducerTrait` -- `TokenProducerTrait`
+
+This trait marks operations that are allowed to produce builtin `token` values
+as operation results or as region entry block arguments.
+
+### TokenConsumerTrait
+
+*   `OpTrait::TokenConsumerTrait` -- `TokenConsumerTrait`
+
+This trait marks operations that are allowed to consume builtin `token` values
+as operands.

mlir/include/mlir/Dialect/Async/IR/Async.h

@@ -50,7 +50,7 @@ namespace async {
 
 /// Returns true if the type is reference counted at runtime.
 inline bool isRefCounted(Type type) {
-  return isa<TokenType, ValueType, GroupType>(type);
+  return isa<async::TokenType, ValueType, GroupType>(type);
 }
 
 } // namespace async

mlir/include/mlir/Dialect/Async/IR/AsyncOps.td

@@ -174,7 +174,9 @@ def Async_FuncOp : Async_Op<"func",
     unsigned getNumResults() {return getResultTypes().size();}
 
     /// Is the async func stateful
-    bool isStateful() { return isa<TokenType>(getFunctionType().getResult(0));}
+    bool isStateful() {
+      return isa<async::TokenType>(getFunctionType().getResult(0));
+    }
 
     //===------------------------------------------------------------------===//
     // OpAsmOpInterface Methods

mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td

@@ -698,7 +698,7 @@ def LLVM_ThreadlocalAddressOp : LLVM_OneResultIntrOp<"threadlocal.address", [],
 // Coroutine intrinsics.
 //
 
-def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
+def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [TokenProducerTrait], 1> {
   let arguments = (ins I32:$align,
                        LLVM_AnyPointer:$promise,
                        LLVM_AnyPointer:$coroaddr,
@@ -707,8 +707,9 @@ def LLVM_CoroIdOp : LLVM_IntrOp<"coro.id", [], [], [], 1> {
     " attr-dict `:` functional-type(operands, results)";
 }
 
-def LLVM_CoroBeginOp : LLVM_IntrOp<"coro.begin", [], [], [], 1> {
-  let arguments = (ins LLVM_TokenType:$token,
+def LLVM_CoroBeginOp
+    : LLVM_IntrOp<"coro.begin", [], [], [], 1> {
+  let arguments = (ins Token:$token,
                        LLVM_AnyPointer:$mem);
   let assemblyFormat = "$token `,` $mem attr-dict `:` functional-type(operands, results)";
 }
@@ -721,26 +722,30 @@ def LLVM_CoroAlignOp : LLVM_IntrOp<"coro.align", [0], [], [], 1> {
   let assemblyFormat = "attr-dict `:` type($res)";
 }
 
-def LLVM_CoroSaveOp : LLVM_IntrOp<"coro.save", [], [], [], 1> {
+def LLVM_CoroSaveOp
+    : LLVM_IntrOp<"coro.save", [], [], [TokenProducerTrait], 1> {
   let arguments = (ins LLVM_AnyPointer:$handle);
   let assemblyFormat = "$handle attr-dict `:` functional-type(operands, results)";
 }
 
-def LLVM_CoroSuspendOp : LLVM_IntrOp<"coro.suspend", [], [], [], 1> {
-  let arguments = (ins LLVM_TokenType:$save,
+def LLVM_CoroSuspendOp
+    : LLVM_IntrOp<"coro.suspend", [], [], [], 1> {
+  let arguments = (ins Token:$save,
                        I1:$final);
   let assemblyFormat = "$save `,` $final attr-dict `:` type($res)";
 }
 
-def LLVM_CoroEndOp : LLVM_IntrOp<"coro.end", [], [], [], 1> {
+def LLVM_CoroEndOp
+    : LLVM_IntrOp<"coro.end", [], [], [], 1> {
   let arguments = (ins LLVM_AnyPointer:$handle,
                        I1:$unwind,
-                       LLVM_TokenType:$retvals);
+                       Token:$retvals);
   let assemblyFormat = "$handle `,` $unwind `,` $retvals attr-dict `:` functional-type(operands, results)";
 }
 
-def LLVM_CoroFreeOp : LLVM_IntrOp<"coro.free", [], [], [], 1> {
-  let arguments = (ins LLVM_TokenType:$id,
+def LLVM_CoroFreeOp
+    : LLVM_IntrOp<"coro.free", [], [], [], 1> {
+  let arguments = (ins Token:$id,
                        LLVM_AnyPointer:$handle);
   let assemblyFormat = "$id `,` $handle attr-dict `:` functional-type(operands, results)";
 }

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

@@ -29,12 +29,6 @@ def LLVM_Type : DialectType<LLVM_Dialect,
                             CPred<"::mlir::LLVM::isCompatibleOuterType($_self)">,
                             "LLVM dialect-compatible type">;
 
-// Type constraint accepting LLVM token type.
-def LLVM_TokenType : Type<
-  CPred<"::llvm::isa<::mlir::LLVM::LLVMTokenType>($_self)">,
-  "LLVM token type">,
-  BuildableType<"::mlir::LLVM::LLVMTokenType::get($_builder.getContext())">;
-
 // Type constraint accepting LLVM primitive types, i.e. all types except void
 // and function.
 def LLVM_PrimitiveType : Type<

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

@@ -2138,15 +2138,14 @@ def LLVM_NoneTokenOp
     : LLVM_Op<"mlir.none", [Pure]> {
   let summary = "Defines a value containing an empty token to LLVM type.";
   let description = [{
-    Unlike LLVM IR, MLIR does not have first-class token values. They must be
-    explicitly created as SSA values using `llvm.mlir.none`. This operation has
-    no operands or attributes, and returns a none token value of a wrapped LLVM IR
-    pointer type.
+    MLIR does not have a way to spell the LLVM IR `none` token literal. This
+    operation produces a builtin `!token` SSA value that lowers to
+    `llvm::ConstantTokenNone` in LLVM IR.
 
     Examples:
 
     ```mlir
-    %0 = llvm.mlir.none : !llvm.token
+    %0 = llvm.mlir.none : !token
     ```
   }];
 
@@ -2154,7 +2153,7 @@ def LLVM_NoneTokenOp
     $res = llvm::ConstantTokenNone::get(builder.getContext());
   }];
 
-  let results = (outs LLVM_TokenType:$res);
+  let results = (outs Token:$res);
   let builders = [LLVM_OneResultOpBuilder];
   let assemblyFormat = "attr-dict `:` type($res)";
 }

mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h

@@ -67,7 +67,6 @@ namespace LLVM {
   }
 
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMVoidType, "llvm.void");
-DEFINE_TRIVIAL_LLVM_TYPE(LLVMTokenType, "llvm.token");
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMLabelType, "llvm.label");
 DEFINE_TRIVIAL_LLVM_TYPE(LLVMMetadataType, "llvm.metadata");
 

mlir/include/mlir/IR/BuiltinDialectBytecode.td

@@ -324,6 +324,8 @@ def UnrankedTensorType : DialectType<(type
   Type:$elementType
 )>;
 
+def TokenType : DialectType<(type)>;
+
 let cType = "VectorType" in {
 def VectorType : DialectType<(type
   Array<SignedVarIntList>:$shape,
@@ -413,7 +415,8 @@ def BuiltinDialectTypes : DialectTypes<"Builtin"> {
     Float4E2M1FNType,
     Float6E2M3FNType,
     Float6E3M2FNType,
-    Float8E8M0FNUType
+    Float8E8M0FNUType,
+    TokenType
   ];
 }
 

mlir/include/mlir/IR/BuiltinTypes.td

@@ -1237,6 +1237,28 @@ def Builtin_RankedTensor : Builtin_Type<"RankedTensor", "tensor", [
   let genVerifyDecl = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// TokenType
+//===----------------------------------------------------------------------===//
+
+def Builtin_Token : Builtin_Type<"Token", "token"> {
+  let summary = "Token type";
+  let description = [{
+    Syntax:
+
+    ```
+    token-type ::= `token`
+    ```
+
+    A use of a token SSA value is a pointer to an operation (in case of an
+    OpResult) or a pointer to a region (in case of an entry block argument).
+    A token carries no runtime data and cannot be forwarded. Tokens are
+    excluded from the `AnyType` type constraint. Operations must define
+    `TokenProducerTrait` to produce token results or token region entry block
+    arguments, and must define `TokenConsumerTrait` to consume token operands.
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // TupleType
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/CommonTypeConstraints.td

@@ -165,8 +165,17 @@ class SameBuildabilityAs<Type type, code builder> {
   code builderCall = !if(!empty(type.builderCall), "", builder);
 }
 
-// Any type at all.
-def AnyType : Type<CPred<"true">, "any type">;
+// Whether a type is the builtin `TokenType`.
+def IsTokenTypePred : CPred<"::llvm::isa<::mlir::TokenType>($_self)">;
+
+// Any non-token type. Tokens are excluded by default to prevent ops that
+// accept arbitrary types from accidentally accepting tokens as operands /
+// results, since a token must not be value-forwarded.
+def AnyType : Type<Neg<IsTokenTypePred>, "any non-token type">;
+
+// The builtin token type.
+def Token : Type<IsTokenTypePred, "token", "::mlir::TokenType">,
+            BuildableType<"$_builder.getType<::mlir::TokenType>()">;
 
 // None type
 def NoneType : Type<CPred<"::llvm::isa<::mlir::NoneType>($_self)">, "none type",

mlir/include/mlir/IR/OpBase.td

@@ -98,6 +98,10 @@ def SameOperandsAndResultElementType :
   NativeOpTrait<"SameOperandsAndResultElementType">;
 // Op is a terminator.
 def Terminator : NativeOpTrait<"IsTerminator">;
+// Op produces builtin token values.
+def TokenProducerTrait : NativeOpTrait<"TokenProducerTrait">;
+// Op consumes builtin token values.
+def TokenConsumerTrait : NativeOpTrait<"TokenConsumerTrait">;
 // Op can be safely normalized in the presence of MemRefs with
 // non-identity maps.
 def MemRefsNormalizable : NativeOpTrait<"MemRefsNormalizable">;

mlir/include/mlir/IR/OpDefinition.h

@@ -778,6 +778,18 @@ public:
   }
 };
 
+/// This trait marks operations that are allowed to produce builtin token
+/// values.
+template <typename ConcreteType>
+class TokenProducerTrait : public TraitBase<ConcreteType, TokenProducerTrait> {
+};
+
+/// This trait marks operations that are allowed to consume builtin token
+/// values.
+template <typename ConcreteType>
+class TokenConsumerTrait : public TraitBase<ConcreteType, TokenConsumerTrait> {
+};
+
 /// This class provides verification for ops that are known to have zero
 /// successors.
 template <typename ConcreteType>

mlir/lib/AsmParser/TokenKinds.def

@@ -127,6 +127,7 @@ TOK_KEYWORD(symbol)
 TOK_KEYWORD(tensor)
 TOK_KEYWORD(tf32)
 TOK_KEYWORD(to)
+TOK_KEYWORD(token)
 TOK_KEYWORD(true)
 TOK_KEYWORD(tuple)
 TOK_KEYWORD(type)

mlir/lib/AsmParser/TypeParser.cpp

@@ -58,6 +58,7 @@ OptionalParseResult Parser::parseOptionalType(Type &type) {
   case Token::kw_f128:
   case Token::kw_index:
   case Token::kw_none:
+  case Token::kw_token:
   case Token::exclamation_identifier:
     return failure(!(type = parseType()));
 
@@ -371,6 +372,11 @@ Type Parser::parseNonFunctionType() {
     consumeToken(Token::kw_none);
     return builder.getNoneType();
 
+  // token-type
+  case Token::kw_token:
+    consumeToken(Token::kw_token);
+    return builder.getType<TokenType>();
+
   // extended type
   case Token::exclamation_identifier:
     return parseExtendedType();

mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp

@@ -78,8 +78,8 @@ struct AsyncAPI {
     return LLVM::LLVMPointerType::get(ctx);
   }
 
-  static LLVM::LLVMTokenType tokenType(MLIRContext *ctx) {
-    return LLVM::LLVMTokenType::get(ctx);
+  static mlir::TokenType tokenType(MLIRContext *ctx) {
+    return mlir::TokenType::get(ctx);
   }
 
   static FunctionType addOrDropRefFunctionType(MLIRContext *ctx) {
@@ -89,7 +89,7 @@ struct AsyncAPI {
   }
 
   static FunctionType createTokenFunctionType(MLIRContext *ctx) {
-    return FunctionType::get(ctx, {}, {TokenType::get(ctx)});
+    return FunctionType::get(ctx, {}, {async::TokenType::get(ctx)});
   }
 
   static FunctionType createValueFunctionType(MLIRContext *ctx) {
@@ -109,7 +109,7 @@ struct AsyncAPI {
   }
 
   static FunctionType emplaceTokenFunctionType(MLIRContext *ctx) {
-    return FunctionType::get(ctx, {TokenType::get(ctx)}, {});
+    return FunctionType::get(ctx, {async::TokenType::get(ctx)}, {});
   }
 
   static FunctionType emplaceValueFunctionType(MLIRContext *ctx) {
@@ -118,7 +118,7 @@ struct AsyncAPI {
   }
 
   static FunctionType setTokenErrorFunctionType(MLIRContext *ctx) {
-    return FunctionType::get(ctx, {TokenType::get(ctx)}, {});
+    return FunctionType::get(ctx, {async::TokenType::get(ctx)}, {});
   }
 
   static FunctionType setValueErrorFunctionType(MLIRContext *ctx) {
@@ -128,7 +128,7 @@ struct AsyncAPI {
 
   static FunctionType isTokenErrorFunctionType(MLIRContext *ctx) {
     auto i1 = IntegerType::get(ctx, 1);
-    return FunctionType::get(ctx, {TokenType::get(ctx)}, {i1});
+    return FunctionType::get(ctx, {async::TokenType::get(ctx)}, {i1});
   }
 
   static FunctionType isValueErrorFunctionType(MLIRContext *ctx) {
@@ -143,7 +143,7 @@ struct AsyncAPI {
   }
 
   static FunctionType awaitTokenFunctionType(MLIRContext *ctx) {
-    return FunctionType::get(ctx, {TokenType::get(ctx)}, {});
+    return FunctionType::get(ctx, {async::TokenType::get(ctx)}, {});
   }
 
   static FunctionType awaitValueFunctionType(MLIRContext *ctx) {
@@ -162,13 +162,14 @@ struct AsyncAPI {
 
   static FunctionType addTokenToGroupFunctionType(MLIRContext *ctx) {
     auto i64 = IntegerType::get(ctx, 64);
-    return FunctionType::get(ctx, {TokenType::get(ctx), GroupType::get(ctx)},
-                             {i64});
+    return FunctionType::get(
+        ctx, {async::TokenType::get(ctx), GroupType::get(ctx)}, {i64});
   }
 
   static FunctionType awaitTokenAndExecuteFunctionType(MLIRContext *ctx) {
     auto ptrType = opaquePointerType(ctx);
-    return FunctionType::get(ctx, {TokenType::get(ctx), ptrType, ptrType}, {});
+    return FunctionType::get(
+        ctx, {async::TokenType::get(ctx), ptrType, ptrType}, {});
   }
 
   static FunctionType awaitValueAndExecuteFunctionType(MLIRContext *ctx) {
@@ -291,7 +292,7 @@ public:
   }
 
   static std::optional<Type> convertAsyncTypes(Type type) {
-    if (isa<TokenType, GroupType, ValueType>(type))
+    if (isa<async::TokenType, GroupType, ValueType>(type))
       return AsyncAPI::opaquePointerType(type.getContext());
 
     if (isa<CoroIdType, CoroStateType>(type))
@@ -583,7 +584,7 @@ public:
     Type resultType = op->getResultTypes()[0];
 
     // Tokens creation maps to a simple function call.
-    if (isa<TokenType>(resultType)) {
+    if (isa<async::TokenType>(resultType)) {
       rewriter.replaceOpWithNewOp<func::CallOp>(
           op, kCreateToken, converter->convertType(resultType));
       return success();
@@ -659,7 +660,7 @@ public:
                   ConversionPatternRewriter &rewriter) const override {
     StringRef apiFuncName =
         TypeSwitch<Type, StringRef>(op.getOperand().getType())
-            .Case<TokenType>([](Type) { return kEmplaceToken; })
+            .Case<async::TokenType>([](Type) { return kEmplaceToken; })
             .Case<ValueType>([](Type) { return kEmplaceValue; });
 
     rewriter.replaceOpWithNewOp<func::CallOp>(op, apiFuncName, TypeRange(),
@@ -685,7 +686,7 @@ public:
                   ConversionPatternRewriter &rewriter) const override {
     StringRef apiFuncName =
         TypeSwitch<Type, StringRef>(op.getOperand().getType())
-            .Case<TokenType>([](Type) { return kSetTokenError; })
+            .Case<async::TokenType>([](Type) { return kSetTokenError; })
             .Case<ValueType>([](Type) { return kSetValueError; });
 
     rewriter.replaceOpWithNewOp<func::CallOp>(op, apiFuncName, TypeRange(),
@@ -710,7 +711,7 @@ public:
                   ConversionPatternRewriter &rewriter) const override {
     StringRef apiFuncName =
         TypeSwitch<Type, StringRef>(op.getOperand().getType())
-            .Case<TokenType>([](Type) { return kIsTokenError; })
+            .Case<async::TokenType>([](Type) { return kIsTokenError; })
             .Case<GroupType>([](Type) { return kIsGroupError; })
             .Case<ValueType>([](Type) { return kIsValueError; });
 
@@ -735,7 +736,7 @@ public:
                   ConversionPatternRewriter &rewriter) const override {
     StringRef apiFuncName =
         TypeSwitch<Type, StringRef>(op.getOperand().getType())
-            .Case<TokenType>([](Type) { return kAwaitToken; })
+            .Case<async::TokenType>([](Type) { return kAwaitToken; })
             .Case<ValueType>([](Type) { return kAwaitValue; })
             .Case<GroupType>([](Type) { return kAwaitGroup; });
 
@@ -763,7 +764,7 @@ public:
                   ConversionPatternRewriter &rewriter) const override {
     StringRef apiFuncName =
         TypeSwitch<Type, StringRef>(op.getOperand().getType())
-            .Case<TokenType>([](Type) { return kAwaitTokenAndExecute; })
+            .Case<async::TokenType>([](Type) { return kAwaitTokenAndExecute; })
             .Case<ValueType>([](Type) { return kAwaitValueAndExecute; })
             .Case<GroupType>([](Type) { return kAwaitAllAndExecute; });
 
@@ -906,7 +907,7 @@ public:
   matchAndRewrite(RuntimeAddToGroupOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Currently we can only add tokens to the group.
-    if (!isa<TokenType>(op.getOperand().getType()))
+    if (!isa<async::TokenType>(op.getOperand().getType()))
       return rewriter.notifyMatchFailure(op, "only token type is supported");
 
     // Replace with a runtime API function call.
@@ -1151,7 +1152,7 @@ public:
 void mlir::populateAsyncStructuralTypeConversionsAndLegality(
     TypeConverter &typeConverter, RewritePatternSet &patterns,
     ConversionTarget &target) {
-  typeConverter.addConversion([&](TokenType type) { return type; });
+  typeConverter.addConversion([&](async::TokenType type) { return type; });
   typeConverter.addConversion([&](ValueType type) {
     Type converted = typeConverter.convertType(type.getValueType());
     return converted ? ValueType::get(converted) : converted;

mlir/lib/Dialect/Async/IR/Async.cpp

@@ -84,7 +84,7 @@ void ExecuteOp::build(OpBuilder &builder, OperationState &result,
 
   // First result is always a token, and then `resultTypes` wrapped into
   // `async.value`.
-  result.addTypes({TokenType::get(result.getContext())});
+  result.addTypes({async::TokenType::get(result.getContext())});
   for (Type type : resultTypes)
     result.addTypes(ValueType::get(type));
 
@@ -139,7 +139,7 @@ ParseResult ExecuteOp::parse(OpAsmParser &parser, OperationState &result) {
   // Sizes of parsed variadic operands, will be updated below after parsing.
   int32_t numDependencies = 0;
 
-  auto tokenTy = TokenType::get(ctx);
+  auto tokenTy = async::TokenType::get(ctx);
 
   // Parse dependency tokens.
   if (succeeded(parser.parseOptionalLSquare())) {
@@ -280,7 +280,7 @@ LogicalResult AwaitOp::verify() {
   Type argType = getOperand().getType();
 
   // Awaiting on a token does not have any results.
-  if (llvm::isa<TokenType>(argType) && !getResultTypes().empty())
+  if (llvm::isa<async::TokenType>(argType) && !getResultTypes().empty())
     return emitOpError("awaiting on a token must have empty result");
 
   // Awaiting on a value unwraps the async value type.
@@ -345,12 +345,12 @@ LogicalResult FuncOp::verify() {
 
   for (unsigned i = 0, e = resultTypes.size(); i != e; ++i) {
     auto type = resultTypes[i];
-    if (!llvm::isa<TokenType>(type) && !llvm::isa<ValueType>(type))
+    if (!llvm::isa<async::TokenType>(type) && !llvm::isa<ValueType>(type))
       return emitOpError() << "result type must be async value type or async "
                               "token type, but got "
                            << type;
     // We only allow AsyncToken appear as the first return value
-    if (llvm::isa<TokenType>(type) && i != 0) {
+    if (llvm::isa<async::TokenType>(type) && i != 0) {
       return emitOpError()
              << " results' (optional) async token type is expected "
                 "to appear as the 1st return value, but got "

mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp

@@ -526,7 +526,7 @@ void AsyncRuntimePolicyBasedRefCountingPass::initializeDefaultPolicy() {
     Operation *op = operand.getOwner();
     Type type = operand.get().getType();
 
-    bool isToken = isa<TokenType>(type);
+    bool isToken = isa<async::TokenType>(type);
     bool isGroup = isa<GroupType>(type);
     bool isValue = isa<ValueType>(type);
 

mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp

@@ -188,11 +188,12 @@ static CoroMachinery setupCoroMachinery(func::FuncOp func) {
 
   // We treat TokenType as state update marker to represent side-effects of
   // async computations
-  bool isStateful = isa<TokenType>(func.getResultTypes().front());
+  bool isStateful = isa<async::TokenType>(func.getResultTypes().front());
 
   std::optional<Value> retToken;
   if (isStateful)
-    retToken.emplace(RuntimeCreateOp::create(builder, TokenType::get(ctx)));
+    retToken.emplace(
+        RuntimeCreateOp::create(builder, async::TokenType::get(ctx)));
 
   llvm::SmallVector<Value, 4> retValues;
   ArrayRef<Type> resValueTypes =
@@ -673,8 +674,9 @@ private:
 };
 
 /// Lowering for `async.await` with a token operand.
-class AwaitTokenOpLowering : public AwaitOpLoweringBase<AwaitOp, TokenType> {
-  using Base = AwaitOpLoweringBase<AwaitOp, TokenType>;
+class AwaitTokenOpLowering
+    : public AwaitOpLoweringBase<AwaitOp, async::TokenType> {
+  using Base = AwaitOpLoweringBase<AwaitOp, async::TokenType>;
 
 public:
   using Base::Base;

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

@@ -2605,8 +2605,8 @@ static bool isZeroAttribute(Attribute value) {
 
 LogicalResult GlobalOp::verify() {
   bool validType = isCompatibleOuterType(getType())
-                       ? !llvm::isa<LLVMVoidType, LLVMTokenType,
-                                    LLVMMetadataType, LLVMLabelType>(getType())
+                       ? !llvm::isa<LLVMVoidType, TokenType, LLVMMetadataType,
+                                    LLVMLabelType>(getType())
                        : llvm::isa<PointerElementTypeInterface>(getType());
   if (!validType)
     return emitOpError(
@@ -2826,8 +2826,8 @@ ParseResult AliasOp::parse(OpAsmParser &parser, OperationState &result) {
 
 LogicalResult AliasOp::verify() {
   bool validType = isCompatibleOuterType(getType())
-                       ? !llvm::isa<LLVMVoidType, LLVMTokenType,
-                                    LLVMMetadataType, LLVMLabelType>(getType())
+                       ? !llvm::isa<LLVMVoidType, TokenType, LLVMMetadataType,
+                                    LLVMLabelType>(getType())
                        : llvm::isa<PointerElementTypeInterface>(getType());
   if (!validType)
     return emitOpError(
@@ -4462,7 +4462,6 @@ void LLVMDialect::initialize() {
 
   // clang-format off
   addTypes<LLVMVoidType,
-           LLVMTokenType,
            LLVMLabelType,
            LLVMMetadataType>();
   // clang-format on

mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp

@@ -36,7 +36,6 @@ static StringRef getTypeKeyword(Type type) {
   return TypeSwitch<Type, StringRef>(type)
       .Case<LLVMVoidType>([&](Type) { return "void"; })
       .Case<LLVMPPCFP128Type>([&](Type) { return "ppc_fp128"; })
-      .Case<LLVMTokenType>([&](Type) { return "token"; })
       .Case<LLVMLabelType>([&](Type) { return "label"; })
       .Case<LLVMMetadataType>([&](Type) { return "metadata"; })
       .Case<LLVMFunctionType>([&](Type) { return "func"; })
@@ -266,7 +265,6 @@ static Type dispatchParse(AsmParser &parser, bool allowAny = true) {
   return StringSwitch<function_ref<Type()>>(key)
       .Case("void", [&] { return LLVMVoidType::get(ctx); })
       .Case("ppc_fp128", [&] { return LLVMPPCFP128Type::get(ctx); })
-      .Case("token", [&] { return LLVMTokenType::get(ctx); })
       .Case("label", [&] { return LLVMLabelType::get(ctx); })
       .Case("metadata", [&] { return LLVMMetadataType::get(ctx); })
       .Case("func", [&] { return LLVMFunctionType::parse(parser); })

mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp

@@ -150,7 +150,7 @@ generatedTypeParser(AsmParser &parser, StringRef *mnemonic, Type &value);
 
 bool LLVMArrayType::isValidElementType(Type type) {
   return !llvm::isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
-                    LLVMFunctionType, LLVMTokenType>(type);
+                    LLVMFunctionType, TokenType>(type);
 }
 
 LLVMArrayType LLVMArrayType::get(Type elementType, uint64_t numElements) {
@@ -435,7 +435,7 @@ LogicalResult LLVMPointerType::verifyEntries(DataLayoutEntryListRef entries,
 
 bool LLVMStructType::isValidElementType(Type type) {
   return !llvm::isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
-                    LLVMFunctionType, LLVMTokenType>(type);
+                    LLVMFunctionType, TokenType>(type);
 }
 
 LLVMStructType LLVMStructType::getIdentified(MLIRContext *context,
@@ -743,10 +743,10 @@ bool mlir::LLVM::isCompatibleOuterType(Type type) {
       LLVMPPCFP128Type,
       LLVMPointerType,
       LLVMStructType,
-      LLVMTokenType,
       LLVMTargetExtType,
       LLVMVoidType,
-      LLVMX86AMXType
+      LLVMX86AMXType,
+      TokenType
     >(type)) {
     // clang-format on
     return true;
@@ -803,9 +803,9 @@ static bool isCompatibleImpl(Type type, DenseSet<Type> &compatibleTypes) {
             LLVMLabelType,
             LLVMMetadataType,
             LLVMPPCFP128Type,
-            LLVMTokenType,
             LLVMVoidType,
-            LLVMX86AMXType
+            LLVMX86AMXType,
+            TokenType
           >([](Type) { return true; })
           // clang-format on
           .Case<PtrLikeTypeInterface>(
@@ -917,11 +917,11 @@ llvm::TypeSize mlir::LLVM::getPrimitiveTypeSizeInBits(Type type) {
                               elementSize.isScalable());
       })
       .Default([](Type ty) {
-        assert((llvm::isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType,
-                          LLVMTokenType, LLVMStructType, LLVMArrayType,
-                          LLVMPointerType, LLVMFunctionType, LLVMTargetExtType>(
-                   ty)) &&
-               "unexpected missing support for primitive type");
+        assert(
+            (llvm::isa<LLVMVoidType, LLVMLabelType, LLVMMetadataType, TokenType,
+                       LLVMStructType, LLVMArrayType, LLVMPointerType,
+                       LLVMFunctionType, LLVMTargetExtType>(ty)) &&
+            "unexpected missing support for primitive type");
         return llvm::TypeSize::getFixed(0);
       });
 }

mlir/lib/IR/AsmPrinter.cpp

@@ -2907,6 +2907,7 @@ void AsmPrinter::Impl::printTypeImpl(Type type) {
         os << '>';
       })
       .Case<NoneType>([&](Type) { os << "none"; })
+      .Case<TokenType>([&](Type) { os << "token"; })
       .Case([&](GraphType graphTy) {
         os << '(';
         interleaveComma(graphTy.getInputs(), [&](Type ty) { printType(ty); });

mlir/lib/IR/Verifier.cpp

@@ -26,6 +26,7 @@
 
 #include "mlir/IR/Verifier.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
@@ -42,8 +43,8 @@ class OperationVerifier {
 public:
   /// If `verifyRecursively` is true, then this will also recursively verify
   /// nested operations.
-  explicit OperationVerifier(bool verifyRecursively)
-      : verifyRecursively(verifyRecursively) {}
+  OperationVerifier(MLIRContext *ctx, bool verifyRecursively)
+      : tokenType(TokenType::get(ctx)), verifyRecursively(verifyRecursively) {}
 
   /// Verify the given operation.
   LogicalResult verifyOpAndDominance(Operation &op);
@@ -58,6 +59,12 @@ private:
   /// upon exit from the subtree, i.e. when we visit a node for the second time.
   LogicalResult verifyOnEntrance(Block &block);
   LogicalResult verifyOnEntrance(Operation &op);
+  LogicalResult
+  verifyTokenValue(Operation &producer, Value value,
+                   function_ref<InFlightDiagnostic()> emitProducerError);
+  LogicalResult verifyTokenValues(Operation &op);
+  LogicalResult verifyTokenBlockArgument(Block &block, BlockArgument arg,
+                                         unsigned idx);
 
   LogicalResult verifyOnExit(Block &block);
   LogicalResult verifyOnExit(Operation &op);
@@ -70,6 +77,9 @@ private:
   LogicalResult verifyDominanceOfContainedRegions(Operation &op,
                                                   DominanceInfo &domInfo);
 
+  /// The cached instance of the builtin token type.
+  TokenType tokenType;
+
   /// A flag indicating if this verifier should recursively verify nested
   /// operations.
   bool verifyRecursively;
@@ -109,6 +119,77 @@ static bool mayBeValidWithoutTerminator(Block *block) {
   return !op || op->mightHaveTrait<OpTrait::NoTerminator>();
 }
 
+LogicalResult OperationVerifier::verifyTokenValue(
+    Operation &producer, Value value,
+    function_ref<InFlightDiagnostic()> emitProducerError) {
+  if (value.getType() != tokenType)
+    return success();
+
+  if (!producer.mightHaveTrait<OpTrait::TokenProducerTrait>())
+    return emitProducerError();
+
+  for (OpOperand &use : value.getUses()) {
+    Operation *user = use.getOwner();
+    if (user->mightHaveTrait<OpTrait::TokenConsumerTrait>())
+      continue;
+
+    return user->emitOpError()
+           << "consumes token operand #" << use.getOperandNumber()
+           << " but does not have the TokenConsumerTrait";
+  }
+
+  return success();
+}
+
+LogicalResult OperationVerifier::verifyTokenValues(Operation &op) {
+  for (auto resultIt : llvm::enumerate(op.getResults())) {
+    unsigned idx = resultIt.index();
+    OpResult result = resultIt.value();
+    if (failed(verifyTokenValue(op, result, [&]() {
+          return op.emitOpError()
+                 << "produces token result #" << idx
+                 << " but does not have the TokenProducerTrait";
+        })))
+      return failure();
+  }
+
+  for (Region &region : op.getRegions()) {
+    if (region.empty())
+      continue;
+
+    Block &entryBlock = region.front();
+    for (auto argIt : llvm::enumerate(entryBlock.getArguments())) {
+      unsigned idx = argIt.index();
+      BlockArgument arg = argIt.value();
+      if (failed(verifyTokenValue(op, arg, [&]() {
+            return emitError(arg.getLoc(), "token entry block argument #")
+                   << idx << " requires the parent operation to have the "
+                   << "TokenProducerTrait";
+          })))
+        return failure();
+    }
+  }
+
+  return success();
+}
+
+LogicalResult OperationVerifier::verifyTokenBlockArgument(Block &block,
+                                                          BlockArgument arg,
+                                                          unsigned idx) {
+  if (arg.getType() != tokenType)
+    return success();
+
+  // The producer-trait check on the parent op (and the token consumer check
+  // on the uses) is performed by `verifyTokenValues` when it iterates the
+  // entry block arguments of an op's regions. Here we only enforce that
+  // tokens are not used as non-entry block arguments.
+  if (!block.getParent() || !block.isEntryBlock())
+    return emitError(arg.getLoc(), "token block argument #")
+           << idx << " is only allowed in a region entry block";
+
+  return success();
+}
+
 LogicalResult OperationVerifier::verifyOnEntrance(Block &block) {
   // Get the parent op and context for cross-context checks. Both are available
   // whenever the block lives inside a region that has a parent operation.
@@ -133,6 +214,8 @@ LogicalResult OperationVerifier::verifyOnEntrance(Block &block) {
                << " type from a different MLIRContext than its "
                   "parent operation";
     }
+    if (failed(verifyTokenBlockArgument(block, arg, idx)))
+      return failure();
   }
 
   // Verify that this block has a terminator.
@@ -232,6 +315,9 @@ LogicalResult OperationVerifier::verifyOnEntrance(Operation &op) {
   if (registeredInfo && failed(registeredInfo->verifyInvariants(&op)))
     return failure();
 
+  if (failed(verifyTokenValues(op)))
+    return failure();
+
   unsigned numRegions = op.getNumRegions();
   if (!numRegions)
     return success();
@@ -478,6 +564,6 @@ OperationVerifier::verifyDominanceOfContainedRegions(Operation &op,
 //===----------------------------------------------------------------------===//
 
 LogicalResult mlir::verify(Operation *op, bool verifyRecursively) {
-  OperationVerifier verifier(verifyRecursively);
+  OperationVerifier verifier(op->getContext(), verifyRecursively);
   return verifier.verifyOpAndDominance(*op);
 }

mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp

@@ -73,7 +73,7 @@ private:
     if (type->isMetadataTy())
       return LLVM::LLVMMetadataType::get(&context);
     if (type->isTokenTy())
-      return LLVM::LLVMTokenType::get(&context);
+      return TokenType::get(&context);
     llvm_unreachable("not a primitive type");
   }
 

mlir/lib/Target/LLVMIR/TypeToLLVM.cpp

@@ -58,9 +58,7 @@ public:
             .Case([this](LLVM::LLVMPPCFP128Type) {
               return llvm::Type::getPPC_FP128Ty(context);
             })
-            .Case([this](LLVM::LLVMTokenType) {
-              return llvm::Type::getTokenTy(context);
-            })
+            .Case([this](TokenType) { return llvm::Type::getTokenTy(context); })
             .Case([this](LLVM::LLVMLabelType) {
               return llvm::Type::getLabelTy(context);
             })

mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir

@@ -4,7 +4,7 @@
 func.func @coro_id() {
   // CHECK: %0 = llvm.mlir.constant(0 : i32) : i32
   // CHECK: %1 = llvm.mlir.zero : !llvm.ptr
-  // CHECK: %2 = llvm.intr.coro.id %0, %1, %1, %1 : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  // CHECK: %2 = llvm.intr.coro.id %0, %1, %1, %1 : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   %0 = async.coro.id
   return
 }

mlir/test/Dialect/ArmSME/invalid.mlir

@@ -132,7 +132,7 @@ func.func @arm_sme_tile_load__pad_but_no_mask(%src : memref<?x?xf64>, %pad : f64
 
 func.func @arm_sme_tile_load__bad_memref_rank(%src : memref<?xf64>, %pad : f64) {
   %c0 = arith.constant 0 : index
-  // expected-error@+1 {{op operand #0 must be 2D memref of any type values, but got 'memref<?xf64>'}}
+  // expected-error@+1 {{op operand #0 must be 2D memref of any non-token type values, but got 'memref<?xf64>'}}
   %tile = arm_sme.tile_load %src[%c0], %pad, : memref<?xf64>, vector<[2]x[2]xf64>
   return
 }
@@ -186,7 +186,7 @@ func.func @arm_sme_tile_store__bad_mask_type(%tile : vector<[16]x[16]xi8>, %mask
 
 func.func @arm_sme_tile_store__bad_memref_rank(%tile : vector<[16]x[16]xi8>, %dest : memref<?xi8>) {
   %c0 = arith.constant 0 : index
-  // expected-error@+1 {{op operand #1 must be 2D memref of any type values, but got 'memref<?xi8>'}}
+  // expected-error@+1 {{op operand #1 must be 2D memref of any non-token type values, but got 'memref<?xi8>'}}
   arm_sme.tile_store %tile, %dest[%c0] : memref<?xi8>, vector<[16]x[16]xi8>
   return
 }

mlir/test/Dialect/Builtin/Bytecode/builtin_fixed.mlir

@@ -282,6 +282,7 @@ module @TestBasicTypes attributes {
   // CHECK-DAG: bytecode.ui64 = ui64
   // CHECK-DAG: bytecode.index = index
   // CHECK-DAG: bytecode.none = none
+  // CHECK-DAG: bytecode.token = token
   bytecode.i1 = i1,
   bytecode.i8 = i8,
   bytecode.i32 = i32,
@@ -289,7 +290,8 @@ module @TestBasicTypes attributes {
   bytecode.si32 = si32,
   bytecode.ui64 = ui64,
   bytecode.index = index,
-  bytecode.none = none
+  bytecode.none = none,
+  bytecode.token = token
 } {} loc(unknown)
 
 //===----------------------------------------------------------------------===//

mlir/test/Dialect/Builtin/Bytecode/types.mlir

@@ -169,3 +169,13 @@ module @TestVector attributes {
   bytecode.test = vector<8x8x128xi8>,
   bytecode.test1 = vector<8x[8]xf32>
 } {}
+
+//===----------------------------------------------------------------------===//
+// TokenType
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: @TestToken
+module @TestToken attributes {
+  // CHECK: bytecode.test = token
+  bytecode.test = token
+} {}

mlir/test/Dialect/LLVMIR/types.mlir

@@ -6,12 +6,12 @@ func.func @primitive() {
   "some.op"() : () -> !llvm.void
   // CHECK: !llvm.ppc_fp128
   "some.op"() : () -> !llvm.ppc_fp128
-  // CHECK: !llvm.token
-  "some.op"() : () -> !llvm.token
   // CHECK: !llvm.label
   "some.op"() : () -> !llvm.label
   // CHECK: !llvm.metadata
   "some.op"() : () -> !llvm.metadata
+  // CHECK: token
+  "some.op"() : () -> token
   return
 }
 

mlir/test/Dialect/Linalg/invalid.mlir

@@ -415,7 +415,7 @@ func.func @illegal_fill_memref_with_tensor_return
 func.func @illegal_fill_tensor_with_memref_return
   (%arg0 : tensor<?x?xf32>, %arg1 : f32) -> memref<?x?xf32>
 {
-  // expected-error @+1 {{result #0 must be variadic of ranked tensor of any type values, but got 'memref<?x?xf32>'}}
+  // expected-error @+1 {{result #0 must be variadic of ranked tensor of any non-token type values, but got 'memref<?x?xf32>'}}
   %0 = linalg.fill ins(%arg1 : f32) outs(%arg0 : tensor<?x?xf32>) -> memref<?x?xf32>
   return %0 : memref<?x?xf32>
 }
@@ -468,7 +468,7 @@ func.func @invalid_scalar_input_matmul(%arg0: f32, %arg1: memref<3x4xf32>, %arg2
 // -----
 
 func.func @invalid_scalar_output_matmul(%arg0: memref<2x3xf32>, %arg1: memref<3x4xf32>, %arg2: f32) {
-  // expected-error @+1 {{'linalg.matmul' op operand #2 must be variadic of shaped of any type values, but got 'f32'}}
+  // expected-error @+1 {{'linalg.matmul' op operand #2 must be variadic of shaped of any non-token type values, but got 'f32'}}
   linalg.matmul ins(%arg0, %arg1 : memref<2x3xf32>, memref<3x4xf32>)
                 outs(%arg2 : f32)
   return

mlir/test/Dialect/MemRef/invalid.mlir

@@ -1037,7 +1037,7 @@ func.func @test_alloc_memref_map_rank_mismatch() {
 // -----
 
 func.func @rank(%0: f32) {
-  // expected-error@+1 {{'memref.rank' op operand #0 must be ranked or unranked memref of any type values}}
+  // expected-error@+1 {{'memref.rank' op operand #0 must be ranked or unranked memref of any non-token type values}}
   "memref.rank"(%0): (f32)->index
   return
 }
@@ -1172,7 +1172,7 @@ func.func @memref_realloc_type(%src : memref<256xf32>) -> memref<?xi32>{
 
 // Asking the dimension of a 0-D shape doesn't make sense.
 func.func @dim_0_ranked(%arg : memref<f32>, %arg1 : index) {
-  memref.dim %arg, %arg1 : memref<f32> // expected-error {{'memref.dim' op operand #0 must be unranked.memref of any type values or non-0-ranked.memref of any type values, but got 'memref<f32>'}}
+  memref.dim %arg, %arg1 : memref<f32> // expected-error {{'memref.dim' op operand #0 must be unranked.memref of any non-token type values or non-0-ranked.memref of any non-token type values, but got 'memref<f32>'}}
   return
 }
 

mlir/test/Dialect/SparseTensor/invalid.mlir

@@ -1,7 +1,7 @@
 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 func.func @invalid_new_dense(%arg0: !llvm.ptr) -> tensor<32xf32> {
-  // expected-error@+1 {{'sparse_tensor.new' op result #0 must be sparse tensor of any type values, but got 'tensor<32xf32>'}}
+  // expected-error@+1 {{'sparse_tensor.new' op result #0 must be sparse tensor of any non-token type values, but got 'tensor<32xf32>'}}
   %0 = sparse_tensor.new %arg0 : !llvm.ptr to tensor<32xf32>
   return %0 : tensor<32xf32>
 }
@@ -96,7 +96,7 @@ func.func @invalid_unpack_mis_position(%sp: tensor<2x100xf64, #CSR>, %values: te
 // -----
 
 func.func @invalid_positions_dense(%arg0: tensor<128xf64>) -> memref<?xindex> {
-  // expected-error@+1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<128xf64>'}}
   %0 = sparse_tensor.positions %arg0 { level = 0 : index } : tensor<128xf64> to memref<?xindex>
   return %0 : memref<?xindex>
 }
@@ -104,7 +104,7 @@ func.func @invalid_positions_dense(%arg0: tensor<128xf64>) -> memref<?xindex> {
 // -----
 
 func.func @invalid_positions_unranked(%arg0: tensor<*xf64>) -> memref<?xindex> {
-  // expected-error@+1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any type values, but got 'tensor<*xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<*xf64>'}}
   %0 = "sparse_tensor.positions"(%arg0) { level = 0 : index } : (tensor<*xf64>) -> (memref<?xindex>)
   return %0 : memref<?xindex>
 }
@@ -132,7 +132,7 @@ func.func @positions_oob(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xinde
 // -----
 
 func.func @invalid_indices_dense(%arg0: tensor<10x10xi32>) -> memref<?xindex> {
-  // expected-error@+1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any type values, but got 'tensor<10x10xi32>'}}
+  // expected-error@+1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<10x10xi32>'}}
   %0 = sparse_tensor.coordinates %arg0 { level = 1 : index } : tensor<10x10xi32> to memref<?xindex>
   return %0 : memref<?xindex>
 }
@@ -140,7 +140,7 @@ func.func @invalid_indices_dense(%arg0: tensor<10x10xi32>) -> memref<?xindex> {
 // -----
 
 func.func @invalid_indices_unranked(%arg0: tensor<*xf64>) -> memref<?xindex> {
-  // expected-error@+1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any type values, but got 'tensor<*xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<*xf64>'}}
   %0 = "sparse_tensor.coordinates"(%arg0) { level = 0 : index } : (tensor<*xf64>) -> (memref<?xindex>)
   return %0 : memref<?xindex>
 }
@@ -168,7 +168,7 @@ func.func @indices_oob(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex>
 // -----
 
 func.func @invalid_values_dense(%arg0: tensor<1024xf32>) -> memref<?xf32> {
-  // expected-error@+1 {{'sparse_tensor.values' op operand #0 must be sparse tensor of any type values, but got 'tensor<1024xf32>'}}
+  // expected-error@+1 {{'sparse_tensor.values' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<1024xf32>'}}
   %0 = sparse_tensor.values %arg0 : tensor<1024xf32> to memref<?xf32>
   return %0 : memref<?xf32>
 }
@@ -186,7 +186,7 @@ func.func @indices_buffer_noncoo(%arg0: tensor<128xf64, #SparseVector>) -> memre
 // -----
 
 func.func @indices_buffer_dense(%arg0: tensor<1024xf32>) -> memref<?xindex> {
-  // expected-error@+1 {{must be sparse tensor of any type values}}
+  // expected-error@+1 {{must be sparse tensor of any non-token type values}}
   %0 = sparse_tensor.coordinates_buffer %arg0 : tensor<1024xf32> to memref<?xindex>
   return %0 : memref<?xindex>
 }
@@ -283,7 +283,7 @@ func.func @sparse_get_md(%arg0: !sparse_tensor.storage_specifier<#COO>) -> index
 // -----
 
 func.func @sparse_unannotated_load(%arg0: tensor<16x32xf64>) -> tensor<16x32xf64> {
-  // expected-error@+1 {{'sparse_tensor.load' op operand #0 must be sparse tensor of any type values, but got 'tensor<16x32xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.load' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<16x32xf64>'}}
   %0 = sparse_tensor.load %arg0 : tensor<16x32xf64>
   return %0 : tensor<16x32xf64>
 }
@@ -308,7 +308,7 @@ func.func @sparse_push_back_n(%arg0: index, %arg1: memref<?xf32>, %arg2: f32) ->
 // -----
 
 func.func @sparse_unannotated_expansion(%arg0: tensor<128xf64>) {
-  // expected-error@+1 {{'sparse_tensor.expand' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.expand' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<128xf64>'}}
   %values, %filled, %added, %count = sparse_tensor.expand %arg0
     : tensor<128xf64> to memref<?xf64>, memref<?xi1>, memref<?xindex>
   return
@@ -322,7 +322,7 @@ func.func @sparse_unannotated_compression(%arg0: memref<?xf64>,
                                           %arg3: index,
                                           %arg4: tensor<8x8xf64>,
                                           %arg5: index) {
-  // expected-error@+1 {{'sparse_tensor.compress' op operand #4 must be sparse tensor of any type values, but got 'tensor<8x8xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.compress' op operand #4 must be sparse tensor of any non-token type values, but got 'tensor<8x8xf64>'}}
   sparse_tensor.compress %arg0, %arg1, %arg2, %arg3 into %arg4[%arg5]
     : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64>
   return
@@ -375,7 +375,7 @@ func.func @sparse_convert_dim_mismatch(%arg0: tensor<10x?xf32>) -> tensor<10x10x
 // -----
 
 func.func @invalid_out_dense(%arg0: tensor<10xf64>, %arg1: !llvm.ptr) {
-  // expected-error@+1 {{'sparse_tensor.out' op operand #0 must be sparse tensor of any type values, but got 'tensor<10xf64>'}}
+  // expected-error@+1 {{'sparse_tensor.out' op operand #0 must be sparse tensor of any non-token type values, but got 'tensor<10xf64>'}}
   sparse_tensor.out %arg0, %arg1 : tensor<10xf64>, !llvm.ptr
   return
 }
@@ -1022,7 +1022,7 @@ func.func @sparse_reinterpret_map(%t0 : tensor<6x12xi32, #BSR>) -> tensor<3x4x2x
 #CSR = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed, d1 : compressed)}>
 
 func.func @sparse_print(%arg0: tensor<10x10xf64>) {
-  // expected-error@+1 {{'sparse_tensor.print' op operand #0 must be sparse tensor of any type values}}
+  // expected-error@+1 {{'sparse_tensor.print' op operand #0 must be sparse tensor of any non-token type values}}
   sparse_tensor.print %arg0 : tensor<10x10xf64>
   return
 }

mlir/test/Dialect/Tensor/invalid.mlir

@@ -404,7 +404,7 @@ func.func @illegal_collapsing_reshape_mixed_tensor_2(%arg0 : tensor<?x4x5xf32>)
 // -----
 
 func.func @rank(%0: f32) {
-  // expected-error@+1 {{'tensor.rank' op operand #0 must be tensor of any type values}}
+  // expected-error@+1 {{'tensor.rank' op operand #0 must be tensor of any non-token type values}}
   "tensor.rank"(%0): (f32)->index
   return
 }

mlir/test/Dialect/Vector/invalid.mlir

@@ -106,7 +106,7 @@ func.func @shuffle_index_out_of_range(%arg0: vector<2xf32>, %arg1: vector<2xf32>
 // -----
 
 func.func @shuffle_scalable_vec(%arg0: vector<[2]xf32>, %arg1: vector<[2]xf32>) {
-  // expected-error@+1 {{'vector.shuffle' op operand #0 must be fixed-length vector of any type values}}
+  // expected-error@+1 {{'vector.shuffle' op operand #0 must be fixed-length vector of any non-token type values}}
   %1 = vector.shuffle %arg0, %arg1 [0, 1, 2, 3] : vector<[2]xf32>, vector<[2]xf32>
 }
 
@@ -1460,7 +1460,7 @@ func.func @maskedstore_memref_mismatch(%base: memref<?xf32>, %mask: vector<16xi1
 func.func @gather_from_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
                                 %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
   %c0 = arith.constant 0 : index
-  // expected-error@+1 {{'vector.gather' op operand #0 must be Tensor or MemRef of any type values, but got 'vector<16xf32>'}}
+  // expected-error@+1 {{'vector.gather' op operand #0 must be Tensor or MemRef of any non-token type values, but got 'vector<16xf32>'}}
   %0 = vector.gather %base[%c0][%indices], %mask, %pass_thru
     : vector<16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32>
 }
@@ -1557,7 +1557,7 @@ func.func @gather_tensor_alignment(%base: tensor<16xf32>, %indices: vector<16xi3
 func.func @scatter_to_vector(%base: vector<16xf32>, %indices: vector<16xi32>,
                              %mask: vector<16xi1>, %pass_thru: vector<16xf32>) {
   %c0 = arith.constant 0 : index
-  // expected-error@+1 {{'vector.scatter' op operand #0 must be Tensor or MemRef of any type values, but got 'vector<16xf32>'}}
+  // expected-error@+1 {{'vector.scatter' op operand #0 must be Tensor or MemRef of any non-token type values, but got 'vector<16xf32>'}}
   vector.scatter %base[%c0][%indices], %mask, %pass_thru
     : vector<16xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32>
 }
@@ -1959,7 +1959,7 @@ func.func @invalid_outerproduct1(%src : memref<?xf32>, %lhs : vector<[4]x[4]xf32
 // -----
 
 func.func @deinterleave_zero_dim_fail(%vec : vector<f32>) {
-  // expected-error @+1 {{'vector.deinterleave' op operand #0 must be vector of any type values, but got 'vector<f32>}}
+  // expected-error @+1 {{'vector.deinterleave' op operand #0 must be vector of any non-token type values, but got 'vector<f32>}}
   %0, %1 = vector.deinterleave %vec : vector<f32> -> vector<f32>
   return
 }
@@ -2048,7 +2048,7 @@ func.func @from_elements_wrong_operand_type(%a: f32, %b: i32) {
 // -----
 
 func.func @invalid_from_elements_scalable(%a: f32, %b: i32) {
-  // expected-error @+1 {{'dest' must be fixed-length vector of any type values, but got 'vector<[2]xf32>'}}
+  // expected-error @+1 {{'dest' must be fixed-length vector of any non-token type values, but got 'vector<[2]xf32>'}}
   vector.from_elements %a, %b : vector<[2]xf32>
   return
 }

mlir/test/Dialect/traits.mlir

@@ -58,7 +58,7 @@ func.func @broadcast_tensor_tensor_tensor(tensor<8x1x?x1xi32>, tensor<7x1x5xi32>
 // Check incompatible vector and tensor result type
 func.func @broadcast_scalar_vector_vector(tensor<4xf32>, tensor<4xf32>) -> vector<4xf32> {
 ^bb0(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>):
-  // expected-error @+1 {{op result #0 must be tensor of any type values, but got 'vector<4xf32>'}}
+  // expected-error @+1 {{op result #0 must be tensor of any non-token type values, but got 'vector<4xf32>'}}
   %0 = "test.broadcastable"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> vector<4xf32>
   return %0 : vector<4xf32>
 }

mlir/test/IR/operand.mlir

@@ -13,7 +13,7 @@ func.func @correct_variadic_operand(%arg0: tensor<f32>, %arg1: f32) {
 // -----
 
 func.func @error_in_first_variadic_operand(%arg0: tensor<f32>, %arg1: f32) {
-  // expected-error @+1 {{operand #1 must be variadic of tensor of any type}}
+  // expected-error @+1 {{operand #1 must be variadic of tensor of any non-token type}}
   "test.mixed_normal_variadic_operand"(%arg0, %arg1, %arg0, %arg0, %arg0) : (tensor<f32>, f32, tensor<f32>, tensor<f32>, tensor<f32>) -> ()
   return
 }
@@ -21,7 +21,7 @@ func.func @error_in_first_variadic_operand(%arg0: tensor<f32>, %arg1: f32) {
 // -----
 
 func.func @error_in_normal_operand(%arg0: tensor<f32>, %arg1: f32) {
-  // expected-error @+1 {{operand #2 must be tensor of any type}}
+  // expected-error @+1 {{operand #2 must be tensor of any non-token type}}
   "test.mixed_normal_variadic_operand"(%arg0, %arg0, %arg1, %arg0, %arg0) : (tensor<f32>, tensor<f32>, f32, tensor<f32>, tensor<f32>) -> ()
   return
 }
@@ -29,7 +29,7 @@ func.func @error_in_normal_operand(%arg0: tensor<f32>, %arg1: f32) {
 // -----
 
 func.func @error_in_second_variadic_operand(%arg0: tensor<f32>, %arg1: f32) {
-  // expected-error @+1 {{operand #3 must be variadic of tensor of any type}}
+  // expected-error @+1 {{operand #3 must be variadic of tensor of any non-token type}}
   "test.mixed_normal_variadic_operand"(%arg0, %arg0, %arg0, %arg1, %arg0) : (tensor<f32>, tensor<f32>, tensor<f32>, f32, tensor<f32>) -> ()
   return
 }

mlir/test/IR/result.mlir

@@ -13,7 +13,7 @@ func.func @correct_variadic_result() -> tensor<f32> {
 // -----
 
 func.func @error_in_first_variadic_result() -> tensor<f32> {
-  // expected-error @+1 {{result #1 must be variadic of tensor of any type}}
+  // expected-error @+1 {{result #1 must be variadic of tensor of any non-token type}}
   %0:5 = "test.mixed_normal_variadic_result"() : () -> (tensor<f32>, f32, tensor<f32>, tensor<f32>, tensor<f32>)
   return %0#4 : tensor<f32>
 }
@@ -21,7 +21,7 @@ func.func @error_in_first_variadic_result() -> tensor<f32> {
 // -----
 
 func.func @error_in_normal_result() -> tensor<f32> {
-  // expected-error @+1 {{result #2 must be tensor of any type}}
+  // expected-error @+1 {{result #2 must be tensor of any non-token type}}
   %0:5 = "test.mixed_normal_variadic_result"() : () -> (tensor<f32>, tensor<f32>, f32, tensor<f32>, tensor<f32>)
   return %0#4 : tensor<f32>
 }
@@ -29,7 +29,7 @@ func.func @error_in_normal_result() -> tensor<f32> {
 // -----
 
 func.func @error_in_second_variadic_result() -> tensor<f32> {
-  // expected-error @+1 {{result #3 must be variadic of tensor of any type}}
+  // expected-error @+1 {{result #3 must be variadic of tensor of any non-token type}}
   %0:5 = "test.mixed_normal_variadic_result"() : () -> (tensor<f32>, tensor<f32>, tensor<f32>, f32, tensor<f32>)
   return %0#4 : tensor<f32>
 }

mlir/test/IR/token-type.mlir

@@ -0,0 +1,143 @@
+// RUN: mlir-opt %s -verify-diagnostics -split-input-file | FileCheck %s
+
+// Tests for the builtin `token` type, the token producer/consumer operation
+// traits, and the `Token`, `AnyType` ODS predicates. The default `AnyType`
+// predicate excludes tokens.
+
+// CHECK-LABEL: @token_produce_consume
+func.func @token_produce_consume() {
+  // CHECK: %[[T:.*]] = test.token.produce
+  %t = test.token.produce
+  // CHECK: test.token.consume %[[T]]
+  test.token.consume %t
+  return
+}
+
+// -----
+
+// Region entry block arguments may produce tokens when the parent op opts in.
+// CHECK-LABEL: @token_region_entry_block_arg
+func.func @token_region_entry_block_arg() {
+  // CHECK: "test.token.region"
+  "test.token.region"() ({
+  ^bb0(%arg0: token):
+    // CHECK: test.token.consume
+    test.token.consume %arg0
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// `AnyType` accepts arbitrary non-token types.
+// CHECK-LABEL: @any_type_with_non_token
+func.func @any_type_with_non_token(%arg0: i32) {
+  // CHECK: test.token.any_type %{{.*}} : i32
+  test.token.any_type %arg0 : i32
+  return
+}
+
+// -----
+
+// `AnyType` rejects tokens by default.
+func.func @any_type_rejects_token() {
+  %t = test.token.produce
+  // expected-error @below {{operand #0 must be any non-token type}}
+  test.token.any_type %t : token
+  return
+}
+
+// -----
+
+// Token-producing ops must have the TokenProducerTrait.
+func.func @token_result_requires_producer_trait() {
+  // expected-error @below {{'test.token.produce_without_trait' op produces token result #0 but does not have the TokenProducerTrait}}
+  %t = test.token.produce_without_trait : token
+  return
+}
+
+// -----
+
+// Token-consuming ops must have the TokenConsumerTrait.
+func.func @token_operand_requires_consumer_trait() {
+  %t = test.token.produce
+  // expected-error @below {{'test.token.consume_without_trait' op consumes token operand #0 but does not have the TokenConsumerTrait}}
+  test.token.consume_without_trait %t : token
+  return
+}
+
+// -----
+
+// Token entry block arguments require the parent op to have the
+// TokenProducerTrait.
+func.func @token_entry_block_arg_requires_parent_producer_trait() {
+  "test.token.region_without_trait"() ({
+  // expected-error @below {{token entry block argument #0 requires the parent operation to have the TokenProducerTrait}}
+  ^bb0(%arg0: token):
+    test.token.consume %arg0
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// A region with a parent op still cannot have token entry block arguments unless
+// the parent op has the TokenProducerTrait.
+func.func @token_entry_block_arg_requires_parent_producer_trait_without_uses() {
+  "test.token.region_without_trait"() ({
+  // expected-error @below {{token entry block argument #0 requires the parent operation to have the TokenProducerTrait}}
+  ^bb0(%arg0: token):
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// Token entry block arguments still require consumers to have the
+// TokenConsumerTrait.
+func.func @token_entry_block_arg_use_requires_consumer_trait() {
+  "test.token.region"() ({
+  ^bb0(%arg0: token):
+    // expected-error @below {{'test.token.consume_without_trait' op consumes token operand #0 but does not have the TokenConsumerTrait}}
+    test.token.consume_without_trait %arg0 : token
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// Tokens cannot be non-entry block arguments.
+func.func @token_non_entry_block_arg_is_rejected() {
+  "test.token.region"() ({
+    "test.finish"() : () -> ()
+  // expected-error @below {{token block argument #0 is only allowed in a region entry block}}
+  ^bb1(%arg0: token):
+    "test.finish"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// -----
+
+// Function entry blocks do not opt in to producing builtin tokens.
+// expected-error @below {{token entry block argument #0 requires the parent operation to have the TokenProducerTrait}}
+func.func @token_region_arg(%arg0: token) {
+  test.token.consume %arg0
+  return
+}
+
+// -----
+
+// `Token` rejects non-token types. The op's operand type is fixed to the
+// builtin `token` (it's a `BuildableType`), so passing a non-token SSA value
+// fails at parse time with an SSA type mismatch.
+// expected-note @below {{prior use here}}
+func.func @token_rejects_non_token(%arg0: i32) {
+  // expected-error @below {{use of value '%arg0' expects different type than prior uses: 'token' vs 'i32'}}
+  test.token.consume %arg0
+  return
+}

mlir/test/Target/LLVMIR/Import/intrinsic.ll

@@ -793,16 +793,16 @@ define void @threadlocal_test() {
 ; CHECK-LABEL:  llvm.func @coro_id
 define void @coro_id() {
   %a = alloca [16 x i8]
-  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   %3 = call token @llvm.coro.id(i32 0, ptr %a, ptr null, ptr null)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @coro_begin
 define void @coro_begin(ptr %0) {
-  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   %3 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
-  ; CHECK: llvm.intr.coro.begin %{{.*}}, %{{.*}} : (!llvm.token, !llvm.ptr) -> !llvm.ptr
+  ; CHECK: llvm.intr.coro.begin %{{.*}}, %{{.*}} : (token, !llvm.ptr) -> !llvm.ptr
   %4 = call ptr @llvm.coro.begin(token %3, ptr %0)
   ret void
 }
@@ -826,14 +826,14 @@ define void @coro_align() {
 
 ; CHECK-LABEL:  llvm.func @coro_save
 define void @coro_save(ptr %0) {
-  ; CHECK: llvm.intr.coro.save %{{.*}} : (!llvm.ptr) -> !llvm.token
+  ; CHECK: llvm.intr.coro.save %{{.*}} : (!llvm.ptr) -> token
   %2 = call token @llvm.coro.save(ptr %0)
   ret void
 }
 
 ; CHECK-LABEL:  llvm.func @coro_suspend
 define void @coro_suspend(i1 %0) {
-  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   %4 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
   ; CHECK: llvm.intr.coro.suspend %{{.*}}, %{{.*}} : i8
   %5 = call i8 @llvm.coro.suspend(token %4, i1 %0)
@@ -849,9 +849,9 @@ define void @coro_end(ptr %0, i1 %1) {
 
 ; CHECK-LABEL:  llvm.func @coro_free
 define void @coro_free(ptr %0) {
-  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  ; CHECK: llvm.intr.coro.id %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   %3 = call token @llvm.coro.id(i32 0, ptr null, ptr null, ptr null)
-  ; CHECK: llvm.intr.coro.free %{{.*}}, %{{.*}} : (!llvm.token, !llvm.ptr) -> !llvm.ptr
+  ; CHECK: llvm.intr.coro.free %{{.*}}, %{{.*}} : (token, !llvm.ptr) -> !llvm.ptr
   %4 = call ptr @llvm.coro.free(token %3, ptr %0)
   ret void
 }

mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir

@@ -816,7 +816,7 @@ llvm.func @coro_id() {
   %a = llvm.alloca %c x i8 : (i64) -> !llvm.ptr
   // CHECK: call token @llvm.coro.id
   %null = llvm.mlir.zero : !llvm.ptr
-  llvm.intr.coro.id %zero, %a, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  llvm.intr.coro.id %zero, %a, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   llvm.return
 }
 
@@ -824,9 +824,9 @@ llvm.func @coro_id() {
 llvm.func @coro_begin(%arg0: !llvm.ptr) {
   %zero = llvm.mlir.constant(0 : i32) : i32
   %null = llvm.mlir.zero : !llvm.ptr
-  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   // CHECK: call ptr @llvm.coro.begin
-  llvm.intr.coro.begin %token, %arg0 : (!llvm.token, !llvm.ptr) -> !llvm.ptr
+  llvm.intr.coro.begin %token, %arg0 : (token, !llvm.ptr) -> !llvm.ptr
   llvm.return
 }
 
@@ -851,7 +851,7 @@ llvm.func @coro_align() {
 // CHECK-LABEL: @coro_save
 llvm.func @coro_save(%arg0: !llvm.ptr) {
   // CHECK: call token @llvm.coro.save
-  %0 = llvm.intr.coro.save %arg0 : (!llvm.ptr) -> !llvm.token
+  %0 = llvm.intr.coro.save %arg0 : (!llvm.ptr) -> token
   llvm.return
 }
 
@@ -859,7 +859,7 @@ llvm.func @coro_save(%arg0: !llvm.ptr) {
 llvm.func @coro_suspend(%arg0 : i1) {
   %zero = llvm.mlir.constant(0 : i32) : i32
   %null = llvm.mlir.zero : !llvm.ptr
-  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   // CHECK: call i8 @llvm.coro.suspend
   %0 = llvm.intr.coro.suspend %token, %arg0 : i8
   llvm.return
@@ -867,9 +867,9 @@ llvm.func @coro_suspend(%arg0 : i1) {
 
 // CHECK-LABEL: @coro_end
 llvm.func @coro_end(%arg0: !llvm.ptr, %arg1 : i1) {
-  %none = llvm.mlir.none : !llvm.token
+  %none = llvm.mlir.none : token
   // CHECK: call void @llvm.coro.end
-  llvm.intr.coro.end %arg0, %arg1, %none : (!llvm.ptr, i1, !llvm.token) -> !llvm.void
+  llvm.intr.coro.end %arg0, %arg1, %none : (!llvm.ptr, i1, token) -> !llvm.void
   llvm.return
 }
 
@@ -877,9 +877,9 @@ llvm.func @coro_end(%arg0: !llvm.ptr, %arg1 : i1) {
 llvm.func @coro_free(%arg0 : !llvm.ptr) {
   %zero = llvm.mlir.constant(0 : i32) : i32
   %null = llvm.mlir.zero : !llvm.ptr
-  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> !llvm.token
+  %token = llvm.intr.coro.id %zero, %null, %null, %null : (i32, !llvm.ptr, !llvm.ptr, !llvm.ptr) -> token
   // CHECK: call ptr @llvm.coro.free
-  %0 = llvm.intr.coro.free %token, %arg0 : (!llvm.token, !llvm.ptr) -> !llvm.ptr
+  %0 = llvm.intr.coro.free %token, %arg0 : (token, !llvm.ptr) -> !llvm.ptr
   llvm.return
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

@@ -110,6 +110,61 @@ def SignlessLikeVariadic : TEST_Op<"signless_like_variadic"> {
   let arguments = (ins Variadic<SignlessIntegerLike>:$x);
 }
 
+//===----------------------------------------------------------------------===//
+// Test Token Type
+//===----------------------------------------------------------------------===//
+
+// Produce a builtin `!token` value.
+def TestTokenProduceOp : TEST_Op<"token.produce"> {
+  let results = (outs Token:$token);
+  let assemblyFormat = "attr-dict";
+}
+
+// Produce a builtin `!token` value via an opt-in result type without the
+// producer trait. ODS only auto-derives the trait when the declared type is
+// statically `Token`, so this op (whose result type accepts any type
+// including tokens) exercises the runtime trait check.
+def TestTokenProduceWithoutTraitOp
+    : TEST_Op<"token.produce_without_trait"> {
+  let results = (outs Type<CPred<"true">>:$token);
+  let assemblyFormat = "attr-dict `:` type($token)";
+}
+
+// Consume a builtin `!token` value (token-only operand). The consumer trait
+// is auto-derived by ODS from the statically-typed `Token` operand.
+def TestTokenConsumeOp : TEST_Op<"token.consume"> {
+  let arguments = (ins Token:$token);
+  let assemblyFormat = "$token attr-dict";
+}
+
+// Consume a builtin `!token` value via an opt-in operand type without the
+// consumer trait. ODS only auto-derives the trait when the declared type is
+// statically `Token`, so this op (whose operand type accepts any type
+// including tokens) exercises the runtime trait check.
+def TestTokenConsumeWithoutTraitOp
+    : TEST_Op<"token.consume_without_trait"> {
+  let arguments = (ins Type<CPred<"true">>:$token);
+  let assemblyFormat = "$token attr-dict `:` type($token)";
+}
+
+// Op that uses the default `AnyType` predicate. Tokens are excluded by
+// default and should be rejected by the verifier when passed here.
+def TestTokenAnyTypeOp : TEST_Op<"token.any_type", [TokenConsumerTrait]> {
+  let arguments = (ins AnyType:$value);
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
+// Op whose region entry blocks may produce builtin `!token` values.
+def TestTokenRegionOp : TEST_Op<"token.region", [TokenProducerTrait]> {
+  let regions = (region AnyRegion:$body);
+}
+
+// Op whose regions do not opt in to producing builtin `!token` values.
+def TestTokenRegionWithoutTraitOp
+    : TEST_Op<"token.region_without_trait"> {
+  let regions = (region AnyRegion:$body);
+}
+
 //===----------------------------------------------------------------------===//
 // Test Symbols
 //===----------------------------------------------------------------------===//

mlir/test/mlir-tblgen/predicate.td

@@ -27,9 +27,9 @@ def OpA : NS_Op<"op_for_CPred_containing_multiple_same_placeholder", []> {
 // CHECK-NOT.        << " must be 32-bit integer or floating-point type, but got " << type;
 
 // CHECK: static ::llvm::LogicalResult [[$TENSOR_CONSTRAINT:__mlir_ods_local_type_constraint.*]](
-// CHECK:       if (!(((::llvm::isa<::mlir::TensorType>(type))) && ([](::mlir::Type elementType) { return (true); }(::llvm::cast<::mlir::ShapedType>(type).getElementType())))) {
+// CHECK:       if (!(((::llvm::isa<::mlir::TensorType>(type))) && ([](::mlir::Type elementType) { return !((::llvm::isa<::mlir::TokenType>(elementType))); }(::llvm::cast<::mlir::ShapedType>(type).getElementType())))) {
 // CHECK-NEXT:    return op->emitOpError(valueKind) << " #" << valueIndex
-// CHECK-NEXT:        << " must be tensor of any type values, but got " << type;
+// CHECK-NEXT:        << " must be tensor of any non-token type values, but got " << type;
 
 // CHECK: static ::llvm::LogicalResult [[$TENSOR_INTEGER_FLOAT_CONSTRAINT:__mlir_ods_local_type_constraint.*]](
 // CHECK:       if (!(((::llvm::isa<::mlir::TensorType>(type))) && ([](::mlir::Type elementType) { return ((elementType.isF32())) || ((elementType.isSignlessInteger(32))); }(::llvm::cast<::mlir::ShapedType>(type).getElementType())))) {

mlir/test/mlir-tblgen/types.mlir

@@ -204,7 +204,7 @@ func.func @ranked_tensor_success(%arg0: tensor<i8>, %arg1: tensor<1xi32>, %arg2:
 // -----
 
 func.func @ranked_tensor_success(%arg0: tensor<*xf32>) {
-  // expected-error @+1 {{must be ranked tensor of any type values}}
+  // expected-error @+1 {{must be ranked tensor of any non-token type values}}
   "test.ranked_tensor_op"(%arg0) : (tensor<*xf32>) -> ()
   return
 }
@@ -212,7 +212,7 @@ func.func @ranked_tensor_success(%arg0: tensor<*xf32>) {
 // -----
 
 func.func @ranked_tensor_success(%arg0: vector<2xf32>) {
-  // expected-error @+1 {{must be ranked tensor of any type values}}
+  // expected-error @+1 {{must be ranked tensor of any non-token type values}}
   "test.ranked_tensor_op"(%arg0) : (vector<2xf32>) -> ()
   return
 }
@@ -510,7 +510,7 @@ func.func @does_not_have_i32(%arg0: tensor<1x2xi32>, %arg1: none) {
 // -----
 
 func.func @does_not_have_static_memref(%arg0: memref<?xi32>) {
-  // expected-error@+1 {{'test.takes_static_memref' op operand #0 must be statically shaped memref of any type values}}
+  // expected-error@+1 {{'test.takes_static_memref' op operand #0 must be statically shaped memref of any non-token type values}}
   "test.takes_static_memref"(%arg0) : (memref<?xi32>) -> ()
 }
 

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

@@ -4154,6 +4154,19 @@ void OpEmitter::genTraits() {
       opClass.addTrait(opTrait->getFullyQualifiedTraitName());
     }
   }
+
+  // Auto-derive the builtin token producer/consumer traits whenever the op
+  // statically declares a Token operand or result.
+  constexpr llvm::StringLiteral kTokenCppType = "::mlir::TokenType";
+  auto hasStaticTokenType = [&](auto &&values) {
+    return llvm::any_of(values, [&](const tblgen::NamedTypeConstraint &v) {
+      return v.constraint.getCppType() == kTokenCppType;
+    });
+  };
+  if (hasStaticTokenType(op.getOperands()))
+    opClass.addTrait("::mlir::OpTrait::TokenConsumerTrait");
+  if (hasStaticTokenType(op.getResults()))
+    opClass.addTrait("::mlir::OpTrait::TokenProducerTrait");
 }
 
 void OpEmitter::genOpNameGetter() {