[Mlir-commits] [mlir] 67aaa03 - [mlir][IR] Add builtin `TokenType` (#195640)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jun 2 03:33:51 PDT 2026
Author: Matthias Springer
Date: 2026-06-02T12:33:46+02:00
New Revision: 67aaa03dac582b5ccf65776e449eb6c7517eb11c
URL: https://github.com/llvm/llvm-project/commit/67aaa03dac582b5ccf65776e449eb6c7517eb11c
DIFF: https://github.com/llvm/llvm-project/commit/67aaa03dac582b5ccf65776e449eb6c7517eb11c.diff
LOG: [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 at gmail.com>
Added:
mlir/docs/Tokens.md
mlir/test/IR/token-type.mlir
Modified:
mlir/docs/Dialects/LLVM.md
mlir/docs/LangRef.md
mlir/docs/Traits/_index.md
mlir/include/mlir/Dialect/Async/IR/Async.h
mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
mlir/include/mlir/IR/BuiltinDialectBytecode.td
mlir/include/mlir/IR/BuiltinTypes.td
mlir/include/mlir/IR/CommonTypeConstraints.td
mlir/include/mlir/IR/OpBase.td
mlir/include/mlir/IR/OpDefinition.h
mlir/lib/AsmParser/TokenKinds.def
mlir/lib/AsmParser/TypeParser.cpp
mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
mlir/lib/Dialect/Async/IR/Async.cpp
mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp
mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
mlir/lib/IR/AsmPrinter.cpp
mlir/lib/IR/Verifier.cpp
mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp
mlir/lib/Target/LLVMIR/TypeToLLVM.cpp
mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
mlir/test/Dialect/ArmSME/invalid.mlir
mlir/test/Dialect/Builtin/Bytecode/builtin_fixed.mlir
mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc
mlir/test/Dialect/Builtin/Bytecode/types.mlir
mlir/test/Dialect/LLVMIR/types.mlir
mlir/test/Dialect/Linalg/invalid.mlir
mlir/test/Dialect/MemRef/invalid.mlir
mlir/test/Dialect/SparseTensor/invalid.mlir
mlir/test/Dialect/Tensor/invalid.mlir
mlir/test/Dialect/Vector/invalid.mlir
mlir/test/Dialect/traits.mlir
mlir/test/IR/operand.mlir
mlir/test/IR/result.mlir
mlir/test/Target/LLVMIR/Import/intrinsic.ll
mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
mlir/test/lib/Dialect/Test/TestOps.td
mlir/test/mlir-tblgen/predicate.td
mlir/test/mlir-tblgen/types.mlir
mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
Removed:
################################################################################
diff --git a/mlir/docs/Dialects/LLVM.md b/mlir/docs/Dialects/LLVM.md
index 4b5d518ca4eab..f8419d726fc33 100644
--- a/mlir/docs/Dialects/LLVM.md
+++ b/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
diff --git a/mlir/docs/LangRef.md b/mlir/docs/LangRef.md
index 5e53df83997e2..0e6fb006da48b 100644
--- a/mlir/docs/LangRef.md
+++ b/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
diff erent
+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
diff --git a/mlir/docs/Tokens.md b/mlir/docs/Tokens.md
new file mode 100644
index 0000000000000..20df548ca9c7e
--- /dev/null
+++ b/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.
diff --git a/mlir/docs/Traits/_index.md b/mlir/docs/Traits/_index.md
index 866716b9f5193..7d01f8517477a 100644
--- a/mlir/docs/Traits/_index.md
+++ b/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.
diff --git a/mlir/include/mlir/Dialect/Async/IR/Async.h b/mlir/include/mlir/Dialect/Async/IR/Async.h
index f16e87e71373a..fc0b086126f52 100644
--- a/mlir/include/mlir/Dialect/Async/IR/Async.h
+++ b/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
diff --git a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td b/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
index 2cebeac767f29..058f58bda6433 100644
--- a/mlir/include/mlir/Dialect/Async/IR/AsyncOps.td
+++ b/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
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index 688bc19cbf18a..5acfec359e074 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/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)";
}
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
index bd59319c79ad3..5e5b9e9faa070 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td
+++ b/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<
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index d7c8cf236f0da..c9ddfa335c4f2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/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)";
}
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h b/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
index a1506497dc85c..a54f83660ca9c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMTypes.h
+++ b/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");
diff --git a/mlir/include/mlir/IR/BuiltinDialectBytecode.td b/mlir/include/mlir/IR/BuiltinDialectBytecode.td
index 207b99164d0d6..213ba260cc4fd 100644
--- a/mlir/include/mlir/IR/BuiltinDialectBytecode.td
+++ b/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
];
}
diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td
index 20c41c5f79729..40ccaefa6de3f 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.td
+++ b/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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index 57caaae08462f..d8615a4730c32 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/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",
diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 1e34959d0d557..0d0669e90c3f7 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/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">;
diff --git a/mlir/include/mlir/IR/OpDefinition.h b/mlir/include/mlir/IR/OpDefinition.h
index b93dfc0258207..a0a36f2bd53c4 100644
--- a/mlir/include/mlir/IR/OpDefinition.h
+++ b/mlir/include/mlir/IR/OpDefinition.h
@@ -778,6 +778,18 @@ class IsTerminator : public TraitBase<ConcreteType, IsTerminator> {
}
};
+/// 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>
diff --git a/mlir/lib/AsmParser/TokenKinds.def b/mlir/lib/AsmParser/TokenKinds.def
index fe7c53753e156..f5e5c25832a30 100644
--- a/mlir/lib/AsmParser/TokenKinds.def
+++ b/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)
diff --git a/mlir/lib/AsmParser/TypeParser.cpp b/mlir/lib/AsmParser/TypeParser.cpp
index a461ebed967a8..2cdec14d65fa6 100644
--- a/mlir/lib/AsmParser/TypeParser.cpp
+++ b/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();
diff --git a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp b/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
index 29e6552231f9c..46e53e71d35f5 100644
--- a/mlir/lib/Conversion/AsyncToLLVM/AsyncToLLVM.cpp
+++ b/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 @@ class AsyncRuntimeTypeConverter : public TypeConverter {
}
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 @@ class RuntimeCreateOpLowering : public ConvertOpToLLVMPattern<RuntimeCreateOp> {
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 @@ class RuntimeSetAvailableOpLowering
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 @@ class RuntimeSetErrorOpLowering
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 @@ class RuntimeIsErrorOpLowering : public OpConversionPattern<RuntimeIsErrorOp> {
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 @@ class RuntimeAwaitOpLowering : public OpConversionPattern<RuntimeAwaitOp> {
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 @@ class RuntimeAwaitAndResumeOpLowering
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 @@ class RuntimeAddToGroupOpLowering
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 @@ class ConvertYieldOpTypes : public OpConversionPattern<async::YieldOp> {
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;
diff --git a/mlir/lib/Dialect/Async/IR/Async.cpp b/mlir/lib/Dialect/Async/IR/Async.cpp
index 71be1d275280e..1713da07da60d 100644
--- a/mlir/lib/Dialect/Async/IR/Async.cpp
+++ b/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 "
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
index 91e37dd9ac36e..2a726f3fd2999 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncRuntimeRefCounting.cpp
+++ b/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);
diff --git a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp b/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
index 6ed50671fb3b3..f4d37bd0e8f34 100644
--- a/mlir/lib/Dialect/Async/Transforms/AsyncToAsyncRuntime.cpp
+++ b/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 @@ class AwaitOpLoweringBase : public OpConversionPattern<AwaitType> {
};
/// 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;
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 63bd9f8a3d625..aac100198509a 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/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
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp
index 705d07d3e6c42..c498f0746169d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp
+++ b/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); })
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
index 2b3ba1b8b5a35..2c29b38c26f08 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMTypes.cpp
+++ b/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);
});
}
diff --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index ec270db189081..ca5c2d2a88ee5 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/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); });
diff --git a/mlir/lib/IR/Verifier.cpp b/mlir/lib/IR/Verifier.cpp
index 33da2cd867f42..11771e78d5f20 100644
--- a/mlir/lib/IR/Verifier.cpp
+++ b/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 @@ class OperationVerifier {
/// 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 @@ class OperationVerifier {
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 ®ion : 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
diff erent 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);
}
diff --git a/mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp b/mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp
index 5d9345d707a44..018f66802fc47 100644
--- a/mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp
+++ b/mlir/lib/Target/LLVMIR/TypeFromLLVM.cpp
@@ -73,7 +73,7 @@ class TypeFromLLVMIRTranslatorImpl {
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");
}
diff --git a/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp b/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp
index 807a94c61f0c8..61997c691e35d 100644
--- a/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp
+++ b/mlir/lib/Target/LLVMIR/TypeToLLVM.cpp
@@ -58,9 +58,7 @@ class TypeToLLVMIRTranslatorImpl {
.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);
})
diff --git a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir b/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
index a398bc5710a86..7874aca75dc4c 100644
--- a/mlir/test/Conversion/AsyncToLLVM/convert-coro-to-llvm.mlir
+++ b/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
}
diff --git a/mlir/test/Dialect/ArmSME/invalid.mlir b/mlir/test/Dialect/ArmSME/invalid.mlir
index 8c5a098a0c785..f00945e18cc1f 100644
--- a/mlir/test/Dialect/ArmSME/invalid.mlir
+++ b/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 at +1 {{op operand #0 must be 2D memref of any type values, but got 'memref<?xf64>'}}
+ // expected-error at +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 at +1 {{op operand #1 must be 2D memref of any type values, but got 'memref<?xi8>'}}
+ // expected-error at +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
}
diff --git a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed.mlir b/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed.mlir
index 638689406ab17..31061a6fa8653 100644
--- a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed.mlir
+++ b/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)
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc b/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc
index e2f58a42751a1..f08188cf1b8c0 100644
Binary files a/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc and b/mlir/test/Dialect/Builtin/Bytecode/builtin_fixed_0.mlirbc
diff er
diff --git a/mlir/test/Dialect/Builtin/Bytecode/types.mlir b/mlir/test/Dialect/Builtin/Bytecode/types.mlir
index 5e421e2bf75bf..91d4512998b99 100644
--- a/mlir/test/Dialect/Builtin/Bytecode/types.mlir
+++ b/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
+} {}
diff --git a/mlir/test/Dialect/LLVMIR/types.mlir b/mlir/test/Dialect/LLVMIR/types.mlir
index b87c3dd6f2d7a..69582546cb429 100644
--- a/mlir/test/Dialect/LLVMIR/types.mlir
+++ b/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
}
diff --git a/mlir/test/Dialect/Linalg/invalid.mlir b/mlir/test/Dialect/Linalg/invalid.mlir
index 06f3fcb41190b..a446cfcc4eec1 100644
--- a/mlir/test/Dialect/Linalg/invalid.mlir
+++ b/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
diff --git a/mlir/test/Dialect/MemRef/invalid.mlir b/mlir/test/Dialect/MemRef/invalid.mlir
index 2f061a1bb773e..ecffd683a98c2 100644
--- a/mlir/test/Dialect/MemRef/invalid.mlir
+++ b/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 at +1 {{'memref.rank' op operand #0 must be ranked or unranked memref of any type values}}
+ // expected-error at +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
}
diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir
index ae706b9b148a6..d14229b011f11 100644
--- a/mlir/test/Dialect/SparseTensor/invalid.mlir
+++ b/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 at +1 {{'sparse_tensor.new' op result #0 must be sparse tensor of any type values, but got 'tensor<32xf32>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.positions' op operand #0 must be sparse tensor of any type values, but got 'tensor<*xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any type values, but got 'tensor<10x10xi32>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.coordinates' op operand #0 must be sparse tensor of any type values, but got 'tensor<*xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.values' op operand #0 must be sparse tensor of any type values, but got 'tensor<1024xf32>'}}
+ // expected-error at +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 at +1 {{must be sparse tensor of any type values}}
+ // expected-error at +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 at +1 {{'sparse_tensor.load' op operand #0 must be sparse tensor of any type values, but got 'tensor<16x32xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.expand' op operand #0 must be sparse tensor of any type values, but got 'tensor<128xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.compress' op operand #4 must be sparse tensor of any type values, but got 'tensor<8x8xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.out' op operand #0 must be sparse tensor of any type values, but got 'tensor<10xf64>'}}
+ // expected-error at +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 at +1 {{'sparse_tensor.print' op operand #0 must be sparse tensor of any type values}}
+ // expected-error at +1 {{'sparse_tensor.print' op operand #0 must be sparse tensor of any non-token type values}}
sparse_tensor.print %arg0 : tensor<10x10xf64>
return
}
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index 6ee2f9911663f..a526d7ed61722 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/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 at +1 {{'tensor.rank' op operand #0 must be tensor of any type values}}
+ // expected-error at +1 {{'tensor.rank' op operand #0 must be tensor of any non-token type values}}
"tensor.rank"(%0): (f32)->index
return
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 662e4e8b5b561..2fed3002596a3 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/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 at +1 {{'vector.shuffle' op operand #0 must be fixed-length vector of any type values}}
+ // expected-error at +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 at +1 {{'vector.gather' op operand #0 must be Tensor or MemRef of any type values, but got 'vector<16xf32>'}}
+ // expected-error at +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 at +1 {{'vector.scatter' op operand #0 must be Tensor or MemRef of any type values, but got 'vector<16xf32>'}}
+ // expected-error at +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
}
diff --git a/mlir/test/Dialect/traits.mlir b/mlir/test/Dialect/traits.mlir
index 4d583435adeee..ae48cadbf370f 100644
--- a/mlir/test/Dialect/traits.mlir
+++ b/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>
}
diff --git a/mlir/test/IR/operand.mlir b/mlir/test/IR/operand.mlir
index 507e37c775c0b..1ac12dc4b9556 100644
--- a/mlir/test/IR/operand.mlir
+++ b/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
}
diff --git a/mlir/test/IR/result.mlir b/mlir/test/IR/result.mlir
index 1e4eb3bede4c5..cdeae4202f0ff 100644
--- a/mlir/test/IR/result.mlir
+++ b/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>
}
diff --git a/mlir/test/IR/token-type.mlir b/mlir/test/IR/token-type.mlir
new file mode 100644
index 0000000000000..0218ae6e81ac5
--- /dev/null
+++ b/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
diff erent type than prior uses: 'token' vs 'i32'}}
+ test.token.consume %arg0
+ return
+}
diff --git a/mlir/test/Target/LLVMIR/Import/intrinsic.ll b/mlir/test/Target/LLVMIR/Import/intrinsic.ll
index f79d09aa3d633..959a04fff6dca 100644
--- a/mlir/test/Target/LLVMIR/Import/intrinsic.ll
+++ b/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
}
diff --git a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir b/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
index 11882a0a1d4c6..5865e046aa5ac 100644
--- a/mlir/test/Target/LLVMIR/llvmir-intrinsics.mlir
+++ b/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
}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index a1529e3020c82..d008343e43a39 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/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
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
index 41e041f171213..ae436885b421f 100644
--- a/mlir/test/mlir-tblgen/predicate.td
+++ b/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())))) {
diff --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index c2acce0903bf4..30aea48e3e369 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/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 at +1 {{'test.takes_static_memref' op operand #0 must be statically shaped memref of any type values}}
+ // expected-error at +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>) -> ()
}
diff --git a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp b/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
index 2cb47d084ce69..90ada40302296 100644
--- a/mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
+++ b/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() {
More information about the Mlir-commits
mailing list