Update Wall-X to 1.1.0 (#104)

wall_x/model/core/ops/base.py

@@ -0,0 +1,143 @@
+"""Base class for operator proxies with lazy backend resolution."""
+
+import importlib
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+
+class OpsProxy:
+    """Callable proxy that lazily resolves to the best available backend.
+
+    Three-level fallback:
+      Level 1: external_accel (fastest, requires external_accel package)
+      Level 2: CUDA inline kernel (JIT compiled, requires CUDA)
+      Level 3: Pure PyTorch (always available)
+
+    Subclasses define _external_accel_name, _pytorch_fallback, and optionally _cuda_kernel.
+
+    Set _external_accel_name to None to skip external_accel resolution (e.g., when the
+    external_accel API is incompatible).
+    """
+
+    def __init__(self):
+        self._resolved_fn = None
+        self._backend = None
+        self._resolve_lock = threading.Lock()
+
+    def _resolve(self):
+        """Resolve the best available backend. Called once on first use (thread-safe)."""
+        if self._resolved_fn is not None:
+            return
+        with self._resolve_lock:
+            # Double-check after acquiring lock
+            if self._resolved_fn is not None:
+                return
+
+            # Level 1: try external_accel (skip if _external_accel_name is None)
+            fn = self._probe_external_accel()
+            if fn is not None:
+                self._backend = "external_accel"
+                self._resolved_fn = fn
+                return
+
+            # Level 2: try CUDA inline kernel (subclass override)
+            cuda_fn = self._get_cuda_kernel()
+            if cuda_fn is not None:
+                self._backend = "cuda_inline"
+                self._resolved_fn = cuda_fn
+                return
+
+            # Level 3: PyTorch fallback
+            self._backend = "pytorch"
+            self._resolved_fn = self._pytorch_fallback
+            logger.info(f"{self.__class__.__name__}: using PyTorch fallback")
+
+    def _get_cuda_kernel(self):
+        """Override in subclass to provide Level 2 CUDA kernel. Returns None if unavailable."""
+        return None
+
+    @property
+    def _external_accel_name(self):
+        """Return external_accel attribute name, or None to skip external_accel resolution.
+
+        Base returns None (skip external_accel). Subclasses override to return the
+        attribute name when external_accel integration is desired.
+        """
+        return None
+
+    def _pytorch_fallback(self, *args, **kwargs):
+        raise NotImplementedError(f"{self.__class__.__name__} has no PyTorch fallback")
+
+    def __call__(self, *args, **kwargs):
+        if self._resolved_fn is None:
+            self._resolve()
+        return self._resolved_fn(*args, **kwargs)
+
+    @property
+    def backend(self) -> str:
+        if self._resolved_fn is None:
+            self._resolve()
+        return self._backend
+
+    # ------------------------------------------------------------------
+    # Multi-backend API (for testing / benchmarking)
+    # ------------------------------------------------------------------
+
+    def _probe_external_accel(self):
+        """Return the external_accel function for this op, or None if unavailable."""
+        if self._external_accel_name is None:
+            return None
+        try:
+            mod = importlib.import_module("external_accel")
+            return getattr(mod, self._external_accel_name)
+        except (ImportError, AttributeError):
+            return None
+
+    def available_backends(self):
+        """Return list of available backend names for this operator."""
+        backends = ["pytorch"]
+        if self._probe_external_accel() is not None:
+            backends.append("external_accel")
+        if self._get_cuda_kernel() is not None:
+            backends.append("cuda_inline")
+        return backends
+
+    def _get_backend_fn(self, backend):
+        """Return the callable for a specific backend, or None if unavailable."""
+        if backend == "pytorch":
+            return self._pytorch_fallback
+        elif backend == "external_accel":
+            return self._probe_external_accel()
+        elif backend == "cuda_inline":
+            return self._get_cuda_kernel()
+        else:
+            raise ValueError(f"Unknown backend: {backend}")
+
+    def call_with_backend(self, backend, *args, **kwargs):
+        """Call this operator using a specific backend.
+
+        This does NOT change the default backend used by ``__call__``.
+        The ``backend`` property still reflects the auto-resolved default.
+
+        Args:
+            backend: One of "external_accel", "cuda_inline", "pytorch".
+
+        Returns:
+            Operator output from the specified backend.
+
+        Raises:
+            RuntimeError: If the requested backend is not available.
+        """
+        fn = self._get_backend_fn(backend)
+        if fn is None:
+            raise RuntimeError(
+                f"{self.__class__.__name__}: backend '{backend}' not available. "
+                f"Available: {self.available_backends()}"
+            )
+        return fn(*args, **kwargs)
+
+    def __repr__(self):
+        backend = self._backend or "unresolved"
+        return f"<{self.__class__.__name__} backend={backend}>"