From c9c895cd30f07c527bc82913e2446e22e83b7aad Mon Sep 17 00:00:00 2001 From: Nicolai Stawinoga Date: Tue, 21 Apr 2026 12:26:30 +0200 Subject: [PATCH] irdl: add custom directive for attr assembly format Add AttrCustomDirective ABC and irdl_attr_custom_directive decorator for user-defined assembly format directives in attribute/type formats, mirroring the existing CustomDirective system for operations. --- .../test_attr_declarative_assembly_format.py | 60 ++++++++++++++++++- xdsl/irdl/declarative_assembly_format.py | 34 +++++++++++ .../declarative_assembly_format_parser.py | 24 ++++++++ 3 files changed, 117 insertions(+), 1 deletion(-) diff --git a/tests/irdl/test_attr_declarative_assembly_format.py b/tests/irdl/test_attr_declarative_assembly_format.py index 1d00096507..d902acdd50 100644 --- a/tests/irdl/test_attr_declarative_assembly_format.py +++ b/tests/irdl/test_attr_declarative_assembly_format.py @@ -16,7 +16,14 @@ ) from xdsl.ir import Attribute, ParametrizedAttribute, TypeAttribute from xdsl.irdl import ParamAttrDef, irdl_attr_definition, param_def -from xdsl.irdl.declarative_assembly_format import AttrFormatProgram +from xdsl.irdl.declarative_assembly_format import ( + AttrCustomDirective, + AttrFormatProgram, + AttrParsingState, + ParameterVariable, + PrintingState, + irdl_attr_custom_directive, +) from xdsl.parser import AttrParser, Parser from xdsl.printer import Printer from xdsl.utils.exceptions import ParseError, PyRDLAttrDefinitionError @@ -686,6 +693,57 @@ class StructOptType(ParametrizedAttribute, TypeAttribute): check_roundtrip(StructOptType, "req = i32", ctx) +# ============================================================================ +# Integration tests — custom directive +# ============================================================================ + + +def test_custom_directive(): + """Custom directive with two parameters.""" + + @irdl_attr_custom_directive + class SwapDirective(AttrCustomDirective): + """Prints two parameters in reversed order.""" + + first: ParameterVariable + second: ParameterVariable + + def parse(self, parser: AttrParser, state: AttrParsingState) -> bool: + self.second.parse(parser, state) + parser.parse_punctuation(",") + self.first.parse(parser, state) + return True + + def print( + self, + printer: Printer, + state: PrintingState, + attr: ParametrizedAttribute, + /, + ) -> None: + self.second.print(printer, state, attr) + printer.print_string(", ") + state.should_emit_space = False + state.last_was_punctuation = True + self.first.print(printer, state, attr) + + @irdl_attr_definition + class CustomType(ParametrizedAttribute, TypeAttribute): + name = "test_af.custom" + a: IntegerType = param_def() + b: IntegerType = param_def() + assembly_format = "custom($a, $b)" + custom_directives = (SwapDirective,) + + ctx = Context(allow_unregistered=True) + ctx.load_attr_or_type(CustomType) + attr = CustomType(i32, i64) + printed = print_attr(attr) + assert printed == "!test_af.custom" + parsed = parse_type(printed, ctx) + assert parsed == attr + + # ============================================================================ # Integration tests — printing correctness # ============================================================================ diff --git a/xdsl/irdl/declarative_assembly_format.py b/xdsl/irdl/declarative_assembly_format.py index 5fc1a0f090..2ffcaa7263 100644 --- a/xdsl/irdl/declarative_assembly_format.py +++ b/xdsl/irdl/declarative_assembly_format.py @@ -1698,6 +1698,40 @@ class QualifiedParameterVariable(ParameterVariable): """ +class AttrCustomDirective(AttrFormatDirective, ABC): + """A user defined assembly format directive for attribute/type formats. + + Mirrors CustomDirective for operations. Custom directives can have + multiple parameters, whose types should be declared in the + `parameters` field. + """ + + parameters: ClassVar[dict[str, type[AttrFormatDirective]]] + + +AttrCustomDirectiveInvT = TypeVar("AttrCustomDirectiveInvT", bound=AttrCustomDirective) + + +def irdl_attr_custom_directive( + cls: type[AttrCustomDirectiveInvT], +) -> type[AttrCustomDirectiveInvT]: + """Decorator used on custom attr directives to define the `parameters` + class variable.""" + + cls.parameters = {} + param_types = inspect.get_annotations(cls, eval_str=True) + for field_name, ty in param_types.items(): + if is_const_classvar(field_name, ty, PyRDLError): + continue + if not issubclass(ty, AttrFormatDirective): + raise PyRDLError( + f"Custom attr directive {cls.__name__} has parameter " + f"{field_name} which is not an attr format directive." + ) + cls.parameters[field_name] = ty + return dataclass(frozen=True)(cls) + + @dataclass(frozen=True) class ParamsDirective(AttrFormatDirective): """Captures all parameters of an attribute, printed comma-separated.""" diff --git a/xdsl/irdl/declarative_assembly_format_parser.py b/xdsl/irdl/declarative_assembly_format_parser.py index f85ca9d29b..752f7c37de 100644 --- a/xdsl/irdl/declarative_assembly_format_parser.py +++ b/xdsl/irdl/declarative_assembly_format_parser.py @@ -954,6 +954,8 @@ def parse_format_directive(self) -> AttrFormatDirective: return self.parse_params_directive() if self.parse_optional_keyword("struct"): return self.parse_struct_directive() + if self.parse_optional_keyword("custom"): + return self.parse_attr_custom_directive() self.raise_error(f"unexpected token '{self._current_token.text}'") def parse_variable(self, inside_ref: bool = False) -> ParameterVariable: @@ -1095,3 +1097,25 @@ def _parse_struct_variables(self) -> tuple[ParameterVariable, ...]: def parse_struct_directive(self) -> StructDirective: return StructDirective(self._parse_struct_variables()) + + def parse_attr_custom_directive(self) -> AttrFormatDirective: + """Parse `custom` `<` name `>` `(` args `)` for attr format.""" + with self.in_angle_brackets(): + name = self.parse_identifier() + if name not in self.attr_def.custom_directives: + self.raise_error(f"Custom attr directive '{name}' cannot be found.") + directive = self.attr_def.custom_directives[name] + param_types = directive.parameters + params: list[AttrFormatDirective] = [] + with self.in_parens(): + for i, (field_name, ty) in enumerate(param_types.items()): + if i: + self.parse_punctuation(",") + param = self.parse_variable() + if not isinstance(param, ty): + self.raise_error( + f"{name}.{field_name} was expected to be of type " + f"{ty.__name__}, but got {type(param).__name__}" + ) + params.append(param) + return directive(*params)