Shard

@irdl_attr_definition
class ShardAxesArrayAttr(ParametrizedAttribute, OpaqueSyntaxAttribute):
    """
    ShardAxesArrayAttr attribute for representing multiple grid axes.

    Reflects [the MLIR attribute](https://github.com/llvm/llvm-project/blob/fef02d48c08db859ef83f84232ed78bd9d1c323a/mlir/include/mlir/Dialect/Shard/IR/ShardBase.td#L83).
    """

    name = "shard.axisarray"

    axes: ArrayAttr[ShardAxesAttr]

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> Sequence[Attribute]:
        """
        Parses a ShardAxesArrayAttr, which has the syntax of a list
        of lists, e.g.:

        [[1, 2, 3], [], [4, 5]]
        """
        axes = parser.parse_comma_separated_list(
            parser.Delimiter.SQUARE,
            lambda: _parse_grid_axes_attr(parser),
        )

        return (ArrayAttr(axes),)

    def print_parameters(self, printer: Printer) -> None:
        """
        Prints a ShardAxesArrayAttr, which has the syntax of a list
        of lists, e.g.:

        [[1, 2, 3], [], [4, 5]]
        """
        with printer.in_square_brackets():
            printer.print_list(
                self.axes.data,
                lambda x: _print_sublist(printer, x),
            )

class ReductionKind(StrEnum):
    "Reduction kind for grid dialect"

    SUM = auto()
    MAX = auto()
    MIN = auto()
    PRODUCT = auto()
    AVERAGE = auto()
    BITWISE_AND = auto()
    BITWISE_OR = auto()
    BITWISE_XOR = auto()
    GENERIC = auto()

@irdl_attr_definition
class ReductionKindAttr(EnumAttribute[ReductionKind], SpacedOpaqueSyntaxAttribute):
    name = "shard.partial"

    assembly_format = "$value"

@irdl_attr_definition
class ShardingType(ParametrizedAttribute, TypeAttribute):
    name = "shard.sharding"

class CollectiveCommunicationOp(IRDLOperation, ABC):
    """
    Base class for collective communication ops.
    """

    grid = prop_def(FlatSymbolRefAttr)
    grid_axes = prop_def(
        ShardAxesAttr, default_value=ShardAxesAttr(i16, BytesAttr(b""))
    )

@irdl_op_definition
class AllReduceOp(CollectiveCommunicationOp):
    """
    All-reduce over a device grid.

    Within each device group reduce the input using the `reduction` method.
    Each device in a group receives a replicated copy of the reduction result.
    The accumulation element type is determined by the result type.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardall_reduce-shardallreduceop).
    """

    name = "shard.all_reduce"

    input = operand_def(MemRefType.constr() | TensorType.constr())
    reduction = prop_def(
        ReductionKindAttr, default_value=ReductionKindAttr(ReductionKind.SUM)
    )

    result = result_def(MemRefType.constr() | TensorType.constr())

    traits = traits_def(Pure())

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "(`reduction` `=` $reduction^)? "
        + "attr-dict `:` type($input) `->` type($result)"
    )

@irdl_op_definition
class BroadcastOp(CollectiveCommunicationOp):
    """
    Broadcast tensor from one device to many devices.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardbroadcast-shardbroadcastop).
    """

    name = "shard.broadcast"

    input = operand_def(TensorType)
    root = prop_def(DenseArrayBase[I64])
    root_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    traits = traits_def(Pure())

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "`root` `=` custom<DynamicIndexList>($root_dynamic, $root) "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class GatherOp(CollectiveCommunicationOp):
    """
    Gather tensor shards from many devices to a single device.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardgather-shardgatherop).
    """

    name = "shard.gather"

    input = operand_def(TensorType)
    gather_axis = prop_def(IntegerAttr.constr(IndexTypeConstr))
    root = prop_def(DenseArrayBase[I64])
    root_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    traits = traits_def(Pure())

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "`gather_axis` `=` $gather_axis "
        + "`root` `=` custom<DynamicIndexList>($root_dynamic, $root) "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class ScatterOp(CollectiveCommunicationOp):
    """
    Scatter tensor over a device shard.

    For each device group split the input tensor on the `root` device along
    axis `scatter_axis` and scatter the parts across the group devices.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardscatter-shardscatterop).
    """

    name = "shard.scatter"

    input = operand_def(TensorType)
    scatter_axis = prop_def(IntegerAttr.constr(IndexTypeConstr))

    root = prop_def(DenseArrayBase[I64])
    root_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    traits = traits_def(
        Pure(),
    )

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "`scatter_axis` `=` $scatter_axis "
        + "`root` `=` custom<DynamicIndexList>($root_dynamic, $root) "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class RecvOp(CollectiveCommunicationOp):
    """
    Receive from a device within a device group.
    """

    name = "shard.recv"

    input = operand_def(TensorType)
    source = opt_prop_def(DenseArrayBase[I64])
    source_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "(`source` `=` custom<DynamicIndexList>($source_dynamic, $source)^)? "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class ReduceOp(CollectiveCommunicationOp):
    """
    Reduce over a device grid.

    Within each device group reduce the input tensor using the `reduction`
    method. The result is returned on the `root` device of each group and is
    undefined on all other devices.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardreduce-shardreduceop).
    """

    name = "shard.reduce"

    input = operand_def(TensorType)
    reduction = prop_def(
        ReductionKindAttr, default_value=ReductionKindAttr(ReductionKind.SUM)
    )
    root = prop_def(DenseArrayBase[I64])
    root_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    traits = traits_def(Pure())

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "(`reduction` `=` $reduction^)? "
        + "`root` `=` custom<DynamicIndexList>($root_dynamic, $root) "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class SendOp(CollectiveCommunicationOp):
    """
    Send from one device to another within a device group.
    """

    name = "shard.send"

    input = operand_def(TensorType)

    destination = prop_def(DenseArrayBase[I64])
    destination_dynamic = var_operand_def(IndexType)

    result = result_def(TensorType)

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "`destination` `=` custom<DynamicIndexList>($destination_dynamic, $destination) "
        + "attr-dict `:` functional-type(operands, results)"
    )

    custom_directives = (DynamicIndexList,)

@irdl_op_definition
class ShiftOp(CollectiveCommunicationOp):
    """
    Shift over a device shard.

    Within each device group shift along `shift_axis` by `offset`. If the
    `rotate` flag is present a rotation is performed instead of a shift.
    """

    name = "shard.shift"

    input = operand_def(TensorType)

    shift_axis = prop_def(IntegerAttr.constr(IndexTypeConstr))
    offset = prop_def(IntegerAttr[I64])
    rotate = prop_def(UnitAttr)

    result = result_def(TensorType)

    traits = traits_def(
        Pure(),
    )

    assembly_format = (
        "$input `on` $grid (`grid_axes` `=` $grid_axes^)? "
        + "`shift_axis` `=` $shift_axis "
        + "`offset` `=` $offset "
        + "(`rotate` $rotate^)? "
        + "attr-dict `:` type($input) `->` type($result)"
    )

@irdl_op_definition
class GridOp(IRDLOperation):
    name = "shard.grid"

    sym_name = prop_def(SymbolNameConstraint())
    shape = prop_def(DenseArrayBase[I64])

    traits = traits_def(SymbolOpInterface(), Pure())

    assembly_format = (
        "$sym_name `(` `shape` `=` custom<DimensionList>($shape) `)` attr-dict"
    )

    custom_directives = (DimensionList,)

    def verify_(self):
        if not self.shape.get_values():
            raise VerifyException(
                "'shard.grid' op rank of grid is expected to be a positive integer"
            )

@irdl_op_definition
class ShardingOp(IRDLOperation):
    """
    Define the sharding of a tensor.

    Note: `halo_sizes` and `sharded_dims_offsets` are mutually exclusive.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardsharding-shardshardingop)
    """

    name = "shard.sharding"

    grid = prop_def(FlatSymbolRefAttr)
    split_axes = prop_def(ShardAxesArrayAttr)
    static_sharded_dims_offsets = prop_def(
        DenseArrayBase[I64], default_value=DenseArrayBase[I64](i64, BytesAttr(b""))
    )
    dynamic_sharded_dims_offsets = var_operand_def(I64)
    static_halo_sizes = prop_def(
        DenseArrayBase[I64], default_value=DenseArrayBase[I64](i64, BytesAttr(b""))
    )
    dynamic_halo_sizes = var_operand_def(I64)

    result = result_def(ShardingType)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    traits = traits_def(
        Pure(),
    )

    assembly_format = (
        "$grid "
        + "`split_axes` `=` $split_axes "
        + "(`halo_sizes` `=` custom<DynamicIndexList>($dynamic_halo_sizes, $static_halo_sizes)^)? "
        + "(`sharded_dims_offsets` `=` "
        + "custom<DynamicIndexList>($dynamic_sharded_dims_offsets, $static_sharded_dims_offsets)^)? "
        + "attr-dict `:` type($result)"
    )

    custom_directives = (DynamicIndexList,)

    def verify_(self) -> None:
        dims_offsets = (
            self.static_sharded_dims_offsets or self.dynamic_sharded_dims_offsets
        )
        halo_sizes = self.static_halo_sizes or self.dynamic_halo_sizes

        if dims_offsets and halo_sizes:
            raise VerifyException(
                "'shard.sharding' op halo sizes and shard offsets are mutually exclusive"
            )

@irdl_op_definition
class ShardOp(IRDLOperation):
    """
    Annotate on how a tensor is sharded across a shard.

    See [external documentation](https://mlir.llvm.org/docs/Dialects/Shard/#shardshard-shardshardop).
    """

    name = "shard.shard"

    T: ClassVar = VarConstraint("T", TensorType.constr())

    src = operand_def(T)
    sharding = operand_def(ShardingType)
    annotate_for_users = opt_prop_def(UnitAttr)

    result = result_def(T)

    traits = traits_def(
        Pure(),
    )

    assembly_format = "$src `to` $sharding (`annotate_for_users` $annotate_for_users^)? attr-dict `:` type($result)"

    def __init__(
        self,
        src: SSAValue,
        sharding: SSAValue,
        annotate_for_users: UnitAttr | None,
    ):
        return super().__init__(
            operands=[src, sharding],
            result_types=[src.type],
            properties={
                "annotate_for_users": annotate_for_users,
            },
        )