[SOT] Guard dict itself in dict.get

[SigureMo]

PR Category

Execute Infrastructure

PR Types

Bug fixes

Description

def get_arg_from_kwargs(x, **kwargs):
    x = x if x is not None else None
    y = kwargs.get("y", None)
    paddle.jit.sot.psdb.breakgraph()
    return x, y

static_fn = paddle.jit.to_static(get_arg_from_kwargs)
static_fn(None)     # (None, None)
static_fn(None, 1)  # (None, None) 错误命中 cache

当 MIN_GRAPH_SIZE=10 时，这里我们生成的字节码如下：

 14           0 LOAD_GLOBAL              6 (paddle_set_eval_frame_fn)
              2 LOAD_CONST               0 (None)
              4 CALL_FUNCTION            1
              6 STORE_FAST               3 (___old_eval_frame)

 15           8 LOAD_FAST                0 (x)
             10 LOAD_CONST               0 (None)
             12 IS_OP                    1
             14 POP_JUMP_IF_FALSE       10 (to 20)
             16 LOAD_FAST                0 (x)
             18 JUMP_FORWARD             1 (to 22)
        >>   20 LOAD_CONST               0 (None)
        >>   22 STORE_FAST               0 (x)

 16          24 LOAD_FAST                1 (kwargs)
             26 LOAD_METHOD              0 (get)
             28 LOAD_CONST               1 ('y')
             30 LOAD_CONST               0 (None)
             32 CALL_METHOD              2
             34 STORE_FAST               2 (y)

 17          36 LOAD_GLOBAL              1 (paddle)
             38 LOAD_ATTR                2 (jit)
             40 LOAD_ATTR                3 (sot)
             42 LOAD_ATTR                4 (psdb)
             44 LOAD_METHOD              5 (breakgraph)
             46 NOP
             48 LOAD_GLOBAL              6 (paddle_set_eval_frame_fn)
             50 LOAD_FAST                3 (___old_eval_frame)
             52 CALL_FUNCTION            1
             54 POP_TOP
             56 STORE_FAST               4 (___graph_fn_saved_orig_0)
             58 STORE_FAST               5 (___graph_fn_saved_orig_1)
             60 LOAD_GLOBAL              7 (___null_var)
             62 LOAD_FAST                4 (___graph_fn_saved_orig_0)
             64 CALL_METHOD              0
             66 LOAD_GLOBAL              8 ($resume_0@fn_af1a0)
             68 ROT_N                    2
             70 LOAD_FAST                0 (x)
             72 LOAD_FAST                0 (x)
             74 CALL_FUNCTION            3
             76 RETURN_VALUE
             78 POP_TOP

 18          80 LOAD_FAST                0 (x)
             82 LOAD_FAST                2 (y)
             84 BUILD_TUPLE              2
             86 RETURN_VALUE

这里连续 load 两次 x，这是因为 x 和 y 都是同一个对象 None，是 LOAD_CONST 产生的，这里没什么问题

但是生成的 guard 却只是 lambda frame: id(type(frame.f_locals['x'])) == 94068900165056 and frame.f_locals['x'] == None

这里明显有一个问题，就是 y 无论是什么都会命中 guard，而且这里生成的代码全是 load x，因此后面当 y 传入和 x 不一样的值时就出问题了

其实这里本质问题是缺失了 kwargs 的 guard，这里 d.get(key, default) 等价于 d[key] if key in d else default，这里控制流 cond key in d 应该加到 guard 里，而这里我们是没有加的

这里我们当然可以构造 key in d 的 Variable 并将其 guard 住，但我们这里使用了一个简单的方式，直接 guard 住 d 和 key，d 是新加的，key 在 getitem 时会自动 guard 住

未来我们可以实现 polyfill 的 dispatch 机制，将 d.get(key, default) 派发到 d[key] if key in d else default，以自动记录这里的 key in d 到 guard

PCard-66972

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[DrRyanHuang]

Suggested change

	# TODO: We should add some tracker to record the key and the dict
	# TODO: We should add some trackers to record the key and the dict

🤪

[Copilot]

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (2)

python/paddle/jit/sot/opcode_translator/executor/variables/container.py:885

• The current implementation of get() still lacks a guard for the kwargs case (i.e. when a default is provided). This may lead to unintended cache hits when different values for kwargs are supplied; consider adding a guard for kwargs as described in the PR details.

self.graph.add_global_guarded_variable(self)

test/sot/test_min_graph_size.py:117

• [nitpick] The new test for get_arg_from_kwargs checks that the function returns (None, None) for both calls, but it does not verify that the cache correctly distinguishes between different kwargs inputs. Consider extending this test with assertions that explicitly validate cache invalidation when kwargs differ.

self.assert_results(get_arg_from_kwargs, None, y=1)