core: introduce param_def API for ParametrizedAttribute declaration

tests/irdl/test_attribute_definition.py

@@ -50,6 +50,7 @@
     base,
     irdl_attr_definition,
     irdl_to_attr_constraint,
+    param_def,
 )
 from xdsl.parser import AttrParser
 from xdsl.printer import Printer
@@ -816,6 +817,26 @@ def test_irdl_definition():
     )
 
 
+@irdl_attr_definition
+class ParamAttrDefAttr2(ParametrizedAttribute):
+    name = "test.param_attr_def_attr"
+
+    arg1 = param_def(Attribute)
+    arg2 = param_def(BoolData)
+
+    # Check that we can define methods in attribute definition
+    def test(self):
+        pass
+
+
+def test_irdl_definition2():
+    """Test that we can get the IRDL definition of a parametrized attribute."""
+
+    assert ParamAttrDefAttr2.get_irdl_definition() == ParamAttrDef(
+        "test.param_attr_def_attr", [("arg1", AnyAttr()), ("arg2", BaseAttr(BoolData))]
+    )
+
+
 class InvalidTypedFieldTestAttr(ParametrizedAttribute):
     name = "test.invalid_typed_field"
 
@@ -899,6 +920,50 @@ def test_generic_attr():
     )
 
 
+@irdl_attr_definition
+class GenericAttr2(Generic[AttributeInvT], ParametrizedAttribute):
+    name = "test.generic_attr"
+
+    param = param_def(AttributeInvT)
+
+
+def test_generic_attr2():
+    """Test the generic parameter of a ParametrizedAttribute."""
+
+    assert GenericAttr2.get_irdl_definition() == ParamAttrDef(
+        "test.generic_attr",
+        [
+            (
+                "param",
+                TypeVarConstraint(
+                    type_var=AttributeInvT,
+                    constraint=AnyAttr(),
+                ),
+            )
+        ],
+    )
+
+    assert base(GenericAttr2[IntAttr]) == ParamAttrConstraint(
+        GenericAttr2, (BaseAttr(IntAttr),)
+    )
+
+
+def test_mixed_param_def_apis():
+    """Test that using both ParamDef and param_def raises an error."""
+    with pytest.raises(
+        ValueError,
+        match="ParametrizedAttribute definitions must not mix `param_def` and ParamDef "
+        "declarations.",
+    ):
+
+        @irdl_attr_definition
+        class InvalidAttr(ParametrizedAttribute):  # pyright: ignore[reportUnusedClass]
+            name = "test.invalid"
+            # Using both styles is invalid
+            param1: ParameterDef[IntegerType]  # Using annotation style
+            param2 = param_def(IntegerType)  # Using param_def style
+
+
 ################################################################################
 # ConstraintVar
 ################################################################################

xdsl/dialects/riscv_snitch.py

@@ -40,6 +40,7 @@
 )
 from xdsl.ir import Attribute, Block, Dialect, Operation, Region, SSAValue
 from xdsl.irdl import (
+    BaseAttr,
     VarConstraint,
     attr_def,
     base,
@@ -457,7 +458,9 @@ def assembly_instruction_name(self) -> str:
 class GetStreamOp(RISCVAsmOperation):
     name = "riscv_snitch.get_stream"
 
-    stream = result_def(stream.StreamType[riscv.FloatRegisterType])
+    stream = result_def(
+        stream.StreamType.constr(element_type=BaseAttr(riscv.FloatRegisterType))
+    )
 
     def __init__(self, result_type: Attribute):
         super().__init__(result_types=[result_type])

xdsl/irdl/attributes.py

@@ -14,12 +14,14 @@
     Annotated,
     Any,
     Generic,
+    Literal,
     TypeVar,
     Union,
     cast,
     get_args,
     get_origin,
     get_type_hints,
+    overload,
 )
 
 from xdsl.ir import (
@@ -81,12 +83,56 @@ def generic_constraint_coercion(args: tuple[Any]) -> AttrConstraint:
 
 _A = TypeVar("_A", bound=Attribute)
 
+
+# Will deprecate soon
 ParameterDef = Annotated[_A, IRDLAnnotations.ParamDefAnnot]
 
 
-def irdl_param_attr_get_param_type_hints(cls: type[_A]) -> list[tuple[str, Any]]:
+# Field definition classes for `@irdl_param_attr_definition`
+class _ParameterDef:
+    param: AttrConstraint | Attribute | type[Attribute] | TypeVar | ConstraintVar
+
+    def __init__(
+        self,
+        param: AttrConstraint | Attribute | type[Attribute] | TypeVar | ConstraintVar,
+    ):
+        self.param = param
+
+
+@overload
+def param_def(
+    constraint: TypeVar,
+    *,
+    default: None = None,
+    resolver: None = None,
+    init: Literal[False] = False,
+) -> Attribute: ...
+
+
+@overload
+def param_def(
+    constraint: type[AttributeInvT] | GenericAttrConstraint[AttributeInvT],
+    *,
+    default: None = None,
+    resolver: None = None,
+    init: Literal[False] = False,
+) -> AttributeInvT: ...
+
+
+def param_def(
+    constraint: type[AttributeInvT] | TypeVar | GenericAttrConstraint[AttributeInvT],
+    *,
+    default: None = None,
+    resolver: None = None,
+    init: Literal[False] = False,
+) -> AttributeInvT:
+    """Defines a property of an operation."""
+    return cast(AttributeInvT, _ParameterDef(constraint))
+
+
+def irdl_param_attr_get_param_type_hints(cls: type[_A]) -> dict[str, Any]:
     """Get the type hints of an IRDL parameter definitions."""
-    res = list[tuple[str, Any]]()
+    res: dict[str, Any] = {}
     for field_name, field_type in get_type_hints(cls, include_extras=True).items():
         if field_name == "name" or field_name == "parameters":
             continue
@@ -100,7 +146,7 @@ def irdl_param_attr_get_param_type_hints(cls: type[_A]) -> list[tuple[str, Any]]
                 + f"type `ParameterDef[<Constraint>]`, got type {field_type}."
             )
 
-        res.append((field_name, field_type))
+        res[field_name] = field_type
     return res
 
 
@@ -133,16 +179,28 @@ def from_pyrdl(
             for key, value in parent_cls.__dict__.items()
             if key not in _PARAMETRIZED_ATTRIBUTE_DICT_KEYS
         }
+        # Fields of the class with type annotations
+        param_hints = irdl_param_attr_get_param_type_hints(pyrdl_def)
+
+        # The resulting parameters
+        parameters = list[tuple[str, AttrConstraint]]()
 
-        # Check that all fields of the attribute definition are either already
-        # in ParametrizedAttribute, or are class functions or methods.
         for field_name, value in clsdict.items():
             if field_name == "name":
+                # Ignore name field
                 continue
             if isinstance(
                 value, FunctionType | PropertyType | classmethod | staticmethod
             ):
+                # Ignore functions
+                continue
+
+            # Parameter def must be a field def
+            if isinstance(value, _ParameterDef):
+                constraint = irdl_to_attr_constraint(value.param, allow_type_var=True)
+                parameters.append((field_name, constraint))
                 continue
+
             # Constraint variables are allowed
             if get_origin(value) is Annotated:
                 if any(isinstance(arg, ConstraintVar) for arg in get_args(value)):
@@ -160,17 +218,13 @@ def from_pyrdl(
 
         name = clsdict["name"]
 
-        param_hints = irdl_param_attr_get_param_type_hints(pyrdl_def)
         if issubclass(pyrdl_def, TypedAttribute):
             pyrdl_def = cast(type[TypedAttribute[Attribute]], pyrdl_def)
-            try:
-                param_names = [name for name, _ in param_hints]
-                type_index = param_names.index("type")
-            except ValueError:
+            if "type" not in param_hints:
                 raise PyRDLAttrDefinitionError(
                     f"TypedAttribute {pyrdl_def.__name__} should have a 'type' parameter."
                 )
-            typed_hint = param_hints[type_index][1]
+            typed_hint = param_hints["type"]
             if get_origin(typed_hint) is Annotated:
                 typed_hint = get_args(typed_hint)[0]
             type_var = get_type_var_mapping(pyrdl_def)[1][AttributeCovT]
@@ -181,8 +235,13 @@ def from_pyrdl(
                     " as the type variable in the TypedAttribute base class."
                 )
 
-        parameters = list[tuple[str, AttrConstraint]]()
-        for param_name, param_type in param_hints:
+        if parameters and param_hints:
+            raise ValueError(
+                "ParametrizedAttribute definitions must not mix `param_def` and "
+                "ParamDef declarations."
+            )
+
+        for param_name, param_type in param_hints.items():
             constraint = irdl_to_attr_constraint(param_type, allow_type_var=True)
             parameters.append((param_name, constraint))
 
@@ -310,7 +369,8 @@ def irdl_to_attr_constraint(
     allow_type_var: bool = False,
 ) -> AttrConstraint:
     if isinstance(irdl, GenericAttrConstraint):
-        return cast(AttrConstraint, irdl)
+        constr: GenericAttrConstraint[Attribute] = irdl
+        return constr
 
     if isinstance(irdl, Attribute):
         return EqAttrConstraint(irdl)
@@ -357,9 +417,8 @@ def irdl_to_attr_constraint(
         args = get_args(irdl)
         if len(args) != 1:
             raise Exception(f"GenericData args must have length 1, got {args}")
-        origin = cast(type[GenericData[Any]], origin)
-        args = cast(tuple[Attribute], args)
-        return AllOf((BaseAttr(origin), origin.generic_constraint_coercion(args)))
+        cls: type[GenericData[Attribute]] = origin
+        return AllOf((BaseAttr(cls), origin.generic_constraint_coercion(args)))
 
     # Generic ParametrizedAttributes case
     # We translate it to constraints over the attribute parameters.
@@ -383,7 +442,10 @@ def irdl_to_attr_constraint(
 
         type_var_mapping = dict(zip(generic_args, args))
 
-        origin_parameters = irdl_param_attr_get_param_type_hints(origin)
+        # Map the constraints in the attribute definition
+        attr_def = origin.get_irdl_definition()
+        origin_parameters = attr_def.parameters
+
         origin_constraints = [
             irdl_to_attr_constraint(param, allow_type_var=True).mapping_type_vars(
                 type_var_mapping

xdsl/transforms/experimental/convert_stencil_to_ll_mlir.py

@@ -49,6 +49,7 @@
     PatternRewriteWalker,
     RewritePattern,
     TypeConversionPattern,
+    attr_constr_rewrite_pattern,
     attr_type_rewrite_pattern,
     op_type_rewrite_pattern,
 )
@@ -649,7 +650,7 @@ def return_target_analysis(module: builtin.ModuleOp):
 
 
 class StencilTypeConversion(TypeConversionPattern):
-    @attr_type_rewrite_pattern
+    @attr_constr_rewrite_pattern(StencilTypeConstr)
     def convert_type(self, typ: StencilType[Attribute]) -> MemRefType[Attribute]:
         return StencilToMemRefType(typ)