python / python/cpython

Encode immediate def-use in the JIT specializer

Open
#148,080 4 comments 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

interpreter-core topic-JIT type-feature
Dominant language
Python
Stars
77.2k
Forks
35.9k
PR merge metrics
PR metrics pending

Description

Feature or enhancement

Proposal:

Our JIT IR naturally encodes liveness in it by virtue of being a stack IR. However, we're missing a crucial piece of info to allow for stronger optimizations: def-use.

Background: Normal SSA (static single assignment form) IR already naturally encodes def-use, and LLVM's SSA flavor also includes use-def as a property.

The idea of def-use is as the name suggests, to include information of where something is defined and where it is used (hence the name def-use).

We need def-use for the following optimizations:

  1. Eliminating redundant ops after we've constant folded combined with refcount elimination. The problem is that the current refcount elimination complicates the stack effect. So naturally the current remove_unneeded_uops pass cannot do its job properly for all of those, as the ref shuffling form of uops blocks it. This issue is becoming a bigger one as we are starting to constant fold more in the JIT as we record more info, which leaves a bunch of redundant ops lying around that we could really do away with. E.g we can constant fold most LOAD_ATTR, but we leave their operand loads lying around even when no guards need it.
  2. To feed info to the future partial evaluator (PE). In my past experiments, partial evaluation is not worth it if the unboxed operand goes straight into a boxed consumer (e.g. a C API call). This mostly wrecks the perf of PE. We need a good heuristic to determine if something is worth to unbox or whether we should keep it boxed. A reminder that the PE pass is planned to take info from the current specializer pass.

Design (def-use):

typedef union {
    uintptr_t bits;
} JitOptRef;

becomes

typedef struct {
    uintptr_t bits;
    _PyUOpInstruction *originator;
} JitOptRef;

At selected uops, we then assign orginator to be the originating instruction this symbol comes from. E.g.

[0xfefefefe] LOAD_FAST x

symbol x's originator will point to 0xfefefefe (LOAD_FAST).

Not that this encodes only the most recent def-use as a property of the reference, which should be safe. We do not need full def-use chains further back, as an interesting property of a stack IR I've observed, but not yet proven, is that local rewrites can generalize to a whole-trace optimization if done correctly. E.g. see how remove_unneeded_uops just by performing local rewrites is able to remove dead code for a larger-than-local chunk of the trace.

Has this already been discussed elsewhere?

No response given

Links to previous discussion of this feature:

No response

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading the JIT specializer's JitOptRef design and the remove_unneeded_uops pass mentioned in the issue. Trace where selected uops create or update references, then determine how originator data could support redundant-op elimination and future partial-evaluation input. Done should include def-use information at the specified uops without breaking existing stack-effect handling.

Written by the indexing model from the issue text.

Assessment

Tech stack
c, python
Domain
compilers, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.