Skip to content

Commit 62de833

Browse files
alexaricealexarice
authored
core: Add successors to assembly format (#3124)
Adds successors to assembly format. Implementation is very similar to regions.
1 parent d9e7918 commit 62de833

3 files changed

+298
-4
lines changed

‎tests/irdl/test_declarative_assembly_format.py

+179
Original file line numberDiff line numberDiff line change
@@ -38,12 +38,15 @@
3838
opt_prop_def,
3939
opt_region_def,
4040
opt_result_def,
41+
opt_successor_def,
4142
prop_def,
4243
region_def,
4344
result_def,
45+
successor_def,
4446
var_operand_def,
4547
var_region_def,
4648
var_result_def,
49+
var_successor_def,
4750
)
4851
from xdsl.parser import Parser
4952
from xdsl.printer import Printer
@@ -1364,6 +1367,182 @@ class OptionalRegionOp(IRDLOperation):
13641367
check_equivalence(program, generic_program, ctx)
13651368

13661369

1370+
################################################################################
1371+
# Successors #
1372+
################################################################################
1373+
1374+
1375+
def test_missing_successor():
1376+
"""Test that successors should be parsed."""
1377+
with pytest.raises(PyRDLOpDefinitionError, match="successor 'successor' not found"):
1378+
1379+
@irdl_op_definition
1380+
class NoSuccessorOp(IRDLOperation): # pyright: ignore[reportUnusedClass]
1381+
name = "test.no_successor_op"
1382+
successor = successor_def()
1383+
1384+
assembly_format = "attr-dict-with-keyword"
1385+
1386+
1387+
def test_successors():
1388+
"""Test the parsing of successors"""
1389+
1390+
program = textwrap.dedent(
1391+
"""\
1392+
"test.op"() ({
1393+
"test.op"() [^0] : () -> ()
1394+
^0:
1395+
test.two_successors ^0 ^0
1396+
}) : () -> ()"""
1397+
)
1398+
1399+
generic_program = textwrap.dedent(
1400+
"""\
1401+
"test.op"() ({
1402+
"test.op"() [^0] : () -> ()
1403+
^0:
1404+
"test.two_successors"() [^0, ^0] : () -> ()
1405+
}) : () -> ()"""
1406+
)
1407+
1408+
@irdl_op_definition
1409+
class TwoSuccessorsOp(IRDLOperation):
1410+
name = "test.two_successors"
1411+
fst = successor_def()
1412+
snd = successor_def()
1413+
1414+
assembly_format = "$fst $snd attr-dict"
1415+
1416+
ctx = MLContext()
1417+
ctx.load_op(TwoSuccessorsOp)
1418+
ctx.load_dialect(Test)
1419+
1420+
check_roundtrip(program, ctx)
1421+
check_equivalence(program, generic_program, ctx)
1422+
1423+
1424+
@pytest.mark.parametrize(
1425+
"program, generic_program",
1426+
[
1427+
(
1428+
'"test.op"() ({\n "test.op"() [^0] : () -> ()\n^0:\n test.var_successor \n}) : () -> ()',
1429+
textwrap.dedent(
1430+
"""\
1431+
"test.op"() ({
1432+
"test.op"() [^0] : () -> ()
1433+
^0:
1434+
"test.var_successor"() : () -> ()
1435+
}) : () -> ()"""
1436+
),
1437+
),
1438+
(
1439+
textwrap.dedent(
1440+
"""\
1441+
"test.op"() ({
1442+
"test.op"() [^0] : () -> ()
1443+
^0:
1444+
test.var_successor ^0
1445+
}) : () -> ()"""
1446+
),
1447+
textwrap.dedent(
1448+
"""\
1449+
"test.op"() ({
1450+
"test.op"() [^0] : () -> ()
1451+
^0:
1452+
"test.var_successor"() [^0] : () -> ()
1453+
}) : () -> ()"""
1454+
),
1455+
),
1456+
(
1457+
textwrap.dedent(
1458+
"""\
1459+
"test.op"() ({
1460+
"test.op"() [^0] : () -> ()
1461+
^0:
1462+
test.var_successor ^0 ^0
1463+
}) : () -> ()"""
1464+
),
1465+
textwrap.dedent(
1466+
"""\
1467+
"test.op"() ({
1468+
"test.op"() [^0] : () -> ()
1469+
^0:
1470+
"test.var_successor"() [^0, ^0] : () -> ()
1471+
}) : () -> ()"""
1472+
),
1473+
),
1474+
],
1475+
)
1476+
def test_variadic_successor(program: str, generic_program: str):
1477+
"""Test the parsing of successors"""
1478+
1479+
@irdl_op_definition
1480+
class VarSuccessorOp(IRDLOperation):
1481+
name = "test.var_successor"
1482+
succ = var_successor_def()
1483+
1484+
assembly_format = "$succ attr-dict"
1485+
1486+
ctx = MLContext()
1487+
ctx.load_op(VarSuccessorOp)
1488+
ctx.load_dialect(Test)
1489+
1490+
check_roundtrip(program, ctx)
1491+
check_equivalence(program, generic_program, ctx)
1492+
1493+
1494+
@pytest.mark.parametrize(
1495+
"program, generic_program",
1496+
[
1497+
(
1498+
'"test.op"() ({\n "test.op"() [^0] : () -> ()\n^0:\n test.opt_successor \n}) : () -> ()',
1499+
textwrap.dedent(
1500+
"""\
1501+
"test.op"() ({
1502+
"test.op"() [^0] : () -> ()
1503+
^0:
1504+
"test.opt_successor"() : () -> ()
1505+
}) : () -> ()"""
1506+
),
1507+
),
1508+
(
1509+
textwrap.dedent(
1510+
"""\
1511+
"test.op"() ({
1512+
"test.op"() [^0] : () -> ()
1513+
^0:
1514+
test.opt_successor ^0
1515+
}) : () -> ()"""
1516+
),
1517+
textwrap.dedent(
1518+
"""\
1519+
"test.op"() ({
1520+
"test.op"() [^0] : () -> ()
1521+
^0:
1522+
"test.opt_successor"() [^0] : () -> ()
1523+
}) : () -> ()"""
1524+
),
1525+
),
1526+
],
1527+
)
1528+
def test_optional_successor(program: str, generic_program: str):
1529+
"""Test the parsing of successors"""
1530+
1531+
@irdl_op_definition
1532+
class OptSuccessorOp(IRDLOperation):
1533+
name = "test.opt_successor"
1534+
succ = opt_successor_def()
1535+
1536+
assembly_format = "$succ attr-dict"
1537+
1538+
ctx = MLContext()
1539+
ctx.load_op(OptSuccessorOp)
1540+
ctx.load_dialect(Test)
1541+
1542+
check_roundtrip(program, ctx)
1543+
check_equivalence(program, generic_program, ctx)
1544+
1545+
13671546
################################################################################
13681547
# Inference #
13691548
################################################################################

‎xdsl/irdl/declarative_assembly_format.py

+78-4
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
IRDLOperationInvT,
2727
OpDef,
2828
OptionalDef,
29+
Successor,
2930
VariadicDef,
3031
VarIRConstruct,
3132
)
@@ -50,19 +51,17 @@ class ParsingState:
5051
operand_types: list[Attribute | None | list[Attribute | None]]
5152
result_types: list[Attribute | None | list[Attribute | None]]
5253
regions: list[Region | None | list[Region]]
54+
successors: list[Successor | None | list[Successor]]
5355
attributes: dict[str, Attribute]
5456
properties: dict[str, Attribute]
5557
constraint_context: ConstraintContext
5658

5759
def __init__(self, op_def: OpDef):
58-
if op_def.successors:
59-
raise NotImplementedError(
60-
"Operation definitions with successors are not yet supported"
61-
)
6260
self.operands = [None] * len(op_def.operands)
6361
self.operand_types = [None] * len(op_def.operands)
6462
self.result_types = [None] * len(op_def.results)
6563
self.regions = [None] * len(op_def.regions)
64+
self.successors = [None] * len(op_def.successors)
6665
self.attributes = {}
6766
self.properties = {}
6867
self.constraint_context = ConstraintContext()
@@ -168,6 +167,7 @@ def parse(
168167
attributes=state.attributes,
169168
properties=properties,
170169
regions=state.regions,
170+
successors=state.successors,
171171
)
172172

173173
def assign_constraint_variables(
@@ -810,6 +810,80 @@ def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> No
810810
state.should_emit_space = True
811811

812812

813+
class SuccessorVariable(VariableDirective, OptionallyParsableDirective):
814+
"""
815+
A successor variable, with the following format:
816+
successor-directive ::= dollar-ident
817+
The directive will request a space to be printed after.
818+
"""
819+
820+
def parse_optional(self, parser: Parser, state: ParsingState) -> bool:
821+
successor = parser.parse_optional_successor()
822+
823+
state.successors[self.index] = successor
824+
825+
return successor is not None
826+
827+
def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None:
828+
if state.should_emit_space or not state.last_was_punctuation:
829+
printer.print(" ")
830+
printer.print_block_name(getattr(op, self.name))
831+
state.last_was_punctuation = False
832+
state.should_emit_space = True
833+
834+
835+
class VariadicSuccessorVariable(VariadicVariable, OptionallyParsableDirective):
836+
"""
837+
A variadic successor variable, with the following format:
838+
successor-directive ::= dollar-ident
839+
The directive will request a space to be printed after.
840+
"""
841+
842+
def parse_optional(self, parser: Parser, state: ParsingState) -> bool:
843+
successors: list[Successor] = []
844+
current_successor = parser.parse_optional_successor()
845+
while current_successor is not None:
846+
successors.append(current_successor)
847+
current_successor = parser.parse_optional_successor()
848+
849+
state.successors[self.index] = successors
850+
851+
return bool(successors)
852+
853+
def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None:
854+
if state.should_emit_space or not state.last_was_punctuation:
855+
printer.print(" ")
856+
successor = getattr(op, self.name)
857+
if successor:
858+
printer.print_list(successor, printer.print_block_name, delimiter=" ")
859+
state.last_was_punctuation = False
860+
state.should_emit_space = True
861+
862+
863+
class OptionalSuccessorVariable(OptionalVariable, OptionallyParsableDirective):
864+
"""
865+
An optional successor variable, with the following format:
866+
successor-directive ::= dollar-ident
867+
The directive will request a space to be printed after.
868+
"""
869+
870+
def parse_optional(self, parser: Parser, state: ParsingState) -> bool:
871+
successor = parser.parse_optional_successor()
872+
if successor is None:
873+
successor = list[Successor]()
874+
state.successors[self.index] = successor
875+
return bool(successor)
876+
877+
def print(self, printer: Printer, state: PrintingState, op: IRDLOperation) -> None:
878+
if state.should_emit_space or not state.last_was_punctuation:
879+
printer.print(" ")
880+
successor = getattr(op, self.name)
881+
if successor:
882+
printer.print_block_name(successor)
883+
state.last_was_punctuation = False
884+
state.should_emit_space = True
885+
886+
813887
@dataclass(frozen=True)
814888
class AttributeVariable(FormatDirective):
815889
"""

0 commit comments

Comments
 (0)