Skip to content

Commit ffd74bd

Browse files
diegorussodiegorusso
committed
Address Brandt's feedback.
1 parent 23e8697 commit ffd74bd

3 files changed

Lines changed: 33 additions & 25 deletions

File tree

‎Tools/jit/_stencils.py‎

Lines changed: 16 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,6 @@
66

77
import _schema
88

9-
# Number of 32-bit words needed to store the bit mask of external symbols
10-
SYMBOL_MASK_SIZE: int = 4
11-
12-
known_symbols: dict[str | None, int] = {}
13-
149

1510
@enum.unique
1611
class HoleValue(enum.Enum):
@@ -252,9 +247,14 @@ class StencilGroup:
252247
default_factory=dict, init=False
253248
)
254249
_got: dict[str, int] = dataclasses.field(default_factory=dict, init=False)
255-
trampolines: set[int] = dataclasses.field(default_factory=set, init=False)
256-
257-
def process_relocations(self, *, alignment: int = 1) -> None:
250+
_trampolines: set[int] = dataclasses.field(default_factory=set, init=False)
251+
252+
def process_relocations(
253+
self,
254+
known_symbols: dict[str | None, int],
255+
*,
256+
alignment: int = 1,
257+
) -> None:
258258
"""Fix up all GOT and internal relocations for this stencil group."""
259259
for hole in self.code.holes.copy():
260260
if (
@@ -268,7 +268,7 @@ def process_relocations(self, *, alignment: int = 1) -> None:
268268
else:
269269
ordinal = len(known_symbols)
270270
known_symbols[hole.symbol] = ordinal
271-
self.trampolines.add(ordinal)
271+
self._trampolines.add(ordinal)
272272
hole.addend = ordinal
273273
hole.symbol = None
274274
self.code.remove_jump(alignment=alignment)
@@ -328,18 +328,17 @@ def _emit_global_offset_table(self) -> None:
328328
def _get_trampoline_mask(self) -> str:
329329
bitmask: int = 0
330330
trampoline_mask: list[str] = []
331-
for ordinal in self.trampolines:
331+
for ordinal in self._trampolines:
332332
bitmask |= 1 << ordinal
333-
if bitmask:
334-
trampoline_mask = [
335-
f"0x{(bitmask >> i*32) & ((1 << 32) - 1):x}"
336-
for i in range(0, SYMBOL_MASK_SIZE)
337-
]
338-
return ", ".join(trampoline_mask)
333+
while bitmask:
334+
word = bitmask & ((1 << 32) - 1)
335+
trampoline_mask.append(f"{word:#04x}")
336+
bitmask >>= 32
337+
return "{" + ", ".join(trampoline_mask) + "}"
339338

340339
def as_c(self, opname: str) -> str:
341340
"""Dump this hole as a StencilGroup initializer."""
342-
return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {{{self._get_trampoline_mask()}}}}}"
341+
return f"{{emit_{opname}, {len(self.code.body)}, {len(self.data.body)}, {self._get_trampoline_mask()}}}"
343342

344343

345344
def symbol_to_value(symbol: str) -> tuple[HoleValue, str | None]:

‎Tools/jit/_targets.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ class _Target(typing.Generic[_S, _R]):
4444
stable: bool = False
4545
debug: bool = False
4646
verbose: bool = False
47+
known_symbols: dict[str | None, int] = dataclasses.field(default_factory=dict)
4748

4849
def _compute_digest(self, out: pathlib.Path) -> str:
4950
hasher = hashlib.sha256()
@@ -95,7 +96,9 @@ async def _parse(self, path: pathlib.Path) -> _stencils.StencilGroup:
9596
if group.data.body:
9697
line = f"0: {str(bytes(group.data.body)).removeprefix('b')}"
9798
group.data.disassembly.append(line)
98-
group.process_relocations(alignment=self.alignment)
99+
group.process_relocations(
100+
known_symbols=self.known_symbols, alignment=self.alignment
101+
)
99102
return group
100103

101104
def _handle_section(self, section: _S, group: _stencils.StencilGroup) -> None:
@@ -231,7 +234,7 @@ def build(
231234
if comment:
232235
file.write(f"// {comment}\n")
233236
file.write("\n")
234-
for line in _writer.dump(stencil_groups):
237+
for line in _writer.dump(stencil_groups, self.known_symbols):
235238
file.write(f"{line}\n")
236239
try:
237240
jit_stencils_new.replace(jit_stencils)

‎Tools/jit/_writer.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,16 @@
22

33
import itertools
44
import typing
5+
import math
56

67
import _stencils
78

89

9-
def _dump_footer(groups: dict[str, _stencils.StencilGroup]) -> typing.Iterator[str]:
10-
yield f"typedef uint32_t SymbolMask[{_stencils.SYMBOL_MASK_SIZE}];"
10+
def _dump_footer(
11+
groups: dict[str, _stencils.StencilGroup], symbols: dict[str | None, int]
12+
) -> typing.Iterator[str]:
13+
symbol_mask_size = math.ceil(len(symbols) / 32)
14+
yield f"typedef uint32_t SymbolMask[{symbol_mask_size}];"
1115
yield ""
1216
yield "typedef struct {"
1317
yield " void (*emit)("
@@ -27,8 +31,8 @@ def _dump_footer(groups: dict[str, _stencils.StencilGroup]) -> typing.Iterator[s
2731
yield f" [{opname}] = {group.as_c(opname)},"
2832
yield "};"
2933
yield ""
30-
yield f"static const void * const symbols_map[{max(len(_stencils.known_symbols), 1)}] = {{"
31-
for symbol, ordinal in _stencils.known_symbols.items():
34+
yield f"static const void * const symbols_map[{max(len(symbols), 1)}] = {{"
35+
for symbol, ordinal in symbols.items():
3236
yield f" [{ordinal}] = &{symbol},"
3337
yield "};"
3438

@@ -66,8 +70,10 @@ def _dump_stencil(opname: str, group: _stencils.StencilGroup) -> typing.Iterator
6670
yield ""
6771

6872

69-
def dump(groups: dict[str, _stencils.StencilGroup]) -> typing.Iterator[str]:
73+
def dump(
74+
groups: dict[str, _stencils.StencilGroup], symbols: dict[str | None, int]
75+
) -> typing.Iterator[str]:
7076
"""Yield a JIT compiler line-by-line as a C header file."""
7177
for opname, group in sorted(groups.items()):
7278
yield from _dump_stencil(opname, group)
73-
yield from _dump_footer(groups)
79+
yield from _dump_footer(groups, symbols)

0 commit comments

Comments
 (0)