Skip to content

Commit f52d3a1

Browse files
alexmalyshevmeta-codesync[bot]
authored andcommitted
Refactor subscript simplifying
Summary: Putting it in its own function because it's complicated. This diff adds a `constexpr bool kFreeThreadedBuild` value for cases like the one in simplifySubscript() where we don't need conditional compilation between FT and non-FT. It also improves the case for unicode subscript with constant values on both sides in multithreaded compilation. Previously it would always early exit, now it'll fall through to the UnicodeSubscr case which will be a little bit faster. Reviewed By: yoney Differential Revision: D103345172 fbshipit-source-id: 6849f4482f038e19819d37b56df08081bab3ea42
1 parent b83ff7a commit f52d3a1

2 files changed

Lines changed: 115 additions & 92 deletions

File tree

cinderx/Common/util.h

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,17 @@ constexpr bool kPyRefDebug =
5858
false;
5959
#endif
6060

61+
// True when CinderX is built against a free-threaded (Py_GIL_DISABLED) Python.
62+
//
63+
// When false, code can assume the GIL is held. When true, it cannot, the GIL
64+
// might still be held at any given moment but that's no longer guaranteed.
65+
constexpr bool kFreeThreadedBuild =
66+
#ifdef Py_GIL_DISABLED
67+
true;
68+
#else
69+
false;
70+
#endif
71+
6172
struct jit_string_deleter {
6273
void operator()(jit_string_t* ss) const {
6374
ss_free(ss);

cinderx/Jit/hir/simplify.cpp

Lines changed: 104 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -720,116 +720,128 @@ Register* simplifyLoadMethod(Env& env, const LoadMethod* load_meth) {
720720
*load_meth->frameState());
721721
}
722722

723-
Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
723+
Register* simplifySubscript(Env& env, const BinaryOp* instr) {
724724
BinaryOpKind op = instr->op();
725725
Register* lhs = instr->left();
726726
Register* rhs = instr->right();
727+
JIT_CHECK(
728+
op == BinaryOpKind::kSubscript,
729+
"simplifySubscript trying to optimize {}",
730+
op);
727731

728-
if (op == BinaryOpKind::kSubscript) {
729-
if (lhs->isA(TDictExact)) {
730-
return env.emit<DictSubscr>(lhs, rhs, *instr->frameState());
732+
if (lhs->isA(TDictExact)) {
733+
return env.emit<DictSubscr>(lhs, rhs, *instr->frameState());
734+
}
735+
if (!rhs->isA(TLongExact)) {
736+
return nullptr;
737+
}
738+
739+
Type lhs_type = lhs->type();
740+
Type rhs_type = rhs->type();
741+
742+
// Constant tuple subscripted by constant long.
743+
if (lhs_type <= TTupleExact && lhs_type.hasObjectSpec() &&
744+
rhs_type.hasObjectSpec()) {
745+
int overflow;
746+
Py_ssize_t index =
747+
PyLong_AsLongAndOverflow(rhs_type.objectSpec(), &overflow);
748+
if (!overflow) {
749+
BorrowedRef<> lhs_obj = lhs_type.objectSpec();
750+
if (index >= 0 && index < PyTuple_GET_SIZE(lhs_obj)) {
751+
BorrowedRef<> item = PyTuple_GET_ITEM(lhs_obj.get(), index);
752+
env.emit<UseType>(lhs, lhs_type);
753+
env.emit<UseType>(rhs, rhs_type);
754+
return env.emit<LoadConst>(
755+
Type::fromObject(env.func.env.addReference(item)));
756+
}
731757
}
732-
if (!rhs->isA(TLongExact)) {
733-
return nullptr;
758+
}
759+
760+
// TODO(T255264263). Enable this for FT builds. See P2169673256.
761+
if (!kFreeThreadedBuild && (lhs->isA(TListExact) || lhs->isA(TTupleExact))) {
762+
// TASK(T93509109): Replace TCInt64 with a less platform-specific
763+
// representation of the type, which should be analagous to Py_ssize_t.
764+
env.emit<UseType>(lhs, lhs->isA(TListExact) ? TListExact : TTupleExact);
765+
env.emit<UseType>(rhs, TLongExact);
766+
Register* right_index = env.emit<IndexUnbox>(rhs);
767+
env.emit<IsNegativeAndErrOccurred>(right_index, *instr->frameState());
768+
Register* adjusted_idx =
769+
env.emit<CheckSequenceBounds>(lhs, right_index, *instr->frameState());
770+
Py_ssize_t offset = offsetof(PyTupleObject, ob_item);
771+
Register* array = lhs;
772+
// Lists carry a nested array of ob_item whereas tuples are variable-sized
773+
// structs.
774+
if (lhs->isA(TListExact)) {
775+
array = env.emit<LoadField>(
776+
lhs, "ob_item", offsetof(PyListObject, ob_item), TCPtr);
777+
offset = 0;
734778
}
735-
Type lhs_type = lhs->type();
736-
Type rhs_type = rhs->type();
737-
if (lhs_type <= TTupleExact && lhs_type.hasObjectSpec() &&
738-
rhs_type.hasObjectSpec()) {
739-
int overflow;
740-
Py_ssize_t index =
741-
PyLong_AsLongAndOverflow(rhs_type.objectSpec(), &overflow);
742-
if (!overflow) {
743-
PyObject* lhs_obj = lhs_type.objectSpec();
744-
if (index >= 0 && index < PyTuple_GET_SIZE(lhs_obj)) {
745-
BorrowedRef<> item = PyTuple_GET_ITEM(lhs_obj, index);
746-
env.emit<UseType>(lhs, lhs_type);
747-
env.emit<UseType>(rhs, rhs_type);
748-
return env.emit<LoadConst>(
749-
Type::fromObject(env.func.env.addReference(item)));
750-
}
751-
// Fallthrough
779+
return env.emit<LoadArrayItem>(array, adjusted_idx, lhs, offset, TObject);
780+
}
781+
782+
// Unicode subscript.
783+
if (lhs_type <= TUnicodeExact && rhs_type <= TLongExact) {
784+
// Constant fold. This isn't safe in multi-threaded compilation because the
785+
// worker doesn't hold the GIL and that's required for creating a new
786+
// string.
787+
if (!getThreadedCompileContext().compileRunning() &&
788+
lhs_type.hasObjectSpec() && rhs_type.hasObjectSpec()) {
789+
Py_ssize_t idx = PyLong_AsSsize_t(rhs_type.objectSpec());
790+
if (idx == -1 && PyErr_Occurred()) {
791+
PyErr_Clear();
792+
return nullptr;
752793
}
753-
// Fallthrough
754-
}
755-
// TODO(T255264263). Enable this again. See P2169673256.
756-
#ifndef Py_GIL_DISABLED
757-
if (lhs->isA(TListExact) || lhs->isA(TTupleExact)) {
758-
// TASK(T93509109): Replace TCInt64 with a less platform-specific
759-
// representation of the type, which should be analagous to Py_ssize_t.
760-
env.emit<UseType>(lhs, lhs->isA(TListExact) ? TListExact : TTupleExact);
761-
env.emit<UseType>(rhs, TLongExact);
762-
Register* right_index = env.emit<IndexUnbox>(rhs);
763-
env.emit<IsNegativeAndErrOccurred>(right_index, *instr->frameState());
764-
Register* adjusted_idx =
765-
env.emit<CheckSequenceBounds>(lhs, right_index, *instr->frameState());
766-
Py_ssize_t offset = offsetof(PyTupleObject, ob_item);
767-
Register* array = lhs;
768-
// Lists carry a nested array of ob_item whereas tuples are variable-sized
769-
// structs.
770-
if (lhs->isA(TListExact)) {
771-
array = env.emit<LoadField>(
772-
lhs, "ob_item", offsetof(PyListObject, ob_item), TCPtr);
773-
offset = 0;
794+
Py_ssize_t n = PyUnicode_GetLength(lhs_type.objectSpec());
795+
if (idx < -n || idx >= n) {
796+
return nullptr;
797+
}
798+
799+
if (idx < 0) {
800+
idx += n;
801+
}
802+
803+
ThreadedCompileSerialize guard;
804+
Py_UCS4 c = PyUnicode_ReadChar(lhs_type.objectSpec(), idx);
805+
PyObject* substr = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, &c, 1);
806+
if (substr == nullptr) {
807+
return nullptr;
774808
}
775-
return env.emit<LoadArrayItem>(array, adjusted_idx, lhs, offset, TObject);
809+
PyUnicode_InternInPlace(&substr);
810+
Ref<> result = Ref<>::steal(substr);
811+
812+
// Use exact types since we're relying on the object specializations.
813+
env.emit<UseType>(lhs, lhs_type);
814+
env.emit<UseType>(rhs, rhs_type);
815+
return env.emit<LoadConst>(
816+
Type::fromObject(env.func.env.addReference(std::move(result))));
776817
}
777-
#endif
778-
if (lhs_type <= TUnicodeExact && rhs_type <= TLongExact) { // Unicode subscr
779-
if (lhs_type.hasObjectSpec() && rhs_type.hasObjectSpec()) {
780-
// This isn't safe in the multi-threaded compilation on 3.12 because
781-
// we don't hold the GIL which is required for
782-
// PyUnicode_InternInPlace.
783-
RETURN_MULTITHREADED_COMPILE(nullptr);
784-
785-
// Constant propagation
786-
Py_ssize_t idx = PyLong_AsSsize_t(rhs_type.objectSpec());
787-
if (idx == -1 && PyErr_Occurred()) {
788-
PyErr_Clear();
789-
return nullptr;
790-
}
791-
Py_ssize_t n = PyUnicode_GetLength(lhs_type.objectSpec());
792818

793-
if (idx < -n || idx >= n) {
794-
return nullptr;
795-
}
819+
env.emit<UseType>(lhs, TUnicodeExact);
820+
env.emit<UseType>(rhs, TLongExact);
821+
Register* unboxed_idx = env.emit<IndexUnbox>(rhs);
822+
env.emit<IsNegativeAndErrOccurred>(unboxed_idx, *instr->frameState());
823+
Register* adjusted_idx =
824+
env.emit<CheckSequenceBounds>(lhs, unboxed_idx, *instr->frameState());
825+
return env.emit<UnicodeSubscr>(lhs, adjusted_idx, *instr->frameState());
826+
}
796827

797-
if (idx < 0) {
798-
idx += n;
799-
}
828+
return nullptr;
829+
}
800830

801-
ThreadedCompileSerialize guard;
802-
Py_UCS4 c = PyUnicode_ReadChar(lhs_type.objectSpec(), idx);
803-
PyObject* substr =
804-
PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, &c, 1);
805-
if (substr == nullptr) {
806-
return nullptr;
807-
}
808-
PyUnicode_InternInPlace(&substr);
809-
Ref<> result = Ref<>::steal(substr);
831+
Register* simplifyBinaryOp(Env& env, const BinaryOp* instr) {
832+
BinaryOpKind op = instr->op();
833+
Register* lhs = instr->left();
834+
Register* rhs = instr->right();
810835

811-
// Use exact types since we're relying on the object specializations.
812-
env.emit<UseType>(lhs, lhs_type);
813-
env.emit<UseType>(rhs, rhs_type);
814-
return env.emit<LoadConst>(
815-
Type::fromObject(env.func.env.addReference(std::move(result))));
816-
} else {
817-
env.emit<UseType>(lhs, TUnicodeExact);
818-
env.emit<UseType>(rhs, TLongExact);
819-
Register* unboxed_idx = env.emit<IndexUnbox>(rhs);
820-
env.emit<IsNegativeAndErrOccurred>(unboxed_idx, *instr->frameState());
821-
Register* adjusted_idx = env.emit<CheckSequenceBounds>(
822-
lhs, unboxed_idx, *instr->frameState());
823-
return env.emit<UnicodeSubscr>(lhs, adjusted_idx, *instr->frameState());
824-
}
825-
}
836+
if (op == BinaryOpKind::kSubscript) {
837+
return simplifySubscript(env, instr);
826838
}
827839

828840
if (lhs->isA(TLongExact) && rhs->isA(TLongExact)) {
829841
// All binary ops on TLong's return mutable so can be freely simplified with
830842
// no explicit check.
831-
if (op == BinaryOpKind::kMatrixMultiply || op == BinaryOpKind::kSubscript) {
832-
// These will generate an error at runtime.
843+
if (op == BinaryOpKind::kMatrixMultiply) {
844+
// This will generate an error at runtime.
833845
return nullptr;
834846
}
835847
env.emit<UseType>(lhs, TLongExact);

0 commit comments

Comments
 (0)