From 4b5ecd8ea80d85d39cc28cb57a44f70d68cec822 Mon Sep 17 00:00:00 2001 From: FlagOS Templates Date: Wed, 7 Oct 2026 19:19:26 +0800 Subject: [PATCH] fix: upgrade legacy operator register contracts --- .github/workflows/ci.yml | 107 ++++++++++++++++-- CHANGELOG.md | 14 +++ ops/_local_scalar_dense/register.py | 45 +++++++- ops/clone/register.py | 45 +++++++- ops/contiguous/register.py | 45 +++++++- ops/copy_/register.py | 45 +++++++- ops/detach/register.py | 45 +++++++- ops/detach_/register.py | 31 ++++- ops/dropout/register.py | 45 +++++++- ops/empty/register.py | 32 +++++- ops/empty_like/register.py | 45 +++++++- ops/empty_strided/register.py | 45 +++++++- ops/item/register.py | 45 +++++++- ops/result_type/register.py | 31 ++++- ops/type_as/register.py | 47 +++++++- .../unit/test_issue_19_register_contracts.py | 104 +++++++++++++++++ 16 files changed, 721 insertions(+), 50 deletions(-) create mode 100644 tests/unit/test_issue_19_register_contracts.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6cea5c9..8f940c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -6,9 +6,6 @@ on: pull_request: jobs: - # wt 2026-10-07-fix 新增 packaging 门: 算子包化回归（CPU 可跑部分） - # ——全算子 __init__.py 导入扫描; NPU 消费段在本机跑, CI 里跳过。 - # # wt ops-packaging: runs-on: ubuntu-latest steps: @@ -70,9 +67,103 @@ jobs: def anchors_of(p): k = str(p) if k not in anchor_cache: - text = p.read_text(encoding='utf-8', errors='replace') - hs = re.findall(r'^#{1,6}\s+(.*)$', text, re.M) - anchor_cache[k] = { - re.sub(r'[^\w\u4e00-\u9fff-]', '', h.strip().lower().replace(' ', '-')) - for h in hs} + try: + anchor_cache[k] = set(re.findall(r' {target}') + elif anchor and anchor not in anchors_of(tp): + errs.append(f'{md}: #{anchor} in {target}') + + print(f'检查完成: {len(errs)} 个错误') + for e in errs: + print(' -', e) + sys.exit(1 if errs else 0) + EOF + + - name: Python 语法检查 + run: | + python3 - <<'EOF' + import ast, pathlib, sys + fails = [] + for p in pathlib.Path('.').rglob('*.py'): + if '.git' in str(p): + continue + try: + ast.parse(p.read_text(), filename=str(p)) + except SyntaxError as e: + fails.append(f'{p}: {e}') + if fails: + for f in fails: + print(' -', f) + sys.exit(1) + print('Python 语法检查通过') + EOF + + - name: YAML 配置检查 + run: | + pip install pyyaml -q + python3 - <<'EOF' + import yaml, glob, sys + fails = [] + for f in glob.glob('configs/devices/*.yaml') + glob.glob('inputs/*.yaml'): + try: + yaml.safe_load(open(f)) + except Exception as e: + fails.append(f'{f}: {e}') + if fails: + for f in fails: + print(' -', f) + sys.exit(1) + print('YAML 检查通过') + EOF + + - name: Intake manifest 契约校验 + run: python3 scripts/intake_validate.py --validate-only + + - name: 性能基线结构检查 + run: | + python3 - <<'EOF' + import json, glob, sys + fails = [] + for f in glob.glob('perf/baselines/*.json') + ['intake/manifest.schema.json']: + try: + d = json.load(open(f)) + except Exception as e: + fails.append(f'{f}: {e}') + continue + if f.startswith('perf/'): + for k in ('schema', 'device', 'created_at', 'updated_at', + 'git_commit', 'env_fingerprint', 'cases'): + if k not in d: + fails.append(f'{f}: 缺字段 {k}') + for cid, c in d.get('cases', {}).items(): + if 'latency_ms' not in c.get('metrics', {}): + fails.append(f'{f}: {cid} 缺 metrics.latency_ms') + if fails: + for x in fails: + print(' -', x) + sys.exit(1) + print('性能基线结构检查通过') + EOF diff --git a/CHANGELOG.md b/CHANGELOG.md index cd6fc38..dacc9e3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,20 @@ ## [Unreleased] +### Changed +- **13 个 MLU 算子 register 签名对齐新库约定**（#19）: + `_local_scalar_dense` / `clone` / `contiguous` / `copy_` / `detach` / + `detach_` / `dropout` / `empty` / `empty_like` / `empty_strided` / + `item` / `result_type` / `type_as` 统一支持 + `counter` 拦截计数与显式 `platform` 参数；单平台 Cambricon 会校验 + platform/dispatch key；aten 实现导入改为包态相对导入 + 单算子脚本 + fallback，三个不可 A1 注册的算子保留原有禁用语义 + +### Fixed +- 恢复 #20 中被截断的 `docs-and-syntax` 完整脚本：Markdown 链接、 + Python 语法、YAML、intake manifest 与性能基线结构检查重新生效； + 保留新增 `ops-packaging` 门 + ### Added - **sdpa_math 第二平台: p800-kunlunxin**（vendor O/LSE + exact-P + A1） - `_profile` / benchmark / framework 消费方接入 P800；A1 使用 diff --git a/ops/_local_scalar_dense/register.py b/ops/_local_scalar_dense/register.py index 703b66d..7a1bf9d 100644 --- a/ops/_local_scalar_dense/register.py +++ b/ops/_local_scalar_dense/register.py @@ -29,14 +29,53 @@ def _bump(): def lsd_aten(self): _bump() - from kernel.torch_level import _local_scalar_dense_torch + try: + from .kernel.torch_level import _local_scalar_dense_torch + except ImportError: # standalone single-operator script mode + from kernel.torch_level import _local_scalar_dense_torch return _local_scalar_dense_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::_local_scalar_dense`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return lsd_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("_local_scalar_dense", lsd_aten, dispatch_key) + _LIB.impl("_local_scalar_dense", aten_impl, dispatch_key) return _LIB diff --git a/ops/clone/register.py b/ops/clone/register.py index 1dad279..2d053c3 100644 --- a/ops/clone/register.py +++ b/ops/clone/register.py @@ -31,14 +31,53 @@ def _bump() -> None: def clone_aten(self, *, memory_format=None): """aten::clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor""" _bump() - from kernel.triton_level import clone_triton + try: + from .kernel.triton_level import clone_triton + except ImportError: # standalone single-operator script mode + from kernel.triton_level import clone_triton return clone_triton(self, memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::clone`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return clone_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("clone", clone_aten, dispatch_key) + _LIB.impl("clone", aten_impl, dispatch_key) return _LIB diff --git a/ops/contiguous/register.py b/ops/contiguous/register.py index 3e95d8e..5918986 100644 --- a/ops/contiguous/register.py +++ b/ops/contiguous/register.py @@ -30,14 +30,53 @@ def _bump() -> None: def contiguous_aten(self, *, memory_format=torch.contiguous_format): """aten::contiguous(Tensor self, *, MemoryFormat memory_format=...)""" _bump() - from kernel.triton_level import contiguous_triton + try: + from .kernel.triton_level import contiguous_triton + except ImportError: # standalone single-operator script mode + from kernel.triton_level import contiguous_triton return contiguous_triton(self, memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::contiguous`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return contiguous_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("contiguous", contiguous_aten, dispatch_key) + _LIB.impl("contiguous", aten_impl, dispatch_key) return _LIB diff --git a/ops/copy_/register.py b/ops/copy_/register.py index 048d336..2980922 100644 --- a/ops/copy_/register.py +++ b/ops/copy_/register.py @@ -33,14 +33,53 @@ def _bump() -> None: def copy_aten(self, src, non_blocking: bool = False): """aten::copy_(Tensor self, Tensor src, bool non_blocking=False) -> Tensor""" _bump() - from kernel.triton_level import copy_triton + try: + from .kernel.triton_level import copy_triton + except ImportError: # standalone single-operator script mode + from kernel.triton_level import copy_triton return copy_triton(self, src, non_blocking) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::copy_`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return copy_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("copy_", copy_aten, dispatch_key) + _LIB.impl("copy_", aten_impl, dispatch_key) return _LIB diff --git a/ops/detach/register.py b/ops/detach/register.py index b6c0a2c..a6beae6 100644 --- a/ops/detach/register.py +++ b/ops/detach/register.py @@ -28,14 +28,53 @@ def _bump(): def detach_aten(self): _bump() - from kernel.torch_level import detach_torch + try: + from .kernel.torch_level import detach_torch + except ImportError: # standalone single-operator script mode + from kernel.torch_level import detach_torch return detach_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::detach`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return detach_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("detach", detach_aten, dispatch_key) + _LIB.impl("detach", aten_impl, dispatch_key) return _LIB diff --git a/ops/detach_/register.py b/ops/detach_/register.py index a82a2d2..f95dd61 100644 --- a/ops/detach_/register.py +++ b/ops/detach_/register.py @@ -4,7 +4,32 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key="PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> None: + """Validate the new registration contract, then keep A1 disabled.""" + _backend_for(dispatch_key, platform) raise NotImplementedError( - "detach_ 不做 A1 注册：其为 autograd 层原地算子，PrivateUse1 IMPL 不命中。" - ) + "detach_ 不做 A1 注册：其为 autograd 层原地算子，PrivateUse1 IMPL 不命中。") diff --git a/ops/dropout/register.py b/ops/dropout/register.py index 48a2aa2..0427403 100644 --- a/ops/dropout/register.py +++ b/ops/dropout/register.py @@ -30,14 +30,53 @@ def _bump() -> None: def dropout_aten(input, p: float, train: bool): """aten::dropout(Tensor input, float p, bool train) -> Tensor""" _bump() - from kernel.triton_level import dropout_triton + try: + from .kernel.triton_level import dropout_triton + except ImportError: # standalone single-operator script mode + from kernel.triton_level import dropout_triton return dropout_triton(input, p, train) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::dropout`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return dropout_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("dropout", dropout_aten, dispatch_key) + _LIB.impl("dropout", aten_impl, dispatch_key) return _LIB diff --git a/ops/empty/register.py b/ops/empty/register.py index 22b13cf..ae6fb8b 100644 --- a/ops/empty/register.py +++ b/ops/empty/register.py @@ -8,8 +8,32 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key: str = "PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> None: + """Validate the new registration contract, then keep A1 disabled.""" + _backend_for(dispatch_key, platform) raise NotImplementedError( - "empty 不做 A1 注册：分配原语注册会与 empty_strided 递归，" - "且全局覆盖 torch.empty 风险高。route=自用/实验。" - ) + "empty 不做 A1 注册：分配原语注册会与 empty_strided 递归，且全局覆盖 torch.empty 风险高。route=自用/实验。") diff --git a/ops/empty_like/register.py b/ops/empty_like/register.py index 9c319fb..0d1192c 100644 --- a/ops/empty_like/register.py +++ b/ops/empty_like/register.py @@ -33,16 +33,55 @@ def empty_like_aten(self, *, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None): """aten::empty_like(Tensor self, *, ...) -> Tensor""" _bump() - from kernel.torch_level import empty_like_torch + try: + from .kernel.torch_level import empty_like_torch + except ImportError: # standalone single-operator script mode + from kernel.torch_level import empty_like_torch return empty_like_torch(self, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory, memory_format=memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::empty_like`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return empty_like_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("empty_like", empty_like_aten, dispatch_key) + _LIB.impl("empty_like", aten_impl, dispatch_key) return _LIB diff --git a/ops/empty_strided/register.py b/ops/empty_strided/register.py index d2300b8..ca382ac 100644 --- a/ops/empty_strided/register.py +++ b/ops/empty_strided/register.py @@ -32,15 +32,54 @@ def empty_strided_aten(size, stride, *, dtype=None, layout=None, device=None, pin_memory=None): """aten::empty_strided(SymInt[] size, SymInt[] stride, *, ...) -> Tensor""" _bump() - from kernel.torch_level import empty_strided_torch + try: + from .kernel.torch_level import empty_strided_torch + except ImportError: # standalone single-operator script mode + from kernel.torch_level import empty_strided_torch return empty_strided_torch(size, stride, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::empty_strided`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return empty_strided_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("empty_strided", empty_strided_aten, dispatch_key) + _LIB.impl("empty_strided", aten_impl, dispatch_key) return _LIB diff --git a/ops/item/register.py b/ops/item/register.py index 2ea0cd8..542dc00 100644 --- a/ops/item/register.py +++ b/ops/item/register.py @@ -28,14 +28,53 @@ def _bump(): def item_aten(self): _bump() - from kernel.torch_level import item_torch + try: + from .kernel.torch_level import item_torch + except ImportError: # standalone single-operator script mode + from kernel.torch_level import item_torch return item_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::item`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return item_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("item", item_aten, dispatch_key) + _LIB.impl("item", aten_impl, dispatch_key) return _LIB diff --git a/ops/result_type/register.py b/ops/result_type/register.py index eed4e3d..17b239b 100644 --- a/ops/result_type/register.py +++ b/ops/result_type/register.py @@ -4,7 +4,32 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key="PrivateUse1"): +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> None: + """Validate the new registration contract, then keep A1 disabled.""" + _backend_for(dispatch_key, platform) raise NotImplementedError( - "result_type 不做 A1 注册：纯 dtype 元数据函数，PrivateUse1 IMPL 不命中。" - ) + "result_type 不做 A1 注册：纯 dtype 元数据函数，PrivateUse1 IMPL 不命中。") diff --git a/ops/type_as/register.py b/ops/type_as/register.py index cfc7bb1..439b111 100644 --- a/ops/type_as/register.py +++ b/ops/type_as/register.py @@ -36,17 +36,54 @@ def _bump() -> None: def type_as_aten(self, other): """aten::type_as(Tensor self, Tensor other) -> Tensor""" _bump() - from kernel.triton_level import type_as_triton + try: + from .kernel.triton_level import type_as_triton + except ImportError: # standalone single-operator script mode + from kernel.triton_level import type_as_triton return type_as_triton(self, other) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: - """注册 type_as 到 aten；Library 对象必须保持引用（否则被回收）。""" +_PLATFORM_ALIASES = {"mlu": "cambricon", "mlu590": "cambricon"} +_SUPPORTED_PLATFORMS = {"cambricon"} +_SUPPORTED_DISPATCH_KEYS = { + "AutogradPrivateUse1", "PrivateUse1", "Autograd", +} + + +def _backend_for(dispatch_key: str, platform: str | None = None) -> str: + """Resolve the sole current backend without silently accepting typos.""" + requested = platform or "cambricon" + canonical = _PLATFORM_ALIASES.get(requested, requested) + if canonical not in _SUPPORTED_PLATFORMS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: " + f"{sorted(_SUPPORTED_PLATFORMS)}") + if dispatch_key not in _SUPPORTED_DISPATCH_KEYS: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: " + f"{sorted(_SUPPORTED_DISPATCH_KEYS)}（或显式传 platform）") + return canonical + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform: str | None = None) -> torch.library.Library: + """Register ``aten::type_as`` with optional interception counter. + + ``platform`` is accepted explicitly so a second backend can later use the + same call contract. These operators currently have one Cambricon backend. + """ + _backend_for(dispatch_key, platform) + + def aten_impl(*args, **kwargs): + if counter is not None: + counter["n"] = counter.get("n", 0) + 1 + return type_as_aten(*args, **kwargs) + global _LIB _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("type_as", type_as_aten, dispatch_key) + _LIB.impl("type_as", aten_impl, dispatch_key) return _LIB -# 统一入口别名（与 A2/B 插件约定对齐，便于框架级注入） register = register_a1 diff --git a/tests/unit/test_issue_19_register_contracts.py b/tests/unit/test_issue_19_register_contracts.py new file mode 100644 index 0000000..21967bc --- /dev/null +++ b/tests/unit/test_issue_19_register_contracts.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +import importlib +import inspect + +import pytest +import torch + + +ISSUE_19_OPS = ( + "_local_scalar_dense", + "clone", + "contiguous", + "copy_", + "detach", + "detach_", + "dropout", + "empty", + "empty_like", + "empty_strided", + "item", + "result_type", + "type_as", +) + +REGISTERABLE = { + "_local_scalar_dense": "lsd_aten", + "clone": "clone_aten", + "contiguous": "contiguous_aten", + "copy_": "copy_aten", + "detach": "detach_aten", + "dropout": "dropout_aten", + "empty_like": "empty_like_aten", + "empty_strided": "empty_strided_aten", + "item": "item_aten", + "type_as": "type_as_aten", +} + + +@pytest.fixture() +def modules(): + return { + name: importlib.import_module(f"ops.{name}.register") + for name in ISSUE_19_OPS + } + + +def test_issue_19_register_a1_signatures(modules): + for name, module in modules.items(): + parameters = inspect.signature(module.register_a1).parameters + assert list(parameters) == ["dispatch_key", "counter", "platform"], name + assert parameters["dispatch_key"].default == "AutogradPrivateUse1", name + assert parameters["counter"].default is None, name + assert parameters["platform"].default is None, name + + +def test_issue_19_rejects_unknown_platform_before_registration(modules): + for name, module in modules.items(): + with pytest.raises(RuntimeError, match="未知 platform"): + module.register_a1( + "PrivateUse1", counter={"n": 0}, platform="unsupported" + ) + + +def test_issue_19_counter_is_captured_by_registered_function(modules, monkeypatch): + class FakeLibrary: + def __init__(self, namespace, kind): + assert namespace == "aten" + assert kind == "IMPL" + self.impl_fn = None + self.dispatch_key = None + + def impl(self, op_name, function, dispatch_key): + self.op_name = op_name + self.impl_fn = function + self.dispatch_key = dispatch_key + + fake_library = FakeLibrary + monkeypatch.setattr(torch.library, "Library", fake_library) + + for name, aten_name in REGISTERABLE.items(): + module = modules[name] + sentinel = object() + monkeypatch.setattr( + module, aten_name, lambda *args, **kwargs: sentinel + ) + counter = {"n": 0} + library = module.register_a1( + "PrivateUse1", counter=counter, platform="cambricon" + ) + assert isinstance(library, FakeLibrary), name + assert library.dispatch_key == "PrivateUse1", name + assert counter == {"n": 0}, name + assert library.impl_fn(*()) is sentinel, name + assert counter == {"n": 1}, name + + +def test_issue_19_registerable_ops_use_package_relative_kernel_imports(): + import ast + from pathlib import Path + + for name in REGISTERABLE: + source = (Path("ops") / name / "register.py").read_text() + assert "from .kernel." in source, name