Star659
445
代码
代码
Issues36
Pull Requests118
Actions
项目讨论
Wiki
分析
Star659
445
已合并
feat: 自定义算子 AnnotatedArgsOp 地址刷新接口 Python 化 #4254
shangdf创建于 22 天前
feat: 自定义算子 AnnotatedArgsOp 地址刷新接口 Python 化 #4254
已合并
shangdf创建于 22 天前
从shangdf/gefeature/address-refresh-pythonization-pr合入到cann/gedevelop
共 27 个文件变更+2579-73
Mapi/python/ge/ge/custom_op/__init__.py+10-0
@@ -14,9 +14,14 @@
14 14 
15__all__ = [15__all__ = [
16 "BaseCustomOp",16 "BaseCustomOp",
17+ "AnnotatedArgsContext",
18+ "AnnotatedKernelArgs",
19+ "AnnotatedKernelLaunchInfo",
17 "EagerExecuteOp",20 "EagerExecuteOp",
18 "EagerOpExecutionContext",21 "EagerOpExecutionContext",
22+ "WorkspaceAddr",
19 "clear_registered_op_impls",23 "clear_registered_op_impls",
24+ "get_declare_launch_args_ctx",
20 "get_execute_ctx",25 "get_execute_ctx",
21 "get_registered_op_impl_by_descriptor_key",26 "get_registered_op_impl_by_descriptor_key",
22 "get_registered_op_impl_dicts",27 "get_registered_op_impl_dicts",
@@ -27,15 +32,20 @@ __all__ = [
27 32 
28_LAZY_EXPORTS = {33_LAZY_EXPORTS = {
29 "BaseCustomOp": ".base",34 "BaseCustomOp": ".base",
35+ "AnnotatedArgsContext": "._native",
36+ "AnnotatedKernelArgs": "._native",
37+ "AnnotatedKernelLaunchInfo": "._native",
30 "EagerExecuteOp": ".base",38 "EagerExecuteOp": ".base",
31 "EagerOpExecutionContext": ".base",39 "EagerOpExecutionContext": ".base",
32 "clear_registered_op_impls": ".registry",40 "clear_registered_op_impls": ".registry",
41+ "get_declare_launch_args_ctx": ".context",
33 "get_execute_ctx": ".context",42 "get_execute_ctx": ".context",
34 "get_registered_op_impl_by_descriptor_key": ".registry",43 "get_registered_op_impl_by_descriptor_key": ".registry",
35 "get_registered_op_impl_dicts": ".registry",44 "get_registered_op_impl_dicts": ".registry",
36 "get_registered_op_impls": ".registry",45 "get_registered_op_impls": ".registry",
37 "register_op": ".proto",46 "register_op": ".proto",
38 "register_op_impl": ".registry",47 "register_op_impl": ".registry",
48+ "WorkspaceAddr": "._native",
39}49}
40 50 
41 51 
Mapi/python/ge/ge/custom_op/_bridge.py+140-33
@@ -18,11 +18,16 @@ import threading
18from dataclasses import dataclass18from dataclasses import dataclass
19from typing import Dict, Optional19from typing import Dict, Optional
20 20 
21-from ._ir_types import AttrType, InputType21+from ._ir_types import InputType, OutputType
22+from ._signature import _get_runtime_attr_spec, _validate_args_signature
22from .base import EagerOpExecutionContext23from .base import EagerOpExecutionContext
23from .bootstrap import get_registered_op_impls, load_custom_op_plugins24from .bootstrap import get_registered_op_impls, load_custom_op_plugins
24-from .context import _execute_ctx_scope25+from .context import _declare_launch_args_ctx_scope, _execute_ctx_scope
25-from .registry import get_registered_op_impl_by_descriptor_key26+from .registry import (
27+ INTERFACE_ANNOTATED_ARGS,
28+ INTERFACE_EAGER_EXECUTE,
29+ get_registered_op_impl_by_descriptor_key,
30+)
26 31 
27 32 
28@dataclass33@dataclass
@@ -35,21 +40,6 @@ class _OpImplHolder:
35_HOLDER_LOCK = threading.RLock()40_HOLDER_LOCK = threading.RLock()
36_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {}41_OP_IMPL_HOLDERS: Dict[str, _OpImplHolder] = {}
37 42 
38-_RUNTIME_ATTR_GETTERS = {
39- AttrType.INT: "get_int",
40- AttrType.FLOAT: "get_float",
41- AttrType.BOOL: "get_bool",
42- AttrType.STRING: "get_str",
43- AttrType.DATA_TYPE: "get_data_type",
44- AttrType.TENSOR: "get_tensor",
45- AttrType.LIST_INT: "get_list_int",
46- AttrType.LIST_FLOAT: "get_list_float",
47- AttrType.LIST_BOOL: "get_list_bool",
48- AttrType.LIST_STRING: "get_list_str",
49- AttrType.LIST_DATA_TYPE: "get_list_data_type",
50- AttrType.LIST_LIST_INT: "get_list_list_int",
51-}
52- 
53 43 
54def load_and_get_op_impl_descriptors() -> list:44def load_and_get_op_impl_descriptors() -> list:
55 load_custom_op_plugins()45 load_custom_op_plugins()
@@ -64,13 +54,17 @@ def _get_holder(instance_id: str) -> _OpImplHolder:
64 return holder54 return holder
65 55 
66 56 
67-def _get_eager_execute_op(instance_id: str) -> object:57+def _get_eager_execute_holder(instance_id: str) -> _OpImplHolder:
68- instance = _get_holder(instance_id).instance58+ holder = _get_holder(instance_id)
69- if not callable(getattr(instance, "execute", None)):59+ if not callable(getattr(holder.instance, "execute", None)):
70 raise TypeError(60 raise TypeError(
71 f"python op impl does not implement callable execute: {instance_id}"61 f"python op impl does not implement callable execute: {instance_id}"
72 )62 )
73- return instance63+ return holder
64+ 
65+ 
66+def _get_eager_execute_op(instance_id: str) -> object:
67+ return _get_eager_execute_holder(instance_id).instance
74 68 
75 69 
76def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool:70def create_op_impl_holder(instance_id: str, descriptor_key: str) -> bool:
@@ -106,19 +100,63 @@ def _is_legacy_execute(method) -> bool:
106 )100 )
107 101 
108 102 
109-def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list:103+def _get_callback_for_signature(cls, method_name: str):
104+ method = inspect.getattr_static(cls, method_name)
105+ if isinstance(method, staticmethod):
106+ return method.__func__
107+ if isinstance(method, classmethod):
108+ return method.__get__(None, cls)
109+ if inspect.isfunction(method):
110+ return method.__get__(object(), cls)
111+ return getattr(cls, method_name)
112+ 
113+ 
114+def validate_op_impl_descriptor(descriptor_key: str, ir_meta: Optional[dict]) -> bool:
115+ descriptor = get_registered_op_impl_by_descriptor_key(descriptor_key)
116+ if descriptor is None:
117+ raise KeyError(f"python op impl descriptor_key not found: {descriptor_key}")
118+ 
119+ if INTERFACE_EAGER_EXECUTE in descriptor.interfaces:
120+ method = _get_callback_for_signature(descriptor.cls, "execute")
121+ if not _is_legacy_execute(method):
122+ if ir_meta is None:
123+ raise RuntimeError(
124+ "canonical IR not found for schema-bound execute: "
125+ f"{descriptor.op_type}"
126+ )
127+ _validate_args_signature(method, ir_meta, descriptor, method_name="execute")
128+ 
129+ if INTERFACE_ANNOTATED_ARGS in descriptor.interfaces:
130+ if ir_meta is None:
131+ raise RuntimeError(
132+ "canonical IR not found for schema-bound declare_launch_args"
133+ )
134+ method = _get_callback_for_signature(descriptor.cls, "declare_launch_args")
135+ _validate_args_signature(
136+ method, ir_meta, descriptor, method_name="declare_launch_args"
137+ )
138+ return True
139+ 
140+ 
141+def _build_inputs(
142+ ir_inputs: list,
143+ get_required_input,
144+ get_optional_input,
145+ get_dynamic_input_num,
146+ get_dynamic_input,
147+) -> list:
110 args = []148 args = []
111 for ir_index, item in enumerate(ir_inputs):149 for ir_index, item in enumerate(ir_inputs):
112 kind = item["kind"]150 kind = item["kind"]
113 if kind == InputType.REQUIRED:151 if kind == InputType.REQUIRED:
114- args.append(ctx.get_required_input_tensor(ir_index))152+ args.append(get_required_input(ir_index))
115 elif kind == InputType.OPTIONAL:153 elif kind == InputType.OPTIONAL:
116- args.append(ctx.get_optional_input_tensor(ir_index))154+ args.append(get_optional_input(ir_index))
117 elif kind == InputType.DYNAMIC:155 elif kind == InputType.DYNAMIC:
118- instance_num = ctx.get_dynamic_input_num(ir_index)156+ instance_num = get_dynamic_input_num(ir_index)
119 args.append(157 args.append(
120 [158 [
121- ctx.get_dynamic_input_tensor(ir_index, relative_index)159+ get_dynamic_input(ir_index, relative_index)
122 for relative_index in range(instance_num)160 for relative_index in range(instance_num)
123 ]161 ]
124 )162 )
@@ -129,12 +167,18 @@ def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list
129 return args167 return args
130 168 
131 169 
170+def _build_execute_inputs(ctx: EagerOpExecutionContext, ir_inputs: list) -> list:
171+ return _build_inputs(
172+ ir_inputs,
173+ ctx.get_required_input_tensor,
174+ ctx.get_optional_input_tensor,
175+ ctx.get_dynamic_input_num,
176+ ctx.get_dynamic_input_tensor,
177+ )
178+ 
179+ 
132def _read_runtime_attr(attrs, index: int, ir_type: str):180def _read_runtime_attr(attrs, index: int, ir_type: str):
133- getter_name = _RUNTIME_ATTR_GETTERS.get(ir_type)181+ getter_name, _ = _get_runtime_attr_spec(ir_type, index)
134- if getter_name is None:
135- raise ValueError(
136- f"unsupported custom op runtime attr type: {ir_type}, attr index: {index}"
137- )
138 return getattr(attrs, getter_name)(index)182 return getattr(attrs, getter_name)(index)
139 183 
140 184 
@@ -148,13 +192,53 @@ def _build_execute_attrs(ctx: EagerOpExecutionContext, ir_attrs: list) -> dict:
148 }192 }
149 193 
150 194 
195+def _build_declare_inputs(ctx, ir_inputs: list) -> list:
196+ return _build_inputs(
197+ ir_inputs,
198+ ctx._get_required_input_tensor,
199+ ctx._get_optional_input_tensor,
200+ ctx._get_dynamic_input_num,
201+ ctx._get_dynamic_input_tensor,
202+ )
203+ 
204+ 
205+def _build_declare_outputs(ctx, ir_outputs: list) -> list:
206+ args = []
207+ for ir_index, item in enumerate(ir_outputs):
208+ kind = item["kind"]
209+ if kind == OutputType.REQUIRED:
210+ args.append(ctx._get_required_output_tensor(ir_index))
211+ elif kind == OutputType.DYNAMIC:
212+ instance_num = ctx._get_dynamic_output_num(ir_index)
213+ args.append(
214+ [
215+ ctx._get_dynamic_output_tensor(ir_index, relative_index)
216+ for relative_index in range(instance_num)
217+ ]
218+ )
219+ else:
220+ raise ValueError(f"unsupported custom op IR output kind: {kind}")
221+ return args
222+ 
223+ 
224+def _build_declare_attrs(ctx, ir_attrs: list) -> dict:
225+ if not ir_attrs:
226+ return {}
227+ attrs = ctx._get_attrs()
228+ return {
229+ item["name"]: _read_runtime_attr(attrs, index, item["type"])
230+ for index, item in enumerate(ir_attrs)
231+ }
232+ 
233+ 
151def call_execute(234def call_execute(
152 instance_id: str,235 instance_id: str,
153 ir_meta: Optional[dict],236 ir_meta: Optional[dict],
154 ctx: EagerOpExecutionContext,237 ctx: EagerOpExecutionContext,
155) -> None:238) -> None:
156 try:239 try:
157- custom_op = _get_eager_execute_op(instance_id)240+ holder = _get_eager_execute_holder(instance_id)
241+ custom_op = holder.instance
158 method = custom_op.execute242 method = custom_op.execute
159 if _is_legacy_execute(method):243 if _is_legacy_execute(method):
160 method(ctx)244 method(ctx)
@@ -172,6 +256,29 @@ def call_execute(
172 ctx._invalidate()256 ctx._invalidate()
173 257 
174 258 
259+def call_declare_launch_args(instance_id: str, ir_meta: Optional[dict], ctx) -> None:
260+ try:
261+ holder = _get_holder(instance_id)
262+ method = getattr(holder.instance, "declare_launch_args", None)
263+ if not callable(method):
264+ raise TypeError(
265+ f"python op impl does not implement declare_launch_args: {instance_id}"
266+ )
267+ if ir_meta is None:
268+ raise RuntimeError(
269+ "canonical IR not found for schema-bound declare_launch_args"
270+ )
271+ args = _build_declare_inputs(ctx, ir_meta["inputs"])
272+ args.extend(_build_declare_outputs(ctx, ir_meta["outputs"]))
273+ kwargs = _build_declare_attrs(ctx, ir_meta["attrs"])
274+ with _declare_launch_args_ctx_scope(ctx):
275+ result = method(*args, **kwargs)
276+ if result is not None:
277+ raise TypeError("declare_launch_args must return None")
278+ finally:
279+ ctx._invalidate()
280+ 
281+ 
175def clear_op_impl_holders() -> None:282def clear_op_impl_holders() -> None:
176 with _HOLDER_LOCK:283 with _HOLDER_LOCK:
177 _OP_IMPL_HOLDERS.clear()284 _OP_IMPL_HOLDERS.clear()
Mapi/python/ge/ge/custom_op/_ge_custom_op_native.pyi+55-0
@@ -16,7 +16,11 @@ from ge.graph.types import DataType
16from ge.runtime import StorageFormat, StorageShape, Tensor16from ge.runtime import StorageFormat, StorageShape, Tensor
17 17 
18__all__: List[str] = [18__all__: List[str] = [
19+ "AnnotatedArgsContext",
20+ "AnnotatedKernelArgs",
21+ "AnnotatedKernelLaunchInfo",
19 "EagerOpExecutionContext",22 "EagerOpExecutionContext",
23+ "WorkspaceAddr",
20]24]
21 25 
22 26 
@@ -156,3 +160,54 @@ class EagerOpExecutionContext:
156 GE bridge only; user custom op code should not call this method.160 GE bridge only; user custom op code should not call this method.
157 """161 """
158 ...162 ...
163+ 
164+ 
165+class WorkspaceAddr:
166+ """Borrowed workspace address allocated by ``AnnotatedArgsContext``."""
167+ 
168+ @property
169+ def index(self) -> int: ...
170+ 
171+ @property
172+ def addr(self) -> int: ...
173+ 
174+ 
175+class AnnotatedKernelLaunchInfo:
176+ """Owned kernel launch metadata used by ``AnnotatedArgsContext.add_launch``."""
177+ 
178+ def __init__(
179+ self,
180+ *,
181+ kernel_name: str,
182+ kernel_bin: bytes,
183+ block_dim: int,
184+ stream_id: int,
185+ ) -> None: ...
186+ 
187+ 
188+class AnnotatedKernelArgs:
189+ """Borrowed builder for one annotated kernel launch's argument sequence."""
190+ 
191+ def append_input(self, instance_index: int, tensor: Tensor) -> None: ...
192+ 
193+ def append_output(self, instance_index: int, tensor: Tensor) -> None: ...
194+ 
195+ def append_workspace(self, workspace: WorkspaceAddr) -> None: ...
196+ 
197+ def append_scalar(self, value: int) -> None: ...
198+ 
199+ 
200+class AnnotatedArgsContext:
201+ """Borrowed declaration context available only in ``declare_launch_args``."""
202+ 
203+ def malloc_workspace(self, size: int) -> WorkspaceAddr: ...
204+ 
205+ def get_stream_id(self) -> int: ...
206+ 
207+ def create_kernel_args(self) -> AnnotatedKernelArgs: ...
208+ 
209+ def add_launch(
210+ self,
211+ launch_info: AnnotatedKernelLaunchInfo,
212+ args: AnnotatedKernelArgs,
213+ ) -> None: ...
Mapi/python/ge/ge/custom_op/_native.py+8-0
@@ -15,7 +15,11 @@
15from __future__ import annotations15from __future__ import annotations
16 16 
17__all__ = [17__all__ = [
18+ "AnnotatedArgsContext",
19+ "AnnotatedKernelArgs",
20+ "AnnotatedKernelLaunchInfo",
18 "EagerOpExecutionContext",21 "EagerOpExecutionContext",
22+ "WorkspaceAddr",
19]23]
20 24 
21from importlib import import_module25from importlib import import_module
@@ -35,3 +39,7 @@ def _load_native_module():
35_native = _load_native_module()39_native = _load_native_module()
36 40 
37EagerOpExecutionContext = _native.EagerOpExecutionContext41EagerOpExecutionContext = _native.EagerOpExecutionContext
42+AnnotatedArgsContext = _native.AnnotatedArgsContext
43+AnnotatedKernelArgs = _native.AnnotatedKernelArgs
44+AnnotatedKernelLaunchInfo = _native.AnnotatedKernelLaunchInfo
45+WorkspaceAddr = _native.WorkspaceAddr
Aapi/python/ge/ge/custom_op/_signature.py+271-0
@@ -0,0 +1,271 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+ 
13+"""Schema-bound callback signature validation and runtime attribute metadata."""
14+ 
15+import inspect
16+import types
17+import typing
18+ 
19+from ge.graph import DataType
20+from ge.runtime import Tensor
21+ 
22+from ._ir_types import AttrType, InputType, OutputType
23+ 
24+ 
25+_POSITIONAL_KINDS = (
26+ inspect.Parameter.POSITIONAL_ONLY,
27+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
28+)
29+_GET_ORIGIN = getattr(
30+ typing, "get_origin", lambda value: getattr(value, "__origin__", None)
31+)
32+_GET_ARGS = getattr(typing, "get_args", lambda value: getattr(value, "__args__", ()))
33+_UNION_ORIGINS = {typing.Union}
34+_PEP604_UNION = getattr(types, "UnionType", None)
35+if _PEP604_UNION is not None:
36+ _UNION_ORIGINS.add(_PEP604_UNION)
37+ 
38+_RUNTIME_ATTR_SPECS = {
39+ AttrType.INT: ("get_int", int),
40+ AttrType.FLOAT: ("get_float", float),
41+ AttrType.BOOL: ("get_bool", bool),
42+ AttrType.STRING: ("get_str", str),
43+ AttrType.DATA_TYPE: ("get_data_type", DataType),
44+ AttrType.TENSOR: ("get_tensor", Tensor),
45+ AttrType.LIST_INT: ("get_list_int", list[int]),
46+ AttrType.LIST_FLOAT: ("get_list_float", list[float]),
47+ AttrType.LIST_BOOL: ("get_list_bool", list[bool]),
48+ AttrType.LIST_STRING: ("get_list_str", list[str]),
49+ AttrType.LIST_DATA_TYPE: ("get_list_data_type", list[DataType]),
50+ AttrType.LIST_LIST_INT: ("get_list_list_int", list[list[int]]),
51+}
52+ 
53+ 
54+def _signature_error(
55+ descriptor, method_name: str, expected: str, actual: str
56+) -> TypeError:
57+ return TypeError(
58+ f"invalid {method_name} signature for op type "
59+ f"{descriptor.op_type}, descriptor key {descriptor.descriptor_key}, "
60+ f"method {method_name}: expected {expected}, actual {actual}"
61+ )
62+ 
63+ 
64+def _normalize_annotation(annotation):
65+ if annotation is None:
66+ return type(None)
67+ origin = _GET_ORIGIN(annotation)
68+ args = _GET_ARGS(annotation)
69+ if origin is list:
70+ return ("list", tuple(_normalize_annotation(arg) for arg in args))
71+ if origin in _UNION_ORIGINS:
72+ return ("union", frozenset(_normalize_annotation(arg) for arg in args))
73+ return annotation
74+ 
75+ 
76+def _get_expected_input_annotation(kind: int):
77+ if kind == InputType.REQUIRED:
78+ return Tensor
79+ if kind == InputType.OPTIONAL:
80+ return typing.Optional[Tensor]
81+ if kind == InputType.DYNAMIC:
82+ return list[Tensor]
83+ raise ValueError(f"unsupported custom op IR input kind: {kind}")
84+ 
85+ 
86+def _get_expected_output_annotation(kind: int):
87+ if kind == OutputType.REQUIRED:
88+ return Tensor
89+ if kind == OutputType.DYNAMIC:
90+ return list[Tensor]
91+ raise ValueError(f"unsupported custom op IR output kind: {kind}")
92+ 
93+ 
94+def _get_runtime_attr_spec(ir_type: str, index: int):
95+ spec = _RUNTIME_ATTR_SPECS.get(ir_type)
96+ if spec is None:
97+ raise ValueError(
98+ f"unsupported custom op runtime attr type: {ir_type}, attr index: {index}"
99+ )
100+ return spec
101+ 
102+ 
103+def _get_type_hints(method, descriptor, method_name: str) -> dict:
104+ try:
105+ if getattr(method, "__no_type_check__", False):
106+ return {}
107+ if method_name == "execute":
108+ target = getattr(method, "__func__", method)
109+ annotations = dict(getattr(target, "__annotations__", {}))
110+ annotations.pop("return", None)
111+ if not annotations:
112+ return {}
113+ 
114+ def annotation_source():
115+ pass
116+ 
117+ annotation_source.__annotations__ = annotations
118+ return typing.get_type_hints(
119+ annotation_source,
120+ globalns=getattr(target, "__globals__", None),
121+ )
122+ return typing.get_type_hints(method)
123+ except (NameError, TypeError, AttributeError) as exc:
124+ raise _signature_error(
125+ descriptor,
126+ method_name,
127+ "resolvable type annotations",
128+ f"type hint resolution failed: {exc}",
129+ ) from exc
130+ 
131+ 
132+def _validate_annotation(
133+ parameter,
134+ expected,
135+ hints: dict,
136+ descriptor,
137+ method_name: str,
138+ position: str,
139+) -> None:
140+ if parameter.annotation is inspect.Parameter.empty:
141+ return
142+ actual = hints.get(parameter.name, parameter.annotation)
143+ if _normalize_annotation(actual) != _normalize_annotation(expected):
144+ raise _signature_error(
145+ descriptor,
146+ method_name,
147+ f"{position} annotation {_normalize_annotation(expected)!r}",
148+ f"{_normalize_annotation(actual)!r}",
149+ )
150+ 
151+ 
152+def _validate_args_signature(
153+ method,
154+ ir_meta: dict,
155+ descriptor,
156+ *,
157+ method_name: str = "declare_launch_args",
158+) -> None:
159+ if method_name not in ("execute", "declare_launch_args"):
160+ raise ValueError(f"unsupported schema callback: {method_name}")
161+ signature = inspect.signature(method)
162+ parameters = list(signature.parameters.values())
163+ for parameter in parameters:
164+ if parameter.kind in (
165+ inspect.Parameter.VAR_POSITIONAL,
166+ inspect.Parameter.VAR_KEYWORD,
167+ ):
168+ raise _signature_error(
169+ descriptor,
170+ method_name,
171+ "no variadic parameters",
172+ f"variadic parameter {parameter.name}",
173+ )
174+ 
175+ ir_inputs = ir_meta["inputs"]
176+ ir_outputs = ir_meta["outputs"] if method_name == "declare_launch_args" else []
177+ ir_attrs = ir_meta["attrs"]
178+ positional_count = len(ir_inputs) + len(ir_outputs)
179+ expected_count = positional_count + len(ir_attrs)
180+ if len(parameters) != expected_count:
181+ raise _signature_error(
182+ descriptor,
183+ method_name,
184+ f"{positional_count} positional "
185+ f"{'input/output' if method_name == 'declare_launch_args' else 'input'} "
186+ "parameters followed by "
187+ f"{len(ir_attrs)} keyword-only attrs",
188+ f"{len(parameters)} parameters",
189+ )
190+ 
191+ hints = _get_type_hints(method, descriptor, method_name)
192+ for index, item in enumerate(ir_inputs):
193+ parameter = parameters[index]
194+ if parameter.kind not in _POSITIONAL_KINDS:
195+ raise _signature_error(
196+ descriptor,
197+ method_name,
198+ f"positional input parameter at index {index}",
199+ f"parameter {parameter.name} kind {parameter.kind.name}",
200+ )
201+ _validate_annotation(
202+ parameter,
203+ _get_expected_input_annotation(item["kind"]),
204+ hints,
205+ descriptor,
206+ method_name,
207+ f"input parameter at index {index}",
208+ )
209+ 
210+ for output_index, item in enumerate(ir_outputs):
211+ parameter_index = len(ir_inputs) + output_index
212+ parameter = parameters[parameter_index]
213+ if parameter.kind not in _POSITIONAL_KINDS:
214+ raise _signature_error(
215+ descriptor,
216+ method_name,
217+ f"positional output parameter at index {output_index}",
218+ f"parameter {parameter.name} kind {parameter.kind.name}",
219+ )
220+ _validate_annotation(
221+ parameter,
222+ _get_expected_output_annotation(item["kind"]),
223+ hints,
224+ descriptor,
225+ method_name,
226+ f"output parameter at index {output_index}",
227+ )
228+ 
229+ for attr_index, item in enumerate(ir_attrs):
230+ parameter = parameters[positional_count + attr_index]
231+ if parameter.kind is not inspect.Parameter.KEYWORD_ONLY:
232+ raise _signature_error(
233+ descriptor,
234+ method_name,
235+ f"keyword-only attr parameter {item['name']}",
236+ f"parameter {parameter.name} kind {parameter.kind.name}",
237+ )
238+ if parameter.name != item["name"]:
239+ raise _signature_error(
240+ descriptor,
241+ method_name,
242+ f"attr name {item['name']} at index {attr_index}",
243+ f"attr name {parameter.name}",
244+ )
245+ _, expected_annotation = _get_runtime_attr_spec(item["type"], attr_index)
246+ _validate_annotation(
247+ parameter,
248+ expected_annotation,
249+ hints,
250+ descriptor,
251+ method_name,
252+ f"attr parameter {item['name']}",
253+ )
254+ 
255+ if method_name == "execute":
256+ return
257+ if signature.return_annotation is inspect.Signature.empty:
258+ raise _signature_error(
259+ descriptor,
260+ method_name,
261+ "None return annotation",
262+ "missing return annotation",
263+ )
264+ return_annotation = hints.get("return", signature.return_annotation)
265+ if _normalize_annotation(return_annotation) is not type(None):
266+ raise _signature_error(
267+ descriptor,
268+ method_name,
269+ "None return annotation",
270+ repr(_normalize_annotation(return_annotation)),
271+ )
Mapi/python/ge/ge/custom_op/context.py+36-1
@@ -12,12 +12,14 @@
12 12 
13"""Execution context access for schema-bound Python custom ops."""13"""Execution context access for schema-bound Python custom ops."""
14 14 
15+from __future__ import annotations
16+ 
15from contextlib import contextmanager17from contextlib import contextmanager
16from contextvars import ContextVar18from contextvars import ContextVar
17from dataclasses import dataclass19from dataclasses import dataclass
18from typing import Iterator, Optional20from typing import Iterator, Optional
19 21 
20-from ._native import EagerOpExecutionContext22+from ._native import AnnotatedArgsContext, EagerOpExecutionContext
21 23 
22 24 
23@dataclass25@dataclass
@@ -51,3 +53,36 @@ def _execute_ctx_scope(ctx: EagerOpExecutionContext) -> Iterator[None]:
51 finally:53 finally:
52 binding.active = False54 binding.active = False
53 _CURRENT_EXECUTE_CONTEXT.reset(token)55 _CURRENT_EXECUTE_CONTEXT.reset(token)
56+ 
57+ 
58+@dataclass
59+class _DeclareLaunchArgsContextBinding:
60+ ctx: AnnotatedArgsContext
61+ active: bool = True
62+ 
63+ 
64+_CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT: ContextVar[
65+ Optional[_DeclareLaunchArgsContextBinding]
66+] = ContextVar("ge_custom_op_declare_launch_args_context", default=None)
67+ 
68+ 
69+def get_declare_launch_args_ctx() -> AnnotatedArgsContext:
70+ """Return the borrowed context of the active declare_launch_args callback."""
71+ 
72+ binding = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.get()
73+ if binding is None or not binding.active:
74+ raise RuntimeError(
75+ "get_declare_launch_args_ctx() is only available inside declare_launch_args"
76+ )
77+ return binding.ctx
78+ 
79+ 
80+@contextmanager
81+def _declare_launch_args_ctx_scope(ctx: AnnotatedArgsContext) -> Iterator[None]:
82+ binding = _DeclareLaunchArgsContextBinding(ctx=ctx)
83+ token = _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.set(binding)
84+ try:
85+ yield
86+ finally:
87+ binding.active = False
88+ _CURRENT_DECLARE_LAUNCH_ARGS_CONTEXT.reset(token)
Mapi/python/ge/ge/custom_op/native_bindings/bindings.h+1-0
@@ -17,6 +17,7 @@ namespace ge {
17namespace python_custom_op_native {17namespace python_custom_op_native {
18 18 
19void BindEagerOpExecutionContext(py::module_ &m);19void BindEagerOpExecutionContext(py::module_ &m);
20+void BindAnnotatedArgsContext(py::module_ &m);
20 21 
21} // namespace python_custom_op_native22} // namespace python_custom_op_native
22} // namespace ge23} // namespace ge
Mapi/python/ge/ge/custom_op/native_bindings/context_binding.cc+283-0
@@ -9,6 +9,7 @@
9 */9 */
10 10 
11#include "bindings.h"11#include "bindings.h"
12+#include "exe_graph/runtime/annotated_args_context.h"
12#include "exe_graph/runtime/continuous_vector.h"13#include "exe_graph/runtime/continuous_vector.h"
13#include "exe_graph/runtime/eager_op_execution_context.h"14#include "exe_graph/runtime/eager_op_execution_context.h"
14#include "exe_graph/runtime/runtime_attrs.h"15#include "exe_graph/runtime/runtime_attrs.h"
@@ -21,6 +22,7 @@
21#include <stdexcept>22#include <stdexcept>
22#include <string>23#include <string>
23#include <utility>24#include <utility>
25+#include <vector>
24 26 
25namespace ge {27namespace ge {
26namespace python_custom_op_native {28namespace python_custom_op_native {
@@ -299,6 +301,253 @@ BorrowedEagerOpExecutionContext BorrowEagerOpExecutionContext(uintptr_t ctx_hand
299 return BorrowedEagerOpExecutionContext(reinterpret_cast<gert::EagerOpExecutionContext *>(ctx_handle));301 return BorrowedEagerOpExecutionContext(reinterpret_cast<gert::EagerOpExecutionContext *>(ctx_handle));
300}302}
301 303 
304+void EnsureActive(const std::shared_ptr<bool> &active) {
305+ if ((active == nullptr) || (!(*active))) {
306+ throw std::runtime_error("Borrowed native object has expired");
307+ }
308+}
309+ 
310+class NativeWorkspaceAddr {
311+ public:
312+ NativeWorkspaceAddr(gert::WorkspaceAddr workspace, std::shared_ptr<bool> active)
313+ : workspace_(workspace), active_(std::move(active)) {}
314+ 
315+ uint32_t GetIndex() const {
316+ return Get().index;
317+ }
318+ 
319+ uintptr_t GetAddr() const {
320+ return reinterpret_cast<uintptr_t>(Get().addr);
321+ }
322+ 
323+ const gert::WorkspaceAddr &Get() const {
324+ EnsureActive(active_);
325+ return workspace_;
326+ }
327+ 
328+ private:
329+ gert::WorkspaceAddr workspace_{};
330+ std::shared_ptr<bool> active_;
331+};
332+ 
333+class NativeAnnotatedKernelLaunchInfo {
334+ public:
335+ NativeAnnotatedKernelLaunchInfo(std::string kernel_name, const py::bytes &kernel_bin, uint32_t block_dim,
336+ uint32_t stream_id)
337+ : kernel_name_(std::move(kernel_name)), block_dim_(block_dim), stream_id_(stream_id) {
338+ const std::string kernel_bin_str = kernel_bin;
339+ kernel_bin_.assign(kernel_bin_str.cbegin(), kernel_bin_str.cend());
340+ if (kernel_name_.empty()) {
341+ throw std::invalid_argument("kernel_name must not be empty");
342+ }
343+ if (kernel_bin_.empty()) {
344+ throw std::invalid_argument("kernel_bin must not be empty");
345+ }
346+ if (block_dim_ == 0U) {
347+ throw std::invalid_argument("block_dim must be greater than zero");
348+ }
349+ }
350+ 
351+ gert::AnnotatedKernelLaunchInfo GetView() const {
352+ return gert::AnnotatedKernelLaunchInfo{kernel_name_.c_str(), kernel_bin_.data(), kernel_bin_.size(), block_dim_,
353+ stream_id_};
354+ }
355+ 
356+ uint32_t GetStreamId() const {
357+ return stream_id_;
358+ }
359+ 
360+ private:
361+ std::string kernel_name_;
362+ std::vector<uint8_t> kernel_bin_;
363+ uint32_t block_dim_{0U};
364+ uint32_t stream_id_{0U};
365+};
366+ 
367+class NativeAnnotatedKernelArgs {
368+ public:
369+ NativeAnnotatedKernelArgs(gert::AnnotatedArgsContext *ctx, std::shared_ptr<bool> active)
370+ : ctx_(ctx), active_(std::move(active)) {}
371+ NativeAnnotatedKernelArgs(const NativeAnnotatedKernelArgs &) = delete;
372+ NativeAnnotatedKernelArgs &operator=(const NativeAnnotatedKernelArgs &) = delete;
373+ NativeAnnotatedKernelArgs(NativeAnnotatedKernelArgs &&) = default;
374+ NativeAnnotatedKernelArgs &operator=(NativeAnnotatedKernelArgs &&) = default;
375+ 
376+ void AppendInput(uint32_t instance_index, const runtime_native::NativeTensor &tensor) {
377+ const auto input_num = GetContext()->GetComputeNodeInputNum();
378+ if (instance_index >= input_num) {
379+ throw std::out_of_range("input instance index " + std::to_string(instance_index) +
380+ " is out of range for input count " + std::to_string(input_num));
381+ }
382+ Append(gert::InputAddr{instance_index, tensor.Get()->GetAddr()}, "append input");
383+ }
384+ 
385+ void AppendOutput(uint32_t instance_index, const runtime_native::NativeTensor &tensor) {
386+ const auto output_num = GetContext()->GetComputeNodeOutputNum();
387+ if (instance_index >= output_num) {
388+ throw std::out_of_range("output instance index " + std::to_string(instance_index) +
389+ " is out of range for output count " + std::to_string(output_num));
390+ }
391+ Append(gert::OutputAddr{instance_index, tensor.Get()->GetAddr()}, "append output");
392+ }
393+ 
394+ void AppendWorkspace(const NativeWorkspaceAddr &workspace) {
395+ Append(workspace.Get(), "append workspace");
396+ }
397+ 
398+ void AppendScalar(uint64_t value) {
399+ Append(value, "append scalar");
400+ }
401+ 
402+ gert::AnnotatedKernelArgs Take() {
403+ (void)GetContext();
404+ consumed_ = true;
405+ return std::move(args_);
406+ }
407+ 
408+ private:
409+ gert::AnnotatedArgsContext *GetContext() const {
410+ EnsureActive(active_);
411+ if ((ctx_ == nullptr) || consumed_) {
412+ throw std::runtime_error("AnnotatedKernelArgs has been consumed");
413+ }
414+ return ctx_;
415+ }
416+ 
417+ template <typename T>
418+ void Append(const T &arg, const char *operation) {
419+ (void)GetContext();
420+ if (args_.AppendArg(arg) != GRAPH_SUCCESS) {
421+ throw std::runtime_error(std::string("Failed to ") + operation);
422+ }
423+ }
424+ 
425+ gert::AnnotatedArgsContext *ctx_{nullptr};
426+ std::shared_ptr<bool> active_;
427+ gert::AnnotatedKernelArgs args_;
428+ bool consumed_{false};
429+};
430+ 
431+class BorrowedAnnotatedArgsContext {
432+ public:
433+ explicit BorrowedAnnotatedArgsContext(gert::AnnotatedArgsContext *ctx)
434+ : ctx_(ctx), active_(std::make_shared<bool>(true)) {}
435+ 
436+ py::object GetRequiredInputTensor(size_t ir_index) const {
437+ return CastRequiredTensor(Get()->GetRequiredInputTensor(ir_index), "Failed to get required input tensor");
438+ }
439+ 
440+ py::object GetOptionalInputTensor(size_t ir_index) const {
441+ const auto *tensor = Get()->GetOptionalInputTensor(ir_index);
442+ return (tensor == nullptr) ? py::none() : CastTensor(tensor);
443+ }
444+ 
445+ size_t GetDynamicInputNum(size_t ir_index) const {
446+ const auto *instance_info = Get()->GetIrInputInstanceInfo(ir_index);
447+ if (instance_info == nullptr) {
448+ throw std::runtime_error("Failed to get dynamic input instance info");
449+ }
450+ return instance_info->GetInstanceNum();
451+ }
452+ 
453+ py::object GetDynamicInputTensor(size_t ir_index, size_t relative_index) const {
454+ return CastRequiredTensor(Get()->GetDynamicInputTensor(ir_index, relative_index),
455+ "Failed to get dynamic input tensor");
456+ }
457+ 
458+ py::object GetRequiredOutputTensor(size_t ir_index) const {
459+ return CastRequiredTensor(Get()->GetRequiredOutputTensor(ir_index), "Failed to get required output tensor");
460+ }
461+ 
462+ size_t GetDynamicOutputNum(size_t ir_index) const {
463+ const auto *instance_info = Get()->GetIrOutputInstanceInfo(ir_index);
464+ if (instance_info == nullptr) {
465+ throw std::runtime_error("Failed to get dynamic output instance info");
466+ }
467+ return instance_info->GetInstanceNum();
468+ }
469+ 
470+ py::object GetDynamicOutputTensor(size_t ir_index, size_t relative_index) const {
471+ return CastRequiredTensor(Get()->GetDynamicOutputTensor(ir_index, relative_index),
472+ "Failed to get dynamic output tensor");
473+ }
474+ 
475+ py::object GetAttrs() const {
476+ const auto *attrs = Get()->GetAttrs();
477+ if (attrs == nullptr) {
478+ throw std::runtime_error("Failed to get runtime attrs");
479+ }
480+ return py::cast(BorrowedRuntimeAttrs(attrs, active_));
481+ }
482+ 
483+ NativeWorkspaceAddr MallocWorkspace(size_t size) const {
484+ if (size == 0U) {
485+ throw std::invalid_argument("workspace size must be greater than zero");
486+ }
487+ const auto workspace = Get()->MallocWorkSpace(size);
488+ if (workspace.addr == nullptr) {
489+ throw std::runtime_error("Failed to malloc workspace");
490+ }
491+ return NativeWorkspaceAddr(workspace, active_);
492+ }
493+ 
494+ uint32_t GetStreamId() const {
495+ return Get()->GetStreamId();
496+ }
497+ 
498+ NativeAnnotatedKernelArgs CreateKernelArgs() const {
499+ return NativeAnnotatedKernelArgs(Get(), active_);
500+ }
501+ 
502+ void AddLaunch(const NativeAnnotatedKernelLaunchInfo &launch_info, NativeAnnotatedKernelArgs &args) const {
503+ auto *ctx = Get();
504+ if (launch_info.GetStreamId() != ctx->GetStreamId()) {
505+ throw std::invalid_argument("launch stream_id does not match current context stream_id");
506+ }
507+ auto native_args = args.Take();
508+ if (ctx->AddLaunch(launch_info.GetView(), std::move(native_args)) != GRAPH_SUCCESS) {
509+ throw std::runtime_error("Failed to add annotated kernel launch");
510+ }
511+ }
512+ 
513+ void Invalidate() {
514+ if (active_ != nullptr) {
515+ *active_ = false;
516+ }
517+ ctx_ = nullptr;
518+ }
519+ 
520+ private:
521+ gert::AnnotatedArgsContext *Get() const {
522+ EnsureActive(active_);
523+ if (ctx_ == nullptr) {
524+ throw std::runtime_error("Borrowed native object has expired");
525+ }
526+ return ctx_;
527+ }
528+ 
529+ py::object CastTensor(const gert::Tensor *tensor) const {
530+ return py::cast(runtime_native::NativeTensor::Borrow(const_cast<gert::Tensor *>(tensor), active_));
531+ }
532+ 
533+ py::object CastRequiredTensor(const gert::Tensor *tensor, const char *message) const {
534+ if (tensor == nullptr) {
535+ throw std::runtime_error(message);
536+ }
537+ return CastTensor(tensor);
538+ }
539+ 
540+ gert::AnnotatedArgsContext *ctx_{nullptr};
541+ std::shared_ptr<bool> active_;
542+};
543+ 
544+BorrowedAnnotatedArgsContext BorrowAnnotatedArgsContext(uintptr_t ctx_handle) {
545+ if (ctx_handle == 0U) {
546+ throw std::invalid_argument("ctx_handle is null");
547+ }
548+ return BorrowedAnnotatedArgsContext(reinterpret_cast<gert::AnnotatedArgsContext *>(ctx_handle));
549+}
550+ 
302} // namespace551} // namespace
303 552 
304void BindEagerOpExecutionContext(py::module_ &m) {553void BindEagerOpExecutionContext(py::module_ &m) {
@@ -338,5 +587,39 @@ void BindEagerOpExecutionContext(py::module_ &m) {
338 m.def("_borrow_eager_op_execution_context", &BorrowEagerOpExecutionContext, py::arg("ctx_handle"));587 m.def("_borrow_eager_op_execution_context", &BorrowEagerOpExecutionContext, py::arg("ctx_handle"));
339}588}
340 589 
590+void BindAnnotatedArgsContext(py::module_ &m) {
591+ py::class_<NativeWorkspaceAddr>(m, "WorkspaceAddr", "Borrowed annotated workspace address")
592+ .def_property_readonly("index", &NativeWorkspaceAddr::GetIndex)
593+ .def_property_readonly("addr", &NativeWorkspaceAddr::GetAddr);
594+ 
595+ py::class_<NativeAnnotatedKernelLaunchInfo>(m, "AnnotatedKernelLaunchInfo", "Owned kernel launch metadata")
596+ .def(py::init<std::string, const py::bytes &, uint32_t, uint32_t>(), py::kw_only(), py::arg("kernel_name"),
597+ py::arg("kernel_bin"), py::arg("block_dim"), py::arg("stream_id"));
598+ 
599+ py::class_<NativeAnnotatedKernelArgs>(m, "AnnotatedKernelArgs", "Borrowed annotated kernel args builder")
600+ .def("append_input", &NativeAnnotatedKernelArgs::AppendInput, py::arg("instance_index"), py::arg("tensor"))
601+ .def("append_output", &NativeAnnotatedKernelArgs::AppendOutput, py::arg("instance_index"), py::arg("tensor"))
602+ .def("append_workspace", &NativeAnnotatedKernelArgs::AppendWorkspace, py::arg("workspace"))
603+ .def("append_scalar", &NativeAnnotatedKernelArgs::AppendScalar, py::arg("value"));
604+ 
605+ py::class_<BorrowedAnnotatedArgsContext>(m, "AnnotatedArgsContext", "Borrowed view of gert::AnnotatedArgsContext")
606+ .def("_get_required_input_tensor", &BorrowedAnnotatedArgsContext::GetRequiredInputTensor, py::arg("ir_index"))
607+ .def("_get_optional_input_tensor", &BorrowedAnnotatedArgsContext::GetOptionalInputTensor, py::arg("ir_index"))
608+ .def("_get_dynamic_input_num", &BorrowedAnnotatedArgsContext::GetDynamicInputNum, py::arg("ir_index"))
609+ .def("_get_dynamic_input_tensor", &BorrowedAnnotatedArgsContext::GetDynamicInputTensor, py::arg("ir_index"),
610+ py::arg("relative_index"))
611+ .def("_get_required_output_tensor", &BorrowedAnnotatedArgsContext::GetRequiredOutputTensor, py::arg("ir_index"))
612+ .def("_get_dynamic_output_num", &BorrowedAnnotatedArgsContext::GetDynamicOutputNum, py::arg("ir_index"))
613+ .def("_get_dynamic_output_tensor", &BorrowedAnnotatedArgsContext::GetDynamicOutputTensor, py::arg("ir_index"),
614+ py::arg("relative_index"))
615+ .def("_get_attrs", &BorrowedAnnotatedArgsContext::GetAttrs)
616+ .def("malloc_workspace", &BorrowedAnnotatedArgsContext::MallocWorkspace, py::arg("size"))
617+ .def("get_stream_id", &BorrowedAnnotatedArgsContext::GetStreamId)
618+ .def("create_kernel_args", &BorrowedAnnotatedArgsContext::CreateKernelArgs)
619+ .def("add_launch", &BorrowedAnnotatedArgsContext::AddLaunch, py::arg("launch_info"), py::arg("args"))
620+ .def("_invalidate", &BorrowedAnnotatedArgsContext::Invalidate);
621+ m.def("_borrow_annotated_args_context", &BorrowAnnotatedArgsContext, py::arg("ctx_handle"));
622+}
623+ 
341} // namespace python_custom_op_native624} // namespace python_custom_op_native
342} // namespace ge625} // namespace ge
Mapi/python/ge/ge/custom_op/native_bindings/module.cc+1-0
@@ -13,5 +13,6 @@
13namespace ge {13namespace ge {
14PYBIND11_MODULE(_ge_custom_op_native, m) {14PYBIND11_MODULE(_ge_custom_op_native, m) {
15 python_custom_op_native::BindEagerOpExecutionContext(m);15 python_custom_op_native::BindEagerOpExecutionContext(m);
16+ python_custom_op_native::BindAnnotatedArgsContext(m);
16}17}
17} // namespace ge18} // namespace ge
Mapi/python/ge/ge/custom_op/registry.py+5-1
@@ -18,7 +18,11 @@ from dataclasses import dataclass, field
18from typing import Any, Dict, List, Optional, Type18from typing import Any, Dict, List, Optional, Type
19 19 
20INTERFACE_EAGER_EXECUTE = "eager_execute"20INTERFACE_EAGER_EXECUTE = "eager_execute"
21-_INTERFACE_SPECS = ((INTERFACE_EAGER_EXECUTE, "execute"),)21+INTERFACE_ANNOTATED_ARGS = "annotated_args"
22+_INTERFACE_SPECS = (
23+ (INTERFACE_EAGER_EXECUTE, "execute"),
24+ (INTERFACE_ANNOTATED_ARGS, "declare_launch_args"),
25+)
22 26 
23 27 
24@dataclass(frozen=True)28@dataclass(frozen=True)
Mcompiler/engines/custom_engine/custom_ops_kernel_builder.cc+22-5
@@ -38,6 +38,8 @@
38#include "exe_graph/runtime/storage_shape.h"38#include "exe_graph/runtime/storage_shape.h"
39#include "graph/buffer.h"39#include "graph/buffer.h"
40#include "graph/custom_op.h"40#include "graph/custom_op.h"
41+#include "graph/custom_op/args_refresh.h"
42+#include "graph/custom_op/cast.h"
41#include "graph/custom_op_factory.h"43#include "graph/custom_op_factory.h"
42#include "graph/debug/ge_attr_define.h"44#include "graph/debug/ge_attr_define.h"
43#include "graph/ge_context.h"45#include "graph/ge_context.h"
@@ -491,7 +493,7 @@ Status GetAnnotatedArgsOp(const OpDescPtr &op_desc, AnnotatedArgsOp *&annotated_
491 auto *const base_custom_op = CustomOpFactory::CreateOrGetCustomOp(AscendString(op_desc->GetTypePtr()));493 auto *const base_custom_op = CustomOpFactory::CreateOrGetCustomOp(AscendString(op_desc->GetTypePtr()));
492 GE_ASSERT_NOTNULL(base_custom_op, "Create custom op failed, op_name:%s, op_type:%s", op_desc->GetNamePtr(),494 GE_ASSERT_NOTNULL(base_custom_op, "Create custom op failed, op_name:%s, op_type:%s", op_desc->GetNamePtr(),
493 op_desc->GetTypePtr());495 op_desc->GetTypePtr());
494- annotated_args_op = dynamic_cast<AnnotatedArgsOp *>(base_custom_op);496+ annotated_args_op = CustomOpCast<AnnotatedArgsOp>(base_custom_op);
C
CChang-an-HW22 天前
已过期

只保留这部分修改即可，其余部分不在交付范围可认为已完成

likedislike
S
shangdf
22 天前 评论：
495 GE_ASSERT_NOTNULL(annotated_args_op, "Custom op does not implement AnnotatedArgsOp, op_name:%s, op_type:%s",497 GE_ASSERT_NOTNULL(annotated_args_op, "Custom op does not implement AnnotatedArgsOp, op_name:%s, op_type:%s",
496 op_desc->GetNamePtr(), op_desc->GetTypePtr());498 op_desc->GetNamePtr(), op_desc->GetTypePtr());
497 return SUCCESS;499 return SUCCESS;
@@ -761,11 +763,14 @@ Status FillBasicCustomKernelTask(const Node &node, domi::TaskDef &task_def) {
761}763}
762 764 
763Status GenerateBasicCustomKernelTask(const Node &node, const OpDescPtr &op_desc, const std::string &soc_version,765Status GenerateBasicCustomKernelTask(const Node &node, const OpDescPtr &op_desc, const std::string &soc_version,
764- std::vector<domi::TaskDef> &tasks) {766+ const CustomTaskArgsMode args_mode, std::vector<domi::TaskDef> &tasks) {
765 GELOGI("Custom op %s(%s) generate basic custom kernel task, soc_version: %s", op_desc->GetNamePtr(),767 GELOGI("Custom op %s(%s) generate basic custom kernel task, soc_version: %s", op_desc->GetNamePtr(),
766 op_desc->GetTypePtr(), soc_version.c_str());768 op_desc->GetTypePtr(), soc_version.c_str());
767 domi::TaskDef task_def = {};769 domi::TaskDef task_def = {};
768 GE_ASSERT_SUCCESS(FillBasicCustomKernelTask(node, task_def));770 GE_ASSERT_SUCCESS(FillBasicCustomKernelTask(node, task_def));
771+ GE_ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE, static_cast<int64_t>(args_mode)),
772+ "Set %s failed for custom op %s(%s).", ATTR_NAME_CUSTOM_TASK_ARGS_MODE.c_str(), op_desc->GetNamePtr(),
773+ op_desc->GetTypePtr());
769 tasks.push_back(task_def);774 tasks.push_back(task_def);
770 return SUCCESS;775 return SUCCESS;
771}776}
@@ -818,15 +823,24 @@ Status CustomOpsKernelBuilder::GenerateTask(const Node &node, RunContext &contex
818 823 
819 const auto *const owner_graph = node.GetOwnerComputeGraphBarePtr();824 const auto *const owner_graph = node.GetOwnerComputeGraphBarePtr();
820 GE_ASSERT_NOTNULL(owner_graph, "Owner graph of custom op %s(%s) is null.", node.GetNamePtr(), node.GetTypePtr());825 GE_ASSERT_NOTNULL(owner_graph, "Owner graph of custom op %s(%s) is null.", node.GetNamePtr(), node.GetTypePtr());
826+ ArgsRefreshStrategy args_refresh_strategy = ArgsRefreshStrategy::kNone;
821 if (!owner_graph->GetGraphUnknownFlag()) {827 if (!owner_graph->GetGraphUnknownFlag()) {
822 const bool is_mobile_omc = IsMobileSocVersion(soc_version);828 const bool is_mobile_omc = IsMobileSocVersion(soc_version);
823- const auto args_refresh_strategy = CustomOpFactory::GetArgsRefreshStrategy(AscendString(op_desc->GetTypePtr()));829+ args_refresh_strategy = CustomOpFactory::GetArgsRefreshStrategy(AscendString(op_desc->GetTypePtr()));
824 if (args_refresh_strategy == ArgsRefreshStrategy::kAnnotatedArgs) {830 if (args_refresh_strategy == ArgsRefreshStrategy::kAnnotatedArgs) {
825 // 端侧场景下的自定义算子只支持单任务的 AnnotatedArgsOp 生成方式，非端侧场景下的自定义算子支持多任务的831 // 端侧场景下的自定义算子只支持单任务的 AnnotatedArgsOp 生成方式，非端侧场景下的自定义算子支持多任务的
826 // AnnotatedArgsOp 生成方式832 // AnnotatedArgsOp 生成方式
827 const auto mode =833 const auto mode =
828 is_mobile_omc ? OfflineLaunchGenMode::kMobileLegacySingleTask : OfflineLaunchGenMode::kStandardOmMultiTask;834 is_mobile_omc ? OfflineLaunchGenMode::kMobileLegacySingleTask : OfflineLaunchGenMode::kStandardOmMultiTask;
829- return GenerateAnnotatedArgsTask(node, op_desc, context, mode, tasks);835+ const auto status = GenerateAnnotatedArgsTask(node, op_desc, context, mode, tasks);
836+ if (status != SUCCESS) {
837+ return status;
838+ }
839+ GE_ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE,
840+ static_cast<int64_t>(CustomTaskArgsMode::kAnnotatedArgs)),
841+ "Set %s failed for custom op %s(%s).", ATTR_NAME_CUSTOM_TASK_ARGS_MODE.c_str(),
842+ op_desc->GetNamePtr(), op_desc->GetTypePtr());
843+ return SUCCESS;
830 }844 }
831 // 端侧场景下的自定义算子必须实现 AnnotatedArgsOp 接口，其他生成方式不支持845 // 端侧场景下的自定义算子必须实现 AnnotatedArgsOp 接口，其他生成方式不支持
832 if (is_mobile_omc) {846 if (is_mobile_omc) {
@@ -837,7 +851,10 @@ Status CustomOpsKernelBuilder::GenerateTask(const Node &node, RunContext &contex
837 }851 }
838 // 全部动态图场景的自定义算子或者在非端侧场景中没有实现 AnnotatedArgsOp 的自定义算子，使用最基础的 CustomKernelTask852 // 全部动态图场景的自定义算子或者在非端侧场景中没有实现 AnnotatedArgsOp 的自定义算子，使用最基础的 CustomKernelTask
839 // 生成方式853 // 生成方式
840- return GenerateBasicCustomKernelTask(node, op_desc, soc_version, tasks);854+ const auto basic_args_mode = (args_refresh_strategy == ArgsRefreshStrategy::kUpdateCallback)
855+ ? CustomTaskArgsMode::kUpdateCallback
856+ : CustomTaskArgsMode::kNone;
857+ return GenerateBasicCustomKernelTask(node, op_desc, soc_version, basic_args_mode, tasks);
841}858}
842} // namespace custom859} // namespace custom
843} // namespace ge860} // namespace ge
Mgraph_metadef/graph/attr/ge_attr_define.cc+1-0
@@ -1273,6 +1273,7 @@ const std::string ATTR_NAME_CUBE_VECTOR_CORE_TYPE = "_cube_vector_core_type";
1273const std::string ATTR_NAME_CACHE_PERSIST = "_cache_persist";1273const std::string ATTR_NAME_CACHE_PERSIST = "_cache_persist";
1274const std::string ATTR_NAME_ALIAS_ENGINE_NAME = "_alias_engine_name";1274const std::string ATTR_NAME_ALIAS_ENGINE_NAME = "_alias_engine_name";
1275const std::string ATTR_NAME_KERNEL_NAMES_PREFIX = "_kernel_names_prefix";1275const std::string ATTR_NAME_KERNEL_NAMES_PREFIX = "_kernel_names_prefix";
1276+const std::string ATTR_NAME_CUSTOM_TASK_ARGS_MODE = "_custom_task_args_mode";
1276const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_SIZE = "_ffts_sub_task_tensor_size";1277const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_SIZE = "_ffts_sub_task_tensor_size";
1277const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_OFFSETS = "_ffts_sub_task_tensor_offsets";1278const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_OFFSETS = "_ffts_sub_task_tensor_offsets";
1278const std::string ATTR_NAME_IS_FFTS_UNSUPPORTED = "_is_ffts_unsupported";1279const std::string ATTR_NAME_IS_FFTS_UNSUPPORTED = "_is_ffts_unsupported";
Minc/graph_metadef/graph/custom_op/args_refresh.h+10-0
@@ -11,12 +11,22 @@
11#ifndef METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_11#ifndef METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_
12#define METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_12#define METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_
13 13 
14+#include <cstdint>
15+ 
14namespace ge {16namespace ge {
15enum class ArgsRefreshStrategy {17enum class ArgsRefreshStrategy {
16 kNone = 0,18 kNone = 0,
17 kAnnotatedArgs,19 kAnnotatedArgs,
18 kUpdateCallback,20 kUpdateCallback,
19};21};
22+ 
23+// Serialized in ATTR_NAME_CUSTOM_TASK_ARGS_MODE. Keep the numeric values stable for OM compatibility.
24+enum class CustomTaskArgsMode : int64_t {
25+ kUnspecified = 0,
26+ kNone = 1,
27+ kAnnotatedArgs = 2,
28+ kUpdateCallback = 3,
29+};
20} // namespace ge30} // namespace ge
21 31 
22#endif // METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_32#endif // METADEF_CXX_INC_GRAPH_CUSTOM_OP_ARGS_REFRESH_H_
Minc/graph_metadef/graph/debug/ge_attr_define.h+1-0
@@ -1275,6 +1275,7 @@ GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAM
1275GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_CACHE_PERSIST;1275GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_CACHE_PERSIST;
1276GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_ALIAS_ENGINE_NAME;1276GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_ALIAS_ENGINE_NAME;
1277GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_KERNEL_NAMES_PREFIX;1277GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_KERNEL_NAMES_PREFIX;
1278+GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_CUSTOM_TASK_ARGS_MODE;
1278GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_SIZE;1279GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_SIZE;
1279GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_OFFSETS;1280GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_FFTS_SUB_TASK_TENSOR_OFFSETS;
1280GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_IS_FFTS_UNSUPPORTED;1281GE_FUNC_DEV_VISIBILITY GE_FUNC_HOST_VISIBILITY extern const std::string ATTR_NAME_IS_FFTS_UNSUPPORTED;
Mruntime/custom_op/python_custom_op_adapter.cc+13-0
@@ -220,6 +220,19 @@ graphStatus PythonCustomOpAdapter::Execute(gert::EagerOpExecutionContext *ctx) {
220 return holder_->GetCallbacks().execute(holder_->GetHolder(), ctx);220 return holder_->GetCallbacks().execute(holder_->GetHolder(), ctx);
221}221}
222 222 
223+graphStatus PythonCustomOpAdapter::DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) {
224+ if (!HasCapability(CustomOpCapability::kAnnotatedArgs)) {
225+ return ReportUnsupported(CustomOpCapability::kAnnotatedArgs, "DeclareLaunchArgs");
226+ }
227+ if ((holder_ == nullptr) || (!holder_->IsValid()) || (holder_->GetHolder() == nullptr) ||
228+ (holder_->GetCallbacks().declare_launch_args == nullptr)) {
229+ GELOGE(GRAPH_FAILED, "Python custom op adapter is invalid, descriptor key[%s], op type[%s].",
230+ desc_.descriptor_key.c_str(), desc_.op_type.c_str());
231+ return GRAPH_FAILED;
232+ }
233+ return holder_->GetCallbacks().declare_launch_args(holder_->GetHolder(), &ctx);
234+}
235+ 
223graphStatus PythonCustomOpAdapter::Compile(gert::OpCompileContext *ctx) {236graphStatus PythonCustomOpAdapter::Compile(gert::OpCompileContext *ctx) {
224 (void)ctx;237 (void)ctx;
225 return ReportUnsupported(CustomOpCapability::kCompilable, "Compile");238 return ReportUnsupported(CustomOpCapability::kCompilable, "Compile");
Mruntime/custom_op/python_custom_op_adapter.h+2-0
@@ -56,6 +56,7 @@ class PythonCustomOpHolder {
56};56};
57 57 
58class PythonCustomOpAdapter final : public EagerExecuteOp,58class PythonCustomOpAdapter final : public EagerExecuteOp,
59+ public AnnotatedArgsOp,
59 public CompilableOp,60 public CompilableOp,
60 public ShapeInferOp,61 public ShapeInferOp,
61 public PortableOp,62 public PortableOp,
@@ -69,6 +70,7 @@ class PythonCustomOpAdapter final : public EagerExecuteOp,
69 bool HasCapability(CustomOpCapability capability) const override;70 bool HasCapability(CustomOpCapability capability) const override;
70 71 
71 graphStatus Execute(gert::EagerOpExecutionContext *ctx) override;72 graphStatus Execute(gert::EagerOpExecutionContext *ctx) override;
73+ graphStatus DeclareLaunchArgs(gert::AnnotatedArgsContext &ctx) override;
72 graphStatus Compile(gert::OpCompileContext *ctx) override;74 graphStatus Compile(gert::OpCompileContext *ctx) override;
73 graphStatus InferShape(gert::InferShapeContext *ctx) override;75 graphStatus InferShape(gert::InferShapeContext *ctx) override;
74 graphStatus InferDataType(gert::InferDataTypeContext *ctx) override;76 graphStatus InferDataType(gert::InferDataTypeContext *ctx) override;
Mruntime/custom_op/python_custom_op_bridge_types.h+9-2
@@ -18,8 +18,9 @@
18#include "graph/error_codes.h"18#include "graph/error_codes.h"
19 19 
20namespace gert {20namespace gert {
21+class AnnotatedArgsContext;
21class EagerOpExecutionContext;22class EagerOpExecutionContext;
22-}23+} // namespace gert
23 24 
24namespace ge {25namespace ge {
25namespace custom_op {26namespace custom_op {
@@ -35,14 +36,17 @@ struct PythonCustomOpDescriptor {
35using PythonCustomOpHolderCreateFn = void *(*)(const PythonCustomOpDescriptor *desc);36using PythonCustomOpHolderCreateFn = void *(*)(const PythonCustomOpDescriptor *desc);
36using PythonCustomOpHolderDestroyFn = void (*)(void *holder);37using PythonCustomOpHolderDestroyFn = void (*)(void *holder);
37using PythonCustomOpExecuteFn = graphStatus (*)(const void *holder, gert::EagerOpExecutionContext *ctx);38using PythonCustomOpExecuteFn = graphStatus (*)(const void *holder, gert::EagerOpExecutionContext *ctx);
39+using PythonCustomOpDeclareLaunchArgsFn = graphStatus (*)(const void *holder, gert::AnnotatedArgsContext *ctx);
38 40 
39struct PythonCustomOpCallbacks {41struct PythonCustomOpCallbacks {
40 PythonCustomOpHolderCreateFn create{nullptr};42 PythonCustomOpHolderCreateFn create{nullptr};
41 PythonCustomOpHolderDestroyFn destroy{nullptr};43 PythonCustomOpHolderDestroyFn destroy{nullptr};
42 PythonCustomOpExecuteFn execute{nullptr};44 PythonCustomOpExecuteFn execute{nullptr};
45+ PythonCustomOpDeclareLaunchArgsFn declare_launch_args{nullptr};
43 46 
44 bool IsValid(CustomOpCapabilityMask capabilities) const {47 bool IsValid(CustomOpCapabilityMask capabilities) const {
45- const auto supported_capabilities = static_cast<CustomOpCapabilityMask>(CustomOpCapability::kEagerExecute);48+ const auto supported_capabilities = static_cast<CustomOpCapabilityMask>(CustomOpCapability::kEagerExecute) |
49+ static_cast<CustomOpCapabilityMask>(CustomOpCapability::kAnnotatedArgs);
46 if ((capabilities == 0U) || ((capabilities & (~supported_capabilities)) != 0U)) {50 if ((capabilities == 0U) || ((capabilities & (~supported_capabilities)) != 0U)) {
47 return false;51 return false;
48 }52 }
@@ -52,6 +56,9 @@ struct PythonCustomOpCallbacks {
52 if (HasCustomOpCapability(capabilities, CustomOpCapability::kEagerExecute) && (execute == nullptr)) {56 if (HasCustomOpCapability(capabilities, CustomOpCapability::kEagerExecute) && (execute == nullptr)) {
53 return false;57 return false;
54 }58 }
59+ if (HasCustomOpCapability(capabilities, CustomOpCapability::kAnnotatedArgs) && (declare_launch_args == nullptr)) {
60+ return false;
61+ }
55 return true;62 return true;
56 }63 }
57};64};
Mruntime/custom_op/python_custom_op_pybind_bridge.cc+67-1
@@ -47,6 +47,7 @@ constexpr const char *kBridgeModuleName = "ge.custom_op._bridge";
47constexpr const char *kCustomOpModuleName = "ge.custom_op";47constexpr const char *kCustomOpModuleName = "ge.custom_op";
48constexpr const char *kCustomOpNativeModuleName = "ge.custom_op._ge_custom_op_native";48constexpr const char *kCustomOpNativeModuleName = "ge.custom_op._ge_custom_op_native";
49constexpr const char *kEnvCustomOppPath = "ASCEND_CUSTOM_OPP_PATH";49constexpr const char *kEnvCustomOppPath = "ASCEND_CUSTOM_OPP_PATH";
50+constexpr const char *kInterfaceAnnotatedArgs = "annotated_args";
50constexpr const char *kInterfaceEagerExecute = "eager_execute";51constexpr const char *kInterfaceEagerExecute = "eager_execute";
51constexpr const char *kGetRegisteredIrDefSymbol = "GetRegisteredIrDef";52constexpr const char *kGetRegisteredIrDefSymbol = "GetRegisteredIrDef";
52constexpr const char *kRunnerLibraryNames[] = {53constexpr const char *kRunnerLibraryNames[] = {
@@ -205,6 +206,8 @@ CustomOpCapabilityMask ParseInterfaces(const py::object &interfaces_obj) {
205 const std::string interface_name = py::str(item);206 const std::string interface_name = py::str(item);
206 if (interface_name == kInterfaceEagerExecute) {207 if (interface_name == kInterfaceEagerExecute) {
207 AddCustomOpCapability(capabilities, CustomOpCapability::kEagerExecute);208 AddCustomOpCapability(capabilities, CustomOpCapability::kEagerExecute);
209+ } else if (interface_name == kInterfaceAnnotatedArgs) {
210+ AddCustomOpCapability(capabilities, CustomOpCapability::kAnnotatedArgs);
208 }211 }
209 }212 }
210 return capabilities;213 return capabilities;
@@ -248,6 +251,8 @@ class PythonCustomOpPybindBridge {
248 }251 }
249 252 
250 const py::list descriptor_list = descriptors_obj.cast<py::list>();253 const py::list descriptor_list = descriptors_obj.cast<py::list>();
254+ std::vector<PythonCustomOpDescriptor> descriptors;
255+ descriptors.reserve(descriptor_list.size());
251 for (const auto &item : descriptor_list) {256 for (const auto &item : descriptor_list) {
252 PythonCustomOpDescriptor desc;257 PythonCustomOpDescriptor desc;
253 const auto parse_ret = ParseDescriptor(item.cast<py::dict>(), desc);258 const auto parse_ret = ParseDescriptor(item.cast<py::dict>(), desc);
@@ -255,7 +260,26 @@ class PythonCustomOpPybindBridge {
255 GELOGE(parse_ret, "Parse python custom op descriptor failed.");260 GELOGE(parse_ret, "Parse python custom op descriptor failed.");
256 return parse_ret;261 return parse_ret;
257 }262 }
258- const auto callbacks = GetCallbacks();263+ auto ir_meta = CollectPythonCustomOpIrMeta(desc.op_type);
264+ try {
265+ const bool validated =
266+ bridge_module_.attr("validate_op_impl_descriptor")(desc.descriptor_key, BuildPythonIrMeta(ir_meta.get()))
267+ .cast<bool>();
268+ if (!validated) {
269+ GELOGE(FAILED, "Validate python custom op[%s] descriptor[%s] failed.", desc.op_type.c_str(),
270+ desc.descriptor_key.c_str());
271+ return FAILED;
272+ }
273+ } catch (const py::error_already_set &err) {
274+ GELOGE(FAILED, "Validate python custom op[%s] descriptor[%s] failed: %s", desc.op_type.c_str(),
275+ desc.descriptor_key.c_str(), err.what());
276+ return FAILED;
277+ }
278+ descriptors.emplace_back(std::move(desc));
279+ }
280+ 
281+ const auto callbacks = GetCallbacks();
282+ for (const auto &desc : descriptors) {
259 if ((registrar.register_custom_op == nullptr) || (!registrar.register_custom_op(&desc, &callbacks))) {283 if ((registrar.register_custom_op == nullptr) || (!registrar.register_custom_op(&desc, &callbacks))) {
260 GELOGE(FAILED, "Register python custom op[%s] failed.", desc.op_type.c_str());284 GELOGE(FAILED, "Register python custom op[%s] failed.", desc.op_type.c_str());
261 return FAILED;285 return FAILED;
@@ -367,6 +391,39 @@ class PythonCustomOpPybindBridge {
367 }391 }
368 }392 }
369 393 
394+ graphStatus DeclareLaunchArgs(const PythonCustomOpBridgeHolder *holder, gert::AnnotatedArgsContext *ctx) {
395+ if ((holder == nullptr) || (ctx == nullptr)) {
396+ GELOGE(GRAPH_FAILED, "Python custom op bridge holder or context is null.");
397+ return GRAPH_FAILED;
398+ }
399+ const auto prepare_ret = EnsureBridgeReady();
400+ if (prepare_ret != SUCCESS) {
401+ GELOGE(prepare_ret, "Prepare python custom op bridge failed.");
402+ return GRAPH_FAILED;
403+ }
404+ py::gil_scoped_acquire gil;
405+ try {
406+ const bool created =
407+ bridge_module_.attr("create_op_impl_holder")(holder->instance_id, holder->descriptor_key).cast<bool>();
408+ if (!created) {
409+ GELOGE(GRAPH_FAILED, "Ensure python custom op holder failed, descriptor key[%s], instance id[%s].",
410+ holder->descriptor_key.c_str(), holder->instance_id.c_str());
411+ return GRAPH_FAILED;
412+ }
413+ py::object result = bridge_module_.attr("call_declare_launch_args")(
414+ holder->instance_id, BuildPythonIrMeta(holder->ir_meta.get()), BuildPythonAnnotatedArgsContext(ctx));
415+ return TranslateStatusLike(result);
416+ } catch (const py::error_already_set &err) {
417+ GELOGE(GRAPH_FAILED, "DeclareLaunchArgs python custom op failed, descriptor key[%s], instance id[%s]: %s",
418+ holder->descriptor_key.c_str(), holder->instance_id.c_str(), err.what());
419+ return GRAPH_FAILED;
420+ } catch (const std::exception &err) {
421+ GELOGE(GRAPH_FAILED, "DeclareLaunchArgs python custom op failed, descriptor key[%s], instance id[%s]: %s",
422+ holder->descriptor_key.c_str(), holder->instance_id.c_str(), err.what());
423+ return GRAPH_FAILED;
424+ }
425+ }
426+ 
370 private:427 private:
371 Status EnsureBridgeReady() {428 Status EnsureBridgeReady() {
372 std::lock_guard<std::mutex> lock(mutex_);429 std::lock_guard<std::mutex> lock(mutex_);
@@ -536,6 +593,11 @@ class PythonCustomOpPybindBridge {
536 return native_module.attr("_borrow_eager_op_execution_context")(py::int_(reinterpret_cast<uintptr_t>(ctx)));593 return native_module.attr("_borrow_eager_op_execution_context")(py::int_(reinterpret_cast<uintptr_t>(ctx)));
537 }594 }
538 595 
596+ static py::object BuildPythonAnnotatedArgsContext(gert::AnnotatedArgsContext *ctx) {
597+ py::module_ native_module = py::module_::import(kCustomOpNativeModuleName);
598+ return native_module.attr("_borrow_annotated_args_context")(py::int_(reinterpret_cast<uintptr_t>(ctx)));
599+ }
600+ 
539 static graphStatus TranslateStatusLike(const py::object &result) {601 static graphStatus TranslateStatusLike(const py::object &result) {
540 if (result.is_none()) {602 if (result.is_none()) {
541 return GRAPH_SUCCESS;603 return GRAPH_SUCCESS;
@@ -566,6 +628,10 @@ class PythonCustomOpPybindBridge {
566 return PythonCustomOpPybindBridge::GetInstance().Execute(static_cast<const PythonCustomOpBridgeHolder *>(holder),628 return PythonCustomOpPybindBridge::GetInstance().Execute(static_cast<const PythonCustomOpBridgeHolder *>(holder),
567 ctx);629 ctx);
568 };630 };
631+ callbacks.declare_launch_args = [](const void *holder, gert::AnnotatedArgsContext *ctx) -> graphStatus {
632+ return PythonCustomOpPybindBridge::GetInstance().DeclareLaunchArgs(
633+ static_cast<const PythonCustomOpBridgeHolder *>(holder), ctx);
634+ };
569 return callbacks;635 return callbacks;
570 }636 }
571 637 
Mruntime/v1/graph/load/model_manager/task_info/ge/custom_task_info.cc+49-2
@@ -174,6 +174,54 @@ Status GetCustomKernelBinaryMagic(const OpDescPtr &op_desc, int32_t &binary_magi
174std::string MakeCustomKernelBinName(const uint32_t model_id, const OpDescPtr &op_desc, const std::string &kernel_name) {174std::string MakeCustomKernelBinName(const uint32_t model_id, const OpDescPtr &op_desc, const std::string &kernel_name) {
175 return std::to_string(model_id) + "_" + op_desc->GetName() + "_" + kernel_name;175 return std::to_string(model_id) + "_" + op_desc->GetName() + "_" + kernel_name;
176}176}
177+ 
178+ArgsRefreshStrategy GetLegacyCustomTaskArgsRefreshStrategy(const AscendString &op_type,
179+ const domi::KernelContext &context,
180+ const CustomOpRegistryPtr &custom_op_registry) {
181+ const auto registry_strategy = custom_op_registry->GetArgsRefreshStrategy(op_type);
182+ if ((registry_strategy == ArgsRefreshStrategy::kNone) && (!context.args_format().empty())) {
183+ return ArgsRefreshStrategy::kAnnotatedArgs;
184+ }
185+ return registry_strategy;
186+}
187+ 
188+Status GetCustomTaskArgsRefreshStrategy(const OpDescPtr &op_desc, const domi::KernelContext &context,
189+ const CustomOpRegistryPtr &custom_op_registry,
190+ ArgsRefreshStrategy &args_refresh_strategy) {
191+ GE_ASSERT_NOTNULL(op_desc);
192+ GE_ASSERT_NOTNULL(custom_op_registry);
193+ const AscendString op_type(op_desc->GetTypePtr());
194+ if (!op_desc->HasAttr(ATTR_NAME_CUSTOM_TASK_ARGS_MODE)) {
195+ args_refresh_strategy = GetLegacyCustomTaskArgsRefreshStrategy(op_type, context, custom_op_registry);
196+ return SUCCESS;
197+ }
198+ 
199+ int64_t args_mode = 0;
200+ GE_ASSERT_TRUE(AttrUtils::GetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE, args_mode),
201+ "[CUSTOM OP] get %s failed for op %s(%s).", ATTR_NAME_CUSTOM_TASK_ARGS_MODE.c_str(),
202+ op_desc->GetNamePtr(), op_desc->GetTypePtr());
203+ switch (static_cast<CustomTaskArgsMode>(args_mode)) {
204+ case CustomTaskArgsMode::kUnspecified:
205+ args_refresh_strategy = GetLegacyCustomTaskArgsRefreshStrategy(op_type, context, custom_op_registry);
206+ return SUCCESS;
207+ case CustomTaskArgsMode::kNone:
208+ args_refresh_strategy = ArgsRefreshStrategy::kNone;
209+ return SUCCESS;
210+ case CustomTaskArgsMode::kAnnotatedArgs:
211+ args_refresh_strategy = ArgsRefreshStrategy::kAnnotatedArgs;
212+ return SUCCESS;
213+ case CustomTaskArgsMode::kUpdateCallback:
214+ GE_ASSERT_TRUE(custom_op_registry->GetArgsRefreshStrategy(op_type) == ArgsRefreshStrategy::kUpdateCallback,
215+ "[CUSTOM OP] update callback is not registered for op %s(%s).", op_desc->GetNamePtr(),
216+ op_desc->GetTypePtr());
217+ args_refresh_strategy = ArgsRefreshStrategy::kUpdateCallback;
218+ return SUCCESS;
219+ default:
220+ GELOGE(PARAM_INVALID, "[CUSTOM OP] invalid %s value %" PRId64 " for op %s(%s).",
221+ ATTR_NAME_CUSTOM_TASK_ARGS_MODE.c_str(), args_mode, op_desc->GetNamePtr(), op_desc->GetTypePtr());
222+ return PARAM_INVALID;
223+ }
224+}
177} // namespace225} // namespace
178 226 
179void CustomTaskInfo::SetCustomDumpInfo(const DumpProperties &dump_properties, DumpOp &dump_op) const {227void CustomTaskInfo::SetCustomDumpInfo(const DumpProperties &dump_properties, DumpOp &dump_op) const {
@@ -281,11 +329,10 @@ Status CustomTaskInfo::ParseTaskRunParam(const domi::TaskDef &task_def, DavinciM
281 workspace_addrs_ = ModelUtils::GetWorkspaceDataAddrsValue(rts_param, op_desc_, workspace_mem_types_);329 workspace_addrs_ = ModelUtils::GetWorkspaceDataAddrsValue(rts_param, op_desc_, workspace_mem_types_);
282 GE_ASSERT_SUCCESS(ValidateIoWorkspaceAddrAndMemTypeSizes());330 GE_ASSERT_SUCCESS(ValidateIoWorkspaceAddrAndMemTypeSizes());
283 331 
284- AscendString op_type(op_desc_->GetType().c_str());
285 const auto &custom_op_registry = davinci_model->GetCustomOpRegistry();332 const auto &custom_op_registry = davinci_model->GetCustomOpRegistry();
286 GE_ASSERT_NOTNULL(custom_op_registry, "[CUSTOM OP] custom op registry is nullptr for op %s.",333 GE_ASSERT_NOTNULL(custom_op_registry, "[CUSTOM OP] custom op registry is nullptr for op %s.",
287 op_desc_->GetName().c_str());334 op_desc_->GetName().c_str());
288- args_refresh_strategy_ = custom_op_registry->GetArgsRefreshStrategy(op_type);335+ GE_ASSERT_SUCCESS(GetCustomTaskArgsRefreshStrategy(op_desc_, context, custom_op_registry, args_refresh_strategy_));
289 is_args_refreshable_ = args_refresh_strategy_ != ArgsRefreshStrategy::kNone;336 is_args_refreshable_ = args_refresh_strategy_ != ArgsRefreshStrategy::kNone;
290 337 
291 if (args_refresh_strategy_ == ArgsRefreshStrategy::kAnnotatedArgs) {338 if (args_refresh_strategy_ == ArgsRefreshStrategy::kAnnotatedArgs) {
Mscripts/build_fwk.sh+4-0
@@ -253,6 +253,8 @@ run_pyge_pytests() {
253 PYGE_SRC_PATH=${BASEPATH}/api/python/ge253 PYGE_SRC_PATH=${BASEPATH}/api/python/ge
254 PYGE_COVERAGE_RCFILE=${BUILD_PATH}/pyge.coveragerc254 PYGE_COVERAGE_RCFILE=${BUILD_PATH}/pyge.coveragerc
255 ORIGINAL_LD_LIBRARY_PATH=$LD_LIBRARY_PATH255 ORIGINAL_LD_LIBRARY_PATH=$LD_LIBRARY_PATH
256+ ORIGINAL_LD_PRELOAD=$LD_PRELOAD
257+ PYGE_USE_LIBSTDCXX=$(LD_PRELOAD= gcc -print-file-name=libstdc++.so.6)
GengChao
GengChaoGengChao16 天前

这个修改是为啥？

likedislike
S
shangdf
16 天前 评论：
GengChao
GengChao
16 天前 评论：
S
shangdf
15 天前 评论：
256 cat > ${PYGE_COVERAGE_RCFILE} <<EOF258 cat > ${PYGE_COVERAGE_RCFILE} <<EOF
257[run]259[run]
258source =260source =
@@ -268,8 +270,10 @@ EOF
268 # 先走安装目录，确保 ge.passes 能导到安装产物里的 _ge_pass_native.so；270 # 先走安装目录，确保 ge.passes 能导到安装产物里的 _ge_pass_native.so；
269 # 覆盖率通过 coveragerc 的 [paths] 映射回源码目录。271 # 覆盖率通过 coveragerc 的 [paths] 映射回源码目录。
270 export PYTHONPATH=${PYGE_INSTALL_PATH}:${PYGE_SRC_PATH}:$PYTHON_ORIGINAL_PATH272 export PYTHONPATH=${PYGE_INSTALL_PATH}:${PYGE_SRC_PATH}:$PYTHON_ORIGINAL_PATH
273+ export LD_PRELOAD=${USE_ASAN}:${PYGE_USE_LIBSTDCXX}
271 ASAN_OPTIONS=detect_leaks=0:detect_container_overflow=0 ${HI_PYTHON} -m coverage run --rcfile=${PYGE_COVERAGE_RCFILE} --data-file=coverage_pyge -m pytest ${BASEPATH}/tests/ge/ut/ge/graph/pyge_tests/*_test.py -vv -rs -s274 ASAN_OPTIONS=detect_leaks=0:detect_container_overflow=0 ${HI_PYTHON} -m coverage run --rcfile=${PYGE_COVERAGE_RCFILE} --data-file=coverage_pyge -m pytest ${BASEPATH}/tests/ge/ut/ge/graph/pyge_tests/*_test.py -vv -rs -s
272 export LD_LIBRARY_PATH=${ORIGINAL_LD_LIBRARY_PATH}275 export LD_LIBRARY_PATH=${ORIGINAL_LD_LIBRARY_PATH}
276+ export LD_PRELOAD=${ORIGINAL_LD_PRELOAD}
273 export PYTHONPATH=$PYTHON_ORIGINAL_PATH277 export PYTHONPATH=$PYTHON_ORIGINAL_PATH
274}278}
275 279 
Mtests/ge/st/testcase/test_custom_op.cc+406-12
@@ -13,6 +13,7 @@
13#include <cstdlib>13#include <cstdlib>
14#include <dlfcn.h>14#include <dlfcn.h>
15#include <fstream>15#include <fstream>
16+#include <map>
16#include <memory>17#include <memory>
17#include <mutex>18#include <mutex>
18#include <numeric>19#include <numeric>
@@ -29,6 +30,8 @@
29#include "graph/utils/graph_utils_ex.h"30#include "graph/utils/graph_utils_ex.h"
30#include "graph/utils/op_desc_utils.h"31#include "graph/utils/op_desc_utils.h"
31#include "graph/load/model_manager/model_utils.h"32#include "graph/load/model_manager/model_utils.h"
33+#include "graph/load/model_manager/davinci_model.h"
34+#include "graph/load/model_manager/task_info/ge/custom_task_info.h"
32#include "ge_graph_dsl/assert/graph_assert.h"35#include "ge_graph_dsl/assert/graph_assert.h"
33#include "utils/mock_ops_kernel_builder.h"36#include "utils/mock_ops_kernel_builder.h"
34#include "utils/taskdef_builder.h"37#include "utils/taskdef_builder.h"
@@ -50,12 +53,33 @@
50#include "hcom/hcom_topo_info.h"53#include "hcom/hcom_topo_info.h"
51#include "common/opskernel/ops_kernel_info_types.h"54#include "common/opskernel/ops_kernel_info_types.h"
52#include "engines/custom_engine/custom_graph_optimizer.h"55#include "engines/custom_engine/custom_graph_optimizer.h"
56+#include "engines/custom_engine/custom_ops_kernel_builder.h"
57+#include "graph/compute_graph.h"
58+#include "graph/custom_op/cast.h"
53#include "graph/custom_op_factory.h"59#include "graph/custom_op_factory.h"
54#include "graph/custom_op.h"60#include "graph/custom_op.h"
61+#include "graph/ge_tensor.h"
62+#include "graph/operator_reg.h"
63+#include "graph/op_desc.h"
64+#include "graph/utils/args_format_desc_utils.h"
55#include "common/python_runtime/ge_python_runtime_manager.h"65#include "common/python_runtime/ge_python_runtime_manager.h"
56#include "runtime/custom_op/custom_op_loader.h"66#include "runtime/custom_op/custom_op_loader.h"
57#include "runtime/custom_op/python_custom_op_bridge_loader.h"67#include "runtime/custom_op/python_custom_op_bridge_loader.h"
58 68 
69+namespace ge {
70+REG_OP(StPythonAnnotatedArgsCustomOp)
71+ .INPUT(x, TensorType::ALL())
72+ .OUTPUT(z, TensorType::ALL())
73+ .REQUIRED_ATTR(alpha, Int)
74+ .OP_END_FACTORY_REG(StPythonAnnotatedArgsCustomOp);
75+ 
76+REG_OP(StPythonAnnotatedArgsBadAttrCustomOp)
77+ .INPUT(x, TensorType::ALL())
78+ .OUTPUT(z, TensorType::ALL())
79+ .REQUIRED_ATTR(alpha, Int)
80+ .OP_END_FACTORY_REG(StPythonAnnotatedArgsBadAttrCustomOp);
81+} // namespace ge
82+ 
59namespace ge {83namespace ge {
60using namespace gert;84using namespace gert;
61namespace {85namespace {
@@ -149,8 +173,117 @@ void MockGenerateTask() {
149void *output_addr = nullptr;173void *output_addr = nullptr;
150void **args_table = nullptr;174void **args_table = nullptr;
151constexpr const char *kPythonCustomOpTypeForSt = "StPythonPybindRemoveCoverageCustomOp";175constexpr const char *kPythonCustomOpTypeForSt = "StPythonPybindRemoveCoverageCustomOp";
176+constexpr const char *kPythonAnnotatedArgsOpTypeForSt = "StPythonAnnotatedArgsCustomOp";
177+constexpr const char *kPythonAnnotatedArgsBadAttrOpTypeForSt = "StPythonAnnotatedArgsBadAttrCustomOp";
152constexpr const char *kEnvPythonCustomOpPath = "ASCEND_CUSTOM_OPP_PATH";178constexpr const char *kEnvPythonCustomOpPath = "ASCEND_CUSTOM_OPP_PATH";
153constexpr const char *kEnvPythonPath = "PYTHONPATH";179constexpr const char *kEnvPythonPath = "PYTHONPATH";
180+constexpr char kSharedPybindCustomOpPreambleForSt[] = R"PY(from pathlib import Path
181+from ge.custom_op import (
182+ AnnotatedKernelLaunchInfo,
183+ EagerExecuteOp,
184+ get_declare_launch_args_ctx,
185+ register_op_impl,
186+)
187+from ge.runtime import Tensor
188+ 
189+MARKER_FILE = r')PY";
190+constexpr char kSharedPybindEagerCustomOpForSt[] = R"PY('
191+ 
192+@register_op_impl(op_type=')PY";
193+constexpr char kSharedPybindAnnotatedArgsPrefixForSt[] = R"PY(')
194+class StPythonPybindRemoveCoverageCustomOp(EagerExecuteOp):
195+ def execute(self, ctx):
196+ Path(MARKER_FILE).write_text('executed', encoding='utf-8')
197+ 
198+@register_op_impl(op_type=')PY";
199+constexpr char kSharedPybindAnnotatedArgsBodyForSt[] = R"PY(')
200+class StPythonAnnotatedArgsCustomOp:
201+ def __init__(self):
202+ self.saved = None
203+ 
204+ def declare_launch_args(self, x: Tensor, z: Tensor, *, alpha: int) -> None:
205+ ctx = get_declare_launch_args_ctx()
206+ args = ctx.create_kernel_args()
207+ if alpha == 1:
208+ args.append_input(2, x)
209+ if alpha == 2:
210+ args.append_input(-1, x)
211+ if alpha == 3:
212+ args.append_scalar(-1)
213+ if alpha == 9:
214+ try:
215+ args.append_input(2, x)
216+ except IndexError as error:
217+ if 'index 2' not in str(error):
218+ raise AssertionError('input index is missing from error')
219+ else:
220+ raise AssertionError('invalid input index was accepted')
221+ try:
222+ args.append_output(3, z)
223+ except IndexError as error:
224+ if 'index 3' not in str(error):
225+ raise AssertionError('output index is missing from error')
226+ else:
227+ raise AssertionError('invalid output index was accepted')
228+ if alpha == 7:
229+ saved_tensor, saved_workspace, saved_args = self.saved
230+ for access in (
231+ lambda: saved_tensor.addr,
232+ lambda: saved_workspace.index,
233+ lambda: saved_args.append_scalar(0),
234+ ):
235+ try:
236+ access()
237+ except RuntimeError:
238+ pass
239+ else:
240+ raise AssertionError('borrowed DLA object did not expire')
241+ args.append_input(0, x)
242+ args.append_output(0, z)
243+ args.append_scalar(alpha)
244+ workspace = None
245+ if alpha == 6:
246+ workspace = ctx.malloc_workspace(64)
247+ for name in ('index', 'addr'):
248+ try:
249+ setattr(workspace, name, 0)
250+ except AttributeError:
251+ pass
252+ else:
253+ raise AssertionError('workspace property is writable')
254+ args.append_workspace(workspace)
255+ stream_id = ctx.get_stream_id()
256+ if alpha == 4:
257+ stream_id += 1
258+ ctx.add_launch(
259+ AnnotatedKernelLaunchInfo(
260+ kernel_name='st_python_dla',
261+ kernel_bin=b'\x01\x02',
262+ block_dim=1,
263+ stream_id=stream_id,
264+ ),
265+ args,
266+ )
267+ if alpha == 5:
268+ args.append_scalar(0)
269+ if alpha == 6:
270+ self.saved = (x, workspace, args)
271+)PY";
272+constexpr char kSharedPybindBadAttrCustomOpForSt[] = R"PY(')
273+class StPythonAnnotatedArgsBadAttrCustomOp:
274+ def declare_launch_args(self, x: Tensor, z: Tensor, *, beta: int) -> None:
275+ pass
276+)PY";
277+constexpr char kInvalidSignaturePybindPreambleForSt[] = R"PY(from ge.custom_op import register_op_impl
278+from ge.runtime import Tensor
279+ 
280+@register_op_impl(op_type=')PY";
281+constexpr char kValidBeforeInvalidPybindCustomOpForSt[] = R"PY(')
282+class StPythonValidBeforeInvalidCustomOp:
283+ def execute(self, ctx):
284+ pass
285+ 
286+@register_op_impl(op_type=')PY";
154 287 
155class ScopedTempDirForCustomOpSt {288class ScopedTempDirForCustomOpSt {
156 public:289 public:
@@ -210,6 +343,76 @@ class ScopedEnvVarForCustomOpSt {
210 bool has_old_value_{false};343 bool has_old_value_{false};
211};344};
212 345 
346+class ScopedGraphOptionsForCustomOpSt {
347+ public:
348+ explicit ScopedGraphOptionsForCustomOpSt(const std::map<std::string, std::string> &options)
349+ : old_options_(GetThreadLocalContext().GetAllGraphOptions()) {
350+ GetThreadLocalContext().SetGraphOption(options);
351+ }
352+ 
353+ ~ScopedGraphOptionsForCustomOpSt() {
354+ GetThreadLocalContext().SetGraphOption(old_options_);
355+ }
356+ 
357+ private:
358+ std::map<std::string, std::string> old_options_;
359+};
360+ 
361+Status GeneratePythonAnnotatedArgsTaskForSt(const char *const op_type, const int64_t alpha,
362+ const std::string &soc_version, std::vector<domi::TaskDef> &tasks,
363+ const bool is_unknown_shape = false) {
364+ const std::map<std::string, std::string> graph_options = {{SOC_VERSION, soc_version}};
365+ ScopedGraphOptionsForCustomOpSt scoped_graph_options(graph_options);
366+ auto graph = std::make_shared<ComputeGraph>(std::string("st_python_annotated_args_") + std::to_string(alpha));
367+ GE_ASSERT_NOTNULL(graph);
368+ graph->SetGraphUnknownFlag(is_unknown_shape);
369+ auto op_desc = std::make_shared<OpDesc>("st_python_annotated_args_node", op_type);
370+ GE_ASSERT_NOTNULL(op_desc);
371+ op_desc->SetId(7);
372+ op_desc->SetStreamId(3);
373+ op_desc->AppendIrInput("x", kIrInputRequired);
374+ op_desc->AppendIrOutput("z", kIrOutputRequired);
375+ op_desc->AppendIrAttrName("alpha");
376+ GE_ASSERT_TRUE(AttrUtils::SetInt(op_desc, "alpha", alpha));
377+ GeTensorDesc input_desc(GeShape({1, 16}), FORMAT_ND, DT_FLOAT16);
378+ input_desc.SetOriginShape(GeShape({1, 16}));
379+ GeTensorDesc output_desc(GeShape({1, 16}), FORMAT_ND, DT_FLOAT16);
380+ output_desc.SetOriginShape(GeShape({1, 16}));
381+ GE_ASSERT_GRAPH_SUCCESS(op_desc->AddInputDesc("x", input_desc));
382+ GE_ASSERT_GRAPH_SUCCESS(op_desc->AddOutputDesc("z", output_desc));
383+ op_desc->SetInputOffset({1024});
384+ op_desc->SetOutputOffset({2048});
385+ const auto node = graph->AddNode(op_desc);
386+ GE_ASSERT_NOTNULL(node);
387+ 
388+ RunContext run_context = {};
389+ run_context.dataMemBase = reinterpret_cast<uint8_t *>(0x80000000UL);
390+ run_context.dataMemSize = 4096U;
391+ custom::CustomOpsKernelBuilder builder;
392+ return builder.GenerateTask(*node, run_context, tasks);
393+}
394+ 
395+Status AddPythonAnnotatedArgsOpToModelForSt(DavinciModel &model, const uint32_t op_index, OpDescPtr &op_desc) {
396+ op_desc = std::make_shared<OpDesc>("st_python_annotated_args_loaded", kPythonAnnotatedArgsOpTypeForSt);
397+ GE_ASSERT_NOTNULL(op_desc);
398+ op_desc->SetId(op_index);
399+ op_desc->SetStreamId(3);
400+ op_desc->AppendIrInput("x", kIrInputRequired);
401+ op_desc->AppendIrOutput("z", kIrOutputRequired);
402+ GeTensorDesc input_desc(GeShape({1, 16}), FORMAT_ND, DT_FLOAT16);
403+ input_desc.SetOriginShape(GeShape({1, 16}));
404+ GeTensorDesc output_desc(GeShape({1, 16}), FORMAT_ND, DT_FLOAT16);
405+ output_desc.SetOriginShape(GeShape({1, 16}));
406+ GE_ASSERT_GRAPH_SUCCESS(op_desc->AddInputDesc("x", input_desc));
407+ GE_ASSERT_GRAPH_SUCCESS(op_desc->AddOutputDesc("z", output_desc));
408+ op_desc->SetInputOffset({1024});
409+ op_desc->SetOutputOffset({2048});
410+ GE_ASSERT_TRUE(AttrUtils::SetStr(op_desc, TVM_ATTR_NAME_MAGIC, "RT_DEV_BINARY_MAGIC_ELF_AIVEC"));
411+ model.op_list_[op_index] = op_desc;
412+ model.SetCustomOpRegistry(CustomOpFactory::GetGlobalRegistryPtr());
413+ return SUCCESS;
414+}
415+ 
213void WriteTextFileForCustomOpSt(const std::string &file_path, const std::string &content) {416void WriteTextFileForCustomOpSt(const std::string &file_path, const std::string &content) {
214 std::ofstream file(file_path, std::ios::out | std::ios::trunc);417 std::ofstream file(file_path, std::ios::out | std::ios::trunc);
215 ASSERT_TRUE(file.is_open());418 ASSERT_TRUE(file.is_open());
@@ -235,21 +438,30 @@ const std::string &GetSharedPybindCustomOpMarkerFilePathForSt() {
235 return path;438 return path;
236}439}
237 440 
441+const std::string &GetInvalidSignaturePybindCustomOpFilePathForSt() {
442+ static ScopedTempDirForCustomOpSt dir;
443+ static const std::string path = dir.CreateFilePath("pybind_invalid_signature_custom_op.py");
444+ return path;
445+}
446+ 
238void EnsureSharedPybindCustomOpFileForSt() {447void EnsureSharedPybindCustomOpFileForSt() {
239 static std::once_flag once;448 static std::once_flag once;
240 std::call_once(once, []() {449 std::call_once(once, []() {
241- WriteTextFileForCustomOpSt(GetSharedPybindCustomOpFilePathForSt(),450+ const auto python_file = std::string(kSharedPybindCustomOpPreambleForSt) +
242- "from pathlib import Path\n"451+ GetSharedPybindCustomOpMarkerFilePathForSt() + kSharedPybindEagerCustomOpForSt +
243- "from ge.custom_op import EagerExecuteOp, register_op_impl\n\n"452+ kPythonCustomOpTypeForSt + kSharedPybindAnnotatedArgsPrefixForSt +
244- "MARKER_FILE = r'" +453+ kPythonAnnotatedArgsOpTypeForSt + kSharedPybindAnnotatedArgsBodyForSt;
245- GetSharedPybindCustomOpMarkerFilePathForSt() +454+ WriteTextFileForCustomOpSt(GetSharedPybindCustomOpFilePathForSt(), python_file);
246- "'\n\n"455+ });
247- "@register_op_impl(op_type='" +456+}
248- std::string(kPythonCustomOpTypeForSt) +457+ 
249- "')\n"458+void EnsureInvalidSignaturePybindCustomOpFileForSt() {
250- "class StPythonPybindRemoveCoverageCustomOp(EagerExecuteOp):\n"459+ static std::once_flag once;
251- " def execute(self, ctx):\n"460+ std::call_once(once, []() {
252- " Path(MARKER_FILE).write_text('executed', encoding='utf-8')\n");461+ const auto python_file = std::string(kInvalidSignaturePybindPreambleForSt) + kPythonCustomOpTypeForSt +
462+ kValidBeforeInvalidPybindCustomOpForSt + kPythonAnnotatedArgsBadAttrOpTypeForSt +
463+ kSharedPybindBadAttrCustomOpForSt;
464+ WriteTextFileForCustomOpSt(GetInvalidSignaturePybindCustomOpFilePathForSt(), python_file);
253 });465 });
254}466}
255 467 
@@ -1540,6 +1752,29 @@ TEST_F(CustomOpRefreshTest, eager_only_op_with_malloc_read_only_dev_args) {
1540 ReInitGe();1752 ReInitGe();
1541}1753}
1542 1754 
1755+/**
1756+ * 用例描述：测试Python自定义算子在注册阶段拒绝与IR不匹配的回调签名。
1757+ * 预置条件：
1758+ * 1. 构造一个合法legacy实现和一个属性名与REG_OP定义不一致的Python实现。
1759+ * 测试步骤：
1760+ * 1. 通过LoadPythonCustomOps加载Python实现。
1761+ * 2. 查询CustomOpFactory中是否存在该Python自定义算子creator。
1762+ * 预期结果：
1763+ * 1. 注册阶段签名校验失败，LoadPythonCustomOps返回FAILED。
1764+ * 2. CustomOpFactory中不存在合法或非法Python自定义算子creator。
1765+ */
1766+TEST_F(CustomOpFactoryStTest, PythonCustomOpLoaderRejectsInvalidSignatureDuringRegistration) {
1767+ EnsureInvalidSignaturePybindCustomOpFileForSt();
1768+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath,
1769+ GetInvalidSignaturePybindCustomOpFilePathForSt());
1770+ 
1771+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1772+ EXPECT_EQ(custom_op::LoadPythonCustomOps(), FAILED);
1773+ EXPECT_EQ(CustomOpFactory::CreateOrGetCustomOp(AscendString(kPythonCustomOpTypeForSt)), nullptr);
1774+ EXPECT_EQ(CustomOpFactory::CreateOrGetCustomOp(AscendString(kPythonAnnotatedArgsBadAttrOpTypeForSt)), nullptr);
1775+ custom_op::UnloadPythonCustomOps();
1776+}
1777+ 
1543/**1778/**
1544 * 用例描述：测试Python自定义算子通过loader注册、执行后，可以按类型移除注册信息和已创建实例。1779 * 用例描述：测试Python自定义算子通过loader注册、执行后，可以按类型移除注册信息和已创建实例。
1545 * 预置条件：1780 * 预置条件：
@@ -1581,6 +1816,165 @@ TEST_F(CustomOpFactoryStTest, remove_python_custom_ops_clears_creator_and_create
1581 EXPECT_EQ(CustomOpFactory::CreateOrGetCustomOp(op_type), nullptr);1816 EXPECT_EQ(CustomOpFactory::CreateOrGetCustomOp(op_type), nullptr);
1582}1817}
1583 1818 
1819+TEST_F(CustomOpFactoryStTest, PythonAnnotatedArgsCustomOpLoaderGeneratesTaskDef) {
1820+ EnsureSharedPybindCustomOpFileForSt();
1821+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath, GetSharedPybindCustomOpFilePathForSt());
1822+ 
1823+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1824+ ASSERT_EQ(custom_op::LoadPythonCustomOps(), SUCCESS);
1825+ ScopedLoadedPythonCustomOpsForSt loaded_python_custom_ops;
1826+ 
1827+ const AscendString op_type(kPythonAnnotatedArgsOpTypeForSt);
1828+ ASSERT_TRUE(CustomOpFactory::IsExistOp(op_type));
1829+ auto *const base_op = CustomOpFactory::CreateOrGetCustomOp(op_type);
1830+ ASSERT_NE(base_op, nullptr);
1831+ EXPECT_NE(CustomOpCast<AnnotatedArgsOp>(base_op), nullptr);
1832+ EXPECT_EQ(CustomOpCast<EagerExecuteOp>(base_op), nullptr);
1833+ 
1834+ std::vector<domi::TaskDef> tasks;
1835+ ASSERT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 0, "Ascend910B", tasks), SUCCESS);
1836+ 
1837+ ASSERT_EQ(tasks.size(), 1U);
1838+ const auto &task = tasks[0];
1839+ EXPECT_EQ(task.stream_id(), 3U);
1840+ const auto &kernel = task.kernel();
1841+ EXPECT_EQ(kernel.kernel_name(), "st_python_dla");
1842+ EXPECT_EQ(kernel.stub_func(), "st_python_dla");
1843+ EXPECT_EQ(kernel.block_dim(), 1U);
1844+ EXPECT_EQ(kernel.args_size(), 24U);
1845+ std::vector<ArgDesc> arg_descs;
1846+ ASSERT_EQ(ArgsFormatDescUtils::Parse(kernel.context().args_format(), arg_descs), GRAPH_SUCCESS);
1847+ ASSERT_EQ(arg_descs.size(), 3U);
1848+ EXPECT_EQ(arg_descs[0].addr_type, AddrType::INPUT_INSTANCE);
1849+ EXPECT_EQ(arg_descs[0].ir_idx, 0);
1850+ EXPECT_EQ(arg_descs[1].addr_type, AddrType::OUTPUT_INSTANCE);
1851+ EXPECT_EQ(arg_descs[1].ir_idx, 0);
1852+ EXPECT_EQ(arg_descs[2].addr_type, AddrType::CUSTOM_VALUE);
1853+ 
1854+ custom_op::UnloadPythonCustomOps();
1855+ loaded_python_custom_ops.Dismiss();
1856+ EXPECT_FALSE(CustomOpFactory::IsExistOp(op_type));
1857+}
1858+ 
1859+TEST_F(CustomOpFactoryStTest, PythonAnnotatedArgsUnknownShapeGeneratesBasicTaskDef) {
1860+ EnsureSharedPybindCustomOpFileForSt();
1861+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath, GetSharedPybindCustomOpFilePathForSt());
1862+ 
1863+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1864+ ASSERT_EQ(custom_op::LoadPythonCustomOps(), SUCCESS);
1865+ ScopedLoadedPythonCustomOpsForSt loaded_python_custom_ops;
1866+ 
1867+ std::vector<domi::TaskDef> tasks;
1868+ ASSERT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 0, "Ascend910B", tasks, true),
1869+ SUCCESS);
1870+ 
1871+ ASSERT_EQ(tasks.size(), 1U);
1872+ const auto &task = tasks[0];
1873+ EXPECT_EQ(task.stream_id(), 3U);
1874+ EXPECT_EQ(task.type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
1875+ EXPECT_EQ(task.sqe_num(), 5U);
1876+ EXPECT_EQ(task.kernel().context().op_index(), 0U);
1877+ EXPECT_TRUE(task.kernel().kernel_name().empty());
1878+ EXPECT_TRUE(task.kernel().context().args_format().empty());
1879+}
1880+ 
1881+TEST_F(CustomOpFactoryStTest, PythonAnnotatedArgsRealCallbackRejectsInvalidNativeUsage) {
1882+ EnsureSharedPybindCustomOpFileForSt();
1883+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath, GetSharedPybindCustomOpFilePathForSt());
1884+ 
1885+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1886+ ASSERT_EQ(custom_op::LoadPythonCustomOps(), SUCCESS);
1887+ ScopedLoadedPythonCustomOpsForSt loaded_python_custom_ops;
1888+ 
1889+ for (const int64_t alpha : {1, 2, 3, 4, 5}) {
1890+ std::vector<domi::TaskDef> tasks;
1891+ EXPECT_NE(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, alpha, "Ascend910B", tasks),
1892+ SUCCESS)
1893+ << "alpha=" << alpha;
1894+ EXPECT_TRUE(tasks.empty()) << "alpha=" << alpha;
1895+ }
1896+ 
1897+ std::vector<domi::TaskDef> index_message_tasks;
1898+ EXPECT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 9, "Ascend910B", index_message_tasks),
1899+ SUCCESS);
1900+ EXPECT_EQ(index_message_tasks.size(), 1U);
1901+}
1902+ 
1903+TEST_F(CustomOpFactoryStTest, PythonAnnotatedArgsRealCallbackEnforcesBorrowedLifetime) {
1904+ EnsureSharedPybindCustomOpFileForSt();
1905+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath, GetSharedPybindCustomOpFilePathForSt());
1906+ 
1907+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1908+ ASSERT_EQ(custom_op::LoadPythonCustomOps(), SUCCESS);
1909+ ScopedLoadedPythonCustomOpsForSt loaded_python_custom_ops;
1910+ 
1911+ std::vector<domi::TaskDef> capture_tasks;
1912+ ASSERT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 6, "Ascend910B", capture_tasks),
1913+ SUCCESS);
1914+ ASSERT_EQ(capture_tasks.size(), 1U);
1915+ EXPECT_EQ(capture_tasks[0].kernel().args_size(), 32U);
1916+ 
1917+ std::vector<domi::TaskDef> verify_tasks;
1918+ ASSERT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 7, "Ascend910B", verify_tasks),
1919+ SUCCESS);
1920+ ASSERT_EQ(verify_tasks.size(), 1U);
1921+ EXPECT_EQ(verify_tasks[0].kernel().args_size(), 24U);
1922+}
1923+ 
1924+TEST_F(CustomOpFactoryStTest, PythonAnnotatedArgsMobileOmcAcceptsSingleLaunch) {
1925+ EnsureSharedPybindCustomOpFileForSt();
1926+ ScopedEnvVarForCustomOpSt scoped_custom_opp_path(kEnvPythonCustomOpPath, GetSharedPybindCustomOpFilePathForSt());
1927+ 
1928+ ASSERT_EQ(GePythonRuntimeManager::Instance().EnsureReady(), SUCCESS);
1929+ ASSERT_EQ(custom_op::LoadPythonCustomOps(), SUCCESS);
1930+ ScopedLoadedPythonCustomOpsForSt loaded_python_custom_ops;
1931+ 
1932+ std::vector<domi::TaskDef> tasks;
1933+ ASSERT_EQ(GeneratePythonAnnotatedArgsTaskForSt(kPythonAnnotatedArgsOpTypeForSt, 8, "KirinX90", tasks), SUCCESS);
1934+ ASSERT_EQ(tasks.size(), 1U);
1935+ EXPECT_EQ(tasks[0].type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
1936+ EXPECT_EQ(tasks[0].kernel().kernel_name(), "st_python_dla");
1937+ EXPECT_EQ(tasks[0].kernel().args_size(), 24U);
1938+ const auto op_index = tasks[0].kernel().context().op_index();
1939+ 
1940+ DavinciModel model(0, nullptr);
1941+ std::vector<uint8_t> feature_mem(4096U, 0U);
1942+ model.runtime_param_.mem_base = reinterpret_cast<uintptr_t>(feature_mem.data());
1943+ model.runtime_param_.mem_size = feature_mem.size();
1944+ OpDescPtr op_desc;
1945+ ASSERT_EQ(AddPythonAnnotatedArgsOpToModelForSt(model, op_index, op_desc), SUCCESS);
1946+ ASSERT_EQ(model.GetOpByIndex(op_index), op_desc);
1947+ 
1948+ auto ge_model = MakeShared<GeModel>();
1949+ ASSERT_NE(ge_model, nullptr);
1950+ std::vector<char> kernel_bin = {0x01, 0x02};
1951+ ge_model->GetTBEKernelStore().AddTBEKernel(MakeShared<OpKernelBin>("st_python_dla", std::move(kernel_bin)));
1952+ model.ge_model_ = ge_model;
1953+ 
1954+ const uint64_t input_addr = model.runtime_param_.mem_base + 1024U;
1955+ const uint64_t output_addr = model.runtime_param_.mem_base + 2048U;
1956+ model.logical_mem_allocations_.push_back({0U, input_addr, 32U, MemAllocation::INPUT, 0U, 0U, 0U, 32U});
1957+ model.logical_mem_allocations_.push_back({1U, output_addr, 32U, MemAllocation::OUTPUT, 0U, 0U, 0U, 0U});
1958+ model.reusable_stream_allocator_ = ReusableStreamAllocator::Create();
1959+ std::vector<rtStream_t> streams(4U, nullptr);
1960+ for (auto &stream : streams) {
1961+ ASSERT_EQ(model.reusable_stream_allocator_->GetOrCreateRtStream(stream, 0U, 0, 0U), SUCCESS);
1962+ }
1963+ model.stream_list_ = streams;
1964+ 
1965+ CustomTaskInfo task_info;
1966+ TaskRunParam task_run_param;
1967+ ASSERT_EQ(task_info.ParseTaskRunParam(tasks[0], &model, task_run_param), SUCCESS);
1968+ std::vector<uint8_t> loaded_args(tasks[0].kernel().args().cbegin(), tasks[0].kernel().args().cend());
1969+ PisToArgs args;
1970+ args[static_cast<size_t>(ArgsPlacement::kArgsPlacementHbm)].dev_addr = reinterpret_cast<uint64_t>(loaded_args.data());
1971+ IowAddrs iow_addrs;
1972+ iow_addrs.input_logic_addrs = {{input_addr, static_cast<uint64_t>(MemoryAppType::kMemoryTypeFeatureMap)}};
1973+ iow_addrs.output_logic_addrs = {{output_addr, static_cast<uint64_t>(MemoryAppType::kMemoryTypeFeatureMap)}};
1974+ ASSERT_EQ(task_info.Init(tasks[0], &model, args, {}, iow_addrs), SUCCESS);
1975+ EXPECT_EQ(task_info.Distribute(), SUCCESS);
1976+}
1977+ 
1584TEST_F(CustomOpFactoryStTest, load_python_custom_ops_if_needed_fails_for_missing_python_file) {1978TEST_F(CustomOpFactoryStTest, load_python_custom_ops_if_needed_fails_for_missing_python_file) {
1585 ScopedTempDirForCustomOpSt temp_dir;1979 ScopedTempDirForCustomOpSt temp_dir;
1586 const auto missing_python_file = temp_dir.FilePath("missing_custom_op.py");1980 const auto missing_python_file = temp_dir.FilePath("missing_custom_op.py");
Mtests/ge/ut/ge/common/custom_ops_kernel_info_store_unittest.cc+124-0
@@ -24,11 +24,13 @@
24#include "exe_graph/runtime/annotated_args_context.h"24#include "exe_graph/runtime/annotated_args_context.h"
25#include "graph/compute_graph.h"25#include "graph/compute_graph.h"
26#include "graph/custom_op_factory.h"26#include "graph/custom_op_factory.h"
27+#include "graph/custom_op/args_refresh.h"
27#include "graph/ascend_string.h"28#include "graph/ascend_string.h"
28#include "graph/debug/ge_attr_define.h"29#include "graph/debug/ge_attr_define.h"
29#include "graph/ge_tensor.h"30#include "graph/ge_tensor.h"
30#include "graph/op_kernel_bin.h"31#include "graph/op_kernel_bin.h"
31#include "graph/op_desc.h"32#include "graph/op_desc.h"
33+#include "graph/operator_reg.h"
32#include "graph/custom_op.h"34#include "graph/custom_op.h"
33#include "graph/ge_context.h"35#include "graph/ge_context.h"
34#include "graph/ge_local_context.h"36#include "graph/ge_local_context.h"
@@ -37,8 +39,17 @@
37#include "graph/args_format_desc.h"39#include "graph/args_format_desc.h"
38#include "graph/utils/attr_utils.h"40#include "graph/utils/attr_utils.h"
39#include "graph/utils/tensor_utils.h"41#include "graph/utils/tensor_utils.h"
42+#include "runtime/custom_op/python_custom_op_adapter.h"
40#include "securec.h"43#include "securec.h"
41 44 
45+namespace ge {
46+REG_OP(TestPythonAnnotatedArgsCustomOp_BuilderTest)
47+ .INPUT(x, TensorType::ALL())
48+ .INPUT(w, TensorType::ALL())
49+ .OUTPUT(y, TensorType::ALL())
50+ .OP_END_FACTORY_REG(TestPythonAnnotatedArgsCustomOp_BuilderTest);
51+} // namespace ge
52+ 
42namespace ge {53namespace ge {
43namespace custom {54namespace custom {
44 55 
@@ -92,6 +103,41 @@ class MockPortableCustomOp : public PortableOp {
92 }103 }
93};104};
94 105 
106+std::atomic_uint32_t g_python_annotated_args_declare_count{0U};
107+ 
108+struct MockPythonAnnotatedArgsHolder {};
109+ 
110+void *CreateMockPythonAnnotatedArgsHolder(const custom_op::PythonCustomOpDescriptor *desc) {
111+ return (desc == nullptr) ? nullptr : new (std::nothrow) MockPythonAnnotatedArgsHolder();
112+}
113+ 
114+void DestroyMockPythonAnnotatedArgsHolder(void *holder) {
115+ delete static_cast<MockPythonAnnotatedArgsHolder *>(holder);
116+}
117+ 
118+graphStatus DeclareMockPythonAnnotatedArgs(const void *holder, gert::AnnotatedArgsContext *ctx) {
119+ if ((holder == nullptr) || (ctx == nullptr)) {
120+ return GRAPH_FAILED;
121+ }
122+ const auto *x = ctx->GetInputTensor(0U);
123+ const auto *w = ctx->GetInputTensor(1U);
124+ const auto *y = ctx->GetOutputTensor(0U);
125+ if ((x == nullptr) || (w == nullptr) || (y == nullptr)) {
126+ return GRAPH_FAILED;
127+ }
128+ ++g_python_annotated_args_declare_count;
129+ const auto workspace = ctx->MallocWorkSpace(64U);
130+ if (workspace.addr == nullptr) {
131+ return GRAPH_FAILED;
132+ }
133+ gert::AnnotatedKernelArgs args(gert::InputAddr{0U, x->GetAddr()}, gert::InputAddr{1U, w->GetAddr()},
134+ gert::OutputAddr{0U, y->GetAddr()}, workspace, uint64_t{7U});
135+ static const uint8_t kBin[] = {0x91U, 0x92U};
136+ return ctx->AddLaunch(
137+ gert::AnnotatedKernelLaunchInfo{"python_annotated_args_ut", kBin, sizeof(kBin), 1U, ctx->GetStreamId()},
138+ std::move(args));
139+}
140+ 
95std::atomic_bool g_compile_context_output_called{false};141std::atomic_bool g_compile_context_output_called{false};
96constexpr uintptr_t kLogicDataMemBase = 0x80000000UL;142constexpr uintptr_t kLogicDataMemBase = 0x80000000UL;
97constexpr uintptr_t kLogicWeightMemBase = 0x90000000UL;143constexpr uintptr_t kLogicWeightMemBase = 0x90000000UL;
@@ -108,6 +154,29 @@ uint16_t ReadUint16Slot(const std::string &args, const size_t index) {
108 return value;154 return value;
109}155}
110 156 
157+void ExpectPythonAnnotatedArgsKernel(const domi::KernelDef &kernel) {
158+ EXPECT_EQ(kernel.kernel_name(), "python_annotated_args_ut");
159+ EXPECT_EQ(kernel.args_size(), 40U);
160+ EXPECT_EQ(ReadUint64Slot(kernel.args(), 0U), static_cast<uint64_t>(kLogicDataMemBase + 1024U));
161+ EXPECT_EQ(ReadUint64Slot(kernel.args(), 1U), static_cast<uint64_t>(kLogicWeightMemBase + 4096U));
162+ EXPECT_EQ(ReadUint64Slot(kernel.args(), 2U), static_cast<uint64_t>(kLogicDataMemBase + 2048U));
163+ EXPECT_EQ(ReadUint64Slot(kernel.args(), 3U), static_cast<uint64_t>(kLogicDataMemBase + 4096U));
164+ EXPECT_EQ(ReadUint64Slot(kernel.args(), 4U), 7U);
165+ 
166+ std::vector<ArgDesc> arg_descs;
167+ ASSERT_EQ(ArgsFormatDescUtils::Parse(kernel.context().args_format(), arg_descs), GRAPH_SUCCESS);
168+ ASSERT_EQ(arg_descs.size(), 5U);
169+ EXPECT_EQ(arg_descs[0].addr_type, AddrType::INPUT_INSTANCE);
170+ EXPECT_EQ(arg_descs[0].ir_idx, 0);
171+ EXPECT_EQ(arg_descs[1].addr_type, AddrType::INPUT_INSTANCE);
172+ EXPECT_EQ(arg_descs[1].ir_idx, 1);
173+ EXPECT_EQ(arg_descs[2].addr_type, AddrType::OUTPUT_INSTANCE);
174+ EXPECT_EQ(arg_descs[2].ir_idx, 0);
175+ EXPECT_EQ(arg_descs[3].addr_type, AddrType::WORKSPACE);
176+ EXPECT_EQ(arg_descs[3].ir_idx, 0);
177+ EXPECT_EQ(arg_descs[4].addr_type, AddrType::CUSTOM_VALUE);
178+}
179+ 
111class MockCompileContextOutputOp : public EagerExecuteOp, public CompilableOp {180class MockCompileContextOutputOp : public EagerExecuteOp, public CompilableOp {
112 public:181 public:
113 graphStatus Execute(gert::EagerOpExecutionContext *ctx) override {182 graphStatus Execute(gert::EagerOpExecutionContext *ctx) override {
@@ -991,6 +1060,9 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskDeclaresAnnotatedArgsAndFillsK
991 EXPECT_EQ(parsed_arg_descs[3].addr_type, AddrType::CUSTOM_VALUE);1060 EXPECT_EQ(parsed_arg_descs[3].addr_type, AddrType::CUSTOM_VALUE);
992 1061 
993 auto op_desc = node->GetOpDesc();1062 auto op_desc = node->GetOpDesc();
1063+ int64_t task_args_mode = -1;
1064+ ASSERT_TRUE(AttrUtils::GetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE, task_args_mode));
1065+ EXPECT_EQ(task_args_mode, static_cast<int64_t>(CustomTaskArgsMode::kAnnotatedArgs));
994 std::vector<std::string> prefixes;1066 std::vector<std::string> prefixes;
995 EXPECT_FALSE(AttrUtils::GetListStr(op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, prefixes));1067 EXPECT_FALSE(AttrUtils::GetListStr(op_desc, ATTR_NAME_KERNEL_NAMES_PREFIX, prefixes));
996 const std::string prefixed_attr = "_custom_launch_0_";1068 const std::string prefixed_attr = "_custom_launch_0_";
@@ -1016,6 +1088,46 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskDeclaresAnnotatedArgsAndFillsK
1016 EXPECT_EQ(std::memcmp(tbe_kernel->GetBinData(), expected_bin, sizeof(expected_bin)), 0);1088 EXPECT_EQ(std::memcmp(tbe_kernel->GetBinData(), expected_bin, sizeof(expected_bin)), 0);
1017}1089}
1018 1090 
1091+TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskUsesPythonAnnotatedArgsAdapter) {
1092+ GetThreadLocalContext().SetGraphOption({{ge::SOC_VERSION, "Ascend910B"}});
1093+ const std::string kTestOpType = "TestPythonAnnotatedArgsCustomOp_BuilderTest";
1094+ custom_op::PythonCustomOpDescriptor desc;
1095+ desc.descriptor_key = "python_annotated_args_custom_engine";
1096+ desc.op_type = kTestOpType;
1097+ AddCustomOpCapability(desc.capabilities, CustomOpCapability::kAnnotatedArgs);
1098+ 
1099+ custom_op::PythonCustomOpCallbacks callbacks;
1100+ callbacks.create = CreateMockPythonAnnotatedArgsHolder;
1101+ callbacks.destroy = DestroyMockPythonAnnotatedArgsHolder;
1102+ callbacks.declare_launch_args = DeclareMockPythonAnnotatedArgs;
1103+ ASSERT_TRUE(custom_op::PythonCustomOpRuntimeRegistry::Register(desc, callbacks));
1104+ const auto creator = [desc]() -> std::unique_ptr<BaseCustomOp> {
1105+ auto adapter = std::make_unique<custom_op::PythonCustomOpAdapter>(desc);
1106+ if (!adapter->IsValid()) {
1107+ return nullptr;
1108+ }
1109+ return adapter;
1110+ };
1111+ ASSERT_EQ(CustomOpFactory::RegisterCustomOpCreator(AscendString(kTestOpType.c_str()), creator), GRAPH_SUCCESS);
1112+ 
1113+ ComputeGraphPtr graph;
1114+ auto node = BuildStaticCustomNodeWithConstInput(kTestOpType, graph);
1115+ ASSERT_NE(node, nullptr);
1116+ EXPECT_EQ(CustomOpFactory::GetArgsRefreshStrategy(AscendString(kTestOpType.c_str())),
1117+ ArgsRefreshStrategy::kAnnotatedArgs);
1118+ 
1119+ g_python_annotated_args_declare_count.store(0U);
1120+ std::vector<domi::TaskDef> tasks;
1121+ ASSERT_EQ(GenerateTaskForNode(node, tasks), SUCCESS);
1122+ ASSERT_EQ(tasks.size(), 1U);
1123+ EXPECT_EQ(g_python_annotated_args_declare_count.load(), 1U);
1124+ 
1125+ ExpectPythonAnnotatedArgsKernel(tasks[0].kernel());
1126+ 
1127+ CustomOpFactory::RemoveCustomOps({AscendString(kTestOpType.c_str())});
1128+ EXPECT_TRUE(custom_op::PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
1129+}
1130+ 
1019TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocDeclaresAnnotatedArgsOp) {1131TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocDeclaresAnnotatedArgsOp) {
1020 GetThreadLocalContext().SetGraphOption({{ge::SOC_VERSION, "Ascend910B"}});1132 GetThreadLocalContext().SetGraphOption({{ge::SOC_VERSION, "Ascend910B"}});
1021 1133 
@@ -1346,6 +1458,9 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocFillsBasicCustom
1346 EXPECT_EQ(task.type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));1458 EXPECT_EQ(task.type(), static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
1347 EXPECT_EQ(task.sqe_num(), 5U);1459 EXPECT_EQ(task.sqe_num(), 5U);
1348 EXPECT_EQ(task.kernel().context().op_index(), 7);1460 EXPECT_EQ(task.kernel().context().op_index(), 7);
1461+ int64_t task_args_mode = -1;
1462+ ASSERT_TRUE(AttrUtils::GetInt(node->GetOpDesc(), ATTR_NAME_CUSTOM_TASK_ARGS_MODE, task_args_mode));
1463+ EXPECT_EQ(task_args_mode, static_cast<int64_t>(CustomTaskArgsMode::kNone));
1349}1464}
1350 1465 
1351TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocUsesUpdateCallbackWhenBothRefreshInterfacesExist) {1466TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocUsesUpdateCallbackWhenBothRefreshInterfacesExist) {
@@ -1375,6 +1490,9 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnNonMobileSocUsesUpdateCallba
1375 EXPECT_TRUE(task.kernel().kernel_name().empty());1490 EXPECT_TRUE(task.kernel().kernel_name().empty());
1376 EXPECT_TRUE(task.kernel().context().args_format().empty());1491 EXPECT_TRUE(task.kernel().context().args_format().empty());
1377 EXPECT_EQ(g_both_refresh_interfaces_declare_count.load(), 0U);1492 EXPECT_EQ(g_both_refresh_interfaces_declare_count.load(), 0U);
1493+ int64_t task_args_mode = -1;
1494+ ASSERT_TRUE(AttrUtils::GetInt(node->GetOpDesc(), ATTR_NAME_CUSTOM_TASK_ARGS_MODE, task_args_mode));
1495+ EXPECT_EQ(task_args_mode, static_cast<int64_t>(CustomTaskArgsMode::kUpdateCallback));
1378}1496}
1379 1497 
1380TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnMobileSocRejectsUpdateCallbackWhenBothRefreshInterfacesExist) {1498TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnMobileSocRejectsUpdateCallbackWhenBothRefreshInterfacesExist) {
@@ -1852,6 +1970,9 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnUnknownGraphDoesNotValidateE
1852 EXPECT_EQ(task.kernel().context().op_index(), 7);1970 EXPECT_EQ(task.kernel().context().op_index(), 7);
1853 EXPECT_TRUE(task.kernel().kernel_name().empty());1971 EXPECT_TRUE(task.kernel().kernel_name().empty());
1854 EXPECT_TRUE(task.kernel().context().args_format().empty());1972 EXPECT_TRUE(task.kernel().context().args_format().empty());
1973+ int64_t task_args_mode = -1;
1974+ ASSERT_TRUE(AttrUtils::GetInt(node->GetOpDesc(), ATTR_NAME_CUSTOM_TASK_ARGS_MODE, task_args_mode));
1975+ EXPECT_EQ(task_args_mode, static_cast<int64_t>(CustomTaskArgsMode::kNone));
1855}1976}
1856 1977 
1857TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnUnknownGraphUsesBasicTaskForAnnotatedEagerOp) {1978TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnUnknownGraphUsesBasicTaskForAnnotatedEagerOp) {
@@ -1891,6 +2012,9 @@ TEST_F(UtestCustomOpsKernelInfoStore, GenerateTaskOnUnknownGraphUsesBasicTaskFor
1891 EXPECT_EQ(task.kernel().context().op_index(), 7);2012 EXPECT_EQ(task.kernel().context().op_index(), 7);
1892 EXPECT_TRUE(task.kernel().kernel_name().empty());2013 EXPECT_TRUE(task.kernel().kernel_name().empty());
1893 EXPECT_TRUE(task.kernel().context().args_format().empty());2014 EXPECT_TRUE(task.kernel().context().args_format().empty());
2015+ int64_t task_args_mode = -1;
2016+ ASSERT_TRUE(AttrUtils::GetInt(node->GetOpDesc(), ATTR_NAME_CUSTOM_TASK_ARGS_MODE, task_args_mode));
2017+ EXPECT_EQ(task_args_mode, static_cast<int64_t>(CustomTaskArgsMode::kNone));
1894 }2018 }
1895}2019}
1896 2020 
Mtests/ge/ut/ge/graph/load/custom_task_info_unittest.cc+121-15
@@ -28,6 +28,7 @@
28#include "depends/runtime/src/runtime_stub.h"28#include "depends/runtime/src/runtime_stub.h"
29#include "depends/ascendcl/src/ascendcl_stub.h"29#include "depends/ascendcl/src/ascendcl_stub.h"
30#include "graph/custom_op.h"30#include "graph/custom_op.h"
31+#include "graph/custom_op/args_refresh.h"
31#include "graph/custom_op_factory.h"32#include "graph/custom_op_factory.h"
32#include "graph/custom_op_registry.h"33#include "graph/custom_op_registry.h"
33#include "exe_graph/runtime/kernel_args.h"34#include "exe_graph/runtime/kernel_args.h"
@@ -1061,6 +1062,8 @@ TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_ArgsUpdater_SupportRefreshTrue)
1061 1062 
1062 DavinciModel model(0, nullptr);1063 DavinciModel model(0, nullptr);
1063 const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);1064 const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1065+ ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE,
1066+ static_cast<int64_t>(CustomTaskArgsMode::kUpdateCallback)));
1064 SetUpMinimalDavinciModel(model, op_desc);1067 SetUpMinimalDavinciModel(model, op_desc);
1065 1068 
1066 domi::TaskDef task_def;1069 domi::TaskDef task_def;
@@ -1095,7 +1098,15 @@ TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_AnnotatedArgsOp_UsesAnnotatedAr
1095 1098 
1096 DavinciModel model(0, nullptr);1099 DavinciModel model(0, nullptr);
1097 const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);1100 const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1101+ ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE,
1102+ static_cast<int64_t>(CustomTaskArgsMode::kAnnotatedArgs)));
1098 SetUpMinimalDavinciModel(model, op_desc);1103 SetUpMinimalDavinciModel(model, op_desc);
1104+ auto registry = std::make_shared<CustomOpRegistry>();
1105+ ASSERT_EQ(registry->RegisterCreator(
1106+ op_type.c_str(),
1107+ []() -> std::unique_ptr<BaseCustomOp> { return std::make_unique<TestArgsUpdaterCustomOp>(); }),
1108+ GRAPH_SUCCESS);
1109+ model.SetCustomOpRegistry(registry);
1099 1110 
1100 domi::TaskDef task_def;1111 domi::TaskDef task_def;
1101 FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});1112 FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});
@@ -1112,6 +1123,107 @@ TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_AnnotatedArgsOp_UsesAnnotatedAr
1112 model.runtime_param_.mem_base = 0U;1123 model.runtime_param_.mem_base = 0U;
1113}1124}
1114 1125 
1126+TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_AnnotatedArgsTaskDef_UsesTaskMetadataWithoutRegistryCreator) {
1127+ const std::string op_type = GenerateUniqueOpType();
1128+ DavinciModel model(0, nullptr);
1129+ const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1130+ SetUpMinimalDavinciModel(model, op_desc);
1131+ model.SetCustomOpRegistry(std::make_shared<CustomOpRegistry>());
1132+ 
1133+ domi::TaskDef task_def;
1134+ FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});
1135+ 
1136+ CustomTaskInfo task_info;
1137+ TaskRunParam task_run_param;
1138+ EXPECT_EQ(task_info.ParseTaskRunParam(task_def, &model, task_run_param), SUCCESS);
1139+ EXPECT_EQ(task_info.GetArgsRefreshStrategy(), ArgsRefreshStrategy::kAnnotatedArgs);
1140+ EXPECT_FALSE(task_info.NeedReserveArgsTable());
1141+ ASSERT_FALSE(task_run_param.parsed_input_addrs.empty());
1142+ EXPECT_TRUE(task_run_param.parsed_input_addrs[0].support_refresh);
1143+ 
1144+ model.runtime_param_.mem_base = 0U;
1145+}
1146+ 
1147+TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_ExplicitNoneMode_IgnoresLegacyArgsFormat) {
1148+ const std::string op_type = GenerateUniqueOpType();
1149+ DavinciModel model(0, nullptr);
1150+ const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1151+ ASSERT_TRUE(
1152+ AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE, static_cast<int64_t>(CustomTaskArgsMode::kNone)));
1153+ SetUpMinimalDavinciModel(model, op_desc);
1154+ model.SetCustomOpRegistry(std::make_shared<CustomOpRegistry>());
1155+ 
1156+ domi::TaskDef task_def;
1157+ FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});
1158+ 
1159+ CustomTaskInfo task_info;
1160+ TaskRunParam task_run_param;
1161+ EXPECT_EQ(task_info.ParseTaskRunParam(task_def, &model, task_run_param), SUCCESS);
1162+ EXPECT_EQ(task_info.GetArgsRefreshStrategy(), ArgsRefreshStrategy::kNone);
1163+ ASSERT_FALSE(task_run_param.parsed_input_addrs.empty());
1164+ EXPECT_FALSE(task_run_param.parsed_input_addrs[0].support_refresh);
1165+ 
1166+ model.runtime_param_.mem_base = 0U;
1167+}
1168+ 
1169+TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_UnspecifiedMode_UsesLegacyArgsFormatFallback) {
1170+ const std::string op_type = GenerateUniqueOpType();
1171+ DavinciModel model(0, nullptr);
1172+ const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1173+ ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE,
1174+ static_cast<int64_t>(CustomTaskArgsMode::kUnspecified)));
1175+ SetUpMinimalDavinciModel(model, op_desc);
1176+ model.SetCustomOpRegistry(std::make_shared<CustomOpRegistry>());
1177+ 
1178+ domi::TaskDef task_def;
1179+ FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});
1180+ 
1181+ CustomTaskInfo task_info;
1182+ TaskRunParam task_run_param;
1183+ EXPECT_EQ(task_info.ParseTaskRunParam(task_def, &model, task_run_param), SUCCESS);
1184+ EXPECT_EQ(task_info.GetArgsRefreshStrategy(), ArgsRefreshStrategy::kAnnotatedArgs);
1185+ 
1186+ model.runtime_param_.mem_base = 0U;
1187+}
1188+ 
1189+TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_ExplicitUpdateCallbackWithoutRegistryCreator_Fails) {
1190+ const std::string op_type = GenerateUniqueOpType();
1191+ DavinciModel model(0, nullptr);
1192+ const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1193+ ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE,
1194+ static_cast<int64_t>(CustomTaskArgsMode::kUpdateCallback)));
1195+ SetUpMinimalDavinciModel(model, op_desc);
1196+ model.SetCustomOpRegistry(std::make_shared<CustomOpRegistry>());
1197+ 
1198+ domi::TaskDef task_def;
1199+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
1200+ task_def.mutable_kernel()->mutable_context()->set_op_index(op_desc->GetId());
1201+ 
1202+ CustomTaskInfo task_info;
1203+ TaskRunParam task_run_param;
1204+ EXPECT_NE(task_info.ParseTaskRunParam(task_def, &model, task_run_param), SUCCESS);
1205+ 
1206+ model.runtime_param_.mem_base = 0U;
1207+}
1208+ 
1209+TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_InvalidExplicitMode_Fails) {
1210+ const std::string op_type = GenerateUniqueOpType();
1211+ DavinciModel model(0, nullptr);
1212+ const auto op_desc = CreateOpDesc(op_type, op_type, 1, 1);
1213+ ASSERT_TRUE(AttrUtils::SetInt(op_desc, ATTR_NAME_CUSTOM_TASK_ARGS_MODE, 99));
1214+ SetUpMinimalDavinciModel(model, op_desc);
1215+ model.SetCustomOpRegistry(std::make_shared<CustomOpRegistry>());
1216+ 
1217+ domi::TaskDef task_def;
1218+ FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), BuildInputOutputCustomArgDescs(), {0ULL, 0x40ULL, 0x1234ULL});
1219+ 
1220+ CustomTaskInfo task_info;
1221+ TaskRunParam task_run_param;
1222+ EXPECT_NE(task_info.ParseTaskRunParam(task_def, &model, task_run_param), SUCCESS);
1223+ 
1224+ model.runtime_param_.mem_base = 0U;
1225+}
1226+ 
1115TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_EagerOnly_SupportRefreshFalse) {1227TEST_F(UtestCustomTaskInfoE2E, ParseTaskRunParam_EagerOnly_SupportRefreshFalse) {
1116 std::string op_type = GenerateUniqueOpType();1228 std::string op_type = GenerateUniqueOpType();
1117 CustomOpFactory::RegisterCustomOpCreator(1229 CustomOpFactory::RegisterCustomOpCreator(
@@ -1873,12 +1985,10 @@ TEST_F(UtestCustomTaskInfoE2E, Distribute_NonTensorInput_D2HToHost) {
1873 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();1985 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();
1874 model.SetSpaceRegistries(space_registries);1986 model.SetSpaceRegistries(space_registries);
1875 1987 
1876- std::vector<ArgDesc> arg_descs;
1877- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 0);
1878- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 1);
1879- ArgsFormatDescUtils::Append(arg_descs, AddrType::OUTPUT, 0);
1880 domi::TaskDef task_def;1988 domi::TaskDef task_def;
1881- FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), arg_descs, {0ULL, 0x40ULL, 0x80ULL});1989+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
1990+ task_def.set_stream_id(0);
1991+ task_def.mutable_kernel()->mutable_context()->set_op_index(op_desc->GetId());
1882 1992 
1883 CustomTaskInfo task_info;1993 CustomTaskInfo task_info;
1884 TaskRunParam task_run_param;1994 TaskRunParam task_run_param;
@@ -1917,12 +2027,10 @@ TEST_F(UtestCustomTaskInfoE2E, Distribute_WithoutInputKinds_AllDeviceHbm) {
1917 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();2027 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();
1918 model.SetSpaceRegistries(space_registries);2028 model.SetSpaceRegistries(space_registries);
1919 2029 
1920- std::vector<ArgDesc> arg_descs;
1921- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 0);
1922- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 1);
1923- ArgsFormatDescUtils::Append(arg_descs, AddrType::OUTPUT, 0);
1924 domi::TaskDef task_def;2030 domi::TaskDef task_def;
1925- FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), arg_descs, {0ULL, 0x40ULL, 0x80ULL});2031+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
2032+ task_def.set_stream_id(0);
2033+ task_def.mutable_kernel()->mutable_context()->set_op_index(op_desc->GetId());
1926 2034 
1927 CustomTaskInfo task_info;2035 CustomTaskInfo task_info;
1928 TaskRunParam task_run_param;2036 TaskRunParam task_run_param;
@@ -1958,12 +2066,10 @@ TEST_F(UtestCustomTaskInfoE2E, Distribute_WithExplicitNonTensorKindBase_D2HOnlyF
1958 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();2066 auto space_registries = gert::SpaceRegistryFaker().BuildMainSpaceRegistryArray();
1959 model.SetSpaceRegistries(space_registries);2067 model.SetSpaceRegistries(space_registries);
1960 2068 
1961- std::vector<ArgDesc> arg_descs;
1962- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 0);
1963- ArgsFormatDescUtils::Append(arg_descs, AddrType::INPUT, 1);
1964- ArgsFormatDescUtils::Append(arg_descs, AddrType::OUTPUT, 0);
1965 domi::TaskDef task_def;2069 domi::TaskDef task_def;
1966- FillAnnotatedArgsTaskDef(task_def, op_desc->GetId(), arg_descs, {0ULL, 0x40ULL, 0x80ULL});2070+ task_def.set_type(static_cast<uint32_t>(ModelTaskType::MODEL_TASK_CUSTOM_KERNEL));
2071+ task_def.set_stream_id(0);
2072+ task_def.mutable_kernel()->mutable_context()->set_op_index(op_desc->GetId());
1967 2073 
1968 CustomTaskInfo task_info;2074 CustomTaskInfo task_info;
1969 TaskRunParam task_run_param;2075 TaskRunParam task_run_param;
Atests/ge/ut/ge/graph/pyge_tests/python_custom_op_annotated_args_test.py+531-0
@@ -0,0 +1,531 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# -----------------------------------------------------------------------------------------------------------
4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6+# CANN Open Software License Agreement Version 2.0 (the "License").
7+# Please refer to the License for details. You may not use this file except in compliance with the License.
8+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10+# See LICENSE in the root of the software repository for the full text of the License.
11+# -----------------------------------------------------------------------------------------------------------
12+ 
13+"""Pytest coverage for schema-bound Python declare_launch_args callbacks."""
14+ 
15+import contextvars
16+import importlib
17+from pathlib import Path
18+from types import SimpleNamespace
19+from typing import List, Optional
20+ 
21+import pytest
22+ 
23+try:
24+ bridge = importlib.import_module("ge.custom_op._bridge")
25+ context = importlib.import_module("ge.custom_op.context")
26+ custom_op = importlib.import_module("ge.custom_op")
27+ from ge.graph import DataType
28+ from ge.runtime import Tensor
29+except ImportError as exc:
30+ pytest.skip(f"无法导入 Python custom op 相关模块: {exc}", allow_module_level=True)
31+ 
32+ 
33+@pytest.fixture(autouse=True)
34+def clear_python_custom_op_runtime():
35+ custom_op.clear_registered_op_impls()
36+ bridge.clear_op_impl_holders()
37+ yield
38+ bridge.clear_op_impl_holders()
39+ custom_op.clear_registered_op_impls()
40+ 
41+ 
42+class _FakeRuntimeAttrs:
43+ def __init__(self):
44+ self.calls = []
45+ 
46+ def __getattr__(self, name):
47+ if not name.startswith("get_"):
48+ raise AttributeError(name)
49+ return lambda index: self._record(name, index)
50+ 
51+ def _record(self, name, index):
52+ self.calls.append((name, index))
53+ return (name, index)
54+ 
55+ 
56+class _FakeDlaContext:
57+ def __init__(self):
58+ self.invalidated = False
59+ self.attrs_requested = False
60+ self.attrs = _FakeRuntimeAttrs()
61+ self.calls = []
62+ 
63+ def _get_required_input_tensor(self, ir_index):
64+ self.calls.append(("required_input", ir_index))
65+ return ("required_input", ir_index)
66+ 
67+ def _get_optional_input_tensor(self, ir_index):
68+ self.calls.append(("optional_input", ir_index))
69+ return None
70+ 
71+ def _get_dynamic_input_num(self, ir_index):
72+ self.calls.append(("dynamic_input_num", ir_index))
73+ return 2
74+ 
75+ def _get_dynamic_input_tensor(self, ir_index, relative_index):
76+ self.calls.append(("dynamic_input", ir_index, relative_index))
77+ return ("dynamic_input", ir_index, relative_index)
78+ 
79+ def _get_required_output_tensor(self, ir_index):
80+ self.calls.append(("required_output", ir_index))
81+ return ("required_output", ir_index)
82+ 
83+ def _get_dynamic_output_num(self, ir_index):
84+ self.calls.append(("dynamic_output_num", ir_index))
85+ return 2
86+ 
87+ def _get_dynamic_output_tensor(self, ir_index, relative_index):
88+ self.calls.append(("dynamic_output", ir_index, relative_index))
89+ return ("dynamic_output", ir_index, relative_index)
90+ 
91+ def _get_attrs(self):
92+ self.attrs_requested = True
93+ return self.attrs
94+ 
95+ def _invalidate(self):
96+ self.invalidated = True
97+ 
98+ 
99+def _create_holder(instance_id: str, op_type: str) -> None:
100+ descriptors = {
101+ item["op_type"]: item for item in bridge.load_and_get_op_impl_descriptors()
102+ }
103+ assert bridge.create_op_impl_holder(
104+ instance_id, descriptors[op_type]["descriptor_key"]
105+ )
106+ 
107+ 
108+def _get_descriptor_key(op_type: str) -> str:
109+ return next(
110+ item["descriptor_key"]
111+ for item in bridge.load_and_get_op_impl_descriptors()
112+ if item["op_type"] == op_type
113+ )
114+ 
115+ 
116+def _full_ir_meta(op_type: str) -> dict:
117+ return {
118+ "op_type": op_type,
119+ "inputs": [
120+ {"name": "first", "kind": 0},
121+ {"name": "maybe", "kind": 1},
122+ {"name": "many", "kind": 2},
123+ ],
124+ "attrs": [{"name": "alpha", "type": "VT_INT"}],
125+ "outputs": [
126+ {"name": "result", "kind": 0},
127+ {"name": "result_many", "kind": 1},
128+ ],
129+ }
130+ 
131+ 
132+def test_plain_class_declares_annotated_args_capability():
133+ @custom_op.register_op_impl(op_type="DlaOnlyCustom")
134+ class DlaOnlyCustom:
135+ def declare_launch_args(self) -> None:
136+ pass
137+ 
138+ assert DlaOnlyCustom.__ge_op_impl_descriptor__.interfaces == ["annotated_args"]
139+ 
140+ 
141+def test_plain_execute_and_declare_registers_both_capabilities():
142+ @custom_op.register_op_impl(op_type="PlainExecuteAndDlaCustom")
143+ class PlainExecuteAndDlaCustom:
144+ def execute(self):
145+ pass
146+ 
147+ def declare_launch_args(self) -> None:
148+ pass
149+ 
150+ assert PlainExecuteAndDlaCustom.__ge_op_impl_descriptor__.interfaces == [
151+ "eager_execute",
152+ "annotated_args",
153+ ]
154+ 
155+ 
156+def test_dual_capability_is_registered_in_stable_order():
157+ @custom_op.register_op_impl(op_type="DualCapabilityCustom")
158+ class DualCapabilityCustom(custom_op.EagerExecuteOp):
159+ def execute(self, ctx):
160+ pass
161+ 
162+ def declare_launch_args(self) -> None:
163+ pass
164+ 
165+ assert DualCapabilityCustom.__ge_op_impl_descriptor__.interfaces == [
166+ "eager_execute",
167+ "annotated_args",
168+ ]
169+ 
170+ 
171+def test_non_callable_declare_launch_args_does_not_satisfy_registration():
172+ with pytest.raises(TypeError, match="must implement at least one supported method"):
173+ 
174+ @custom_op.register_op_impl(op_type="BadDlaCustom")
175+ class BadDlaCustom:
176+ declare_launch_args = 1
177+ 
178+ 
179+def test_non_callable_declare_launch_args_is_ignored_with_valid_execute():
180+ @custom_op.register_op_impl(op_type="ExecuteWithBadDlaCustom")
181+ class ExecuteWithBadDlaCustom:
182+ declare_launch_args = 1
183+ 
184+ def execute(self):
185+ pass
186+ 
187+ assert ExecuteWithBadDlaCustom.__ge_op_impl_descriptor__.interfaces == [
188+ "eager_execute"
189+ ]
190+ 
191+ 
192+def test_get_declare_launch_args_ctx_is_unavailable_outside_callback():
193+ with pytest.raises(RuntimeError, match="only available inside declare_launch_args"):
194+ custom_op.get_declare_launch_args_ctx()
195+ 
196+ 
197+def test_declare_launch_args_context_scope_restores_and_invalidates():
198+ outer = object()
199+ inner = object()
200+ copied_contexts = []
201+ context_scope = getattr(context, "_declare_launch_args_ctx_scope")
202+ 
203+ with context_scope(outer):
204+ assert custom_op.get_declare_launch_args_ctx() is outer
205+ with context_scope(inner):
206+ assert custom_op.get_declare_launch_args_ctx() is inner
207+ copied_contexts.append(contextvars.copy_context())
208+ assert custom_op.get_declare_launch_args_ctx() is outer
209+ 
210+ with pytest.raises(RuntimeError, match="only available inside declare_launch_args"):
211+ custom_op.get_declare_launch_args_ctx()
212+ with pytest.raises(RuntimeError, match="only available inside declare_launch_args"):
213+ copied_contexts[0].run(custom_op.get_declare_launch_args_ctx)
214+ 
215+ 
216+def test_call_declare_launch_args_binds_schema_arguments_and_context():
217+ seen = []
218+ 
219+ @custom_op.register_op_impl(op_type="SchemaDlaCustom")
220+ class SchemaDlaCustom:
221+ def declare_launch_args(
222+ self,
223+ first: Tensor,
224+ maybe: Optional[Tensor],
225+ many: List[Tensor],
226+ result: Tensor,
227+ result_many: list[Tensor],
228+ *,
229+ alpha: int,
230+ ) -> None:
231+ _ = self
232+ seen.append(
233+ (
234+ first,
235+ maybe,
236+ many,
237+ result,
238+ result_many,
239+ alpha,
240+ custom_op.get_declare_launch_args_ctx(),
241+ )
242+ )
243+ 
244+ ctx = _FakeDlaContext()
245+ descriptor_key = _get_descriptor_key("SchemaDlaCustom")
246+ assert (
247+ bridge.validate_op_impl_descriptor(
248+ descriptor_key, _full_ir_meta("SchemaDlaCustom")
249+ )
250+ is True
251+ )
252+ _create_holder("SchemaDlaCustom#1", "SchemaDlaCustom")
253+ 
254+ assert (
255+ bridge.call_declare_launch_args(
256+ "SchemaDlaCustom#1", _full_ir_meta("SchemaDlaCustom"), ctx
257+ )
258+ is None
259+ )
260+ assert seen == [
261+ (
262+ ("required_input", 0),
263+ None,
264+ [("dynamic_input", 2, 0), ("dynamic_input", 2, 1)],
265+ ("required_output", 0),
266+ [("dynamic_output", 1, 0), ("dynamic_output", 1, 1)],
267+ ("get_int", 0),
268+ ctx,
269+ )
270+ ]
271+ assert ctx.calls == [
272+ ("required_input", 0),
273+ ("optional_input", 1),
274+ ("dynamic_input_num", 2),
275+ ("dynamic_input", 2, 0),
276+ ("dynamic_input", 2, 1),
277+ ("required_output", 0),
278+ ("dynamic_output_num", 1),
279+ ("dynamic_output", 1, 0),
280+ ("dynamic_output", 1, 1),
281+ ]
282+ assert ctx.attrs.calls == [("get_int", 0)]
283+ assert ctx.invalidated is True
284+ 
285+ 
286+@pytest.mark.parametrize(
287+ "method_body, expected",
288+ [
289+ (
290+ "def declare_launch_args(self, first, *, alpha) -> None: pass",
291+ "expected 5 positional",
292+ ),
293+ (
294+ "def declare_launch_args(self, first, maybe, many, result, result_many, alpha) -> None: pass",
295+ "keyword-only",
296+ ),
297+ (
298+ "def declare_launch_args(self, first, maybe, many, result, result_many, *, beta) -> None: pass",
299+ "expected attr name",
300+ ),
301+ (
302+ "def declare_launch_args(self, first, maybe, many, result, result_many, *args, alpha) -> None: pass",
303+ "variadic",
304+ ),
305+ (
306+ "def declare_launch_args(self, first, maybe, many, result, result_many, *, alpha) -> int: pass",
307+ "expected None",
308+ ),
309+ ],
310+)
311+def test_validate_op_impl_descriptor_rejects_invalid_declare_signature(
312+ method_body, expected
313+):
314+ namespace = {}
315+ exec(method_body, namespace)
316+ method = namespace.get("declare_launch_args")
317+ assert method is not None
318+ 
319+ @custom_op.register_op_impl(op_type=f"InvalidSignature{abs(hash(method_body))}")
320+ class InvalidSignature:
321+ declare_launch_args = method
322+ 
323+ op_type = InvalidSignature.__ge_op_impl_descriptor__.op_type
324+ descriptor_key = _get_descriptor_key(op_type)
325+ 
326+ with pytest.raises(TypeError) as exc_info:
327+ bridge.validate_op_impl_descriptor(descriptor_key, _full_ir_meta(op_type))
328+ 
329+ message = str(exc_info.value)
330+ assert op_type in message
331+ assert InvalidSignature.__ge_op_impl_descriptor__.descriptor_key in message
332+ assert "declare_launch_args" in message
333+ assert "expected" in message
334+ assert "actual" in message
335+ assert expected in message
336+ 
337+ 
338+def test_call_declare_launch_args_does_not_validate_signature_at_runtime(
339+ monkeypatch,
340+):
341+ called = []
342+ 
343+ @custom_op.register_op_impl(op_type="RegistrationValidatedDlaCustom")
344+ class RegistrationValidatedDlaCustom:
345+ @staticmethod
346+ def declare_launch_args(
347+ first: Tensor,
348+ maybe: Optional[Tensor],
349+ many: List[Tensor],
350+ result: Tensor,
351+ result_many: list[Tensor],
352+ *,
353+ alpha: int,
354+ ) -> None:
355+ called.append((first, maybe, many, result, result_many, alpha))
356+ 
357+ op_type = "RegistrationValidatedDlaCustom"
358+ descriptor_key = _get_descriptor_key(op_type)
359+ ir_meta = _full_ir_meta(op_type)
360+ assert bridge.validate_op_impl_descriptor(descriptor_key, ir_meta) is True
361+ assert bridge.create_op_impl_holder(f"{op_type}#1", descriptor_key) is True
362+ 
363+ def fail_on_runtime_validation(*args, **kwargs):
364+ _ = (args, kwargs)
365+ pytest.fail("declare_launch_args must not validate its signature at runtime")
366+ 
367+ monkeypatch.setattr(bridge, "_validate_args_signature", fail_on_runtime_validation)
368+ ctx = _FakeDlaContext()
369+ bridge.call_declare_launch_args(f"{op_type}#1", ir_meta, ctx)
370+ 
371+ assert len(called) == 1
372+ assert ctx.invalidated is True
373+ 
374+ 
375+def test_validate_op_impl_descriptor_requires_canonical_ir_for_declare():
376+ @custom_op.register_op_impl(op_type="MissingRegistrationDlaSchemaCustom")
377+ class MissingRegistrationDlaSchemaCustom:
378+ def declare_launch_args(self, x: Tensor) -> None:
379+ pass
380+ 
381+ with pytest.raises(
382+ RuntimeError,
383+ match="canonical IR not found for schema-bound declare_launch_args",
384+ ):
385+ bridge.validate_op_impl_descriptor(
386+ _get_descriptor_key("MissingRegistrationDlaSchemaCustom"), None
387+ )
388+ 
389+ 
390+def test_call_declare_launch_args_rejects_missing_schema_and_non_none_result():
391+ @custom_op.register_op_impl(op_type="MissingSchemaDlaCustom")
392+ class MissingSchemaDlaCustom:
393+ def declare_launch_args(self) -> None:
394+ _ = self
395+ return True
396+ 
397+ ctx = _FakeDlaContext()
398+ _create_holder("MissingSchemaDlaCustom#1", "MissingSchemaDlaCustom")
399+ with pytest.raises(
400+ RuntimeError, match="canonical IR not found.*declare_launch_args"
401+ ):
402+ bridge.call_declare_launch_args("MissingSchemaDlaCustom#1", None, ctx)
403+ assert ctx.invalidated is True
404+ 
405+ 
406+@pytest.mark.parametrize(
407+ ("ir_type", "annotation"),
408+ [
409+ ("VT_INT", int),
410+ ("VT_FLOAT", float),
411+ ("VT_BOOL", bool),
412+ ("VT_STRING", str),
413+ ("VT_DATA_TYPE", DataType),
414+ ("VT_TENSOR", Tensor),
415+ ("VT_LIST_INT", List[int]),
416+ ("VT_LIST_FLOAT", list[float]),
417+ ("VT_LIST_BOOL", List[bool]),
418+ ("VT_LIST_STRING", list[str]),
419+ ("VT_LIST_DATA_TYPE", List[DataType]),
420+ ("VT_LIST_LIST_INT", list[list[int]]),
421+ ],
422+)
423+def test_signature_accepts_canonical_attr_annotations(ir_type, annotation):
424+ def declare_launch_args(*, alpha) -> None:
425+ pass
426+ 
427+ declare_launch_args.__annotations__ = {"alpha": annotation, "return": None}
428+ descriptor = SimpleNamespace(
429+ op_type="AttrAnnotationCustom", descriptor_key="attr-annotation"
430+ )
431+ validate_args_signature = getattr(bridge, "_validate_args_signature")
432+ validate_args_signature(
433+ declare_launch_args,
434+ {
435+ "op_type": "AttrAnnotationCustom",
436+ "inputs": [],
437+ "attrs": [{"name": "alpha", "type": ir_type}],
438+ "outputs": [],
439+ },
440+ descriptor,
441+ method_name="declare_launch_args",
442+ )
443+ 
444+ 
445+def test_signature_rejects_noncanonical_tensor_annotation():
446+ @custom_op.register_op_impl(op_type="WrongTensorAnnotationCustom")
447+ class WrongTensorAnnotationCustom:
448+ def declare_launch_args(self, x: list[Tensor]) -> None:
449+ pass
450+ 
451+ descriptor_key = _get_descriptor_key("WrongTensorAnnotationCustom")
452+ with pytest.raises(TypeError) as exc_info:
453+ bridge.validate_op_impl_descriptor(
454+ descriptor_key,
455+ {
456+ "op_type": "WrongTensorAnnotationCustom",
457+ "inputs": [{"name": "x", "kind": 0}],
458+ "attrs": [],
459+ "outputs": [],
460+ },
461+ )
462+ assert "WrongTensorAnnotationCustom" in str(exc_info.value)
463+ assert "expected" in str(exc_info.value)
464+ assert "actual" in str(exc_info.value)
465+ 
466+ @custom_op.register_op_impl(op_type="NonNoneResultDlaCustom")
467+ class NonNoneResultDlaCustom:
468+ def declare_launch_args(self) -> None:
469+ _ = self
470+ return True
471+ 
472+ ctx = _FakeDlaContext()
473+ _create_holder("NonNoneResultDlaCustom#1", "NonNoneResultDlaCustom")
474+ with pytest.raises(TypeError, match="declare_launch_args must return None"):
475+ bridge.call_declare_launch_args(
476+ "NonNoneResultDlaCustom#1",
477+ {
478+ "op_type": "NonNoneResultDlaCustom",
479+ "inputs": [],
480+ "attrs": [],
481+ "outputs": [],
482+ },
483+ ctx,
484+ )
485+ assert ctx.invalidated is True
486+ 
487+ 
488+def test_native_module_exposes_annotated_args_types():
489+ native_module = getattr(importlib.import_module("ge.custom_op._native"), "_native")
490+ for type_name in (
491+ "AnnotatedArgsContext",
492+ "AnnotatedKernelArgs",
493+ "AnnotatedKernelLaunchInfo",
494+ "WorkspaceAddr",
495+ ):
496+ assert hasattr(native_module, type_name)
497+ with pytest.raises(TypeError):
498+ native_module.AnnotatedArgsContext()
499+ with pytest.raises(TypeError):
500+ native_module.AnnotatedKernelArgs()
501+ with pytest.raises(TypeError):
502+ native_module.WorkspaceAddr()
503+ 
504+ 
505+def test_public_stub_hides_bridge_private_attr_readers():
506+ stub_path = Path(custom_op.__file__).with_name("_ge_custom_op_native.pyi")
507+ stub_text = stub_path.read_text(encoding="utf-8")
508+ annotated_context_stub = stub_text.split("class AnnotatedArgsContext:", 1)[1]
509+ assert "def _get_attrs" not in annotated_context_stub
510+ assert "def get_attrs" not in annotated_context_stub
511+ assert "def get_attr" not in annotated_context_stub
512+ assert "get_attrs" not in custom_op.__all__
513+ 
514+ 
515+def test_python_artifact_uses_bridge_abi_v1():
516+ artifact_utils = importlib.import_module("ge.custom_op._artifact_utils")
517+ assert artifact_utils.BRIDGE_ABI_VERSION == 1
518+ 
519+ 
520+def test_launch_info_owns_and_validates_fields():
521+ info = custom_op.AnnotatedKernelLaunchInfo(
522+ kernel_name="add_custom",
523+ kernel_bin=b"\x01\x02",
524+ block_dim=1,
525+ stream_id=0,
526+ )
527+ assert info is not None
528+ with pytest.raises(ValueError, match="kernel_name"):
529+ custom_op.AnnotatedKernelLaunchInfo(
530+ kernel_name="", kernel_bin=b"\x01", block_dim=1, stream_id=0
531+ )
Mtests/ge/ut/ge/graph/pyge_tests/python_custom_op_test.py+210-1
@@ -24,6 +24,7 @@ try:
24 bridge = importlib.import_module("ge.custom_op._bridge")24 bridge = importlib.import_module("ge.custom_op._bridge")
25 custom_op = importlib.import_module("ge.custom_op")25 custom_op = importlib.import_module("ge.custom_op")
26 ir_types = importlib.import_module("ge.custom_op._ir_types")26 ir_types = importlib.import_module("ge.custom_op._ir_types")
27+ from ge.runtime import Tensor
27except ImportError as exc:28except ImportError as exc:
28 pytest.skip(f"无法导入 Python custom op 相关模块: {exc}", allow_module_level=True)29 pytest.skip(f"无法导入 Python custom op 相关模块: {exc}", allow_module_level=True)
29 30 
@@ -125,6 +126,18 @@ def _create_holder(instance_id: str, op_type: str) -> None:
125 )126 )
126 127 
127 128 
129+def _get_descriptor_key(op_type: str) -> str:
130+ descriptors = {
131+ item["op_type"]: item for item in bridge.load_and_get_op_impl_descriptors()
132+ }
133+ return descriptors[op_type]["descriptor_key"]
134+ 
135+ 
136+def test_bridge_validate_op_impl_descriptor_rejects_unknown_descriptor():
137+ with pytest.raises(KeyError, match="descriptor_key not found"):
138+ bridge.validate_op_impl_descriptor("missing-descriptor", None)
139+ 
140+ 
128def test_register_op_impl_exports_descriptor_dict():141def test_register_op_impl_exports_descriptor_dict():
129 @custom_op.register_op_impl(op_type="AddCustom")142 @custom_op.register_op_impl(op_type="AddCustom")
130 class AddCustom(custom_op.EagerExecuteOp):143 class AddCustom(custom_op.EagerExecuteOp):
@@ -179,6 +192,61 @@ def test_ctx_execute_keeps_context_argument():
179 assert instance.seen_ctx is ctx192 assert instance.seen_ctx is ctx
180 193 
181 194 
195+def test_bridge_validates_descriptor_signature_without_constructing_instance():
196+ constructed = []
197+ 
198+ @custom_op.register_op_impl(op_type="RegistrationValidatedCustom")
199+ class RegistrationValidatedCustom(custom_op.EagerExecuteOp):
200+ def __init__(self):
201+ constructed.append(True)
202+ 
203+ def execute(self, x: Tensor, *, alpha: float):
204+ pass
205+ 
206+ assert (
207+ bridge.validate_op_impl_descriptor(
208+ _get_descriptor_key("RegistrationValidatedCustom"),
209+ {
210+ "op_type": "RegistrationValidatedCustom",
211+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
212+ "attrs": [{"name": "alpha", "type": ir_types.AttrType.FLOAT}],
213+ "outputs": [],
214+ },
215+ )
216+ is True
217+ )
218+ assert constructed == []
219+ 
220+ 
221+def test_bridge_validates_legacy_execute_without_canonical_ir():
222+ @custom_op.register_op_impl(op_type="RegistrationLegacyCustom")
223+ class RegistrationLegacyCustom(custom_op.EagerExecuteOp):
224+ def execute(self, ctx):
225+ pass
226+ 
227+ assert (
228+ bridge.validate_op_impl_descriptor(
229+ _get_descriptor_key("RegistrationLegacyCustom"), None
230+ )
231+ is True
232+ )
233+ 
234+ 
235+def test_bridge_requires_canonical_ir_to_validate_schema_execute():
236+ @custom_op.register_op_impl(op_type="MissingRegistrationSchemaCustom")
237+ class MissingRegistrationSchemaCustom(custom_op.EagerExecuteOp):
238+ def execute(self, x):
239+ pass
240+ 
241+ with pytest.raises(
242+ RuntimeError,
243+ match="canonical IR not found for schema-bound execute",
244+ ):
245+ bridge.validate_op_impl_descriptor(
246+ _get_descriptor_key("MissingRegistrationSchemaCustom"), None
247+ )
248+ 
249+ 
182def test_register_op_impl_rejects_duplicate_op_type():250def test_register_op_impl_rejects_duplicate_op_type():
183 @custom_op.register_op_impl(op_type="AddCustom")251 @custom_op.register_op_impl(op_type="AddCustom")
184 class AddCustom(custom_op.EagerExecuteOp):252 class AddCustom(custom_op.EagerExecuteOp):
@@ -220,7 +288,7 @@ def test_register_op_impl_supports_plain_class_with_execute():
220def test_register_op_impl_rejects_class_without_supported_method():288def test_register_op_impl_rejects_class_without_supported_method():
221 with pytest.raises(289 with pytest.raises(
222 TypeError,290 TypeError,
223- match=r"BaseOnlyCustom' must implement at least one supported method: execute",291+ match=r"BaseOnlyCustom' must implement at least one supported method: execute, declare_launch_args",
224 ):292 ):
225 293 
226 @custom_op.register_op_impl(op_type="BaseOnlyCustom")294 @custom_op.register_op_impl(op_type="BaseOnlyCustom")
@@ -331,6 +399,7 @@ def test_bridge_call_execute_binds_schema_inputs_and_attrs(method_kind):
331 "outputs": [],399 "outputs": [],
332 }400 }
333 401 
402+ assert bridge.validate_op_impl_descriptor(descriptor_key, ir_meta) is True
334 assert bridge.create_op_impl_holder(instance_id, descriptor_key) is True403 assert bridge.create_op_impl_holder(instance_id, descriptor_key) is True
335 assert bridge.call_execute(instance_id, ir_meta, ctx) is None404 assert bridge.call_execute(instance_id, ir_meta, ctx) is None
336 assert called == [405 assert called == [
@@ -357,6 +426,111 @@ def test_bridge_call_execute_binds_schema_inputs_and_attrs(method_kind):
357 assert ctx.invalidated is True426 assert ctx.invalidated is True
358 427 
359 428 
429+def test_bridge_validate_op_impl_descriptor_rejects_mismatched_attr_name():
430+ @custom_op.register_op_impl(op_type="WrongExecuteAttrNameCustom")
431+ class WrongExecuteAttrNameCustom(custom_op.EagerExecuteOp):
432+ def execute(self, x, *, beta):
433+ pass
434+ 
435+ with pytest.raises(
436+ TypeError,
437+ match=r"invalid execute signature.*expected attr name alpha.*actual attr name beta",
438+ ):
439+ bridge.validate_op_impl_descriptor(
440+ _get_descriptor_key("WrongExecuteAttrNameCustom"),
441+ {
442+ "op_type": "WrongExecuteAttrNameCustom",
443+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
444+ "attrs": [{"name": "alpha", "type": ir_types.AttrType.FLOAT}],
445+ "outputs": [],
446+ },
447+ )
448+ 
449+ 
450+def test_bridge_validate_op_impl_descriptor_rejects_mismatched_input_annotation():
451+ @custom_op.register_op_impl(op_type="WrongExecuteAnnotationCustom")
452+ class WrongExecuteAnnotationCustom(custom_op.EagerExecuteOp):
453+ def execute(self, x: list[Tensor], *, alpha: float):
454+ pass
455+ 
456+ with pytest.raises(
457+ TypeError,
458+ match=r"invalid execute signature.*input parameter at index 0 annotation",
459+ ):
460+ bridge.validate_op_impl_descriptor(
461+ _get_descriptor_key("WrongExecuteAnnotationCustom"),
462+ {
463+ "op_type": "WrongExecuteAnnotationCustom",
464+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
465+ "attrs": [{"name": "alpha", "type": ir_types.AttrType.FLOAT}],
466+ "outputs": [],
467+ },
468+ )
469+ 
470+ 
471+def test_bridge_call_execute_ignores_ir_outputs_and_return_annotation():
472+ called = []
473+ 
474+ @custom_op.register_op_impl(op_type="CompatibleExecuteSignatureCustom")
475+ class CompatibleExecuteSignatureCustom(custom_op.EagerExecuteOp):
476+ def execute(
477+ self, x: Tensor, *, alpha: float
478+ ) -> "CompatibleExecuteSignatureCustom":
479+ called.append((x, alpha))
480+ return True
481+ 
482+ instance_id = "CompatibleExecuteSignatureCustom#1"
483+ ctx = _FakeEagerContext()
484+ _create_holder(instance_id, "CompatibleExecuteSignatureCustom")
485+ 
486+ assert (
487+ bridge.call_execute(
488+ instance_id,
489+ {
490+ "op_type": "CompatibleExecuteSignatureCustom",
491+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
492+ "attrs": [{"name": "alpha", "type": ir_types.AttrType.FLOAT}],
493+ "outputs": [{"name": "y", "kind": ir_types.OutputType.REQUIRED}],
494+ },
495+ ctx,
496+ )
497+ is None
498+ )
499+ 
500+ assert called == [(("required", 0), ("get_float", 0))]
501+ assert ctx.invalidated is True
502+ 
503+ 
504+def test_bridge_call_execute_does_not_validate_signature_at_runtime(monkeypatch):
505+ execute_calls = []
506+ 
507+ @custom_op.register_op_impl(op_type="CachedExecuteSignatureCustom")
508+ class CachedExecuteSignatureCustom(custom_op.EagerExecuteOp):
509+ def execute(self, x: Tensor):
510+ execute_calls.append(x)
511+ 
512+ descriptor_key = _get_descriptor_key("CachedExecuteSignatureCustom")
513+ instance_id = "CachedExecuteSignatureCustom#1"
514+ ir_meta = {
515+ "op_type": "CachedExecuteSignatureCustom",
516+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
517+ "attrs": [],
518+ "outputs": [],
519+ }
520+ 
521+ assert bridge.validate_op_impl_descriptor(descriptor_key, ir_meta) is True
522+ assert bridge.create_op_impl_holder(instance_id, descriptor_key) is True
523+ 
524+ def fail_on_runtime_validation(*args, **kwargs):
525+ _ = (args, kwargs)
526+ pytest.fail("schema execute must not validate its signature at runtime")
527+ 
528+ monkeypatch.setattr(bridge, "_validate_args_signature", fail_on_runtime_validation)
529+ bridge.call_execute(instance_id, ir_meta, _FakeEagerContext())
530+ 
531+ assert execute_calls == [("required", 0)]
532+ 
533+ 
360def test_schema_bound_execute_can_get_current_context():534def test_schema_bound_execute_can_get_current_context():
361 called = []535 called = []
362 536 
@@ -411,6 +585,41 @@ def test_get_execute_ctx_is_not_bound_for_legacy_execute():
411 assert ctx.invalidated is True585 assert ctx.invalidated is True
412 586 
413 587 
588+def test_bridge_call_execute_rechecks_legacy_signature_after_method_change():
589+ calls = []
590+ 
591+ @custom_op.register_op_impl(op_type="ReplaceableExecuteCustom")
592+ class ReplaceableExecuteCustom(custom_op.EagerExecuteOp):
593+ def execute(self, ctx):
594+ calls.append(("legacy", ctx))
595+ 
596+ instance_id = "ReplaceableExecuteCustom#1"
597+ legacy_ctx = _FakeEagerContext()
598+ _create_holder(instance_id, "ReplaceableExecuteCustom")
599+ 
600+ bridge.call_execute(instance_id, None, legacy_ctx)
601+ 
602+ def schema_execute(self, x):
603+ calls.append(("schema", x))
604+ 
605+ ReplaceableExecuteCustom.execute = schema_execute
606+ schema_ctx = _FakeEagerContext()
607+ bridge.call_execute(
608+ instance_id,
609+ {
610+ "op_type": "ReplaceableExecuteCustom",
611+ "inputs": [{"name": "x", "kind": ir_types.InputType.REQUIRED}],
612+ "attrs": [],
613+ "outputs": [],
614+ },
615+ schema_ctx,
616+ )
617+ 
618+ assert calls == [("legacy", legacy_ctx), ("schema", ("required", 0))]
619+ assert legacy_ctx.invalidated is True
620+ assert schema_ctx.invalidated is True
621+ 
622+ 
414def test_schema_execute_context_is_deactivated_after_exception():623def test_schema_execute_context_is_deactivated_after_exception():
415 copied_contexts = []624 copied_contexts = []
416 625 
Mtests/ge/ut/ge/graph_ir/ge_custom_op_factory_unittest.cc+54-0
@@ -106,6 +106,11 @@ graphStatus ExecuteMockPythonCustomOp(const void *holder, gert::EagerOpExecution
106 return GRAPH_SUCCESS;106 return GRAPH_SUCCESS;
107}107}
108 108 
109+graphStatus DeclareMockPythonCustomOp(const void *holder, gert::AnnotatedArgsContext *ctx) {
110+ (void)ctx;
111+ return (holder == nullptr) ? GRAPH_FAILED : GRAPH_SUCCESS;
112+}
113+ 
109std::vector<uint8_t> BuildCustomOpPartition(const std::string &name, const std::vector<uint8_t> &bin) {114std::vector<uint8_t> BuildCustomOpPartition(const std::string &name, const std::vector<uint8_t> &bin) {
110 ge::CustomKernelItemHeader header{ge::kCustomKernelItemMagic, static_cast<uint32_t>(name.size()),115 ge::CustomKernelItemHeader header{ge::kCustomKernelItemMagic, static_cast<uint32_t>(name.size()),
111 static_cast<uint32_t>(bin.size())};116 static_cast<uint32_t>(bin.size())};
@@ -237,6 +242,7 @@ TEST(UtestCustomOpCast, filters_python_adapter_by_capability) {
237 242 
238 BaseCustomOp *base = &adapter;243 BaseCustomOp *base = &adapter;
239 EXPECT_NE(nullptr, CustomOpCast<EagerExecuteOp>(base));244 EXPECT_NE(nullptr, CustomOpCast<EagerExecuteOp>(base));
245+ EXPECT_EQ(nullptr, CustomOpCast<AnnotatedArgsOp>(base));
240 EXPECT_EQ(nullptr, CustomOpCast<CompilableOp>(base));246 EXPECT_EQ(nullptr, CustomOpCast<CompilableOp>(base));
241 EXPECT_EQ(nullptr, CustomOpCast<ShapeInferOp>(base));247 EXPECT_EQ(nullptr, CustomOpCast<ShapeInferOp>(base));
242 EXPECT_EQ(nullptr, CustomOpCast<PortableOp>(base));248 EXPECT_EQ(nullptr, CustomOpCast<PortableOp>(base));
@@ -250,6 +256,54 @@ TEST(UtestCustomOpCast, filters_python_adapter_by_capability) {
250 EXPECT_FALSE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));256 EXPECT_FALSE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
251}257}
252 258 
259+TEST(UtestCustomOpCast, filters_python_adapter_annotated_args_by_capability) {
260+ PythonCustomOpDescriptor desc;
261+ desc.descriptor_key = "python_adapter_annotated_args_only";
262+ desc.op_type = "PythonAdapterAnnotatedArgsOnly";
263+ AddCustomOpCapability(desc.capabilities, CustomOpCapability::kAnnotatedArgs);
264+ 
265+ PythonCustomOpCallbacks callbacks;
266+ callbacks.create = CreateMockPythonCustomOpHolder;
267+ callbacks.destroy = DestroyMockPythonCustomOpHolder;
268+ callbacks.declare_launch_args = DeclareMockPythonCustomOp;
269+ 
270+ ASSERT_TRUE(PythonCustomOpRuntimeRegistry::Register(desc, callbacks));
271+ {
272+ PythonCustomOpAdapter adapter(desc);
273+ EXPECT_TRUE(adapter.IsValid());
274+ 
275+ BaseCustomOp *base = &adapter;
276+ EXPECT_EQ(nullptr, CustomOpCast<EagerExecuteOp>(base));
277+ EXPECT_NE(nullptr, CustomOpCast<AnnotatedArgsOp>(base));
278+ EXPECT_EQ(nullptr, CustomOpCast<CompilableOp>(base));
279+ }
280+ EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
281+}
282+ 
283+TEST(UtestCustomOpCast, exposes_each_python_adapter_capability_in_dual_mode) {
284+ PythonCustomOpDescriptor desc;
285+ desc.descriptor_key = "python_adapter_dual_capability";
286+ desc.op_type = "PythonAdapterDualCapability";
287+ AddCustomOpCapability(desc.capabilities, CustomOpCapability::kEagerExecute);
288+ AddCustomOpCapability(desc.capabilities, CustomOpCapability::kAnnotatedArgs);
289+ 
290+ PythonCustomOpCallbacks callbacks;
291+ callbacks.create = CreateMockPythonCustomOpHolder;
292+ callbacks.destroy = DestroyMockPythonCustomOpHolder;
293+ callbacks.execute = ExecuteMockPythonCustomOp;
294+ callbacks.declare_launch_args = DeclareMockPythonCustomOp;
295+ 
296+ ASSERT_TRUE(PythonCustomOpRuntimeRegistry::Register(desc, callbacks));
297+ {
298+ PythonCustomOpAdapter adapter(desc);
299+ EXPECT_TRUE(adapter.IsValid());
300+ BaseCustomOp *base = &adapter;
301+ EXPECT_NE(nullptr, CustomOpCast<EagerExecuteOp>(base));
302+ EXPECT_NE(nullptr, CustomOpCast<AnnotatedArgsOp>(base));
303+ }
304+ EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
305+}
306+ 
253TEST(UtestCustomOpCast, rejects_unsupported_python_adapter_capability) {307TEST(UtestCustomOpCast, rejects_unsupported_python_adapter_capability) {
254 PythonCustomOpDescriptor desc;308 PythonCustomOpDescriptor desc;
255 desc.descriptor_key = "python_adapter_shape_unsupported";309 desc.descriptor_key = "python_adapter_shape_unsupported";
Mtests/ge/ut/ge/runtime/custom_op/python_custom_op_ir_meta_unittest.cc+145-0
@@ -11,12 +11,19 @@
11#include <gtest/gtest.h>11#include <gtest/gtest.h>
12 12 
13#include <new>13#include <new>
14+#include <string>
14 15 
16+#include "common/python_runtime/python_artifact_utils.h"
17+#include "common/python_runtime/python_bridge_loader_utils.h"
15#include "runtime/custom_op/python_custom_op_adapter.h"18#include "runtime/custom_op/python_custom_op_adapter.h"
19+#include "runtime/custom_op/python_custom_op_bridge_c_api.h"
16 20 
17namespace ge {21namespace ge {
18namespace custom_op {22namespace custom_op {
19namespace {23namespace {
24+namespace artifact = ::ge::python_artifact;
25+namespace bridge_loader = ::ge::python_bridge_loader;
26+ 
20struct MockPythonCustomOpHolder {};27struct MockPythonCustomOpHolder {};
21 28 
22void *CreateMockPythonCustomOpHolder(const PythonCustomOpDescriptor *desc) {29void *CreateMockPythonCustomOpHolder(const PythonCustomOpDescriptor *desc) {
@@ -32,6 +39,109 @@ graphStatus ExecuteMockPythonCustomOp(const void *holder, gert::EagerOpExecution
32 EXPECT_NE(holder, nullptr);39 EXPECT_NE(holder, nullptr);
33 return (holder != nullptr) ? GRAPH_SUCCESS : GRAPH_FAILED;40 return (holder != nullptr) ? GRAPH_SUCCESS : GRAPH_FAILED;
34}41}
42+ 
43+graphStatus DeclareMockPythonCustomOp(const void *holder, gert::AnnotatedArgsContext *ctx) {
44+ return ((holder != nullptr) && (ctx != nullptr)) ? GRAPH_SUCCESS : GRAPH_FAILED;
45+}
46+ 
47+struct MockPythonCustomOpBridgeLoadState {
48+ PythonCustomOpBridgeApi api{};
49+ const PythonCustomOpBridgeApi *api_to_return{nullptr};
50+ uint32_t close_count{0U};
51+ uint32_t set_config_count{0U};
52+ uint32_t register_count{0U};
53+};
54+ 
55+int g_mock_python_custom_op_bridge_handle = 0;
56+MockPythonCustomOpBridgeLoadState g_mock_python_custom_op_bridge_load_state;
57+ 
58+Status SetMockPythonCustomOpArtifactConfig(const PythonCustomOpBridgeArtifactConfig *config) {
59+ if ((config == nullptr) || (config->artifact_root == nullptr) || (config->native_module_path == nullptr)) {
60+ return FAILED;
61+ }
62+ ++g_mock_python_custom_op_bridge_load_state.set_config_count;
63+ return SUCCESS;
64+}
65+ 
66+Status RegisterMockPythonCustomOps(const PythonCustomOpRegistrar *registrar) {
67+ (void)registrar;
68+ ++g_mock_python_custom_op_bridge_load_state.register_count;
69+ return SUCCESS;
70+}
71+ 
72+void ResetMockPythonCustomOpBridgeState() {}
73+ 
74+void ShutdownMockPythonCustomOpBridge() {}
75+ 
76+const PythonCustomOpBridgeApi *GetMockPythonCustomOpBridgeApi() {
77+ return g_mock_python_custom_op_bridge_load_state.api_to_return;
78+}
79+ 
80+std::string ResolveMockPythonCustomOpBridgePath(const char *path) {
81+ return (path == nullptr) ? std::string() : std::string("/tmp/libge_python_custom_op_bridge.so");
82+}
83+ 
84+void *OpenMockPythonCustomOpBridge(const char *path, int flags) {
85+ (void)flags;
86+ return (path == nullptr) ? nullptr : static_cast<void *>(&g_mock_python_custom_op_bridge_handle);
87+}
88+ 
89+int CloseMockPythonCustomOpBridge(void *handle) {
90+ if (handle != nullptr) {
91+ ++g_mock_python_custom_op_bridge_load_state.close_count;
92+ }
93+ return 0;
94+}
95+ 
96+void *LookupMockPythonCustomOpBridgeSymbol(void *handle, const char *symbol) {
97+ if ((handle == nullptr) || (symbol == nullptr)) {
98+ return nullptr;
99+ }
100+ return reinterpret_cast<void *>(&GetMockPythonCustomOpBridgeApi);
101+}
102+ 
103+artifact::PythonRuntimeKey ResolveMockPythonRuntimeKey() {
104+ return {};
105+}
106+ 
107+bool IsMockPythonCustomOpBridgeApiValid(const PythonCustomOpBridgeApi *api, const uint32_t expected_abi) {
108+ return (api != nullptr) && (api->abi_version == expected_abi) && (api->set_artifact_config != nullptr) &&
109+ (api->register_custom_ops != nullptr) && (api->reset_bridge_state != nullptr) &&
110+ (api->shutdown_bridge != nullptr);
111+}
112+ 
113+bridge_loader::BridgeLoadDependencies MakeMockPythonCustomOpBridgeLoadDependencies() {
114+ return bridge_loader::BridgeLoadDependencies{
115+ &ResolveMockPythonCustomOpBridgePath, &OpenMockPythonCustomOpBridge,
116+ &CloseMockPythonCustomOpBridge, &LookupMockPythonCustomOpBridgeSymbol,
117+ &ResolveMockPythonRuntimeKey, kPythonCustomOpBridgeGetApiSymbol,
118+ kPythonCustomOpBridgeAbiVersion, 0,
119+ };
120+}
121+ 
122+void ResetMockPythonCustomOpBridgeLoadState(const uint32_t abi_version) {
123+ g_mock_python_custom_op_bridge_load_state = MockPythonCustomOpBridgeLoadState{};
124+ g_mock_python_custom_op_bridge_load_state.api = PythonCustomOpBridgeApi{
125+ abi_version,
126+ &SetMockPythonCustomOpArtifactConfig,
127+ &RegisterMockPythonCustomOps,
128+ &ResetMockPythonCustomOpBridgeState,
129+ &ShutdownMockPythonCustomOpBridge,
130+ };
131+ g_mock_python_custom_op_bridge_load_state.api_to_return = &g_mock_python_custom_op_bridge_load_state.api;
132+}
133+ 
134+bridge_loader::BridgeLoadStatus LoadMockPythonCustomOpBridge(
135+ bridge_loader::LoadedBridgeCandidate<PythonCustomOpBridgeApi> &loaded_bridge) {
136+ const artifact::BridgeLibraryCandidate candidate{
137+ "libge_python_custom_op_bridge.so",
138+ "/tmp/custom_op/python_custom_op_artifacts/cp311-linux-aarch64",
139+ "/tmp/custom_op/python_custom_op_artifacts/cp311-linux-aarch64/_ge_custom_op_native.so",
140+ };
141+ return bridge_loader::TryLoadBridgeCandidate<PythonCustomOpBridgeApi, PythonCustomOpBridgeArtifactConfig>(
142+ artifact::PythonRuntimeKey{}, candidate, MakeMockPythonCustomOpBridgeLoadDependencies(),
143+ &IsMockPythonCustomOpBridgeApiValid, loaded_bridge);
144+}
35} // namespace145} // namespace
36 146 
37TEST(PythonCustomOpAdapter, forwards_execute_without_ir_meta_pod) {147TEST(PythonCustomOpAdapter, forwards_execute_without_ir_meta_pod) {
@@ -73,5 +183,40 @@ TEST(PythonCustomOpAdapter, keeps_legacy_execute_without_registered_ir) {
73 }183 }
74 EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));184 EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
75}185}
186+ 
187+TEST(PythonCustomOpAdapter, validates_annotated_args_callback_by_capability) {
188+ PythonCustomOpDescriptor desc;
189+ desc.descriptor_key = "python_adapter_annotated_args_callback";
190+ desc.op_type = "PythonCustomOpIrMetaUt";
191+ AddCustomOpCapability(desc.capabilities, CustomOpCapability::kAnnotatedArgs);
192+ 
193+ PythonCustomOpCallbacks callbacks;
194+ callbacks.create = CreateMockPythonCustomOpHolder;
195+ callbacks.destroy = DestroyMockPythonCustomOpHolder;
196+ EXPECT_FALSE(callbacks.IsValid(desc.capabilities));
197+ 
198+ callbacks.declare_launch_args = DeclareMockPythonCustomOp;
199+ EXPECT_TRUE(callbacks.IsValid(desc.capabilities));
200+ EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Register(desc, callbacks));
201+ EXPECT_TRUE(PythonCustomOpRuntimeRegistry::Unregister(desc.descriptor_key));
202+}
203+ 
204+TEST(PythonCustomOpBridgeAbi, rejects_mismatched_abi_before_registration_and_accepts_current) {
205+ bridge_loader::LoadedBridgeCandidate<PythonCustomOpBridgeApi> loaded_bridge;
206+ ResetMockPythonCustomOpBridgeLoadState(kPythonCustomOpBridgeAbiVersion + 1U);
207+ 
208+ EXPECT_EQ(LoadMockPythonCustomOpBridge(loaded_bridge), bridge_loader::BridgeLoadStatus::kInvalidApi);
209+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.close_count, 1U);
210+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.set_config_count, 0U);
211+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.register_count, 0U);
212+ 
213+ ResetMockPythonCustomOpBridgeLoadState(kPythonCustomOpBridgeAbiVersion);
214+ ASSERT_EQ(LoadMockPythonCustomOpBridge(loaded_bridge), bridge_loader::BridgeLoadStatus::kSuccess);
215+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.close_count, 0U);
216+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.set_config_count, 1U);
217+ ASSERT_NE(loaded_bridge.api, nullptr);
218+ EXPECT_EQ(loaded_bridge.api->register_custom_ops(nullptr), SUCCESS);
219+ EXPECT_EQ(g_mock_python_custom_op_bridge_load_state.register_count, 1U);
220+}
76} // namespace custom_op221} // namespace custom_op
77} // namespace ge222} // namespace ge