Python中如何编程统计函数的内存访问次数（mems）？

通常大家聊内存相关问题时，关注的都是内存占用量，但这里我们要实现的是Donald Knuth在《The Art of Computer Programming》里提出的mems统计——也就是统计程序执行过程中内存读写操作的总次数，这种度量方式不依赖具体芯片架构和速度，能更精确地衡量算法耗时。

统计规则示例 #

先明确统计逻辑，以下面的代码为例：

mylist = [1, 2, 3]
x = 2
y = mylist[x]

按照约定的统计规则，这段代码总共对应5次mems：

mylist = [1, 2, 3]  # 写入`mylist`，1次mem
x = 2               # 写入`x`，1次mem
y = mylist[x]       # 读取`x`、`mylist`，写入`y`，3次mems

当然也有争议，比如有人认为列表赋值应该按元素数量算多次写入，但当前我们只需要一套合理的统计逻辑，后续再优化细节。

自动统计的实现方案 #

要实现运行函数时自动统计mems，不需要手动加计数代码，可以利用Python的sys.settrace钩子跟踪字节码执行，识别变量的读写操作并计数。

下面是一个简单的装饰器实现：

import sys

def count_mems(func):
    mem_count = 0

    def trace_func(frame, event, arg):
        nonlocal mem_count
        if event == 'opcode':
            # 字节码对应操作：LOAD_NAME/LOAD_FAST是读取变量，STORE_NAME/STORE_FAST是写入变量
            opcode = frame.f_code.co_code[frame.f_lasti]
            # 0x19=LOAD_NAME, 0x1a=LOAD_FAST；0x52=STORE_NAME, 0x53=STORE_FAST
            if opcode in (0x19, 0x1a):
                mem_count += 1
            elif opcode in (0x52, 0x53):
                mem_count += 1
        return trace_func

    def wrapper(*args, **kwargs):
        nonlocal mem_count
        mem_count = 0
        original_trace = sys.gettrace()
        sys.settrace(trace_func)
        try:
            result = func(*args, **kwargs)
            print(f"Total mems: {mem_count}")
            return result
        finally:
            sys.settrace(original_trace)
    return wrapper

使用示例 #

把需要统计的函数用这个装饰器修饰即可：

@count_mems
def test():
    mylist = [1, 2, 3]
    x = 2
    y = mylist[x]

test()
# 输出：Total mems: 5

说明 #

这个实现基于Python的字节码跟踪，识别变量的读取（LOAD_*系列字节码）和写入（STORE_*系列字节码）操作，每一次对应1次mem。当前逻辑和示例中的统计规则一致，把列表赋值算作1次写入；如果后续需要调整统计规则（比如列表元素单独计数），只需要修改trace_func里的判断逻辑即可。

内容的提问来源于stack exchange，提问作者Tyler