fix: 13 算子 register_a1 签名升级（counter + platform，#19)

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 <dashen272@outlook.com>
+_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
 
 

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,

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 <dashen272@outlook.com>
+_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
 
 

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__":

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 <dashen272@outlook.com>
+_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
 
 

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) 别名快路径

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 <dashen272@outlook.com>
+_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