dialects: (stablehlo) add stablehlo.token (#3093)

Adds stablehlo.token type.

Co-authored-by: Erick Ochoa <erick@ceci-nest-pas.me>

tests/filecheck/dialects/stablehlo/attrs.mlir

@@ -6,6 +6,9 @@
     highest = #stablehlo<precision HIGHEST>
 } : () -> ()
 
+%token = "test.op"() : () -> (!stablehlo.token)
+
 // CHECK:       builtin.module {
 // CHECK-NEXT:    "test.op"() {"default" = #stablehlo<precision DEFAULT>, "high" = #stablehlo<precision HIGH>, "highest" = #stablehlo<precision HIGHEST>} : () -> ()
+// CHECK-NEXT:    %token = "test.op"() : () -> !stablehlo.token
 // CHECK-NEXT:  }

xdsl/dialects/stablehlo.py

@@ -14,9 +14,11 @@
     Attribute,
     Dialect,
     EnumAttribute,
+    ParametrizedAttribute,
     SpacedOpaqueSyntaxAttribute,
     SSAValue,
     StrEnum,
+    TypeAttribute,
 )
 from xdsl.irdl import (
     ConstraintVar,
@@ -77,6 +79,22 @@ class PrecisionAttr(EnumAttribute[Precision], SpacedOpaqueSyntaxAttribute):
     name = "stablehlo.precision"
 
 
+@irdl_attr_definition
+class TokenType(TypeAttribute, ParametrizedAttribute):
+    """
+    Token types represent tokens, i.e. opaque values produced and consumed by some operations.
+    Tokens are used for imposing execution order on operations as described in the Execution section.
+
+    E.g.,
+
+      // %input0: !stablehlo.token
+      // %input1: !stablehlo.token
+      %result = "stablehlo.after_all"(%input0, %input1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
+    """
+
+    name = "stablehlo.token"
+
+
 # endregion
 
 
@@ -277,11 +295,12 @@ def verify_(self) -> None:
         AddOp,
         AndOp,
         MultiplyOp,
-        SubtractOp,
         ReturnOp,
+        SubtractOp,
         TransposeOp,
     ],
     [
         PrecisionAttr,
+        TokenType,
     ],
 )