GH-128939: Refactor JIT optimize structs #128940

markshannon · 2025-01-17T12:47:52Z

Converts the ad-hoc set of flags in _Py_UopsSymbol into a tagged union.
This:

Makes the code more robust. Using an enum in a switch allows the C compiler to tell us if we miss any cases
More extensible. New kinds can be added by adding a new tag type
Easier to maintain. There are no (at least much fewer) hidden assumptions in the data
Renames all structs starting to _Py_Uops to start with JitOpt.
Adds a new tuple kind to demonstrate that it can be done easily and reliably.

Issue: Reorganize the JIT optimization passes. #128939

Fidget-Spinner · 2025-01-17T13:29:44Z

Can you add a test that the tuple information is actually propagated? Either through testing type guard elimination happens or something. Unless you plan to refactor the tests as well :).

brandtbucher · 2025-01-18T00:23:01Z

Random comment: JIT is an acronym, and we use the "JIT" capitalization elsewhere. "Jit" just looks... wrong.

brandtbucher

Various thoughts and questions I had while reading through. I think this is going to be a very nice improvement, though!

Other thoughts:

It's probably worth adding a helper for immortality checks. Basically, check if we have an immortal constant, or a known bool type. The bool part is nice, but easy to forget.
Any plans to update the _Py_uop_ naming to be "JIT"-ier?

brandtbucher · 2025-01-18T00:27:39Z

Include/internal/pycore_optimizer.h

+typedef union _jit_opt_symbol {
+    uint8_t tag;
+    JitOptKnownClass cls;
+    JitOptKnownValue value;
+    JitOptKnownVersion version;
+    JitOptTuple tuple;
+} JitOptSymbol;


Could this be:

Suggested change

typedef union _jit_opt_symbol {

uint8_t tag;

JitOptKnownClass cls;

JitOptKnownValue value;

JitOptKnownVersion version;

JitOptTuple tuple;

} JitOptSymbol;

typedef struct {

uint8_t tag;

union {

JitOptKnownClass cls;

JitOptKnownValue value;

JitOptKnownVersion version;

JitOptTuple tuple;

};

} JitOptSymbol;

Then we wouldn't need to repeat the tag in each struct.

Or does the current scheme potentially pack better?

It packs better. Because structs align according to the alignment of their members, we would end up using 8 bytes for the tag if it weren't part of each struct.

brandtbucher · 2025-01-18T00:27:49Z

Include/internal/pycore_optimizer.h

+typedef struct _jit_opt_known_class {
+    uint8_t tag;
+    uint32_t version;
+    PyTypeObject *type;
+} JitOptKnownClass;


We could probably just get the version by doing cls.type->tp_version_tag, right? Having both here just creates opportunities for them to be stale or out-of-sync, I think.

sym->cls.type->tp_version_tag is the version tag of the value's class now. sym->version.version is the version tag it would have should execution reach this point in the code. Not quite the same thing.

I would imagine that a sym->cls.version that's out-of-sync with sym->cls.type->tp_version_tag should result in the bottom type, then, right? It's not possible in future code, since the cached tag is stale.

brandtbucher · 2025-01-18T00:29:03Z

Include/internal/pycore_optimizer.h

+typedef struct _jit_opt_tuple {
+    uint8_t tag;
+    uint8_t length;
+    uint16_t items[6];


I think we have room for 7, right?

Suggested change

uint16_t items[6];

uint16_t items[7];

brandtbucher · 2025-01-18T00:37:30Z

Python/optimizer_symbols.c

+        case JIT_SYM_KNOWN_VALUE_TAG:
+            Py_CLEAR(sym->value.value);
+            sym_set_bottom(ctx, sym);
+            return false;
+        case JIT_SYM_TUPLE_TAG:
+            sym_set_bottom(ctx, sym);
+            return false;


These could potentially have the same version tag, right? Looking up a method or whatever on a constant value is a normal thing to do.

brandtbucher · 2025-01-18T00:39:55Z

Python/optimizer_symbols.c

+        case JIT_SYM_TUPLE_TAG:
+            sym_set_bottom(ctx, sym);
+            return;


This also seems like it could potentially be fine (transitioning from an abstract tuple to a concrete one)... could potentially need to do something recursive for the elements, though.

brandtbucher · 2025-01-18T00:43:04Z

Python/optimizer_symbols.c

-_Py_UopsSymbol *
-_Py_uop_sym_new_unknown(_Py_UOpsContext *ctx)
+JitOptSymbol *
+_Py_uop_sym_new_unknown(JitOptContext *ctx)
 {
    return sym_new(ctx);


Should this be checked for NULL and replaced with out_of_space like the others?

brandtbucher · 2025-01-18T00:48:25Z

Python/optimizer_symbols.c

+_Py_uop_sym_tuple_getitem(JitOptContext *ctx, JitOptSymbol *sym, int item)
+{
+
+    if (sym->tag != JIT_SYM_TUPLE_TAG || item < 0 || item >= sym->tuple.length) {


We can handle negative indexes pretty easily here, right?

Also, this can probably be updated to work correctly for constant tuples.

Negative indexes are errors at this level. If we add a symbolic version of this, we would handle negative indexes there.

brandtbucher · 2025-01-18T00:48:56Z

Python/optimizer_symbols.c

+int
+_Py_uop_sym_tuple_length(JitOptSymbol *sym)
+{
+    if (sym->tag != JIT_SYM_TUPLE_TAG) {


Same here, this can probably be updated to work for constants.

brandtbucher · 2025-01-18T00:50:30Z

Python/optimizer_symbols.c

+    if (size > 6) {
+        res->tag = JIT_SYM_KNOWN_CLASS_TAG;
+        res->cls.type = &PyTuple_Type;
+    }


There are probably cases where this can be represented as a constant value (when all items are known).

Hm, although it could get tricky because nothing owns that constant, so we have to be careful not to promote its uses to constant loads.

A constant tuple holds no more information than a tuple of constants, so there would be no advantage.

Well, it can have more than 7 items... but sure, that probably isn't very common.

markshannon · 2025-01-20T09:19:39Z

Random comment: JIT is an acronym, and we use the "JIT" capitalization elsewhere. "Jit" just looks... wrong.

Readability. Is JITShow "JIT show" or "JITs how"?

markshannon · 2025-01-20T10:30:32Z

Any plans to update the Py_uop naming to be "JIT"-ier?

At some point, yes.

iritkatriel · 2025-01-20T13:30:13Z

Lib/test/test_capi/test_opt.py

+        self.assertIn("_BINARY_OP_ADD_INT", uops)
+        self.assertNotIn("_GUARD_BOTH_INT", uops)
+        self.assertNotIn("_GUARD_NOS_INT", uops)
+        self.assertNotIn("_GUARD_TOS_INT", uops)


This kind of check can easily pass by accident if there's typo or bytecodes change.

Would be good to have a helper that asserts also that the thing not in is an actual uop id.

It would.
I'll do that in another PR though.

brandtbucher · 2025-01-22T02:24:30Z

Random comment: JIT is an acronym, and we use the "JIT" capitalization elsewhere. "Jit" just looks... wrong.

Readability. Is JITShow "JIT show" or "JITs how"?

I don't know, "JIT show" seems like the pretty obvious CamelCase reading. Similar to how "GILState" is "GIL State" and not "GILs tate". Not a hill I'm willing to die on, it just looks kinda weird (like "Ast" or "Ascii" or "Gil" would).

markshannon added 2 commits January 17, 2025 11:39

Refactor optimizer symbols

219db3d

Add support for analysing tuples

e82764d

markshannon requested a review from Fidget-Spinner as a code owner January 17, 2025 12:47

bedevere-app bot mentioned this pull request Jan 17, 2025

Reorganize the JIT optimization passes. #128939

Open

3 tasks

bedevere-app bot added the awaiting core review label Jan 17, 2025

markshannon added the skip news label Jan 17, 2025

markshannon added 3 commits January 17, 2025 17:08

Add basic test for tuple symbols

50803e7

Add test for propagation of symbols through tuples

85da4d5

Merge branch 'main' into refactor-optimize-structs

307ba6c

Fidget-Spinner approved these changes Jan 17, 2025

View reviewed changes

bedevere-app bot added awaiting merge and removed awaiting core review labels Jan 17, 2025

brandtbucher reviewed Jan 18, 2025

View reviewed changes

markshannon added 2 commits January 20, 2025 11:05

Handle symbolic constants in tuple queries

56149ae

Add immortality test to symbols

8b19379

iritkatriel reviewed Jan 20, 2025

View reviewed changes

markshannon merged commit f0f7b97 into python:main Jan 20, 2025
64 checks passed

bedevere-app bot removed the awaiting merge label Jan 20, 2025

srinivasreddy pushed a commit to srinivasreddy/cpython that referenced this pull request Jan 21, 2025

pythonGH-128939: Refactor JIT optimize structs (pythonGH-128940)

beb8d52

markshannon deleted the refactor-optimize-structs branch January 31, 2025 16:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

GH-128939: Refactor JIT optimize structs #128940

GH-128939: Refactor JIT optimize structs #128940

markshannon commented Jan 17, 2025 •

edited

Loading

Fidget-Spinner commented Jan 17, 2025

brandtbucher commented Jan 18, 2025

brandtbucher left a comment

brandtbucher Jan 18, 2025

markshannon Jan 20, 2025

brandtbucher Jan 18, 2025

markshannon Jan 20, 2025

brandtbucher Jan 22, 2025 •

edited

Loading

brandtbucher Jan 18, 2025

brandtbucher Jan 18, 2025

brandtbucher Jan 18, 2025 •

edited

Loading

brandtbucher Jan 18, 2025

brandtbucher Jan 18, 2025

markshannon Jan 20, 2025

brandtbucher Jan 18, 2025

brandtbucher Jan 18, 2025

brandtbucher Jan 18, 2025

markshannon Jan 20, 2025

brandtbucher Jan 22, 2025

markshannon commented Jan 20, 2025

markshannon commented Jan 20, 2025

iritkatriel Jan 20, 2025

markshannon Jan 20, 2025

brandtbucher commented Jan 22, 2025

GH-128939: Refactor JIT optimize structs #128940

GH-128939: Refactor JIT optimize structs #128940

Conversation

markshannon commented Jan 17, 2025 • edited Loading

Fidget-Spinner commented Jan 17, 2025

brandtbucher commented Jan 18, 2025

brandtbucher left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brandtbucher Jan 22, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brandtbucher Jan 18, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

markshannon commented Jan 20, 2025

markshannon commented Jan 20, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brandtbucher commented Jan 22, 2025

markshannon commented Jan 17, 2025 •

edited

Loading

brandtbucher Jan 22, 2025 •

edited

Loading

brandtbucher Jan 18, 2025 •

edited

Loading