Skip to content

Commit 8da097d

Browse files
[mypyc] feat: extend stararg fastpath from #19629 with star2 fastpath (#19630)
This PR further extends the stararg fastpath PRs (#19623 , #19629) with fastpath logic for star2 All 3 PRs were kept separate in order to make them easier to review and to make the changes in the IR more obvious in a diff.
1 parent b99948b commit 8da097d

File tree

5 files changed

+80
-51
lines changed

5 files changed

+80
-51
lines changed

mypyc/irbuild/ll_builder.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,8 @@
128128
from mypyc.primitives.bytes_ops import bytes_compare
129129
from mypyc.primitives.dict_ops import (
130130
dict_build_op,
131+
dict_copy,
132+
dict_copy_op,
131133
dict_new_op,
132134
dict_ssize_t_size_op,
133135
dict_update_in_display_op,
@@ -806,19 +808,40 @@ def _construct_varargs(
806808
return value, self._create_dict([], [], line)
807809
elif len(args) == 2 and args[1][1] == ARG_STAR2:
808810
# fn(*args, **kwargs)
811+
# TODO: extend to cover(*args, **k, **w, **a, **r, **g, **s)
809812
if is_tuple_rprimitive(value.type) or isinstance(value.type, RTuple):
810813
star_result = value
811814
elif is_list_rprimitive(value.type):
812815
star_result = self.primitive_op(list_tuple_op, [value], line)
813816
else:
814817
star_result = self.primitive_op(sequence_tuple_op, [value], line)
815-
continue
818+
819+
star2_arg = args[1]
820+
star2_value = star2_arg[0]
821+
if is_dict_rprimitive(star2_value.type):
822+
star2_fastpath_op = dict_copy_op
823+
else:
824+
star2_fastpath_op = dict_copy
825+
return star_result, self.primitive_op(
826+
star2_fastpath_op, [star2_value], line
827+
)
816828
# elif ...: TODO extend this to optimize fn(*args, k=1, **kwargs) case
817829
# TODO optimize this case using the length utils - currently in review
818830
star_result = self.new_list_op(star_values, line)
819831
self.primitive_op(list_extend_op, [star_result, value], line)
820832
elif kind == ARG_STAR2:
821833
if star2_result is None:
834+
if len(args) == 1:
835+
# early exit with fastpath if the only arg is ARG_STAR2
836+
# TODO: can we maintain an empty tuple in memory and just reuse it again and again?
837+
if is_dict_rprimitive(value.type):
838+
star2_fastpath_op = dict_copy_op
839+
else:
840+
star2_fastpath_op = dict_copy
841+
return self.new_tuple([], line), self.primitive_op(
842+
star2_fastpath_op, [value], line
843+
)
844+
822845
star2_result = self._create_dict(star2_keys, star2_values, line)
823846

824847
self.call_c(dict_update_in_display_op, [star2_result, value], line=line)

mypyc/primitives/dict_ops.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@
5353
)
5454

5555
# Construct a dictionary from another dictionary.
56-
function_op(
56+
dict_copy_op = function_op(
5757
name="builtins.dict",
5858
arg_types=[dict_rprimitive],
5959
return_type=dict_rprimitive,

mypyc/test-data/irbuild-basic.test

Lines changed: 17 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1740,12 +1740,10 @@ def g():
17401740
r6, r7 :: dict
17411741
r8 :: str
17421742
r9 :: object
1743-
r10 :: dict
1744-
r11 :: i32
1745-
r12 :: bit
1746-
r13 :: tuple
1747-
r14 :: object
1748-
r15 :: tuple[int, int, int]
1743+
r10 :: tuple
1744+
r11 :: dict
1745+
r12 :: object
1746+
r13 :: tuple[int, int, int]
17491747
L0:
17501748
r0 = 'a'
17511749
r1 = 'b'
@@ -1757,13 +1755,11 @@ L0:
17571755
r7 = __main__.globals :: static
17581756
r8 = 'f'
17591757
r9 = CPyDict_GetItem(r7, r8)
1760-
r10 = PyDict_New()
1761-
r11 = CPyDict_UpdateInDisplay(r10, r6)
1762-
r12 = r11 >= 0 :: signed
1763-
r13 = PyTuple_Pack(0)
1764-
r14 = PyObject_Call(r9, r13, r10)
1765-
r15 = unbox(tuple[int, int, int], r14)
1766-
return r15
1758+
r10 = PyTuple_Pack(0)
1759+
r11 = PyDict_Copy(r6)
1760+
r12 = PyObject_Call(r9, r10, r11)
1761+
r13 = unbox(tuple[int, int, int], r12)
1762+
return r13
17671763
def h():
17681764
r0, r1 :: str
17691765
r2, r3 :: object
@@ -3670,18 +3666,14 @@ def wrapper_deco_obj.__call__(__mypyc_self__, lst, kwargs):
36703666
r1 :: object
36713667
r2 :: tuple
36723668
r3 :: dict
3673-
r4 :: i32
3674-
r5 :: bit
3675-
r6 :: object
3669+
r4 :: object
36763670
L0:
36773671
r0 = __mypyc_self__.__mypyc_env__
36783672
r1 = r0.fn
36793673
r2 = PyList_AsTuple(lst)
3680-
r3 = PyDict_New()
3681-
r4 = CPyDict_UpdateInDisplay(r3, kwargs)
3682-
r5 = r4 >= 0 :: signed
3683-
r6 = PyObject_Call(r1, r2, r3)
3684-
return r6
3674+
r3 = PyDict_Copy(kwargs)
3675+
r4 = PyObject_Call(r1, r2, r3)
3676+
return r4
36853677
def deco(fn):
36863678
fn :: object
36873679
r0 :: __main__.deco_env
@@ -3777,18 +3769,14 @@ def wrapper_deco_obj.__call__(__mypyc_self__, args, kwargs):
37773769
r1 :: object
37783770
r2 :: tuple
37793771
r3 :: dict
3780-
r4 :: i32
3781-
r5 :: bit
3782-
r6 :: object
3772+
r4 :: object
37833773
L0:
37843774
r0 = __mypyc_self__.__mypyc_env__
37853775
r1 = r0.fn
37863776
r2 = PySequence_Tuple(args)
3787-
r3 = PyDict_New()
3788-
r4 = CPyDict_UpdateInDisplay(r3, kwargs)
3789-
r5 = r4 >= 0 :: signed
3790-
r6 = PyObject_Call(r1, r2, r3)
3791-
return r6
3777+
r3 = PyDict_Copy(kwargs)
3778+
r4 = PyObject_Call(r1, r2, r3)
3779+
return r4
37923780
def deco(fn):
37933781
fn :: object
37943782
r0 :: __main__.deco_env

mypyc/test-data/irbuild-generics.test

Lines changed: 12 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -167,17 +167,13 @@ def execute(func, args, kwargs):
167167
func :: object
168168
args :: tuple
169169
kwargs, r0 :: dict
170-
r1 :: i32
171-
r2 :: bit
172-
r3 :: object
173-
r4 :: int
170+
r1 :: object
171+
r2 :: int
174172
L0:
175-
r0 = PyDict_New()
176-
r1 = CPyDict_UpdateInDisplay(r0, kwargs)
177-
r2 = r1 >= 0 :: signed
178-
r3 = PyObject_Call(func, args, r0)
179-
r4 = unbox(int, r3)
180-
return r4
173+
r0 = PyDict_Copy(kwargs)
174+
r1 = PyObject_Call(func, args, r0)
175+
r2 = unbox(int, r1)
176+
return r2
181177
def f(x):
182178
x :: int
183179
L0:
@@ -703,10 +699,8 @@ def inner_deco_obj.__call__(__mypyc_self__, args, kwargs):
703699
r22, can_iter, r23, can_use_keys, r24, can_use_values :: list
704700
r25 :: object
705701
r26 :: dict
706-
r27 :: i32
707-
r28 :: bit
708-
r29 :: object
709-
r30 :: int
702+
r27 :: object
703+
r28 :: int
710704
L0:
711705
r0 = __mypyc_self__.__mypyc_env__
712706
r1 = var_object_size args
@@ -758,12 +752,10 @@ L9:
758752
r24 = CPyDict_Values(kwargs)
759753
can_use_values = r24
760754
r25 = r0.func
761-
r26 = PyDict_New()
762-
r27 = CPyDict_UpdateInDisplay(r26, kwargs)
763-
r28 = r27 >= 0 :: signed
764-
r29 = PyObject_Call(r25, args, r26)
765-
r30 = unbox(int, r29)
766-
return r30
755+
r26 = PyDict_Copy(kwargs)
756+
r27 = PyObject_Call(r25, args, r26)
757+
r28 = unbox(int, r27)
758+
return r28
767759
def deco(func):
768760
func :: object
769761
r0 :: __main__.deco_env

mypyc/test-data/run-functions.test

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1312,3 +1312,29 @@ from native import f
13121312
print(f(1))
13131313
[out]
13141314
2
1315+
1316+
[case testStarArgFastPaths]
1317+
from typing import Any, Mapping
1318+
def fn(x: str, y: int) -> str:
1319+
return x * y
1320+
def star_tuple(*args: Any) -> str:
1321+
return fn(*args)
1322+
def star_list(args: list[Any]) -> str:
1323+
return fn(*args)
1324+
def star_generic(args: dict[Any, Any]) -> str:
1325+
return fn(*args)
1326+
def star2(**kwargs: Any) -> str:
1327+
return fn(**kwargs)
1328+
def star2_generic(kwargs: Mapping[Any, Any]) -> str:
1329+
return fn(**kwargs)
1330+
1331+
def test_star_fastpath_tuple() -> None:
1332+
assert star_tuple("a", 3) == "aaa"
1333+
def test_star_fastpath_list() -> None:
1334+
assert star_list(["a", 3]) == "aaa"
1335+
def test_star_fastpath_generic() -> None:
1336+
assert star_generic({"a": None, 3: None}) == "aaa"
1337+
def test_star2_fastpath() -> None:
1338+
assert star2(x="a", y=3) == "aaa"
1339+
def test_star2_fastpath_generic() -> None:
1340+
assert star2_generic({"x": "a", "y": 3}) == "aaa"

0 commit comments

Comments
 (0)