Skip to content

AGENTS.md

Latest commit

 

History

History
109 lines (85 loc) · 6.93 KB

AGENTS.md

File metadata and controls

109 lines (85 loc) · 6.93 KB
Raw

Workflow for AI agents

Setup and Installation

  • Run uv sync --all-groups
  • That's it

Before writing code

  • Read existing code in the relevant file(s) first to understand conventions, typing style, and patterns.
  • Check neighboring files for imports, library choices, and naming conventions.
  • Never add comments to code unless the user explicitly asks.
  • This is an educational project — no runtime dependencies are allowed. Only testing/dev tools (pytest, hypothesis, pyrefly, ruff, numpy) may be used.
  • This library should be very easily understood and as unbloated as possible. The exception would be the transpiler/runtime, since they make the kernels accessible, and at the python level.

Code quality checklist

Run all four before submitting changes:

uv run ruff format && uv run ruff check --fix && uv run pyrefly check && uv run pytest --cov --cov-report=term-missing -n auto
  • uv run ruff format — must produce no changes (i.e., already formatted)
  • uv run ruff check --fix — must pass with zero errors
  • uv run pyrefly check — must pass with zero errors (suppressed warnings are OK)
  • uv run pytest --cov — must pass with 100% coverage on all source files

Notes

  • Functions decorated with @compile_fn or passed to register_inline are transpiled to C — their Python body is never executed. They are automatically excluded from code coverage, as the functions will not be directly used (only module.compiled_function will be called)
  • Cache artifacts live in ~/.cache/simplendarray/ — each hash is isolated in its own subdirectory with a per-module .lock.

PythonModule stub system

Compiled C functions are accessible through PythonModule instances via __getattr__ → .so. To make pyrefly see exact signatures, an importable .py stub file is auto-generated:

  • The stub file is named _{stem}_stubs.py next to the source module (e.g. _element_wise_stubs.py alongside element_wise.py).
  • At runtime, methods inside the if TYPE_CHECKING: block are invisible — calls fall through to __getattr__ → C extension. pyrefly reads them statically.
  • PythonModules will set stub_path=__file__ and stub_var={var_name} to get this desired effect. After the first run, the type stub will be autogenerated.

Repository Map

src/simplendarray/
├── __init__.py          # Public API exports
├── array.py             # Array class (from_iterable, arange, reshape, relu, __add__, slicing, transpose)
├── buffer.py            # Buffer (CPU) — wraps array.array
├── buffer_cuda.py       # BufferCuda (GPU) — wraps CUDA device memory
├── dtypes.py            # DType classes, get_dtype(), typecode/ctype/cname maps
├── utils.py             # contiguous_strides, product, ceildiv
├── kernels/
│   ├── __init__.py      # Dispatchers: elem_wise_modules = {"cpu": ..., "gpu": ...}
│   │                    # dispatch_arange, dispatch_element_wise_binary/unary, dispatch_reshape_copy
│   ├── cpu/
│   │   └── element_wise.py        # CPU compiled kernels (all dtypes)
│   └── cuda/
│       ├── __init__.py            # Re-exports buffer_cuda_module, element_wise_module_cuda
│       ├── element_wise_cuda.py   # CUDA compiled kernels (__device__, __global__, pybind)
│       ├── buffer_cuda.py         # CUDA malloc/free/memcpy wrappers
│       └── helpers.py             # dim3, threadIdx, blockIdx, cudaError_t type stubs
└── transpiler/
    ├── __init__.py
    ├── transpiler.py    # Python → C transpiler
    └── runtime.py       # PythonModule class, SpecItem, compile_fn, caching, dispatch
tests/
├── test_array.py        # Array tests, including dispatch error paths with MagicMock
├── test_buffer.py       # Buffer & BufferCuda tests (GPU tests skipif no nvcc)
├── test_transpiler.py
└── test_runtime.py

Architecture

  • Array stores data via self.data: Buffer | BufferCuda, plus shape/strides/offset (typical strided view).
  • device is "cpu" or "gpu", stored on the buffer. Access via Array.device property → self.data.device.
  • Buffer (CPU) wraps array.array. BufferCuda (GPU) allocates CUDA device memory. Both expose .address (int pointer), .typecode (str like "i"/"f"/"d"/"b"), .device, .num_bytes.
  • Dispatch via elem_wise_modules[device].DISPATCH_DICT_{group}[dispatch_key](...). Each compiled module (CPU/GPU) has its own dispatch dicts. Keys are tuple(spec.mapping.items()) — e.g. (("T", "int"), ("Op", "_add_int")).
  • CPU and GPU can have different dispatch keys for the same operation — dispatch_arange uses (("T", dt.ctype),) for CPU and (("T", dt.ctype), ("Kernel", "arange_kernel_{cname}")) for GPU. Check dispatch_reshape_copy for the pattern.
  • spec.mapping maps template param names to C type strings. Keys must match exactly between dispatch function and kernel definition.

DType system

  • get_dtype(dtype) where dtype is a typecode ("i", "f", "d", "b") or a DType class.
  • Each DType has: .typecode (array.array code), .ctype (C type like "int", "float"), .cname (short name like "int", "float", "double", "bool").
  • Kernel specs iterate over all_dtypes (all) or all_float_dtypes (float/double only for math ops).

CUDA kernel pattern (three tiers)

  1. __device__ function (e.g. _add[T: DType](x: T, y: T) -> T) — single operation, single dtype.
  2. __global__ kernel (e.g. element_wise_binary_kernel) — reads threadIdx.x + blockIdx.x * blockDim.x, dispatches to the Op.
  3. pybind=True wrapper (e.g. element_wise_binary) — computes grid dimensions, launches kernel via Kernel[[[blocks, threads]]](args).

Known issues / incomplete integrations

  • dispatch_element_wise_unary raises NotImplementedError for GPU (line 25-26 of kernels/__init__.py). CUDA unary kernels exist but are never dispatched.
  • Array.relu() and Array.__add__() always create CPU Buffer for output — they don't produce GPU results even when inputs are on GPU.
  • GPU dispatch keys differ from CPU for arange and reshape_copy (GPU adds ("Kernel", ...)). dispatch_arange and dispatch_reshape_copy have been fixed; dispatch_element_wise_binary uses matching keys for both.
  • Coverage skips CUDA kernel lines (element_wise_cuda.py at ~87% when nvcc unavailable). Only source files under src/simplendarray/ are measured, excluding *_stubs.py.

Testing patterns

  • GPU dispatch paths are tested with MagicMock + monkeypatch.setitem(_kernels.elem_wise_modules, "gpu", mock_gpu_module) to avoid needing nvcc.
  • Real GPU integration tests use @pytest.mark.skipif(not cuda_available, reason="CUDA not available") where cuda_available = shutil.which("nvcc") is not None.
  • BufferCuda methods tested in test_buffer.py with skipif pattern.
  • Hypothesis property-based tests compare against NumPy for slicing, transpose, reshape.