netcl 0.1.0__py3-none-any.whl

netcl/__init__.py

@@ -0,0 +1,37 @@
+"""
+netcl: PyOpenCL-based experimentation framework.
+This package currently focuses on low-level kernel primitives and helpers.
+"""
+
+from . import core, ops, autograd, distributed, runtime, profiling
+from .ops import (
+    matmul,
+    build_matmul_kernel,
+    elementwise_binary,
+    relu,
+    bias_add,
+    reduce_sum,
+    softmax,
+    conv2d,
+)
+from . import nn, optim, io
+
+__all__ = [
+    "core",
+    "ops",
+    "autograd",
+    "distributed",
+    "runtime",
+    "profiling",
+    "nn",
+    "optim",
+    "io",
+    "matmul",
+    "build_matmul_kernel",
+    "elementwise_binary",
+    "relu",
+    "bias_add",
+    "reduce_sum",
+    "softmax",
+    "conv2d",
+]

netcl/amp.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Sequence
+
+try:
+    import numpy as np  # type: ignore
+except ImportError:  # pragma: no cover
+    np = None
+
+from netcl.core.tensor import Tensor
+from netcl.core.tensor import _np_dtype  # type: ignore
+
+_AUTOCAST_ENABLED = False
+
+
+@dataclass
+class GradScaler:
+    init_scale: float = 2.0**16
+    growth_factor: float = 2.0
+    backoff_factor: float = 0.5
+    growth_interval: int = 2000
+    enabled: bool = True
+
+    def __post_init__(self):
+        self._scale = self.init_scale
+        self._growth_tracker = 0
+
+    @property
+    def scale(self) -> float:
+        return self._scale
+
+    def scale_loss(self, loss: Tensor) -> Tensor:
+        if not self.enabled:
+            return loss
+        from netcl.ops.elementwise import elementwise_binary
+
+        return elementwise_binary(loss, loss, expression=f"MUL(v0, {float(self._scale)})")
+
+    def unscale_grads(self, params: Sequence[Tensor]):
+        if not self.enabled:
+            return False
+        found_inf = False
+        for p in params:
+            if p.grad is None:
+                continue
+            g = p.grad.to_host()
+            if np.any(~np.isfinite(g)):
+                found_inf = True
+                break
+        if not found_inf:
+            inv_scale = 1.0 / self._scale
+            from netcl.ops.elementwise import elementwise_binary
+
+            for p in params:
+                if p.grad is None:
+                    continue
+                p.grad = elementwise_binary(p.grad, p.grad, expression=f"MUL(v0, {inv_scale})")
+        return found_inf
+
+    def step(self, optimizer, params: Sequence[Tensor]):
+        if not self.enabled:
+            optimizer.step()
+            return
+        found_inf = self.unscale_grads(params)
+        if not found_inf:
+            optimizer.step()
+            self._growth_tracker += 1
+            if self._growth_tracker % self.growth_interval == 0:
+                self._scale *= self.growth_factor
+        else:
+            self._scale *= self.backoff_factor
+            self._growth_tracker = 0
+
+    def update(self):
+        # no-op kept for API compatibility
+        pass
+
+
+def supports_fp16(queue) -> bool:
+    """
+    Check device extensions for cl_khr_fp16 support.
+    """
+    try:
+        return "cl_khr_fp16" in queue.device.extensions
+    except Exception:
+        return False
+
+
+def autocast_enabled(profile_supports_fp16: bool) -> bool:
+    return profile_supports_fp16
+
+
+class autocast:
+    """
+    Context manager for autocast. Enables casting only if underlying device supports fp16.
+    """
+
+    def __init__(self, enabled: bool = True, device_queue=None):
+        self.enabled = enabled
+        self.device_queue = device_queue
+        self.prev = False
+        self._capable = True
+
+    def __enter__(self):
+        global _AUTOCAST_ENABLED
+        self.prev = _AUTOCAST_ENABLED
+        if not self.enabled:
+            _AUTOCAST_ENABLED = False
+            return self
+        if self.device_queue is not None:
+            self._capable = supports_fp16(self.device_queue)
+        _AUTOCAST_ENABLED = self.enabled and self._capable
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        global _AUTOCAST_ENABLED
+        _AUTOCAST_ENABLED = self.prev
+        return False
+
+
+def is_autocast_enabled() -> bool:
+    return _AUTOCAST_ENABLED
+
+
+def maybe_cast_tensor(t: Tensor) -> Tensor:
+    if not _AUTOCAST_ENABLED:
+        return t
+    if t.dtype in ("float", "float32"):
+        # only cast if device can handle fp16
+        if supports_fp16(t.queue):
+            arr = t.to_host().astype(np.float16)
+            return Tensor.from_host(t.queue, arr, dtype="float16")
+        return t
+    return t
+
+
+def master_param(param: Tensor) -> Tensor:
+    """
+    Keep master weights in FP32 for optimizers.
+    """
+    if param.dtype in ("float16", "half"):
+        master = Tensor.from_host(param.queue, param.to_host().astype(np.float32), dtype="float32")
+        setattr(master, "_model_param", param)
+        return master
+    setattr(param, "_model_param", param)
+    return param

netcl/autograd/__init__.py

@@ -0,0 +1,86 @@
+"""
+Autograd: Node/Tape plus op wrappers.
+"""
+
+from .engine import Node, Tape, apply_op, no_grad, set_current_tape, get_current_tape
+from .debug import debug_tape
+from .ops import (
+    tensor,
+    add,
+    relu,
+    bias_add,
+    matmul_op,
+    sub,
+    mse_loss,
+    sigmoid,
+    tanh,
+    leaky_relu,
+    gelu,
+    swish,
+    elu,
+    softplus,
+    hard_sigmoid,
+    hard_swish,
+    clamp,
+    hard_tanh,
+    prelu,
+    hinge_loss,
+    l1_loss,
+    l2_loss,
+    depthwise_conv2d,
+    batch_norm2d,
+    layer_norm,
+    pad2d,
+    group_norm,
+    global_avg_pool2d,
+    cross_entropy,
+    conv2d,
+    flatten,
+    max_pool2d,
+    dropout,
+    avg_pool2d,
+)
+
+__all__ = [
+    "Node",
+    "Tape",
+    "apply_op",
+    "no_grad",
+    "tensor",
+    "add",
+    "relu",
+    "bias_add",
+    "matmul_op",
+    "sub",
+    "mse_loss",
+    "sigmoid",
+    "tanh",
+    "leaky_relu",
+    "gelu",
+    "swish",
+    "elu",
+    "softplus",
+    "hard_sigmoid",
+    "hard_swish",
+    "clamp",
+    "hard_tanh",
+    "prelu",
+    "hinge_loss",
+    "l1_loss",
+    "l2_loss",
+    "depthwise_conv2d",
+    "batch_norm2d",
+    "layer_norm",
+    "pad2d",
+    "group_norm",
+    "global_avg_pool2d",
+    "cross_entropy",
+    "conv2d",
+    "flatten",
+    "max_pool2d",
+    "dropout",
+    "avg_pool2d",
+    "debug_tape",
+    "set_current_tape",
+    "get_current_tape",
+]

netcl/autograd/debug.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+
+@contextmanager
+def debug_tape(tape):
+    """
+    Context manager to expose a tape for debugging/inspection.
+    """
+    yield tape

netcl/autograd/engine.py

@@ -0,0 +1,121 @@
+"""
+Minimal autograd backbone (placeholders).
+
+To be extended with full gradient tracking and backward kernels per op.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Callable, List, Optional
+import threading
+
+from netcl.core.tensor import Tensor
+
+
+GradFn = Callable[[Any], List[Optional[Tensor]]]
+
+
+@dataclass
+class Node:
+    value: Tensor
+    grad_fn: Optional[GradFn] = None
+    parents: List["Node"] = field(default_factory=list)
+    grad: Optional[Tensor] = None
+    requires_grad: bool = False
+
+
+class Tape:
+    """
+    Records operations for backward.
+    """
+
+    def __init__(self) -> None:
+        self.nodes: List[Node] = []
+        self.enabled: bool = True
+
+    def record(self, node: Node) -> Node:
+        if self.enabled:
+            self.nodes.append(node)
+        return node
+
+    def backward(self, loss: Node, grad: Optional[Tensor] = None) -> None:
+        if grad is None:
+            # seed with ones
+            loss.grad = ones_like(loss.value)
+        else:
+            loss.grad = grad
+        # Reverse topological order (here: recorded order)
+        for node in reversed(self.nodes):
+            if node.grad is None or node.grad_fn is None:
+                continue
+            grads = node.grad_fn(node.grad)
+            for parent, g in zip(node.parents, grads):
+                if g is None:
+                    continue
+                if parent.grad is None:
+                    parent.grad = g
+                else:
+                    parent.grad = add_inplace(parent.grad, g)
+                # propagate to underlying Tensor for optimizers
+                if parent.value.grad is None:
+                    parent.value.grad = parent.grad
+                else:
+                    parent.value.grad = add_inplace(parent.value.grad, g)
+
+
+# thread-local current tape for tape-free APIs
+_tls = threading.local()
+
+
+def set_current_tape(tape: Optional[Tape]):
+    _tls.current_tape = tape
+
+
+def get_current_tape() -> Optional[Tape]:
+    return getattr(_tls, "current_tape", None)
+
+
+def ones_like(t: Tensor) -> Tensor:
+    from netcl.core.tensor import Tensor as T
+    if t.dtype not in ("float", "float32", "double", "float64"):
+        raise ValueError("ones_like supports float tensors")
+    import numpy as np
+
+    data = np.ones(t.shape, dtype=np.float32 if "32" in t.dtype or t.dtype == "float" else np.float64)
+    return T.from_host(t.queue, data, dtype=t.dtype)
+
+
+def add_inplace(dst: Tensor, src: Tensor) -> Tensor:
+    from netcl.ops.elementwise import elementwise_binary
+
+    return elementwise_binary(dst, src, expression="ADD(v0, v1)", out=dst)
+
+
+def no_grad():
+    """
+    Context manager placeholder to disable gradient tracking.
+    """
+
+    class _NoGrad:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc, tb):
+            return False
+
+    return _NoGrad()
+
+
+def apply_op(fn: Callable[..., Tensor], grad_fn: Optional[GradFn], *args: Node, tape: Optional[Tape] = None) -> Node:
+    tape = tape or get_current_tape()
+    out_value = fn(*[a.value if isinstance(a, Node) else a for a in args])
+    node = Node(
+        value=out_value,
+        grad_fn=grad_fn,
+        parents=[a for a in args if isinstance(a, Node)],
+        requires_grad=any(getattr(a, "requires_grad", False) for a in args),
+    )
+    if tape is not None:
+        tape.record(node)
+    return node