diff --git a/ops/_local_scalar_dense/register.py b/ops/_local_scalar_dense/register.py index 703b66d..e034749 100644 --- a/ops/_local_scalar_dense/register.py +++ b/ops/_local_scalar_dense/register.py @@ -29,14 +29,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..7891925 100644 --- a/ops/clone/register.py +++ b/ops/clone/register.py @@ -31,14 +31,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..4a6ad80 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(), "拦截后仍共享存储" @@ -49,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 3e95d8e..207e89c 100644 --- a/ops/contiguous/register.py +++ b/ops/contiguous/register.py @@ -30,14 +30,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..0a58b6c 100644 --- a/ops/copy_/register.py +++ b/ops/copy_/register.py @@ -33,14 +33,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..9d9edbf 100644 --- a/ops/detach/register.py +++ b/ops/detach/register.py @@ -28,14 +28,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..b83b82a 100644 --- a/ops/dropout/register.py +++ b/ops/dropout/register.py @@ -30,14 +30,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..d94da13 100644 --- a/ops/empty_like/register.py +++ b/ops/empty_like/register.py @@ -33,16 +33,75 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..55ed2d8 100644 --- a/ops/empty_strided/register.py +++ b/ops/empty_strided/register.py @@ -32,15 +32,74 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..1d5a7a3 100644 --- a/ops/item/register.py +++ b/ops/item/register.py @@ -28,14 +28,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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..a463526 100644 --- a/ops/type_as/register.py +++ b/ops/type_as/register.py @@ -36,15 +36,73 @@ 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) -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 = { + "AutogradPrivateUse1": "cambricon", # canonical A1 key + "PrivateUse1": "cambricon", # 兼容 key + } + 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) 方法式调用同样命中