shem-neural 0.1.0__tar.gz

shem_neural-0.1.0/PKG-INFO

@@ -0,0 +1,37 @@
+Metadata-Version: 2.4
+Name: shem-neural
+Version: 0.1.0
+Summary: shem's neural side: an RV32IM CPU compiled into network weights, and model images
+Keywords: shem,containers,riscv,neural-network,language-models
+Author: João Galego
+Author-email: João Galego <jgalego1990@gmail.com>
+License-Expression: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: System :: Emulators
+Requires-Dist: huggingface-hub>=2.4.0
+Requires-Dist: jinja2>=3.1.6
+Requires-Dist: numpy>=2
+Requires-Dist: safetensors>=0.4
+Requires-Dist: tokenizers>=0.23.3
+Requires-Dist: torch>=2.5
+Requires-Python: >=3.12
+Project-URL: Homepage, https://github.com/JGalego/shem
+Project-URL: Documentation, https://github.com/JGalego/shem/tree/main/docs
+Project-URL: Repository, https://github.com/JGalego/shem
+Description-Content-Type: text/markdown
+
+# shem-neural
+
+The neural vCPU for [shem](https://github.com/JGalego/shem): an RV32IM CPU compiled into the weights of a network that runs one instruction per forward pass.
+
+![An RV32 program runs one instruction per forward pass. On ecall, the host handles the write system call, updates the return value and program counter, then resumes the guest.](https://raw.githubusercontent.com/JGalego/shem/main/docs/assets/execution.svg)
+
+```sh
+uv run pytest                                            # unit tests
+uv run python -m shem_neural.server --socket /tmp/shem.sock   # serve vCPUs to the host
+```
+
+See the [neural vCPU documentation](https://github.com/JGalego/shem/blob/main/docs/neural.md) for how the network works and how the host talks to it.

shem_neural-0.1.0/README.md

@@ -0,0 +1,12 @@
+# shem-neural
+
+The neural vCPU for [shem](https://github.com/JGalego/shem): an RV32IM CPU compiled into the weights of a network that runs one instruction per forward pass.
+
+![An RV32 program runs one instruction per forward pass. On ecall, the host handles the write system call, updates the return value and program counter, then resumes the guest.](https://raw.githubusercontent.com/JGalego/shem/main/docs/assets/execution.svg)
+
+```sh
+uv run pytest                                            # unit tests
+uv run python -m shem_neural.server --socket /tmp/shem.sock   # serve vCPUs to the host
+```
+
+See the [neural vCPU documentation](https://github.com/JGalego/shem/blob/main/docs/neural.md) for how the network works and how the host talks to it.

shem_neural-0.1.0/pyproject.toml

@@ -0,0 +1,81 @@
+[project]
+name = "shem-neural"
+version = "0.1.0"
+description = "shem's neural side: an RV32IM CPU compiled into network weights, and model images"
+readme = "README.md"
+license = "MIT"
+keywords = [
+    "shem",
+    "containers",
+    "riscv",
+    "neural-network",
+    "language-models",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: System :: Emulators",
+]
+requires-python = ">=3.12"
+dependencies = [
+    "huggingface-hub>=2.4.0",
+    "jinja2>=3.1.6",
+    "numpy>=2",
+    "safetensors>=0.4",
+    "tokenizers>=0.23.3",
+    "torch>=2.5",
+]
+
+[[project.authors]]
+name = "João Galego"
+email = "jgalego1990@gmail.com"
+
+[project.urls]
+Homepage = "https://github.com/JGalego/shem"
+Documentation = "https://github.com/JGalego/shem/tree/main/docs"
+Repository = "https://github.com/JGalego/shem"
+
+[project.scripts]
+shem-neural-server = "shem_neural.server:main"
+shem-neural-generate = "shem_neural.model.generate:main"
+shem-neural-import = "shem_neural.model.importers:main"
+shem-neural-export = "shem_neural.export:main"
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+    "transformers>=5.22.2",
+]
+
+[tool.uv.sources.torch]
+index = "pytorch-cpu"
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "F",
+    "I",
+    "UP",
+    "B",
+]
+
+[tool.ruff.lint.per-file-ignores]
+"src/shem_neural/export.py" = ["E501"]
+
+[build-system]
+requires = ["uv_build>=0.10.9,<0.11.0"]
+build-backend = "uv_build"

shem_neural-0.1.0/pyproject.toml.orig

@@ -0,0 +1,67 @@
+[project]
+name = "shem-neural"
+version = "0.1.0"
+description = "shem's neural side: an RV32IM CPU compiled into network weights, and model images"
+readme = "README.md"
+license = "MIT"
+authors = [{ name = "João Galego", email = "jgalego1990@gmail.com" }]
+keywords = ["shem", "containers", "riscv", "neural-network", "language-models"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: System :: Emulators",
+]
+requires-python = ">=3.12"
+dependencies = [
+    "huggingface-hub>=2.4.0",
+    "jinja2>=3.1.6",
+    "numpy>=2",
+    "safetensors>=0.4",
+    "tokenizers>=0.23.3",
+    "torch>=2.5",
+]
+
+[project.urls]
+Homepage = "https://github.com/JGalego/shem"
+Documentation = "https://github.com/JGalego/shem/tree/main/docs"
+Repository = "https://github.com/JGalego/shem"
+
+[project.scripts]
+shem-neural-server = "shem_neural.server:main"
+shem-neural-generate = "shem_neural.model.generate:main"
+shem-neural-import = "shem_neural.model.importers:main"
+shem-neural-export = "shem_neural.export:main"
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+    "transformers>=5.22.2",
+]
+
+[tool.uv.sources]
+torch = { index = "pytorch-cpu" }
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[build-system]
+requires = ["uv_build>=0.10.9,<0.11.0"]
+build-backend = "uv_build"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[tool.ruff]
+line-length = 120
+target-version = "py312"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]
+
+[tool.ruff.lint.per-file-ignores]
+# The model card is Markdown, one line per paragraph
+"src/shem_neural/export.py" = ["E501"]

shem_neural-0.1.0/src/shem_neural/__init__.py

@@ -0,0 +1,7 @@
+"""A neural vCPU for shem: an RV32IM CPU compiled into network weights.
+
+`isa` describes the instruction set as data, `compile` turns that description
+into the network's weight tensors, `cpu` is the network, one instruction per
+forward pass, `vcpu` wraps it in the interface the shem host drives, and
+`server` exposes vCPUs to the host over a Unix socket.
+"""

shem_neural-0.1.0/src/shem_neural/compile.py

@@ -0,0 +1,117 @@
+"""Compile the RV32IM ISA table into lookup and routing weights.
+
+Each tensor is derived from `isa.INSTRUCTIONS`; nothing is trained.
+"""
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file, save_file
+
+from . import isa
+
+DTYPE = torch.float64
+
+
+def _matcher(patterns: list[tuple[int, int]], width: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Weights and bias of a layer that counts mismatches: for bits `x`,
+    `x @ w.T + c` is how many masked bits differ from each pattern's match.
+    `relu(1 - that)` is then a one-hot over the patterns."""
+    w = torch.zeros(len(patterns), width, dtype=DTYPE)
+    c = torch.zeros(len(patterns), dtype=DTYPE)
+    for k, (mask, match) in enumerate(patterns):
+        for i in range(width):
+            if mask >> i & 1:
+                bit = match >> i & 1
+                w[k, i] = 1 - 2 * bit
+                c[k] += bit
+    return w, c
+
+
+def _immediate(fmt: str) -> torch.Tensor:
+    """A (32, 32) selection matrix from instruction bits to immediate bits,
+    with sign extension: row j says which instruction bit is immediate bit j."""
+    source: dict[int, int] = {}
+    if fmt == "I":
+        source = {j: 20 + j for j in range(11)}
+        sign = 11
+    elif fmt == "S":
+        source = {j: 7 + j for j in range(5)} | {j: 20 + j for j in range(5, 11)}
+        sign = 11
+    elif fmt == "B":
+        source = {j: 7 + j for j in range(1, 5)} | {j: 20 + j for j in range(5, 11)} | {11: 7}
+        sign = 12
+    elif fmt == "U":
+        source = {j: j for j in range(12, 32)}
+        sign = 32
+    elif fmt == "J":
+        source = {j: 20 + j for j in range(1, 11)} | {11: 20} | {j: j for j in range(12, 20)}
+        sign = 20
+    source |= {j: 31 for j in range(sign, 32)}
+    select = torch.zeros(32, 32, dtype=DTYPE)
+    for out, bit in source.items():
+        select[out, bit] = 1
+    return select
+
+
+def _route(values: list[str | None], names: tuple[str, ...]) -> torch.Tensor:
+    """One row per instruction, one-hot over `names` (all zeros for None)."""
+    route = torch.zeros(len(values), len(names), dtype=DTYPE)
+    for k, value in enumerate(values):
+        if value is not None:
+            route[k, names.index(value)] = 1
+    return route
+
+
+def compile_weights() -> dict[str, torch.Tensor]:
+    """Build every tensor the network uses."""
+    ins = isa.INSTRUCTIONS
+    decode_w, decode_c = _matcher([(i.mask, i.match) for i in ins], 32)
+    index_w, index_c = _matcher([(0x1F, j) for j in range(32)], 5)
+
+    sll = torch.zeros(32, 32, 32, dtype=DTYPE)  # [shift, out bit, in bit]
+    srl = torch.zeros(32, 32, 32, dtype=DTYPE)
+    sra_fill = torch.zeros(32, 32, dtype=DTYPE)  # [shift, out bit]: copies of the sign bit
+    for s in range(32):
+        for i in range(32):
+            if i - s >= 0:
+                sll[s, i, i - s] = 1
+            if i + s < 32:
+                srl[s, i, i + s] = 1
+            else:
+                sra_fill[s, i] = 1
+
+    byte_bits = torch.tensor([[v >> j & 1 for j in range(8)] for v in range(256)], dtype=DTYPE)
+    return {
+        "decode_w": decode_w,
+        "decode_c": decode_c,
+        "index_w": index_w,
+        "index_c": index_c,
+        "imm_bits": torch.stack([_immediate(f) for f in isa.FORMATS]),
+        "format_route": _route([i.format for i in ins], isa.FORMATS),
+        "uses_imm": torch.tensor([float(i.uses_imm) for i in ins], dtype=DTYPE),
+        "result_route": _route([i.result for i in ins], isa.RESULTS),
+        "next_pc_route": _route([i.next_pc for i in ins], isa.NEXT_PC),
+        "branch_route": _route([i.branch for i in ins], isa.BRANCHES),
+        "exit_route": _route([i.exit for i in ins], isa.EXITS),
+        "load_size": torch.tensor([float(i.load) for i in ins], dtype=DTYPE),
+        "store_size": torch.tensor([float(i.store) for i in ins], dtype=DTYPE),
+        "sll": sll,
+        "srl": srl,
+        "sra_fill": sra_fill,
+        "byte_bits": byte_bits,
+        "powers": torch.tensor([2.0**i for i in range(32)], dtype=DTYPE),
+    }
+
+
+def to(weights: dict[str, torch.Tensor], device: torch.device | str) -> dict[str, torch.Tensor]:
+    """The weights on `device`, where the network will run."""
+    return {name: tensor.to(device) for name, tensor in weights.items()}
+
+
+def save(weights: dict[str, torch.Tensor], path: Path) -> None:
+    save_file({name: t.contiguous() for name, t in weights.items()}, str(path))
+
+
+def load(path: Path) -> dict[str, torch.Tensor]:
+    return load_file(str(path))