core: add TypeVarConstraint, and use in irdl_to_attr_constraint

docs/irdl.ipynb

@@ -517,6 +517,8 @@
    "source": [
     "# NBVAL_CHECK_OUTPUT\n",
     "\n",
+    "from typing import TypeVar\n",
+    "\n",
     "from dataclasses import dataclass\n",
     "\n",
     "from xdsl.dialects.builtin import ArrayAttr\n",
@@ -525,6 +527,8 @@
     "from xdsl.utils.exceptions import VerifyException\n",
     "from xdsl.utils.hints import isa\n",
     "\n",
+    "from typing_extensions import Self\n",
+    "\n",
     "\n",
     "@dataclass(frozen=True)\n",
     "class ArrayOfConstraint(AttrConstraint):\n",
@@ -545,6 +549,14 @@
     "        for e in attr.data:\n",
     "            self.elem_constr.verify(e, constraint_context)\n",
     "\n",
+    "    def mapping_type_vars(\n",
+    "        self, type_var_mapping: dict[TypeVar, AttrConstraint]\n",
+    "    ) -> Self:\n",
+    "        mapped_constraints = tuple(\n",
+    "            c.mapping_type_vars(type_var_mapping) for c in self.attr_constrs\n",
+    "        )\n",
+    "        return ArrayOfConstraint(mapped_constraints)\n",
+    "\n",
     "\n",
     "array_constraint = ArrayOfConstraint(IntAttr)\n",
     "\n",

tests/dialects/test_builtin.py

@@ -1,13 +1,16 @@
 from collections.abc import Sequence
+from typing import TypeVar
 
 import pytest
 
 from xdsl.dialects.arith import Constant
 from xdsl.dialects.builtin import (
     AnyTensorType,
     ArrayAttr,
+    ArrayOfConstraint,
     BFloat16Type,
     ComplexType,
+    ContainerOf,
     DenseArrayBase,
     DenseIntOrFPElementsAttr,
     Float16Type,
@@ -33,8 +36,13 @@
     i32,
     i64,
 )
-from xdsl.ir import Attribute
-from xdsl.irdl import ConstraintContext
+from xdsl.ir import Attribute, Data
+from xdsl.irdl import (
+    BaseAttr,
+    ConstraintContext,
+    TypeVarConstraint,
+    irdl_attr_definition,
+)
 from xdsl.utils.exceptions import VerifyException
 
 
@@ -301,3 +309,45 @@ def test_strides():
     assert ShapedType.strides_for_shape((1,), factor=2) == (2,)
     assert ShapedType.strides_for_shape((2, 3)) == (3, 1)
     assert ShapedType.strides_for_shape((4, 5, 6), factor=2) == (60, 12, 2)
+
+
+################################################################################
+# Mapping Type Var
+################################################################################
+
+
+@irdl_attr_definition
+class A(Data[int]):
+    name = "a"
+
+
+@irdl_attr_definition
+class B(Data[int]):
+    name = "b"
+
+
+_A = TypeVar("_A", bound=Attribute)
+_B = TypeVar("_B", bound=Attribute)
+
+
+def test_array_of_constraint():
+    """Test mapping type variables in ArrayOfConstraint."""
+    array_constraint = ArrayOfConstraint(TypeVarConstraint(_A, BaseAttr(A)))
+
+    assert array_constraint.mapping_type_vars({}) == ArrayOfConstraint(BaseAttr(A))
+    assert array_constraint.mapping_type_vars({_B: BaseAttr(B)}) == ArrayOfConstraint(
+        BaseAttr(A)
+    )
+    assert array_constraint.mapping_type_vars({_A: BaseAttr(B)}) == ArrayOfConstraint(
+        BaseAttr(B)
+    )
+
+    container_constraint = ContainerOf(TypeVarConstraint(_A, BaseAttr(A)))
+
+    assert container_constraint.mapping_type_vars({}) == ContainerOf(BaseAttr(A))
+    assert container_constraint.mapping_type_vars({_B: BaseAttr(B)}) == ContainerOf(
+        BaseAttr(A)
+    )
+    assert container_constraint.mapping_type_vars({_A: BaseAttr(B)}) == ContainerOf(
+        BaseAttr(B)
+    )

tests/irdl/test_attribute_definition.py

@@ -22,6 +22,7 @@
 )
 from xdsl.ir import (
     Attribute,
+    AttributeInvT,
     BitEnumAttribute,
     Data,
     EnumAttribute,
@@ -31,15 +32,21 @@
     TypedAttribute,
 )
 from xdsl.irdl import (
+    AllOf,
     AnyAttr,
+    AnyOf,
     AttrConstraint,
     BaseAttr,
     ConstraintContext,
     ConstraintVar,
     GenericData,
     MessageConstraint,
+    ParamAttrConstraint,
     ParamAttrDef,
     ParameterDef,
+    TypeVarConstraint,
+    VarConstraint,
+    base,
     irdl_attr_definition,
     irdl_to_attr_constraint,
 )
@@ -375,6 +382,11 @@ def verify(self, attr: Attribute, constraint_context: ConstraintContext) -> None
         if attr.data <= 0:
             raise VerifyException(f"Expected positive integer, got {attr.data}.")
 
+    def mapping_type_vars(
+        self, type_var_mapping: dict[TypeVar, AttrConstraint]
+    ) -> PositiveIntConstr:
+        return self
+
 
 @irdl_attr_definition
 class PositiveIntAttr(ParametrizedAttribute):
@@ -630,9 +642,6 @@ def test_data_with_generic_missing_generic_data_failure():
     )
 
 
-A = TypeVar("A", bound=Attribute)
-
-
 @dataclass(frozen=True)
 class DataListAttr(AttrConstraint):
     """
@@ -651,13 +660,22 @@ def verify(
         for e in attr.data:
             self.elem_constr.verify(e, constraint_context)
 
+    def mapping_type_vars(
+        self, type_var_mapping: dict[TypeVar, AttrConstraint]
+    ) -> DataListAttr:
+        mapped_constraint = self.elem_constr.mapping_type_vars(type_var_mapping)
+        if mapped_constraint is self.elem_constr:
+            return self
+        else:
+            return DataListAttr(mapped_constraint)
+
 
 @irdl_attr_definition
-class ListData(GenericData[list[A]]):
+class ListData(GenericData[list[AttributeInvT]]):
     name = "test.list"
 
     @classmethod
-    def parse_parameter(cls, parser: AttrParser) -> list[A]:
+    def parse_parameter(cls, parser: AttrParser) -> list[AttributeInvT]:
         raise NotImplementedError()
 
     def print_parameter(self, printer: Printer) -> None:
@@ -672,7 +690,7 @@ def generic_constraint_coercion(args: tuple[Any]) -> AttrConstraint:
         return DataListAttr(irdl_to_attr_constraint(args[0]))
 
     @staticmethod
-    def from_list(data: list[A]) -> ListData[A]:
+    def from_list(data: list[AttributeInvT]) -> ListData[AttributeInvT]:
         return ListData(data)
 
 
@@ -835,6 +853,34 @@ def test_custom_constructor():
     assert OveriddenInitAttr.new([StringData("17")]) == OveriddenInitAttr("17")
 
 
+@irdl_attr_definition
+class GenericAttr(Generic[AttributeInvT], ParametrizedAttribute):
+    name = "test.generic_attr"
+
+    param: ParameterDef[AttributeInvT]
+
+
+def test_generic_attr():
+    """Test the generic parameter of a ParametrizedAttribute."""
+
+    assert GenericAttr.get_irdl_definition() == ParamAttrDef(
+        "test.generic_attr",
+        [
+            (
+                "param",
+                TypeVarConstraint(
+                    type_var=AttributeInvT,
+                    constraint=AnyAttr(),
+                ),
+            )
+        ],
+    )
+
+    assert base(GenericAttr[IntAttr]) == ParamAttrConstraint(
+        GenericAttr, (BaseAttr(IntAttr),)
+    )
+
+
 ################################################################################
 # ConstraintVar
 ################################################################################
@@ -870,3 +916,82 @@ def test_constraint_var_fail_not_satisfy_constraint():
         ConstraintVarAttr.new(
             [IntegerAttr(42, IndexType()), IntegerAttr(17, IndexType())]
         )
+
+
+################################################################################
+# Mapping Type Var
+################################################################################
+
+
+@irdl_attr_definition
+class A(Data[int]):
+    name = "a"
+
+
+@irdl_attr_definition
+class B(Data[int]):
+    name = "b"
+
+
+_A = TypeVar("_A", bound=Attribute)
+_B = TypeVar("_B", bound=Attribute)
+
+
+def test_var_constraint():
+    var_constraint = VarConstraint("var", TypeVarConstraint(_A, BaseAttr(A)))
+
+    assert var_constraint.mapping_type_vars({}) == VarConstraint("var", BaseAttr(A))
+    assert var_constraint.mapping_type_vars({_B: BaseAttr(B)}) == VarConstraint(
+        "var", BaseAttr(A)
+    )
+    assert var_constraint.mapping_type_vars({_A: BaseAttr(B)}) == VarConstraint(
+        "var", BaseAttr(B)
+    )
+
+
+def test_typevar_constraint():
+    typevar_constraint = TypeVarConstraint(_A, BaseAttr(A))
+
+    assert typevar_constraint.mapping_type_vars({}) == BaseAttr(A)
+    assert typevar_constraint.mapping_type_vars({_B: BaseAttr(B)}) == BaseAttr(A)
+    assert typevar_constraint.mapping_type_vars({_A: BaseAttr(B)}) == BaseAttr(B)
+
+
+def test_message_constraint():
+    message_constraint = MessageConstraint(
+        TypeVarConstraint(_A, BaseAttr(A)), "test message"
+    )
+
+    assert message_constraint.mapping_type_vars({}) == MessageConstraint(
+        BaseAttr(A), "test message"
+    )
+    assert message_constraint.mapping_type_vars({_B: BaseAttr(B)}) == MessageConstraint(
+        BaseAttr(A), "test message"
+    )
+    assert message_constraint.mapping_type_vars({_A: BaseAttr(B)}) == MessageConstraint(
+        BaseAttr(B), "test message"
+    )
+
+
+def test_anyof_constraint():
+    anyof_constraint = AnyOf((TypeVarConstraint(_A, BaseAttr(A)), BaseAttr(B)))
+
+    assert anyof_constraint.mapping_type_vars({}) == AnyOf((BaseAttr(A), BaseAttr(B)))
+    assert anyof_constraint.mapping_type_vars({_B: BaseAttr(B)}) == AnyOf(
+        (BaseAttr(A), BaseAttr(B))
+    )
+    assert anyof_constraint.mapping_type_vars({_A: BaseAttr(B)}) == AnyOf(
+        (BaseAttr(B), BaseAttr(B))
+    )
+
+
+def test_allof_constraint():
+    allof_constraint = AllOf((TypeVarConstraint(_A, BaseAttr(A)), BaseAttr(B)))
+
+    assert allof_constraint.mapping_type_vars({}) == AllOf((BaseAttr(A), BaseAttr(B)))
+    assert allof_constraint.mapping_type_vars({_B: BaseAttr(B)}) == AllOf(
+        (BaseAttr(A), BaseAttr(B))
+    )
+    assert allof_constraint.mapping_type_vars({_A: BaseAttr(B)}) == AllOf(
+        (BaseAttr(B), BaseAttr(B))
+    )

tests/test_pyrdl.py

@@ -1,8 +1,10 @@
 """Unit tests for IRDL."""
 
 from dataclasses import dataclass
+from typing import TypeVar
 
 import pytest
+from typing_extensions import Self
 
 from xdsl.ir import Attribute, Data, ParametrizedAttribute
 from xdsl.irdl import (
@@ -142,6 +144,11 @@ def verify(
         if attr.data >= self.bound:
             raise VerifyException(f"{attr} should hold a value less than {self.bound}")
 
+    def mapping_type_vars(
+        self, type_var_mapping: dict[TypeVar, AttrConstraint]
+    ) -> Self:
+        return self
+
 
 @dataclass(frozen=True)
 class GreaterThan(AttrConstraint):
@@ -159,6 +166,11 @@ def verify(
                 f"{attr} should hold a value greater than {self.bound}"
             )
 
+    def mapping_type_vars(
+        self, type_var_mapping: dict[TypeVar, AttrConstraint]
+    ) -> Self:
+        return self
+
 
 def test_anyof_verify():
     """

xdsl/dialects/bufferization.py

@@ -1,6 +1,8 @@
 from collections.abc import Sequence, Set
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any, ClassVar, TypeVar
+
+from typing_extensions import Self
 
 from xdsl.dialects.builtin import (
     AnyMemRefTypeConstr,
@@ -18,6 +20,7 @@
 )
 from xdsl.ir import Attribute, Dialect, Operation, SSAValue
 from xdsl.irdl import (
+    AttrConstraint,
     AttrSizedOperandSegments,
     ConstraintContext,
     GenericAttrConstraint,
@@ -71,6 +74,15 @@ def verify(self, attr: Attribute, constraint_context: ConstraintContext) -> None
             )
         return self.memref_constraint.verify(memref_type, constraint_context)
 
+    def mapping_type_vars(
+        self, type_var_mapping: dict[TypeVar, AttrConstraint]
+    ) -> Self:
+        memref_constraint = self.memref_constraint.mapping_type_vars(type_var_mapping)
+        if memref_constraint is self.memref_constraint:
+            return self
+        else:
+            return type(self)(memref_constraint)
+
 
 @irdl_op_definition
 class AllocTensorOp(IRDLOperation):