Memref

@irdl_op_definition
class LoadOp(IRDLOperation):
    name = "memref.load"

    T: ClassVar = VarConstraint("T", AnyAttr())

    nontemporal = opt_prop_def(BoolAttr, default_value=BoolAttr.from_bool(False))

    memref = operand_def(MemRefType.constr(T))
    indices = var_operand_def(IndexType())
    res = result_def(T)

    traits = traits_def(MemoryReadEffect())

    irdl_options = (ParsePropInAttrDict(),)
    assembly_format = "$memref `[` $indices `]` attr-dict `:` type($memref)"

    # TODO varargs for indexing, which must match the memref dimensions
    # Problem: memref dimensions require variadic type parameters,
    # which is subject to change

    def verify_(self):
        memref_type = self.memref.type
        if not isinstance(memref_type, MemRefType):
            raise VerifyException("expected a memreftype")

        memref_type = cast(MemRefType, memref_type)

        if memref_type.get_num_dims() != len(self.indices):
            raise Exception("expected an index for each dimension")

    @classmethod
    def get(
        cls, ref: SSAValue | Operation, indices: Sequence[SSAValue | Operation]
    ) -> Self:
        ssa_value = SSAValue.get(ref, type=MemRefType)
        return cls(operands=[ref, indices], result_types=[ssa_value.type.element_type])

@irdl_op_definition
class StoreOp(IRDLOperation):
    T: ClassVar = VarConstraint("T", AnyAttr())

    name = "memref.store"

    nontemporal = opt_prop_def(BoolAttr, default_value=BoolAttr.from_bool(False))

    value = operand_def(T)
    memref = operand_def(MemRefType.constr(T))
    indices = var_operand_def(IndexType())

    traits = traits_def(MemoryWriteEffect())

    irdl_options = (ParsePropInAttrDict(),)
    assembly_format = "$value `,` $memref `[` $indices `]` attr-dict `:` type($memref)"

    def verify_(self):
        if not isinstance(memref_type := self.memref.type, MemRefType):
            raise VerifyException("expected a memreftype")

        memref_type = cast(MemRefType, memref_type)

        if memref_type.get_num_dims() != len(self.indices):
            raise Exception("Expected an index for each dimension")

    @classmethod
    def get(
        cls,
        value: Operation | SSAValue,
        ref: Operation | SSAValue,
        indices: Sequence[Operation | SSAValue],
    ) -> Self:
        return cls(operands=[value, ref, indices])

class AllocOpHasCanonicalizationPatterns(HasCanonicalizationPatternsTrait):
    @classmethod
    def get_canonicalization_patterns(cls) -> tuple[RewritePattern, ...]:
        from xdsl.transforms.canonicalization_patterns.memref import ElideUnusedAlloc

        return (ElideUnusedAlloc(),)

@irdl_op_definition
class AllocOp(IRDLOperation):
    name = "memref.alloc"

    dynamic_sizes = var_operand_def(IndexType)
    symbol_operands = var_operand_def(IndexType)

    memref = result_def(MemRefType)

    # TODO how to constraint the IntegerAttr type?
    alignment = opt_prop_def(IntegerAttr)

    irdl_options = (AttrSizedOperandSegments(as_property=True),)

    traits = traits_def(AllocOpHasCanonicalizationPatterns(), MemoryAllocEffect())

    def __init__(
        self,
        dynamic_sizes: Sequence[SSAValue],
        symbol_operands: Sequence[SSAValue],
        result_type: Attribute,
        alignment: Attribute | None = None,
    ):
        super().__init__(
            operands=(dynamic_sizes, symbol_operands),
            result_types=(result_type,),
            properties={"alignment": alignment},
        )

    @classmethod
    def get(
        cls,
        return_type: Attribute,
        alignment: int | IntegerAttr | None = None,
        shape: Iterable[int | IntAttr] | None = None,
        dynamic_sizes: Sequence[SSAValue | Operation] | None = None,
        layout: MemRefLayoutAttr | NoneAttr = NoneAttr(),
        memory_space: Attribute = NoneAttr(),
    ) -> Self:
        if shape is None:
            shape = [1]

        if dynamic_sizes is None:
            dynamic_sizes = []

        if isinstance(alignment, int):
            alignment = IntegerAttr(alignment, 64)

        return cls(
            tuple(SSAValue.get(ds) for ds in dynamic_sizes),
            (),
            MemRefType(return_type, shape, layout, memory_space),
            alignment,
        )

    def verify_(self) -> None:
        memref_type = self.memref.type

        dyn_dims = [x for x in memref_type.shape.data if x.data == DYNAMIC_INDEX]
        if len(dyn_dims) != len(self.dynamic_sizes):
            raise VerifyException(
                "op dimension operand count does not equal memref dynamic dimension count."
            )
        _verify_memref_alloc_alignment(self.alignment)

    def print(self, printer: Printer):
        printer.print_string("(")
        printer.print_list(self.dynamic_sizes, printer.print_ssa_value)
        printer.print_string(")")
        if self.symbol_operands:
            printer.print_string("[")
            printer.print_list(self.symbol_operands, printer.print_ssa_value)
            printer.print_string("]")

        printer.print_op_attributes(
            self.properties | self.attributes,
            print_keyword=False,
            reserved_attr_names="operandSegmentSizes",
        )

        printer.print_string(" : ")
        printer.print_attribute(self.memref.type)

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        #  %alloc = memref.alloc(%a)[%s] {alignment = 64 : i64} : memref<3x2xf32>

        unresolved_dynamic_sizes = parser.parse_comma_separated_list(
            parser.Delimiter.PAREN, parser.parse_unresolved_operand
        )
        unresolved_symbol_operands = parser.parse_optional_comma_separated_list(
            parser.Delimiter.SQUARE, parser.parse_unresolved_operand
        )
        if unresolved_symbol_operands is None:
            unresolved_symbol_operands = []

        attrs = parser.parse_optional_attr_dict()

        parser.parse_punctuation(":")
        res_type = parser.parse_attribute()

        index = IndexType()
        dynamic_sizes = tuple(
            parser.resolve_operand(uop, index) for uop in unresolved_dynamic_sizes
        )
        symbol_operands = tuple(
            parser.resolve_operand(uop, index) for uop in unresolved_symbol_operands
        )

        if "alignment" in attrs:
            alignment = attrs["alignment"]
            del attrs["alignment"]
        else:
            alignment = None

        op = cls(
            dynamic_sizes,
            symbol_operands,
            res_type,
            alignment,
        )

        op.attributes |= attrs

        return op

@irdl_op_definition
class AllocaScopeOp(IRDLOperation):
    name = "memref.alloca_scope"

    res = var_result_def()

    scope = region_def()

    traits = lazy_traits_def(
        lambda: (
            SingleBlockImplicitTerminator(AllocaScopeReturnOp),
            RecursiveMemoryEffect(),
        )
    )

    @classmethod
    def parse(cls, parser: Parser) -> Self:
        result_types: Sequence[Attribute] = ()
        if parser.parse_optional_punctuation("->"):
            types = parser.parse_optional_comma_separated_list(
                parser.Delimiter.PAREN, parser.parse_type
            )
            result_types = [parser.parse_type()] if types is None else types
        region = parser.parse_region()
        attrs = parser.parse_optional_attr_dict()
        op = cls(result_types=(result_types,), regions=(region,), attributes=attrs)
        for trait in op.get_traits_of_type(SingleBlockImplicitTerminator):
            ensure_terminator(op, trait)
        return op

    def print(self, printer: Printer):
        if self.res:
            printer.print_string(" -> (")
            printer.print_list(self.res.types, printer.print_attribute)
            printer.print_string(")")
        printer.print_string(" ")
        printer.print_region(
            self.scope,
            print_entry_block_args=False,
            print_block_terminators=bool(self.res),
        )
        printer.print_op_attributes(self.attributes)

@irdl_op_definition
class AllocaScopeReturnOp(IRDLOperation):
    name = "memref.alloca_scope.return"

    ops = var_operand_def()

    traits = traits_def(IsTerminator(), HasParent(AllocaScopeOp))
    assembly_format = "attr-dict ($ops^ `:` type($ops))?"

    def verify_(self) -> None:
        parent = cast(AllocaScopeOp, self.parent_op())
        if self.ops.types != parent.result_types:
            raise VerifyException(
                "Expected operand types to match parent's return types."
            )

@irdl_op_definition
class AllocaOp(IRDLOperation):
    name = "memref.alloca"

    dynamic_sizes = var_operand_def(IndexType)
    symbol_operands = var_operand_def(IndexType)

    memref = result_def(MemRefType)

    # TODO how to constraint the IntegerAttr type?
    alignment = opt_prop_def(IntegerAttr)

    traits = traits_def(MemoryAllocEffect())

    irdl_options = (AttrSizedOperandSegments(as_property=True), ParsePropInAttrDict())

    assembly_format = """
    `(`$dynamic_sizes`)` (`` `[` $symbol_operands^ `]`)? attr-dict `:` type($memref)
    """

    @staticmethod
    def get(
        return_type: Attribute,
        alignment: int | IntegerAttr | None = None,
        shape: Iterable[int | IntAttr] | None = None,
        dynamic_sizes: Sequence[SSAValue | Operation] | None = None,
        layout: MemRefLayoutAttr | NoneAttr = NoneAttr(),
        memory_space: Attribute = NoneAttr(),
    ) -> AllocaOp:
        if shape is None:
            shape = [1]

        if dynamic_sizes is None:
            dynamic_sizes = []

        if isinstance(alignment, int):
            alignment = IntegerAttr(alignment, 64)

        return AllocaOp.build(
            operands=[dynamic_sizes, []],
            result_types=[MemRefType(return_type, shape, layout, memory_space)],
            properties={
                "alignment": alignment,
            },
        )

    def verify_(self) -> None:
        memref_type = self.memref.type

        dyn_dims = [x for x in memref_type.shape.data if x.data == DYNAMIC_INDEX]
        if len(dyn_dims) != len(self.dynamic_sizes):
            raise VerifyException(
                "op dimension operand count does not equal memref dynamic dimension count."
            )
        _verify_memref_alloc_alignment(self.alignment)

@irdl_op_definition
class AssumeAlignmentOp(IRDLOperation):
    """Assert that a memref has the given byte alignment."""

    name = "memref.assume_alignment"

    T: ClassVar = VarConstraint("T", base(MemRefType))

    memref = operand_def(T)
    alignment = prop_def(IntegerAttr[I32])
    result = result_def(T)

    traits = traits_def(Pure(), SameOperandsAndResultType())

    assembly_format = "$memref `,` $alignment attr-dict `:` type($memref)"

    def __init__(
        self,
        memref: SSAValue | Operation,
        alignment: int | IntegerAttr[I32],
    ) -> None:
        memref = SSAValue.get(memref)
        if isinstance(alignment, int):
            alignment = IntegerAttr(alignment, i32)
        super().__init__(
            operands=[memref],
            result_types=[memref.type],
            properties={"alignment": alignment},
        )

    def verify_(self) -> None:
        alignment = self.alignment.value.data
        if alignment <= 0:
            raise VerifyException(f"Alignment must be positive, got {alignment}")

@irdl_op_definition
class AtomicRMWOp(IRDLOperation):
    name = "memref.atomic_rmw"

    T: ClassVar = VarConstraint("T", AnyFloatConstr | SignlessIntegerConstraint)

    value = operand_def(T)
    memref = operand_def(MemRefType.constr(T))
    indices = var_operand_def(IndexType)

    kind = prop_def(IntegerAttr[I64])

    result = result_def(T)

    traits = traits_def(MemoryWriteEffect(), MemoryReadEffect())

@irdl_op_definition
class DeallocOp(IRDLOperation):
    name = "memref.dealloc"
    memref = operand_def(base(MemRefType) | base(UnrankedMemRefType))

    traits = traits_def(MemoryFreeEffect())

    @staticmethod
    def get(operand: Operation | SSAValue) -> DeallocOp:
        return DeallocOp.build(operands=[operand])

    assembly_format = "$memref attr-dict `:` type($memref)"

@irdl_op_definition
class GetGlobalOp(IRDLOperation):
    name = "memref.get_global"
    memref = result_def(MemRefType)
    name_ = prop_def(SymbolRefAttr, prop_name="name")

    traits = traits_def(NoMemoryEffect())

    assembly_format = "$name `:` type($memref) attr-dict"

    def __init__(self, name: str | SymbolRefAttr, return_type: Attribute):
        if isinstance(name, str):
            name = SymbolRefAttr(name)
        super().__init__(result_types=[return_type], properties={"name": name})

@irdl_op_definition
class GlobalOp(IRDLOperation):
    name = "memref.global"

    sym_name = prop_def(SymbolNameConstraint())
    sym_visibility = prop_def(StringAttr)
    type = prop_def(MemRefType)
    initial_value = prop_def(UnitAttr | DenseIntOrFPElementsAttr)
    constant = opt_prop_def(UnitAttr)
    alignment = opt_prop_def(IntegerAttr[I64])

    traits = traits_def(SymbolOpInterface(), MemoryAllocEffect())

    def verify_(self) -> None:
        if self.alignment is not None:
            alignment_value = self.alignment.value.data
            # Alignment has to be a power of two
            if not (is_power_of_two(alignment_value)):
                raise VerifyException(
                    f"Alignment attribute {alignment_value} is not a power of 2"
                )

    @staticmethod
    def get(
        sym_name: StringAttr,
        sym_type: Attribute,
        initial_value: Attribute,
        sym_visibility: StringAttr = StringAttr("private"),
        constant: UnitAttr | None = None,
        alignment: int | IntegerAttr[IntegerType] | None = None,
    ) -> GlobalOp:
        if isinstance(alignment, int):
            alignment = IntegerAttr(alignment, 64)

        return GlobalOp.build(
            properties={
                "sym_name": sym_name,
                "type": sym_type,
                "initial_value": initial_value,
                "sym_visibility": sym_visibility,
                "constant": constant,
                "alignment": alignment,
            }
        )