From 15e1d7c6a38b720fd0c733fd19bb7094764aa065 Mon Sep 17 00:00:00 2001 From: Dashener2 Date: Fri, 9 Oct 2026 14:51:18 +0800 Subject: [PATCH 1/2] =?UTF-8?q?fix:=2013=20=E7=AE=97=E5=AD=90=20register?= =?UTF-8?q?=5Fa1=20=E7=AD=BE=E5=90=8D=E5=8D=87=E7=BA=A7=EF=BC=88counter=20?= =?UTF-8?q?+=20platform=EF=BC=8C#19)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 按 templates/operator/register.py 与 ops/embedding 的 _BACKENDS + _DEVICE_PROBE 范式，升级 13 个算子的 A1 注册签名： - 签名统一为 register_a1(dispatch_key="AutogradPrivateUse1", counter: dict | None = None, platform=None) - 新增 _BACKENDS/_backend_for（platform 显式优先，dispatch_key 推断； 覆盖 Autograd/PrivateUse1/AutogradPrivateUse1 -> cambricon）与 _DEVICE_PROBE 注册守卫（MLU 可用性探测） - 10 个可注册算子：命中计数支持 counter["n"]；test/op_level.py 同步 传入 counter 并断言 - 3 个不可注册算子（empty/detach_/result_type）同步签名，保持按设计 抛 NotImplementedError 验证（v26.08.0 容器，MLU）: - 13/13 example 三层全绿；pytest tests/unit 24 passed - counter/platform/默认 key 实测通过 - 真实 vLLM A1 框架测试 tests/framework_level/test_a1_aten_ops.py 1 passed - 13 算子包态导入通过（ops.test_packaging.py 依赖 torch_npu，MLU 上未跑） --- ops/_local_scalar_dense/register.py | 60 ++++++++++++++++++++++- ops/_local_scalar_dense/test/op_level.py | 4 +- ops/clone/register.py | 60 ++++++++++++++++++++++- ops/clone/test/op_level.py | 4 +- ops/contiguous/register.py | 60 ++++++++++++++++++++++- ops/contiguous/test/op_level.py | 4 +- ops/copy_/register.py | 60 ++++++++++++++++++++++- ops/copy_/test/op_level.py | 4 +- ops/detach/register.py | 60 ++++++++++++++++++++++- ops/detach/test/op_level.py | 4 +- ops/detach_/register.py | 4 +- ops/dropout/register.py | 60 ++++++++++++++++++++++- ops/dropout/test/op_level.py | 4 +- ops/empty/register.py | 4 +- ops/empty_like/register.py | 60 ++++++++++++++++++++++- ops/empty_like/test/op_level.py | 4 +- ops/empty_strided/register.py | 60 ++++++++++++++++++++++- ops/empty_strided/test/op_level.py | 4 +- ops/item/register.py | 60 ++++++++++++++++++++++- ops/item/test/op_level.py | 4 +- ops/result_type/register.py | 4 +- ops/type_as/register.py | 61 ++++++++++++++++++++++-- ops/type_as/test/op_level.py | 4 +- 23 files changed, 619 insertions(+), 34 deletions(-) diff --git a/ops/_local_scalar_dense/register.py b/ops/_local_scalar_dense/register.py index 703b66d..6cbc620 100644 --- a/ops/_local_scalar_dense/register.py +++ b/ops/_local_scalar_dense/register.py @@ -33,10 +33,66 @@ def lsd_aten(self): return _local_scalar_dense_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 _local_scalar_dense 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = lsd_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return lsd_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("_local_scalar_dense", lsd_aten, dispatch_key) + _LIB.impl("_local_scalar_dense", fn, dispatch_key) return _LIB diff --git a/ops/_local_scalar_dense/test/op_level.py b/ops/_local_scalar_dense/test/op_level.py index a4bce11..5e8a5e6 100644 --- a/ops/_local_scalar_dense/test/op_level.py +++ b/ops/_local_scalar_dense/test/op_level.py @@ -21,12 +21,14 @@ def run(profile): exp = x.item() reg.CALL_COUNT["_local_scalar_dense"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None before = reg.CALL_COUNT["_local_scalar_dense"] x = torch.full((1,), 3.5, dtype=torch.float32, device=dev) got = torch.ops.aten._local_scalar_dense.default(x) assert reg.CALL_COUNT["_local_scalar_dense"] > before, "未拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert got == 3.5, got return {"ok": True, "registered": "aten::_local_scalar_dense@" + profile.dispatch_key, diff --git a/ops/clone/register.py b/ops/clone/register.py index 1dad279..80bfb5b 100644 --- a/ops/clone/register.py +++ b/ops/clone/register.py @@ -35,10 +35,66 @@ def clone_aten(self, *, memory_format=None): return clone_triton(self, memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 clone 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = clone_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return clone_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("clone", clone_aten, dispatch_key) + _LIB.impl("clone", fn, dispatch_key) return _LIB diff --git a/ops/clone/test/op_level.py b/ops/clone/test/op_level.py index 25ad690..7bfff5b 100644 --- a/ops/clone/test/op_level.py +++ b/ops/clone/test/op_level.py @@ -24,13 +24,15 @@ def run(profile): xi_ref = clone_reference(xi) reg.CALL_COUNT["clone"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None # 1) 拦截验证 before = reg.CALL_COUNT["clone"] out = x.clone() assert reg.CALL_COUNT["clone"] > before, "aten::clone 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert torch.equal(out, exp), "拦截后数值不一致" assert out.data_ptr() != x.data_ptr(), "拦截后仍共享存储" diff --git a/ops/contiguous/register.py b/ops/contiguous/register.py index 3e95d8e..299fea8 100644 --- a/ops/contiguous/register.py +++ b/ops/contiguous/register.py @@ -34,10 +34,66 @@ def contiguous_aten(self, *, memory_format=torch.contiguous_format): return contiguous_triton(self, memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 contiguous 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = contiguous_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return contiguous_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("contiguous", contiguous_aten, dispatch_key) + _LIB.impl("contiguous", fn, dispatch_key) return _LIB diff --git a/ops/contiguous/test/op_level.py b/ops/contiguous/test/op_level.py index db0e2bb..1ee4f7d 100644 --- a/ops/contiguous/test/op_level.py +++ b/ops/contiguous/test/op_level.py @@ -24,13 +24,15 @@ def run(profile): cl_ref = contiguous_reference(cl, memory_format=torch.channels_last) reg.CALL_COUNT["contiguous"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None # 1) 拦截验证（非连续输入） before = reg.CALL_COUNT["contiguous"] out = x.contiguous() assert reg.CALL_COUNT["contiguous"] > before, "aten::contiguous 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert out.is_contiguous() and torch.equal(out, exp) # 2) 别名快路径 diff --git a/ops/copy_/register.py b/ops/copy_/register.py index 048d336..7289913 100644 --- a/ops/copy_/register.py +++ b/ops/copy_/register.py @@ -37,10 +37,66 @@ def copy_aten(self, src, non_blocking: bool = False): return copy_triton(self, src, non_blocking) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 copy_ 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = copy_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return copy_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("copy_", copy_aten, dispatch_key) + _LIB.impl("copy_", fn, dispatch_key) return _LIB diff --git a/ops/copy_/test/op_level.py b/ops/copy_/test/op_level.py index 43f1481..56effd4 100644 --- a/ops/copy_/test/op_level.py +++ b/ops/copy_/test/op_level.py @@ -36,13 +36,15 @@ def run(profile): copy_reference(idst_ref, isrc) reg.CALL_COUNT["copy_"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None # 1) 拦截 + 原地返回 + cast before = reg.CALL_COUNT["copy_"] out = dst.copy_(src) assert reg.CALL_COUNT["copy_"] > before, "aten::copy_ 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert out is dst, "copy_ 未返回 self" assert out.dtype == torch.float32 assert torch.equal(out, dst_ref), "拦截后 cast 精度不一致" diff --git a/ops/detach/register.py b/ops/detach/register.py index b6c0a2c..517700b 100644 --- a/ops/detach/register.py +++ b/ops/detach/register.py @@ -32,10 +32,66 @@ def detach_aten(self): return detach_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 detach 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = detach_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return detach_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("detach", detach_aten, dispatch_key) + _LIB.impl("detach", fn, dispatch_key) return _LIB diff --git a/ops/detach/test/op_level.py b/ops/detach/test/op_level.py index 6b67a71..b5e994e 100644 --- a/ops/detach/test/op_level.py +++ b/ops/detach/test/op_level.py @@ -20,11 +20,13 @@ def run(profile): exp_ptr, exp_stride = x.data_ptr(), tuple(x.stride()) reg.CALL_COUNT["detach"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None before = reg.CALL_COUNT["detach"] got = torch.ops.aten.detach.default(x) assert reg.CALL_COUNT["detach"] > before, "未拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert got.data_ptr() == exp_ptr and got.stride() == exp_stride assert got.requires_grad is False return {"ok": True, diff --git a/ops/detach_/register.py b/ops/detach_/register.py index a82a2d2..85808d3 100644 --- a/ops/detach_/register.py +++ b/ops/detach_/register.py @@ -4,7 +4,9 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key="PrivateUse1"): +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None): raise NotImplementedError( "detach_ 不做 A1 注册：其为 autograd 层原地算子，PrivateUse1 IMPL 不命中。" ) diff --git a/ops/dropout/register.py b/ops/dropout/register.py index 48a2aa2..bca2a85 100644 --- a/ops/dropout/register.py +++ b/ops/dropout/register.py @@ -34,10 +34,66 @@ def dropout_aten(input, p: float, train: bool): return dropout_triton(input, p, train) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 dropout 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = dropout_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return dropout_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("dropout", dropout_aten, dispatch_key) + _LIB.impl("dropout", fn, dispatch_key) return _LIB diff --git a/ops/dropout/test/op_level.py b/ops/dropout/test/op_level.py index 2edce24..bdadb40 100644 --- a/ops/dropout/test/op_level.py +++ b/ops/dropout/test/op_level.py @@ -18,7 +18,8 @@ def run(profile): dev = profile.torch_device reg.CALL_COUNT["dropout"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None # 1) 拦截: F.dropout 命中 @@ -26,6 +27,7 @@ def run(profile): before = reg.CALL_COUNT["dropout"] _ = F.dropout(x, 0.0, True) assert reg.CALL_COUNT["dropout"] > before, "F.dropout 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" after = reg.CALL_COUNT["dropout"] # 2) 确定性分支 diff --git a/ops/empty/register.py b/ops/empty/register.py index 22b13cf..337723e 100644 --- a/ops/empty/register.py +++ b/ops/empty/register.py @@ -8,7 +8,9 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key: str = "PrivateUse1"): +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None): raise NotImplementedError( "empty 不做 A1 注册：分配原语注册会与 empty_strided 递归，" "且全局覆盖 torch.empty 风险高。route=自用/实验。" diff --git a/ops/empty_like/register.py b/ops/empty_like/register.py index 9c319fb..aa98230 100644 --- a/ops/empty_like/register.py +++ b/ops/empty_like/register.py @@ -39,10 +39,66 @@ def empty_like_aten(self, *, dtype=None, layout=None, device=None, memory_format=memory_format) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 empty_like 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = empty_like_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return empty_like_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("empty_like", empty_like_aten, dispatch_key) + _LIB.impl("empty_like", fn, dispatch_key) return _LIB diff --git a/ops/empty_like/test/op_level.py b/ops/empty_like/test/op_level.py index cc95266..2e7b9ac 100644 --- a/ops/empty_like/test/op_level.py +++ b/ops/empty_like/test/op_level.py @@ -26,12 +26,14 @@ def run(profile): exp_mf = _meta(torch.empty_like(x, memory_format=torch.contiguous_format)) reg.CALL_COUNT["empty_like"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None before = reg.CALL_COUNT["empty_like"] out = torch.empty_like(x) assert reg.CALL_COUNT["empty_like"] > before, "aten::empty_like 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert _meta(out) == exp_meta, f"元数据变形 {_meta(out)} != {exp_meta}" assert out.data_ptr() != x.data_ptr(), "共享存储" diff --git a/ops/empty_strided/register.py b/ops/empty_strided/register.py index d2300b8..4dfa8ab 100644 --- a/ops/empty_strided/register.py +++ b/ops/empty_strided/register.py @@ -37,10 +37,66 @@ def empty_strided_aten(size, stride, *, dtype=None, layout=None, device=None, device=device, pin_memory=pin_memory) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 empty_strided 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = empty_strided_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return empty_strided_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("empty_strided", empty_strided_aten, dispatch_key) + _LIB.impl("empty_strided", fn, dispatch_key) return _LIB diff --git a/ops/empty_strided/test/op_level.py b/ops/empty_strided/test/op_level.py index aa95e19..9332451 100644 --- a/ops/empty_strided/test/op_level.py +++ b/ops/empty_strided/test/op_level.py @@ -24,13 +24,15 @@ def run(profile): dtype=torch.float32, device=dev)) reg.CALL_COUNT["empty_strided"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None before = reg.CALL_COUNT["empty_strided"] out = torch.empty_strided((128, 256), (1, 128), dtype=torch.float32, device=dev) assert reg.CALL_COUNT["empty_strided"] > before, "未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert _meta(out) == exp, f"元数据变形 {_meta(out)} != {exp}" # 重叠 stride diff --git a/ops/item/register.py b/ops/item/register.py index 2ea0cd8..800ad2d 100644 --- a/ops/item/register.py +++ b/ops/item/register.py @@ -32,10 +32,66 @@ def item_aten(self): return item_torch(self) -def register_a1(dispatch_key="PrivateUse1"): +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 item 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = item_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return item_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("item", item_aten, dispatch_key) + _LIB.impl("item", fn, dispatch_key) return _LIB diff --git a/ops/item/test/op_level.py b/ops/item/test/op_level.py index 3ffdeef..7bbfae2 100644 --- a/ops/item/test/op_level.py +++ b/ops/item/test/op_level.py @@ -21,12 +21,14 @@ def run(profile): exp = x.item() reg.CALL_COUNT["item"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None before = reg.CALL_COUNT["item"] x = torch.full((1,), 3.5, dtype=torch.float32, device=dev) got = torch.ops.aten.item.default(x) assert reg.CALL_COUNT["item"] > before, "未拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert got == 3.5, got return {"ok": True, "registered": "aten::item@" + profile.dispatch_key, diff --git a/ops/result_type/register.py b/ops/result_type/register.py index eed4e3d..4dade96 100644 --- a/ops/result_type/register.py +++ b/ops/result_type/register.py @@ -4,7 +4,9 @@ ROUTE = "自用/实验" -def register_a1(dispatch_key="PrivateUse1"): +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None): raise NotImplementedError( "result_type 不做 A1 注册：纯 dtype 元数据函数，PrivateUse1 IMPL 不命中。" ) diff --git a/ops/type_as/register.py b/ops/type_as/register.py index cfc7bb1..6a9a76d 100644 --- a/ops/type_as/register.py +++ b/ops/type_as/register.py @@ -40,11 +40,66 @@ def type_as_aten(self, other): return type_as_triton(self, other) -def register_a1(dispatch_key: str = "PrivateUse1") -> torch.library.Library: - """注册 type_as 到 aten；Library 对象必须保持引用（否则被回收）。""" +# ── 平台选择（单平台 cambricon/MLU590）── +# 第二平台接入时按 ops/embedding 范式扩成多后端 + PLATFORM.md/MERGE.md。 +# Dashener2 2026-10-09-fix register_a1 签名升级（issue #19）: counter + platform， +# 默认 dispatch_key 对齐模板（AutogradPrivateUse1）；_DEVICE_PROBE 注册守卫。 +# # Dashener2 +_DEVICE_PROBE = ("mlu", lambda: torch.mlu.is_available()) +_BACKENDS = {"cambricon": "cambricon"} + + +def _backend_for(dispatch_key: str, platform=None): + if platform is not None: + if platform not in _BACKENDS: + raise RuntimeError( + f"未知 platform={platform!r}; 已支持: {list(_BACKENDS)}" + "（第二平台接入见 ops/embedding/PLATFORM.md §5）") + return _BACKENDS[platform] + by_key = { + "Autograd": "cambricon", + "PrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", + } + if dispatch_key not in by_key: + raise RuntimeError( + f"未知 dispatch_key={dispatch_key!r}; 已支持: {list(by_key)}" + "（或显式传 platform）") + return by_key[dispatch_key] + + +def _probe_ok() -> bool: + attr, fn = _DEVICE_PROBE + try: + import torch_mlu # noqa: F401 触发 torch.mlu 注册 + except Exception: + return False + return hasattr(torch, attr) and getattr(torch, attr) is not None and fn() + + +def register_a1(dispatch_key: str = "AutogradPrivateUse1", + counter: dict | None = None, + platform=None) -> torch.library.Library: + """注册 type_as 到 aten；Library 对象必须保持引用（否则被回收）。 + + dispatch_key: 默认 AutogradPrivateUse1（对齐模板）；MLU 实际拦截 key + 为 profile 的 PrivateUse1，两 key 均可命中，platform 显式优先。 + counter: 可选 dict，命中计数写入 counter["n"]（op 层拦截验证）。 + platform: 显式平台名（多平台时优先于 dispatch_key 推断）。 + """ global _LIB + backend = _backend_for(dispatch_key, platform) + if backend == "cambricon" and not _probe_ok(): + raise RuntimeError( + f"register_a1 backend={backend!r}, dispatch_key={dispatch_key!r} " + "需要可用的 MLU 设备（torch.mlu.is_available()）") + fn = type_as_aten + if counter is not None: + def fn(*args, **kwargs): + counter["n"] = counter.get("n", 0) + 1 + return type_as_aten(*args, **kwargs) _LIB = torch.library.Library("aten", "IMPL") - _LIB.impl("type_as", type_as_aten, dispatch_key) + _LIB.impl("type_as", fn, dispatch_key) return _LIB diff --git a/ops/type_as/test/op_level.py b/ops/type_as/test/op_level.py index 76cea6f..d31f01e 100644 --- a/ops/type_as/test/op_level.py +++ b/ops/type_as/test/op_level.py @@ -17,7 +17,8 @@ def run(profile): dev = profile.torch_device reg.CALL_COUNT["type_as"] = 0 - lib = reg.register_a1(profile.dispatch_key) + _CALLS = {"n": 0} + lib = reg.register_a1(profile.dispatch_key, counter=_CALLS) assert lib is not None # 1) 拦截验证: torch 宿主调用必须命中我们的实现 @@ -27,6 +28,7 @@ def run(profile): out = torch.ops.aten.type_as.default(x, other) after = reg.CALL_COUNT["type_as"] assert after > before, "aten::type_as 未被拦截" + assert _CALLS["n"] > 0, "counter 未记入（op 层拦截）" assert torch.equal(out, type_as_reference(x, other)), "拦截后精度不一致" # 1b) 方法式调用同样命中 From e28394460e77906bc0b6fe0f3a951060e89bff06 Mon Sep 17 00:00:00 2001 From: Dashener2 Date: Mon, 12 Oct 2026 10:57:47 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20=E5=8C=85=E6=80=81=E6=89=A7=E8=A1=8C?= =?UTF-8?q?=E4=BF=AE=E5=A4=8D=20+=20A1=20key=20=E6=94=B6=E6=95=9B=EF=BC=88?= =?UTF-8?q?#25=20review=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 10 个可注册算子的实现导入改为「包态优先、脚本态 fallback」： try: from .kernel. import / except ImportError: from kernel... 修复 ops..register 包态下顶层 kernel 遮蔽导致的 ModuleNotFoundError（#25 review 可复现缺口） - _backend_for key 收敛为 canonical AutogradPrivateUse1 + 兼容 PrivateUse1， 移除 generic Autograd（多后端同进程覆盖过大，且无真机必要性证据） - clone/test/op_level.py 增加包态实际执行用例（子进程， 经 ops.clone.register -> ops.clone.kernel 执行并断言存储独立） 验证（v26.08.0 容器 / MLU）: - review repro 通过: from ops.clone.register import clone_aten 可执行 - 13/13 example 三层全绿；pytest tests/unit 24 passed - 真实 vLLM A1 框架测试 1 passed --- ops/_local_scalar_dense/register.py | 11 +++++++---- ops/clone/register.py | 11 +++++++---- ops/clone/test/op_level.py | 25 ++++++++++++++++++++++++- ops/contiguous/register.py | 11 +++++++---- ops/copy_/register.py | 11 +++++++---- ops/detach/register.py | 11 +++++++---- ops/dropout/register.py | 11 +++++++---- ops/empty_like/register.py | 11 +++++++---- ops/empty_strided/register.py | 11 +++++++---- ops/item/register.py | 11 +++++++---- ops/type_as/register.py | 11 +++++++---- 11 files changed, 94 insertions(+), 41 deletions(-) diff --git a/ops/_local_scalar_dense/register.py b/ops/_local_scalar_dense/register.py index 6cbc620..e034749 100644 --- a/ops/_local_scalar_dense/register.py +++ b/ops/_local_scalar_dense/register.py @@ -29,7 +29,11 @@ def _bump(): def lsd_aten(self): _bump() - from kernel.torch_level import _local_scalar_dense_torch + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.torch_level import _local_scalar_dense_torch + except ImportError: + from kernel.torch_level import _local_scalar_dense_torch return _local_scalar_dense_torch(self) @@ -50,9 +54,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/clone/register.py b/ops/clone/register.py index 80bfb5b..7891925 100644 --- a/ops/clone/register.py +++ b/ops/clone/register.py @@ -31,7 +31,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.triton_level import clone_triton + except ImportError: + from kernel.triton_level import clone_triton return clone_triton(self, memory_format) @@ -52,9 +56,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/clone/test/op_level.py b/ops/clone/test/op_level.py index 7bfff5b..4a6ad80 100644 --- a/ops/clone/test/op_level.py +++ b/ops/clone/test/op_level.py @@ -51,9 +51,32 @@ def run(profile): # 4) 回退: 不支持 dtype（int32）仍正确 assert torch.equal(xi.clone(), xi_ref), "int32 回退不一致" + # 5) 包态执行（多算子同进程）: 从 ops.clone 包入口进来须能找到并执行 + # ops.clone.kernel，而非顶层 kernel（#25 review 指出的包态运行时缺口） + import os + import subprocess + + root = str(OP_DIR.parent.parent) + child = ( + "import torch, torch_mlu\n" + "from ops.clone.register import clone_aten\n" + f"x = torch.zeros(3, device='{dev}')\n" + "y = clone_aten(x)\n" + "assert tuple(y.shape) == (3,), y.shape\n" + "assert y.data_ptr() != x.data_ptr()\n" + "print('PKG_MODE_OK')\n" + ) + env = os.environ.copy() + env["PYTHONPATH"] = root + os.pathsep + env.get("PYTHONPATH", "") + proc = subprocess.run([sys.executable, "-c", child], capture_output=True, + text=True, cwd=root, env=env) + assert proc.returncode == 0 and "PKG_MODE_OK" in proc.stdout, \ + f"包态执行失败:\n{proc.stdout}\n{proc.stderr}" + return {"ok": True, "registered": f"aten::clone@{profile.dispatch_key}", - "intercepted_calls": reg.CALL_COUNT["clone"]} + "intercepted_calls": reg.CALL_COUNT["clone"], + "package_mode": True} if __name__ == "__main__": diff --git a/ops/contiguous/register.py b/ops/contiguous/register.py index 299fea8..207e89c 100644 --- a/ops/contiguous/register.py +++ b/ops/contiguous/register.py @@ -30,7 +30,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.triton_level import contiguous_triton + except ImportError: + from kernel.triton_level import contiguous_triton return contiguous_triton(self, memory_format) @@ -51,9 +55,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/copy_/register.py b/ops/copy_/register.py index 7289913..0a58b6c 100644 --- a/ops/copy_/register.py +++ b/ops/copy_/register.py @@ -33,7 +33,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.triton_level import copy_triton + except ImportError: + from kernel.triton_level import copy_triton return copy_triton(self, src, non_blocking) @@ -54,9 +58,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/detach/register.py b/ops/detach/register.py index 517700b..9d9edbf 100644 --- a/ops/detach/register.py +++ b/ops/detach/register.py @@ -28,7 +28,11 @@ def _bump(): def detach_aten(self): _bump() - from kernel.torch_level import detach_torch + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.torch_level import detach_torch + except ImportError: + from kernel.torch_level import detach_torch return detach_torch(self) @@ -49,9 +53,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/dropout/register.py b/ops/dropout/register.py index bca2a85..b83b82a 100644 --- a/ops/dropout/register.py +++ b/ops/dropout/register.py @@ -30,7 +30,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.triton_level import dropout_triton + except ImportError: + from kernel.triton_level import dropout_triton return dropout_triton(input, p, train) @@ -51,9 +55,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/empty_like/register.py b/ops/empty_like/register.py index aa98230..d94da13 100644 --- a/ops/empty_like/register.py +++ b/ops/empty_like/register.py @@ -33,7 +33,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.torch_level import empty_like_torch + except ImportError: + 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) @@ -56,9 +60,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/empty_strided/register.py b/ops/empty_strided/register.py index 4dfa8ab..55ed2d8 100644 --- a/ops/empty_strided/register.py +++ b/ops/empty_strided/register.py @@ -32,7 +32,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.torch_level import empty_strided_torch + except ImportError: + from kernel.torch_level import empty_strided_torch return empty_strided_torch(size, stride, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory) @@ -54,9 +58,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/item/register.py b/ops/item/register.py index 800ad2d..1d5a7a3 100644 --- a/ops/item/register.py +++ b/ops/item/register.py @@ -28,7 +28,11 @@ def _bump(): def item_aten(self): _bump() - from kernel.torch_level import item_torch + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.torch_level import item_torch + except ImportError: + from kernel.torch_level import item_torch return item_torch(self) @@ -49,9 +53,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError( diff --git a/ops/type_as/register.py b/ops/type_as/register.py index 6a9a76d..a463526 100644 --- a/ops/type_as/register.py +++ b/ops/type_as/register.py @@ -36,7 +36,11 @@ 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 + # 包态优先、脚本态 fallback（多算子同进程执行，#25 review） + try: + from .kernel.triton_level import type_as_triton + except ImportError: + from kernel.triton_level import type_as_triton return type_as_triton(self, other) @@ -57,9 +61,8 @@ def _backend_for(dispatch_key: str, platform=None): "（第二平台接入见 ops/embedding/PLATFORM.md §5）") return _BACKENDS[platform] by_key = { - "Autograd": "cambricon", - "PrivateUse1": "cambricon", - "AutogradPrivateUse1": "cambricon", + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key } if dispatch_key not in by_key: raise RuntimeError(