From 5ce0e0b00a7c426ee208aa4aa316bcce6cec02df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E5=BC=BA?= Date: Thu, 31 Jul 2025 16:41:36 +0800 Subject: [PATCH] lowering api --- .../npu_bridge/estimator/npu/npu_scope.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tf_adapter/python/npu_bridge/estimator/npu/npu_scope.py b/tf_adapter/python/npu_bridge/estimator/npu/npu_scope.py index 63eb0ce65..20547f587 100644 --- a/tf_adapter/python/npu_bridge/estimator/npu/npu_scope.py +++ b/tf_adapter/python/npu_bridge/estimator/npu/npu_scope.py @@ -218,3 +218,23 @@ def disable_autofuse(): } with ops.get_default_graph()._attr_scope(attrs): yield + + +@contextlib.contextmanager +def limit_core_num_scope(op_aicore_num="0", op_vectorore_num="0"): + """ + Limit the aic abd aiv core num of autofuse operators within the scope. + """ + if not isinstance(op_vectorore_num, str): + raise ValueError("Param op_vectorore_num must be string.") + try: + int_vector_core_num = int(op_vectorore_num) + except ValueError: + raise ValueError("Param op_vectorore_num can not be translated into a valid int number.") + if not (0 < int_vector_core_num < 48): + raise ValueError("Param op_vectorore_num must be in a valid range.") + attrs = { + "_op_vectorcore_num": attr_value_pb2.AttrValue(s=compat.as_bytes(op_vectorore_num)) + } + with ops.get_default_graph()._attr_scope(attrs): + yield -- Gitee