Acc

class DeviceType(StrEnum):
    NONE = "none"
    STAR = "star"
    DEFAULT = "default"
    HOST = "host"
    MULTICORE = "multicore"
    NVIDIA = "nvidia"
    RADEON = "radeon"

class ClauseDefaultValue(StrEnum):
    PRESENT = "present"
    NONE = "none"

class DataClause(StrEnum):
    """
    OpenACC data clause names (decomposed form).

    The original OpenACC `copy`/`copyout`/`copyin(readonly)`/etc. clauses
    are decomposed into individual `acc` ops. This enum keeps track of
    which user-level clause an op was generated for. See upstream
    `mlir::acc::DataClause`.
    """

    ACC_COPYIN = "acc_copyin"
    ACC_COPYIN_READONLY = "acc_copyin_readonly"
    ACC_COPY = "acc_copy"
    ACC_COPYOUT = "acc_copyout"
    ACC_COPYOUT_ZERO = "acc_copyout_zero"
    ACC_PRESENT = "acc_present"
    ACC_CREATE = "acc_create"
    ACC_CREATE_ZERO = "acc_create_zero"
    ACC_DELETE = "acc_delete"
    ACC_ATTACH = "acc_attach"
    ACC_DETACH = "acc_detach"
    ACC_NO_CREATE = "acc_no_create"
    ACC_PRIVATE = "acc_private"
    ACC_FIRSTPRIVATE = "acc_firstprivate"
    ACC_DEVICEPTR = "acc_deviceptr"
    ACC_GETDEVICEPTR = "acc_getdeviceptr"
    ACC_UPDATE_HOST = "acc_update_host"
    ACC_UPDATE_SELF = "acc_update_self"
    ACC_UPDATE_DEVICE = "acc_update_device"
    ACC_USE_DEVICE = "acc_use_device"
    ACC_REDUCTION = "acc_reduction"
    ACC_DECLARE_DEVICE_RESIDENT = "acc_declare_device_resident"
    ACC_DECLARE_LINK = "acc_declare_link"
    ACC_CACHE = "acc_cache"
    ACC_CACHE_READONLY = "acc_cache_readonly"

class DataClauseModifier(StrEnum):
    """
    Bit flags carried by a data clause op (zero / readonly / always*).

    See upstream `mlir::acc::DataClauseModifier`. Modeled as a bit-enum;
    the empty set is rendered as `none`.
    """

    ZERO = "zero"
    READONLY = "readonly"
    ALWAYSIN = "alwaysin"
    ALWAYSOUT = "alwaysout"
    CAPTURE = "capture"

class VariableTypeCategory(StrEnum):
    """
    Bit flags describing the OpenACC type category of a variable.

    See upstream `mlir::acc::VariableTypeCategory`. Used by the
    `MappableTypeInterface`/`PointerLikeTypeInterface` machinery; the
    empty set is rendered as `uncategorized`.
    """

    SCALAR = "scalar"
    ARRAY = "array"
    COMPOSITE = "composite"
    NONSCALAR = "nonscalar"

class ReductionOpKind(StrEnum):
    """
    Built-in reduction operators supported by OpenACC.

    See upstream `mlir::acc::ReductionOperator` (renamed in xDSL to avoid
    clashing with the `Operator` / `Operation` naming used elsewhere).
    Carried on `acc.reduction.recipe` to identify which reduction the recipe
    encodes.
    """

    NONE = "none"
    ADD = "add"
    MUL = "mul"
    MAX = "max"
    MIN = "min"
    IAND = "iand"
    IOR = "ior"
    XOR = "xor"
    EQV = "eqv"
    NEQV = "neqv"
    LAND = "land"
    LOR = "lor"

class LoopParMode(StrEnum):
    """
    Loop parallelism determination mode for `acc.loop` builders.

    See upstream `mlir::acc::LoopParMode`. Used by Python builders to pick
    between the `seq` / `independent` / `auto` attributes; the enum itself
    is not stored on the op.
    """

    SEQ = "loop_seq"
    AUTO = "loop_auto"
    INDEPENDENT = "loop_independent"

class GangArgType(StrEnum):
    """
    Differentiates `num=` / `dim=` / `static=` values inside an
    `acc.loop` `gang(...)` clause.

    See upstream `mlir::acc::GangArgType`. The string spellings match
    upstream's `Num` / `Dim` / `Static` symbol names so the printed
    `#acc.gang_arg_type<...>` attribute round-trips.
    """

    NUM = "Num"
    DIM = "Dim"
    STATIC = "Static"

class CombinedConstructsType(StrEnum):
    """
    Identifies which combined construct an `acc.loop` was decomposed
    from (`kernels loop`, `parallel loop`, `serial loop`).

    See upstream `mlir::acc::CombinedConstructsType`. Compute constructs
    (`acc.parallel` / `acc.serial` / `acc.kernels`) carry just a
    `combined` UnitAttr indicating they were a combined `... loop`
    construct; `acc.loop` carries the typed `CombinedConstructsTypeAttr`
    so the kind is recoverable.
    """

    KERNELS_LOOP = "kernels_loop"
    PARALLEL_LOOP = "parallel_loop"
    SERIAL_LOOP = "serial_loop"

class ParLevel(StrEnum):
    """
    Parallelism level used by `acc.specialized_routine` to mark a
    function variant for a specific OpenACC level.

    See upstream `mlir::acc::ParLevel`. `gang_dim1` is the default gang
    level (equivalent to a bare `gang`); `gang_dim2` / `gang_dim3` cover
    `gang(dim:2)` / `gang(dim:3)`. The enum is consumed by
    `SpecializedRoutineAttr` — it is not stored on any op directly.
    """

    SEQ = "seq"
    GANG_DIM1 = "gang_dim1"
    GANG_DIM2 = "gang_dim2"
    GANG_DIM3 = "gang_dim3"
    WORKER = "worker"
    VECTOR = "vector"

@irdl_attr_definition
class DeviceTypeAttr(EnumAttribute[DeviceType]):
    """
    Device type attribute used to associate values of clauses with a specific
    device_type. Prints using the pretty form `#acc.device_type<value>` to
    match upstream MLIR (which defines `assemblyFormat = "`<` $value `>`"`).
    """

    name = "acc.device_type"

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> DeviceType:
        with parser.in_angle_brackets():
            return parser.parse_str_enum(DeviceType)

@irdl_attr_definition
class ClauseDefaultValueAttr(
    EnumAttribute[ClauseDefaultValue], SpacedOpaqueSyntaxAttribute
):
    """
    Default clause value attribute, selecting either `none` or `present`.
    See upstream `acc.defaultvalue`.
    """

    name = "acc.defaultvalue"

@irdl_attr_definition
class DataClauseAttr(EnumAttribute[DataClause], SpacedOpaqueSyntaxAttribute):
    """
    OpenACC `#acc<data_clause ...>` attribute. Carried on every data-clause op
    so consumers can recover which user clause (e.g. `acc_copy`) the op was
    decomposed from. See upstream `acc.data_clause`.
    """

    name = "acc.data_clause"

@irdl_attr_definition
class DataClauseModifierAttr(
    BitEnumAttribute[DataClauseModifier], SpacedOpaqueSyntaxAttribute
):
    """
    Bit-enum attribute for data-clause modifiers (`zero` / `readonly` /
    `alwaysin` / `alwaysout` / `capture`). The empty set prints as `none`.
    See upstream `acc.data_clause_modifier`.
    """

    name = "acc.data_clause_modifier"
    none_value = "none"
    separator_value = ","
    delimiter_value = AttrParser.Delimiter.NONE

@irdl_attr_definition
class VariableTypeCategoryAttr(
    BitEnumAttribute[VariableTypeCategory], SpacedOpaqueSyntaxAttribute
):
    """
    Bit-enum attribute classifying a variable's type per the OpenACC spec.
    Empty set prints as `uncategorized`. See upstream
    `acc.variable_type_category`.
    """

    name = "acc.variable_type_category"
    none_value = "uncategorized"
    separator_value = ","
    delimiter_value = AttrParser.Delimiter.NONE

@irdl_attr_definition
class ReductionOpKindAttr(EnumAttribute[ReductionOpKind]):
    """
    Reduction operator attribute carried by `acc.reduction.recipe`. Prints
    using the pretty form `#acc.reduction_operator<value>` to match upstream
    MLIR (which defines `assemblyFormat = "`<` $value `>`"`). When referenced
    as `$reductionOperator` in an op assembly format, xDSL prints just the
    `<value>` parameter — matching upstream's inline spelling
    `reduction_operator <add>`.
    """

    name = "acc.reduction_operator"

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> ReductionOpKind:
        with parser.in_angle_brackets():
            return parser.parse_str_enum(ReductionOpKind)

@irdl_attr_definition
class GangArgTypeAttr(EnumAttribute[GangArgType]):
    """
    Gang arg type attribute distinguishing `num=` / `dim=` / `static=` values
    inside `acc.loop`'s `gang(...)` clause. Prints using the pretty form
    `#acc.gang_arg_type<value>` to match upstream MLIR (which defines
    `assemblyFormat = "`<` $value `>`"`).
    """

    name = "acc.gang_arg_type"

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> GangArgType:
        with parser.in_angle_brackets():
            return parser.parse_str_enum(GangArgType)

@irdl_attr_definition
class CombinedConstructsTypeAttr(EnumAttribute[CombinedConstructsType]):
    """
    Combined-constructs attribute carried on `acc.loop` to identify the
    user-level `kernels loop` / `parallel loop` / `serial loop` it was
    decomposed from. Defaulted to no value when the loop stands alone.
    Prints using the pretty form `#acc.combined_constructs<value>` to
    match upstream MLIR (whose `EnumAttr` default
    `assemblyFormat = "`<` $value `>`"` produces the dot form, *not* the
    spaced-opaque form `#acc<combined_constructs ...>`).
    """

    name = "acc.combined_constructs"

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> CombinedConstructsType:
        with parser.in_angle_brackets():
            return parser.parse_str_enum(CombinedConstructsType)

@irdl_attr_definition
class ParLevelAttr(EnumAttribute[ParLevel]):
    """
    Parallelism level attribute consumed by `SpecializedRoutineAttr` to
    record which OpenACC level (`seq` / `gang_dim*` / `worker` /
    `vector`) a function variant was specialized for. Prints using the
    pretty form `#acc.par_level<value>` to match upstream MLIR (whose
    `EnumAttr` default `assemblyFormat = "`<` $value `>`"` produces the
    dot form). When embedded as the `$level` parameter of
    `#acc.specialized_routine`, the surrounding attribute prints just
    the `<value>` slice — matching upstream's inline spelling
    `<@routine, <gang_dim1>, "foo">`.
    """

    name = "acc.par_level"

    def print_parameter(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string(self.data.value)

    @classmethod
    def parse_parameter(cls, parser: AttrParser) -> ParLevel:
        with parser.in_angle_brackets():
            return parser.parse_str_enum(ParLevel)

@irdl_attr_definition
class VarNameAttr(ParametrizedAttribute):
    """
    Variable-name metadata attribute (upstream `#acc.var_name<"x">`).

    Upstream attaches this purely as a *discardable* attribute on
    non-acc-dialect ops (e.g. `memref.alloca`, casts) to preserve
    source-level variable names through transformations. It is **not**
    an op argument on any acc op — the `name` slot on the data-clause
    family is `OptionalAttr<StrAttr>:$name` upstream, so xDSL keeps the
    plain `StringAttr` there too.
    """

    name = "acc.var_name"

    var_name: StringAttr

    def __init__(self, name: str | StringAttr) -> None:
        if isinstance(name, str):
            name = StringAttr(name)
        super().__init__(name)

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_attribute(self.var_name)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        with parser.in_angle_brackets():
            var_name = parser.parse_str_literal()
        return [StringAttr(var_name)]

@irdl_attr_definition
class RoutineInfoAttr(ParametrizedAttribute):
    """
    Records the `acc.routine` declarations associated with a function
    (upstream `#acc.routine_info<[@rt1, @rt2]>`).

    Upstream attaches this as a *discardable* attribute on `func.func`
    declarations whose names were referenced by `acc routine`
    directives. It is **not** an op argument on any acc op — the
    `acc.routine` op itself holds the per-directive clauses, and this
    attribute is the back-edge from the routine target back to its
    declarations.
    """

    name = "acc.routine_info"

    acc_routines: ArrayAttr[SymbolRefAttr]

    def __init__(
        self, acc_routines: ArrayAttr[SymbolRefAttr] | Sequence[SymbolRefAttr]
    ) -> None:
        if not isinstance(acc_routines, ArrayAttr):
            acc_routines = ArrayAttr(acc_routines)
        super().__init__(acc_routines)

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            with printer.in_square_brackets():
                printer.print_list(self.acc_routines.data, printer.print_attribute)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        with parser.in_angle_brackets():
            refs = parser.parse_comma_separated_list(
                AttrParser.Delimiter.SQUARE, parser.parse_attribute
            )
        for ref in refs:
            if not isinstance(ref, SymbolRefAttr):
                parser.raise_error(
                    f"expected symbol reference in #acc.routine_info, got {ref}"
                )
        return [ArrayAttr(cast(list[SymbolRefAttr], refs))]

@irdl_attr_definition
class SpecializedRoutineAttr(ParametrizedAttribute):
    """
    Marks a function as a device-specialized variant of an `acc.routine`
    (upstream `#acc.specialized_routine<@routine, <par_level>, "name">`).

    Upstream attaches this as a *discardable* attribute on the
    specialized `func.func` produced by the routine-specialization pass.
    It captures the parallelism level, a `SymbolRefAttr` reference back
    to the originating `acc.routine`, and the *original* (pre-rename)
    function name. The `par_level` parameter prints inline as
    `<gang_dim1>` (no `#acc.par_level` prefix) to match upstream's
    spelling.
    """

    name = "acc.specialized_routine"

    routine: SymbolRefAttr
    level: ParLevelAttr
    func_name: StringAttr

    def __init__(
        self,
        routine: SymbolRefAttr | str,
        level: ParLevelAttr | ParLevel,
        func_name: StringAttr | str,
    ) -> None:
        if isinstance(routine, str):
            routine = SymbolRefAttr(routine)
        if isinstance(level, ParLevel):
            level = ParLevelAttr(level)
        if isinstance(func_name, str):
            func_name = StringAttr(func_name)
        super().__init__(routine, level, func_name)

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_attribute(self.routine)
            printer.print_string(", ")
            self.level.print_parameter(printer)
            printer.print_string(", ")
            printer.print_attribute(self.func_name)

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        with parser.in_angle_brackets():
            routine = parser.parse_attribute()
            if not isinstance(routine, SymbolRefAttr):
                parser.raise_error(
                    f"expected symbol reference in #acc.specialized_routine, got {routine}"
                )
            parser.parse_punctuation(",")
            level = ParLevelAttr(ParLevelAttr.parse_parameter(parser))
            parser.parse_punctuation(",")
            func_name = parser.parse_str_literal()
        return [routine, level, StringAttr(func_name)]

@irdl_attr_definition
class DeclareAttr(ParametrizedAttribute):
    """
    Records that a variable is captured by an `acc declare` directive
    (upstream `#acc.declare<dataClause = <clause>[, implicit = true]>`).

    Upstream attaches this as a *discardable* attribute to the variable's
    creation site — typically a global op (`memref.global`,
    `llvm.mlir.global`, …) or its allocation site — to advertise which
    user-level declare clause produced the capture. It pairs with
    `acc.declare_enter` / `acc.declare_exit` ops, which describe *how*
    the data action runs at runtime; the attribute itself just preserves
    the source-level clause kind so analyses can find the variable.

    Parameters mirror upstream: a required `DataClauseAttr` and a
    defaulted `implicit` flag (`false` unless an OpenACC routine /
    standalone-declare lowering produced the capture implicitly). When
    `implicit` is `false` the print form omits it, matching upstream's
    `DefaultValuedParameter<"bool", "false">` behaviour.
    """

    name = "acc.declare"

    data_clause: DataClauseAttr
    implicit: BoolAttr

    def __init__(
        self,
        data_clause: DataClauseAttr | DataClause,
        implicit: BoolAttr | bool = False,
    ) -> None:
        if isinstance(data_clause, DataClause):
            data_clause = DataClauseAttr(data_clause)
        if isinstance(implicit, bool):
            implicit = IntegerAttr.from_bool(implicit)
        super().__init__(data_clause, implicit)

    def print_parameters(self, printer: Printer) -> None:
        with printer.in_angle_brackets():
            printer.print_string("dataClause = ")
            printer.print_string(self.data_clause.data.value)
            if self.implicit.value.data:
                printer.print_string(", implicit = true")

    @classmethod
    def parse_parameters(cls, parser: AttrParser) -> list[Attribute]:
        def parse_field() -> tuple[str, Attribute]:
            key = parser.parse_keyword_in(("dataClause", "implicit"))
            parser.parse_punctuation("=")
            if key == "dataClause":
                return key, DataClauseAttr(parser.parse_str_enum(DataClause))
            return key, IntegerAttr.from_bool(parser.parse_boolean())

        seen: dict[str, Attribute] = {}
        for key, value in parser.parse_comma_separated_list(
            AttrParser.Delimiter.ANGLE, parse_field
        ):
            if key in seen:
                parser.raise_error(f"duplicate struct parameter name: {key}")
            seen[key] = value
        if "dataClause" not in seen:
            parser.raise_error("struct is missing required parameter: dataClause")
        return [seen["dataClause"], seen.get("implicit", IntegerAttr.from_bool(False))]