diff --git a/transforms/scalar/cse.py b/transforms/scalar/cse.py index 79141ec..b431642 100644 --- a/transforms/scalar/cse.py +++ b/transforms/scalar/cse.py @@ -115,8 +115,10 @@ """ from collections import defaultdict + +from ...core import BasePass, PassRegistry +from ...utils.hash_utils import hash_tensor_value from ...utils.logger import logger as logging -from ...core import PassRegistry, BasePass def extract_key_attrs(attrs, op_type=None): @@ -165,9 +167,9 @@ def extract_key_attrs(attrs, op_type=None): key_attrs.append((attr_name, "shape", shape_dims)) elif attr_value.HasField("tensor"): # tensor 类型(Const 节点的 value 属性) - # 序列化为字节串确保相同值的常量有相同签名 - tensor_bytes = attr_value.tensor.SerializeToString() - key_attrs.append((attr_name, "tensor", tensor_bytes)) + # 使用哈希值代替完整的序列化字节串,以优化性能 + tensor_hash = hash_tensor_value(attr_value.tensor) + key_attrs.append((attr_name, "tensor", tensor_hash)) elif attr_value.HasField("func"): # 函数引用(如 While 循环的 body/cond) key_attrs.append((attr_name, "func", attr_value.func.name)) @@ -196,7 +198,7 @@ def extract_key_attrs(attrs, op_type=None): key_attrs.append((attr_name, "list_shape", shapes)) elif attr_value.list.tensor: # list of tensor - tensors = tuple(t.SerializeToString() for t in attr_value.list.tensor) + tensors = tuple(hash_tensor_value(t) for t in attr_value.list.tensor) key_attrs.append((attr_name, "list_tensor", tensors)) elif attr_value.list.func: # list of func diff --git a/utils/hash_utils.py b/utils/hash_utils.py new file mode 100644 index 0000000..4eeefa3 --- /dev/null +++ b/utils/hash_utils.py @@ -0,0 +1,16 @@ +import hashlib + +def hash_tensor_value(tensor_proto): + """ + Computes a hash of the tensor's value for efficient comparison. + + Args: + tensor_proto: The TensorProto object. + + Returns: + A string hash of the tensor's value. + """ + # Using MD5 for speed. It's a performance optimization, not a security feature. + hasher = hashlib.md5() + hasher.update(tensor_proto.SerializeToString()) + return hasher.hexdigest()