From b4c077fa413ec65db8bebed8b71d33d5007e1b04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E5=BC=BA?= Date: Mon, 25 Aug 2025 13:36:49 +0000 Subject: [PATCH] =?UTF-8?q?!3005=20autofuse=20api=20for=20set=20aiv=20core?= =?UTF-8?q?num=20Merge=20pull=20request=20!3005=20from=20=E5=88=98?= =?UTF-8?q?=E5=BC=BA/lowering=5Fapi?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../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