[Mlir-commits] [mlir] [mlir][IR] Add builtin `TokenTypeInterface` (PR #195640)

Matthias Springer llvmlistbot at llvm.org
Mon May 4 05:17:25 PDT 2026


https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/195640

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

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


>From 39761ab7c7efdc44d80e65b6c1ced02ef01de999 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Mon, 4 May 2026 12:14:41 +0000
Subject: [PATCH] [mlir][IR] Add builtin `TokenTypeInterface`

---
 mlir/docs/Dialects/Builtin.md                 | 10 +++
 mlir/docs/Tokens.md                           | 89 +++++++++++++++++++
 mlir/include/mlir/IR/BuiltinTypeInterfaces.td | 20 +++++
 mlir/include/mlir/IR/CommonTypeConstraints.td | 21 ++++-
 mlir/test/Dialect/ArmSME/invalid.mlir         |  4 +-
 mlir/test/Dialect/Linalg/invalid.mlir         |  4 +-
 mlir/test/Dialect/MemRef/invalid.mlir         |  4 +-
 mlir/test/Dialect/SparseTensor/invalid.mlir   | 24 ++---
 mlir/test/Dialect/Tensor/invalid.mlir         |  2 +-
 mlir/test/Dialect/Vector/invalid.mlir         | 10 +--
 mlir/test/Dialect/traits.mlir                 |  2 +-
 mlir/test/IR/operand.mlir                     |  6 +-
 mlir/test/IR/result.mlir                      |  6 +-
 mlir/test/IR/token-type-interface.mlir        | 59 ++++++++++++
 mlir/test/lib/Dialect/Test/TestOps.td         | 32 +++++++
 mlir/test/lib/Dialect/Test/TestTypeDefs.td    |  7 ++
 mlir/test/mlir-tblgen/predicate.td            |  4 +-
 mlir/test/mlir-tblgen/types.mlir              |  6 +-
 18 files changed, 272 insertions(+), 38 deletions(-)
 create mode 100644 mlir/docs/Tokens.md
 create mode 100644 mlir/test/IR/token-type-interface.mlir

diff --git a/mlir/docs/Dialects/Builtin.md b/mlir/docs/Dialects/Builtin.md
index 0a9b7ae8919b5..818d6ded486db 100644
--- a/mlir/docs/Dialects/Builtin.md
+++ b/mlir/docs/Dialects/Builtin.md
@@ -65,3 +65,13 @@ marked using one DistinctAttribute instance per alias group.
 ## Type Interfaces
 
 [include "Dialects/BuiltinTypeInterfaces.md"]
+
+## Token Types
+
+A *token type* is any type that implements the builtin `TokenTypeInterface`.
+Tokens are SSA values that exist purely to encode a *static* def–use
+relationship between operations or regions; they carry no runtime data and
+must not be value-forwarded.
+
+See the [Tokens design note](../Tokens.md) for the structural contract,
+ODS predicates (`AnyType` / `AnyTypeOrToken` / `Token`), and examples.
diff --git a/mlir/docs/Tokens.md b/mlir/docs/Tokens.md
new file mode 100644
index 0000000000000..81bcc1203b6aa
--- /dev/null
+++ b/mlir/docs/Tokens.md
@@ -0,0 +1,89 @@
+# 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).
+
+More precisely, a token is an SSA value whose purpose is to encode a
+**static** def–use relationship between operations or regions. It carries
+no runtime data and is not allowed to flow through "regular"
+value-forwarding constructs. A token's provenance cannot be obscured through
+value forwarding.
+
+In MLIR, "token" is not a single concrete builtin type. Instead, any type
+that implements the builtin `TokenTypeInterface` is treated as a token by
+the framework. Dialects can define their own dialect-specific token types.
+
+## `TokenTypeInterface`
+
+`TokenTypeInterface` is a parameterless, methodless marker type interface.
+
+A type opts in by attaching the interface in TableGen:
+
+```tablegen
+def MyDialect_Token : TypeDef<MyDialect, "MyToken", [TokenTypeInterface]> {
+  let mnemonic = "token";
+}
+```
+
+## Structural Contract
+
+A token value is, by construction:
+
+1. **Not value-forwarding.** In particular, a token must not appear as a
+   forwarded value. E.g.:
+    * a forwarded result/operand of a `CallOpInterface` op,
+    * an argument or result type of a `FunctionOpInterface` op (a token
+      block argument *inside* a function body is fine — what is disallowed
+      is forwarding tokens across the call/return boundary),
+    * a successor operand or successor block argument of a
+      `BranchOpInterface` op,
+    * a forwarded operand to/from any region of a `RegionBranchOpInterface`
+      op (iter-args, region results, yielded values), or
+    * the result of any op that selects or merges values it does not
+      understand (e.g. `arith.select`).
+
+2. **Statically resolvable.** Walking the def–use chain from any token use
+   reaches a producing op without crossing a forwarding boundary.
+
+3. **Cannot constant-fold.** No constant of token type exists.
+
+These properties mirror what LLVM IR already documents for its own
+[`token` type](https://llvm.org/docs/LangRef.html#token-type).
+
+## 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.
+Three 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. |
+| `AnyTypeOrToken`   | any type, including tokens           | the op legitimately accepts arbitrary types (including tokens).       |
+| `Token`            | only types implementing `TokenTypeInterface` | the op specifically takes a token operand/result.             |
+
+
+## Examples
+
+### Rejected: tokens in `AnyType` positions
+
+```mlir
+// error: 'scf.if' op result #0 must be variadic of any non-token type,
+//        but got '!my.token'
+%t = scf.if %cond -> !my.token {
+  %a = my.token.produce : !my.token
+  scf.yield %a : !my.token
+} else {
+  %b = my.token.produce : !my.token
+  scf.yield %b : !my.token
+}
+```
+
+`scf.if`'s results are declared with `Variadic<AnyType>` and `scf.yield`'s
+operands likewise use `AnyType`. Because `AnyType` excludes tokens by
+default, yielding (or returning) a token through a `scf.if` (or any other
+op that has not explicitly opted in via `AnyTypeOrToken`) is rejected.
diff --git a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
index 93c8c0694b467..a50d73e21a69a 100644
--- a/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
+++ b/mlir/include/mlir/IR/BuiltinTypeInterfaces.td
@@ -236,6 +236,26 @@ def PtrLikeTypeInterface : TypeInterface<"PtrLikeTypeInterface"> {
   ];
 }
 
+//===----------------------------------------------------------------------===//
+// TokenTypeInterface
+//===----------------------------------------------------------------------===//
+
+def TokenTypeInterface : TypeInterface<"TokenTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let description = [{
+    Intuitively, a *token* value is a pointer to an operation (via an OpResult)
+    or a pointer to a region (via an entry block argument).
+
+    More precisely, a token is an SSA value whose purpose is to encode a
+    static def–use relationship between operations or regions. It carries
+    no runtime data and is not allowed to flow through "regular"
+    value-forwarding constructs. A token's provenance cannot be obscured through
+    value forwarding.
+
+    This interface is a marker. It has no interface methods.
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // ShapedType
 //===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index 57caaae08462f..7898937ca01f5 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/mlir/include/mlir/IR/CommonTypeConstraints.td
@@ -165,8 +165,25 @@ 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 a token (i.e. implements TokenTypeInterface).
+def IsTokenTypePred
+    : CPred<"::llvm::isa<::mlir::TokenTypeInterface>($_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. Ops that legitimately
+// want to accept any type, including tokens, should use `AnyTypeOrToken`
+// instead.
+def AnyType : Type<Neg<IsTokenTypePred>, "any non-token type">;
+
+// Any type at all, including tokens. Used by ops that explicitly opt in to
+// accepting tokens (e.g. ops in interfaces such as `CallOpInterface`,
+// `BranchOpInterface`, etc. that legitimately handle arbitrary types).
+def AnyTypeOrToken : Type<CPred<"true">, "any type">;
+
+// A token type (any type implementing `TokenTypeInterface`).
+def Token : Type<IsTokenTypePred, "token",
+                 "::mlir::TokenTypeInterface">;
 
 // None type
 def NoneType : Type<CPred<"::llvm::isa<::mlir::NoneType>($_self)">, "none type",
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/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 f90312c915334..36c697c78d93d 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>
 }
@@ -1943,7 +1943,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
 }
@@ -2032,7 +2032,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-interface.mlir b/mlir/test/IR/token-type-interface.mlir
new file mode 100644
index 0000000000000..c0632bdcf616b
--- /dev/null
+++ b/mlir/test/IR/token-type-interface.mlir
@@ -0,0 +1,59 @@
+// RUN: mlir-opt %s -verify-diagnostics -split-input-file | FileCheck %s
+
+// Tests for the builtin `TokenTypeInterface` and the
+// `Token` / `AnyType` / `AnyTypeOrToken` ODS predicates.
+//
+// `!test.test_token` is a test-dialect type that implements
+// `TokenTypeInterface`. The default `AnyType` predicate excludes tokens, while
+// `AnyTypeOrToken` and `Token` accept them.
+
+// CHECK-LABEL: @token_produce_consume
+func.func @token_produce_consume() {
+  // CHECK: %[[T:.*]] = test.token.produce : !test.test_token
+  %t = test.token.produce : !test.test_token
+  // CHECK: test.token.consume %[[T]] : !test.test_token
+  test.token.consume %t : !test.test_token
+  // CHECK: test.token.any_or_token %[[T]] : !test.test_token
+  test.token.any_or_token %t : !test.test_token
+  return
+}
+
+// -----
+
+// `AnyTypeOrToken` also accepts non-token types.
+// CHECK-LABEL: @any_or_token_with_non_token
+func.func @any_or_token_with_non_token(%arg0: i32) {
+  // CHECK: test.token.any_or_token %{{.*}} : i32
+  test.token.any_or_token %arg0 : i32
+  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 : !test.test_token
+  // expected-error @below {{operand #0 must be any non-token type}}
+  test.token.any_type %t : !test.test_token
+  return
+}
+
+// -----
+
+// `Token` rejects non-token types. The operand's cppType is
+// `TokenTypeInterface`, so type resolution fails at parse time.
+func.func @token_rejects_non_token(%arg0: i32) {
+  // expected-error @below {{invalid kind of type specified}}
+  test.token.consume %arg0 : i32
+  return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 348ff5d7f4ea0..529fc4b860ad4 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -110,6 +110,38 @@ def SignlessLikeVariadic : TEST_Op<"signless_like_variadic"> {
   let arguments = (ins Variadic<SignlessIntegerLike>:$x);
 }
 
+//===----------------------------------------------------------------------===//
+// Test Token Type
+//===----------------------------------------------------------------------===//
+
+// Produce a token. Demonstrates a type that implements the builtin
+// `TokenTypeInterface`.
+def TestTokenProduceOp : TEST_Op<"token.produce"> {
+  let results = (outs TestTokenType:$token);
+  let assemblyFormat = "attr-dict `:` type($token)";
+}
+
+// Consume a token (token-only operand). Uses the `Token` ODS predicate which
+// only accepts types implementing `TokenTypeInterface`.
+def TestTokenConsumeOp : TEST_Op<"token.consume"> {
+  let arguments = (ins Token:$token);
+  let assemblyFormat = "$token attr-dict `:` type($token)";
+}
+
+// Op that accepts any type, including a token. Uses the `AnyTypeOrToken`
+// opt-in predicate.
+def TestTokenAnyTypeOrTokenOp : TEST_Op<"token.any_or_token"> {
+  let arguments = (ins AnyTypeOrToken:$value);
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
+// 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"> {
+  let arguments = (ins AnyType:$value);
+  let assemblyFormat = "$value attr-dict `:` type($value)";
+}
+
 //===----------------------------------------------------------------------===//
 // Test Symbols
 //===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestTypeDefs.td b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
index 08600ce713a17..5df0eab829a03 100644
--- a/mlir/test/lib/Dialect/Test/TestTypeDefs.td
+++ b/mlir/test/lib/Dialect/Test/TestTypeDefs.td
@@ -173,6 +173,13 @@ def TestMemRefElementType : Test_Type<"TestMemRefElementType",
   let mnemonic = "memref_element";
 }
 
+// A test token type implementing the builtin `TokenTypeInterface`. Used to
+// exercise the default exclusion of tokens from `AnyType` and the explicit
+// `Token` / `AnyTypeOrToken` opt-ins.
+def TestTokenType : Test_Type<"TestToken", [TokenTypeInterface]> {
+  let mnemonic = "test_token";
+}
+
 def TestTypeTrait : NativeTypeTrait<"TestTypeTrait">;
 
 // The definition of a singleton type that has a trait.
diff --git a/mlir/test/mlir-tblgen/predicate.td b/mlir/test/mlir-tblgen/predicate.td
index 41e041f171213..07a5e6f2f261c 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::TokenTypeInterface>(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>) -> ()
 }
 



More information about the Mlir-commits mailing list