dialects: (stablehlo) add stablehlo.after_all

tests/filecheck/dialects/stablehlo/ops.mlir

@@ -8,6 +8,11 @@
 // CHECK: %add = "stablehlo.add"(%t0, %t0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
 %add = "stablehlo.add"(%t0, %t0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
 
+%token0 = "test.op"() : () -> !stablehlo.token
+%token1 = "test.op"() : () -> !stablehlo.token
+// CHECK: %after_all = "stablehlo.after_all"(%token0, %token1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
+%after_all = "stablehlo.after_all"(%token0, %token1) : (!stablehlo.token, !stablehlo.token) -> !stablehlo.token
+
 // CHECK: %multiply = "stablehlo.multiply"(%t0, %t0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
 %multiply = "stablehlo.multiply"(%t0, %t0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
 

xdsl/dialects/stablehlo.py

@@ -7,6 +7,7 @@
 """
 
 import abc
+from collections.abc import Sequence
 from typing import Annotated, TypeAlias, cast
 
 from xdsl.dialects.builtin import AnyTensorType, DenseArrayBase, IntegerType, TensorType
@@ -145,6 +146,23 @@ class AddOp(ElementwiseBinaryOperation):
     name = "stablehlo.add"
 
 
+@irdl_op_definition
+class AfterAllOp(IRDLOperation):
+    """
+    Ensures that the operations producing the inputs are executed before any operations that depend on result.
+    Execution of this operation does nothing, it only exists to establish data dependencies from result to inputs.
+
+    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#after_all
+    """
+
+    name = "stablehlo.after_all"
+    inputs = var_operand_def(TokenType)
+    result = result_def(TokenType)
+
+    def __init__(self, inputs: Sequence[SSAValue]):
+        super().__init__(operands=[inputs], result_types=(TokenType(),))
+
+
 IntegerTensorType: TypeAlias = TensorType[IntegerType]
 
 
@@ -318,6 +336,7 @@ def verify_(self) -> None:
     [
         AbsOp,
         AddOp,
+        AfterAllOp,
         AndOp,
         BitcastConvertOp,
         MultiplyOp,