dialects: (stablehlo) add stablehlo.token

xdslproject · Aug 24, 2024 · 3cd9d19 · 3cd9d19
1 parent b55fadc
commit 3cd9d19
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/tests/filecheck/dialects/stablehlo/attrs.mlir b/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:  }
diff --git a/xdsl/dialects/stablehlo.py b/xdsl/dialects/stablehlo.py
@@ -12,11 +12,14 @@
 from xdsl.dialects.builtin import AnyTensorType, DenseArrayBase, IntegerType, TensorType
 from xdsl.ir import (
     Attribute,
+    Data,
     Dialect,
     EnumAttribute,
     SpacedOpaqueSyntaxAttribute,
     SSAValue,
     StrEnum,
+    ParametrizedAttribute,
+    TypeAttribute,
 )
 from xdsl.irdl import (
     ConstraintVar,
@@ -77,6 +80,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 +296,12 @@ def verify_(self) -> None:
         AddOp,
         AndOp,
         MultiplyOp,
-        SubtractOp,
         ReturnOp,
+        SubtractOp,
         TransposeOp,
     ],
     [
         PrecisionAttr,
+        TokenType,
     ],
 )