Enhance generator for register instructions

Python/bytecodes.c

@@ -186,12 +186,10 @@ dummy_func(
             ERROR_IF(res == NULL, error);
         }
 
-        inst(UNARY_POSITIVE_R, (--)) {
-            PyObject *value = REG(oparg1);
+        register inst(UNARY_POSITIVE_R, (value -- res)) {
             assert(value != NULL);
-            PyObject *res = PyNumber_Positive(value);
+            res = PyNumber_Positive(value);
             ERROR_IF(res == NULL, error);
-            Py_XSETREF(REG(oparg2), res);
         }
 
         inst(UNARY_NEGATIVE, (value -- res)) {

Tools/cases_generator/generate_cases.py

@@ -96,6 +96,8 @@ def assign(self, dst: StackEffect, src: StackEffect):
         cast = self.cast(dst, src)
         if m := re.match(r"^PEEK\((\d+)\)$", dst.name):
             self.emit(f"POKE({m.group(1)}, {cast}{src.name});")
+        elif m := re.match(r"^REG\(oparg(\d+)\)$", dst.name):
+            self.emit(f"Py_XSETREF({dst.name}, {cast}{src.name});")
         else:
             self.emit(f"{dst.name} = {cast}{src.name};")
 
@@ -109,6 +111,7 @@ class Instruction:
 
     # Parts of the underlying instruction definition
     inst: parser.InstDef
+    register: bool
     kind: typing.Literal["inst", "op"]
     name: str
     block: parser.Block
@@ -121,13 +124,16 @@ class Instruction:
     cache_effects: list[parser.CacheEffect]
     input_effects: list[StackEffect]
     output_effects: list[StackEffect]
+    input_registers: list[str]  # Parallel to input_effects
+    output_registers: list[str]  # Etc.
 
     # Set later
     family: parser.Family | None = None
     predicted: bool = False
 
     def __init__(self, inst: parser.InstDef):
         self.inst = inst
+        self.register = inst.register
         self.kind = inst.kind
         self.name = inst.name
         self.block = inst.block
@@ -142,6 +148,14 @@ def __init__(self, inst: parser.InstDef):
         ]
         self.output_effects = inst.outputs  # For consistency/completeness
 
+    def analyze_registers(self, a: "Analyzer") -> None:
+        regs = iter(("REG(oparg1)", "REG(oparg2)", "REG(oparg3)"))
+        try:
+            self.input_registers = [next(regs) for _ in self.input_effects]
+            self.output_registers = [next(regs) for _ in self.output_effects]
+        except StopIteration:  # Running out of registers
+            a.error(f"Instruction {self.name} has too many register effects")
+
     def write(self, out: Formatter) -> None:
         """Write one instruction, sans prologue and epilogue."""
         # Write a static assertion that a family's cache size is correct
@@ -153,10 +167,16 @@ def write(self, out: Formatter) -> None:
                         f'{self.cache_offset}, "incorrect cache size");'
                     )
 
-        # Write input stack effect variable declarations and initializations
-        for i, ieffect in enumerate(reversed(self.input_effects), 1):
-            src = StackEffect(f"PEEK({i})", "")
-            out.declare(ieffect, src)
+        if not self.register:
+            # Write input stack effect variable declarations and initializations
+            for i, ieffect in enumerate(reversed(self.input_effects), 1):
+                src = StackEffect(f"PEEK({i})", "")
+                out.declare(ieffect, src)
+        else:
+            # Write input register variable declarations and initializations
+            for ieffect, reg in zip(self.input_effects, self.input_registers):
+                src = StackEffect(reg, "")
+                out.declare(ieffect, src)
 
         # Write output stack effect variable declarations
         input_names = {ieffect.name for ieffect in self.input_effects}
@@ -170,18 +190,24 @@ def write(self, out: Formatter) -> None:
         if self.always_exits:
             return
 
-        # Write net stack growth/shrinkage
-        diff = len(self.output_effects) - len(self.input_effects)
-        out.stack_adjust(diff)
-
-        # Write output stack effect assignments
-        unmoved_names: set[str] = set()
-        for ieffect, oeffect in zip(self.input_effects, self.output_effects):
-            if ieffect.name == oeffect.name:
-                unmoved_names.add(ieffect.name)
-        for i, oeffect in enumerate(reversed(self.output_effects), 1):
-            if oeffect.name not in unmoved_names:
-                dst = StackEffect(f"PEEK({i})", "")
+        if not self.register:
+            # Write net stack growth/shrinkage
+            diff = len(self.output_effects) - len(self.input_effects)
+            out.stack_adjust(diff)
+
+            # Write output stack effect assignments
+            unmoved_names: set[str] = set()
+            for ieffect, oeffect in zip(self.input_effects, self.output_effects):
+                if ieffect.name == oeffect.name:
+                    unmoved_names.add(ieffect.name)
+            for i, oeffect in enumerate(reversed(self.output_effects), 1):
+                if oeffect.name not in unmoved_names:
+                    dst = StackEffect(f"PEEK({i})", "")
+                    out.assign(dst, oeffect)
+        else:
+            # Write output register assignments
+            for oeffect, reg in zip(self.output_effects, self.output_registers):
+                dst = StackEffect(reg, "")
                 out.assign(dst, oeffect)
 
         # Write cache effect
@@ -218,24 +244,28 @@ def write_body(self, out: Formatter, dedent: int, cache_adjust: int = 0) -> None
                 # ERROR_IF() must pop the inputs from the stack.
                 # The code block is responsible for DECREF()ing them.
                 # NOTE: If the label doesn't exist, just add it to ceval.c.
-                ninputs = len(self.input_effects)
-                # Don't pop common input/output effects at the bottom!
-                # These aren't DECREF'ed so they can stay.
-                for ieff, oeff in zip(self.input_effects, self.output_effects):
-                    if ieff.name == oeff.name:
-                        ninputs -= 1
-                    else:
-                        break
+                if not self.register:
+                    ninputs = len(self.input_effects)
+                    # Don't pop common input/output effects at the bottom!
+                    # These aren't DECREF'ed so they can stay.
+                    for ieff, oeff in zip(self.input_effects, self.output_effects):
+                        if ieff.name == oeff.name:
+                            ninputs -= 1
+                        else:
+                            break
+                else:
+                    ninputs = 0
                 if ninputs:
                     out.write_raw(
                         f"{extra}{space}if ({cond}) goto pop_{ninputs}_{label};\n"
                     )
                 else:
                     out.write_raw(f"{extra}{space}if ({cond}) goto {label};\n")
             elif m := re.match(r"(\s*)DECREF_INPUTS\(\);\s*$", line):
-                space = m.group(1)
-                for ieff in self.input_effects:
-                    out.write_raw(f"{extra}{space}Py_DECREF({ieff.name});\n")
+                if not self.register:
+                    space = m.group(1)
+                    for ieff in self.input_effects:
+                        out.write_raw(f"{extra}{space}Py_DECREF({ieff.name});\n")
             else:
                 out.write_raw(extra + line)
 
@@ -387,6 +417,7 @@ def analyze(self) -> None:
         self.find_predictions()
         self.map_families()
         self.check_families()
+        self.analyze_register_instrs()
         self.analyze_supers_and_macros()
 
     def find_predictions(self) -> None:
@@ -453,6 +484,11 @@ def check_families(self) -> None:
                         family,
                     )
 
+    def analyze_register_instrs(self) -> None:
+        for instr in self.instrs.values():
+            if instr.register:
+                instr.analyze_registers(self)
+
     def analyze_supers_and_macros(self) -> None:
         """Analyze each super- and macro instruction."""
         self.super_instrs = {}

Tools/cases_generator/parser.py

@@ -84,6 +84,7 @@ class OpName(Node):
 
 @dataclass
 class InstHeader(Node):
+    register: bool
     kind: Literal["inst", "op"]
     name: str
     inputs: list[InputEffect]
@@ -92,6 +93,7 @@ class InstHeader(Node):
 
 @dataclass
 class InstDef(Node):
+    register: bool
     kind: Literal["inst", "op"]
     name: str
     inputs: list[InputEffect]
@@ -134,27 +136,28 @@ def definition(self) -> InstDef | Super | Macro | Family | None:
     def inst_def(self) -> InstDef | None:
         if hdr := self.inst_header():
             if block := self.block():
-                return InstDef(hdr.kind, hdr.name, hdr.inputs, hdr.outputs, block)
+                return InstDef(hdr.register, hdr.kind, hdr.name, hdr.inputs, hdr.outputs, block)
             raise self.make_syntax_error("Expected block")
         return None
 
     @contextual
     def inst_header(self) -> InstHeader | None:
         # inst(NAME)
-        #   | inst(NAME, (inputs -- outputs))
-        #   | op(NAME, (inputs -- outputs))
+        #   | [register] inst(NAME, (inputs -- outputs))
+        #   | [register] op(NAME, (inputs -- outputs))
         # TODO: Make INST a keyword in the lexer.
+        register = bool(self.expect(lx.REGISTER))
         if (tkn := self.expect(lx.IDENTIFIER)) and (kind := tkn.text) in ("inst", "op"):
             if self.expect(lx.LPAREN) and (tkn := self.expect(lx.IDENTIFIER)):
                 name = tkn.text
                 if self.expect(lx.COMMA):
                     inp, outp = self.io_effect()
                     if self.expect(lx.RPAREN):
                         if (tkn := self.peek()) and tkn.kind == lx.LBRACE:
-                            return InstHeader(kind, name, inp, outp)
+                            return InstHeader(register, kind, name, inp, outp)
                 elif self.expect(lx.RPAREN) and kind == "inst":
                     # No legacy stack effect if kind is "op".
-                    return InstHeader(kind, name, [], [])
+                    return InstHeader(register, kind, name, [], [])
         return None
 
     def io_effect(self) -> tuple[list[InputEffect], list[OutputEffect]]: