Skip to content

Commit 9711573

Browse files
committed
WIP: LOAD_ATTR_MODULE inline dict access (+32% speedup)
Inline module attribute access by reading dk_version and dict entry index from CPython inline cache. Emits GuardType(module), dk_version guard, then direct DK_UNICODE_ENTRIES access via JITRT_LoadModuleDictEntry helper. Bypasses generic LoadAttr path entirely. 10/10 correctness tests pass on devgpu004. Known issue: Bug 6 (SplitDictDeoptPatcher segfault) affects LOAD_ATTR_INSTANCE_VALUE/SLOT paths but NOT this module path.
1 parent 6a4b2d9 commit 9711573

3 files changed

Lines changed: 80 additions & 2 deletions

File tree

cinderx/Jit/hir/builder.cpp

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
#endif
1515

1616
#include "cinderx/Common/code.h"
17+
#include "cinderx/Common/dict.h"
1718
#include "cinderx/Common/py-portability.h"
1819
#include "cinderx/Common/ref.h"
1920
#include "cinderx/Interpreter/cinder_opcode.h"
@@ -3177,8 +3178,72 @@ void HIRBuilder::emitLoadAttr(
31773178
if (getConfig().specialized_opcodes) {
31783179
switch (bc_instr.specializedOpcode()) {
31793180
case LOAD_ATTR_MODULE: {
3180-
Type type = Type::fromTypeExact(&PyModule_Type);
3181-
tc.emit<GuardType>(receiver, type, receiver, tc.frame);
3181+
// Guard receiver is a module
3182+
Type mod_type = Type::fromTypeExact(&PyModule_Type);
3183+
tc.emit<GuardType>(receiver, mod_type, receiver, tc.frame);
3184+
3185+
// Read dict_version and index from CPython's inline cache
3186+
_Py_CODEUNIT* code_units = codeUnit(code_);
3187+
int instr_idx = bc_instr.opcodeIndex().value();
3188+
const _PyAttrCache* cache =
3189+
reinterpret_cast<const _PyAttrCache*>(
3190+
&code_units[instr_idx + 1]);
3191+
uint32_t dict_version = cache->version[0] |
3192+
(static_cast<uint32_t>(cache->version[1]) << 16);
3193+
uint16_t index = cache->index;
3194+
3195+
if (dict_version != 0) {
3196+
// Inline dict access: module->md_dict->ma_keys->dk_version
3197+
Register* dict = temps_.AllocateStack();
3198+
tc.emit<LoadField>(
3199+
dict, receiver, "md_dict",
3200+
offsetof(PyModuleObject, md_dict), TObject);
3201+
3202+
Register* keys = temps_.AllocateStack();
3203+
tc.emit<LoadField>(
3204+
keys, dict, "ma_keys",
3205+
offsetof(PyDictObject, ma_keys), TCPtr);
3206+
3207+
Register* loaded_version = temps_.AllocateStack();
3208+
tc.emit<LoadField>(
3209+
loaded_version, keys, "dk_version",
3210+
offsetof(PyDictKeysObject, dk_version), TCUInt32);
3211+
3212+
Register* expected_version = temps_.AllocateStack();
3213+
tc.emit<LoadConst>(
3214+
expected_version,
3215+
Type::fromCUInt(dict_version, TCUInt32));
3216+
3217+
Register* version_match = temps_.AllocateStack();
3218+
tc.emit<PrimitiveCompare>(
3219+
version_match, PrimitiveCompareOp::kEqual,
3220+
loaded_version, expected_version);
3221+
tc.emit<Guard>(version_match, tc.frame);
3222+
3223+
// Load entry value via helper (computes DK_UNICODE_ENTRIES)
3224+
Register* index_reg = temps_.AllocateStack();
3225+
tc.emit<LoadConst>(
3226+
index_reg,
3227+
Type::fromCInt(static_cast<int64_t>(index), TCInt64));
3228+
3229+
Register* result = temps_.AllocateStack();
3230+
auto call = tc.emit<CallStatic>(
3231+
2, result,
3232+
reinterpret_cast<void*>(JITRT_LoadModuleDictEntry),
3233+
TOptObject);
3234+
call->SetOperand(0, keys);
3235+
call->SetOperand(1, index_reg);
3236+
3237+
// Deopt if value is NULL (attribute deleted)
3238+
BorrowedRef<> attr_name =
3239+
PyTuple_GET_ITEM(code_->co_names, name_idx);
3240+
auto cf = tc.emit<CheckField>(
3241+
result, result, attr_name, tc.frame);
3242+
cf->setGuiltyReg(receiver);
3243+
3244+
tc.frame.stack.push(result);
3245+
return;
3246+
}
31823247
break;
31833248
}
31843249
case LOAD_ATTR_SLOT:

cinderx/Jit/jit_rt.cpp

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
#include "cinderx/Jit/jit_rt.h"
44

5+
#include "cinderx/Common/dict.h"
6+
57
#include "internal/pycore_call.h"
68
#include "internal/pycore_ceval.h"
79
#include "internal/pycore_object.h"
@@ -2424,3 +2426,10 @@ int JITRT_MatchAndClearException(PyObject* exc_type) {
24242426
PyErr_SetRaisedException(exc);
24252427
return 0;
24262428
}
2429+
2430+
PyObject* JITRT_LoadModuleDictEntry(
2431+
PyDictKeysObject* keys,
2432+
Py_ssize_t index) {
2433+
PyDictUnicodeEntry* ep = DK_UNICODE_ENTRIES(keys) + index;
2434+
return Py_XNewRef(ep->me_value);
2435+
}

cinderx/Jit/jit_rt.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -611,3 +611,7 @@ extern PyObject JITRT_IterDoneSentinel;
611611
* Returns the next value, or JITRT_IterDoneSentinel if the iterator is done.
612612
*/
613613
PyObject* JITRT_InvokeIterNext(PyObject* iterator);
614+
615+
// Return the value at the given index in a unicode dict keys' entries array.
616+
// Used by LOAD_ATTR_MODULE inline specialisation.
617+
PyObject* JITRT_LoadModuleDictEntry(PyDictKeysObject* keys, Py_ssize_t index);

0 commit comments

Comments
 (0)