diff --git a/tests/filecheck/dialects/stablehlo/ops.mlir b/tests/filecheck/dialects/stablehlo/ops.mlir index c0402d73f5..adfd0a96a9 100644 --- a/tests/filecheck/dialects/stablehlo/ops.mlir +++ b/tests/filecheck/dialects/stablehlo/ops.mlir @@ -27,3 +27,6 @@ // [[1,7], [3,9], [5,11]], // [[2,8], [4,10], [6,12]] // ] + +// CHECK: %and = "stablehlo.and"(%t0, %t0) : (tensor, tensor) -> tensor +%and = "stablehlo.and"(%t0, %t0) : (tensor, tensor) -> tensor diff --git a/xdsl/dialects/stablehlo.py b/xdsl/dialects/stablehlo.py index 58b1ce10e9..0b1bdb39de 100644 --- a/xdsl/dialects/stablehlo.py +++ b/xdsl/dialects/stablehlo.py @@ -7,9 +7,9 @@ """ import abc -from typing import Annotated, cast +from typing import Annotated, TypeAlias, cast -from xdsl.dialects.builtin import AnyTensorType, DenseArrayBase, TensorType +from xdsl.dialects.builtin import AnyTensorType, DenseArrayBase, IntegerType, TensorType from xdsl.ir import Attribute, Dialect, SSAValue from xdsl.irdl import ( ConstraintVar, @@ -91,6 +91,37 @@ class AddOp(ElementwiseBinaryOperation): name = "stablehlo.add" +IntegerTensorType: TypeAlias = TensorType[IntegerType] + + +@irdl_op_definition +class AndOp(IRDLOperation): + """ + Performs element-wise AND of two tensors lhs and rhs and produces a result tensor. Depending on the element type, does the following: + + For booleans: logical AND. + For integers: bitwise AND. + + https://github.com/openxla/stablehlo/blob/main/docs/spec.md#and + """ + + name = "stablehlo.and" + + T = Annotated[IntegerTensorType, ConstraintVar("T")] + + lhs = operand_def(T) + rhs = operand_def(T) + + result = result_def(T) + + def __init__( + self, lhs: SSAValue, rhs: SSAValue, result_type: Attribute | None = None + ): + if result_type is None: + result_type = lhs.type + super().__init__(operands=(lhs, rhs), result_types=(result_type,)) + + @irdl_op_definition class MultiplyOp(ElementwiseBinaryOperation): """ @@ -188,6 +219,7 @@ def verify_(self) -> None: [ AbsOp, AddOp, + AndOp, MultiplyOp, SubtractOp, TransposeOp,