From b20bf6a17bbe464a8e7ec7c31b82cb1035c7c2bc Mon Sep 17 00:00:00 2001 From: apalese Date: Fri, 4 Sep 2026 10:23:34 +0200 Subject: [PATCH] In-place ops and caller-provided destinations Every Tensor method was const and allocated its result: one allocation and one free per op per iteration in a training loop, and a peak footprint equal to the whole expression graph rather than the working set. Each op now has three forms, of which exactly one is a real implementation: auto c = a.add(b, s); // allocates the result, then calls add_out a.add_out(b, out, s); // <- the body a.add_(b, s); // == a.add_out(b, a, s) add_out returns Tensor& (the destination), add_ returns *this, so calls chain; the no-stream spellings are one-line delegates. operator+= -= *= /= delegate to the underscore family. Families: add/sub/mul/div (tensor and scalar), apply, apply_binary, relu, sigmoid and fill_ have all three forms. matmul, transpose and permute have _out but no in-place form -- their kernels read elements they do not write, so _check_alias_none rejects an aliased destination rather than returning a plausible wrong answer. The elementwise family may alias because the CPU loop, the contiguous GPU fast path and the rank-specialized kernels all read index i and write index i. That rests on each buffer being one flat run, so _check_alias_elementwise re-checks is_contiguous(): the day a strided view can alias a different region of the same allocation, it throws instead of computing the wrong thing. Consequences: - contiguous.cuh drops __restrict__ from its three kernels. In-place calls them with dst equal to a source, which is precisely the aliasing the qualifier promises away. Measured, it cost nothing: add at 16 M went 230.2-233.4 -> 228.6-230.7 GB/s and relu 232.9-234.6 -> 233.4-235.5, i.e. inside run-to-run spread. The read-only path comes from the explicit __ldg in device_load, not from restrict-driven inference. - _out is where operand validation now lives, so a GPU add between mismatched shapes throws instead of running off the end of a buffer. - fill_ takes a stream and no longer routes through the legacy _fill dispatch struct, which now has no callers at all. - apply_binary gained the stream overload roadmap 4.4 was asking for, since apply_binary_out takes one by construction. Numbers (benchmark_report.md section 9, harness scripts/bench_inplace.py): 1.2-1.6x below ~64 K on both backends where allocation is a large fraction of the op, shrinking to 1.04x at 16 M where the kernel dominates and HostPool/cudaMallocAsync already recycle the block. That is the honest ceiling for the time argument. The memory argument does not shrink, and the tests assert it: data_ptr unchanged across a 100-step loop, the property a correct-but-reallocating implementation would fail while passing every value check. Python gets the whole family through int-returning C entry points (nothing is allocated, no ownership changes hands), with __iadd__ and friends returning self so `a += b` does not silently rebind the name. tests/test_inplace.cpp (25 tests) and python/tests/test_inplace.py (35). All 13 correctness suites and 174 Python tests pass; compute-sanitizer memcheck clean over test_inplace, test_streams, test_allocator_stream and test_stress. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 54 ++- benchmark_report.md | 74 ++++ compile_commands.json | 6 + docs/roadmap.md | 39 +- headers/ops/kernels/contiguous.cuh | 18 +- headers/tensor.cuh | 148 +++++++- headers/tensor.inl | 576 ++++++++++++++++++++++++++--- python/README.md | 47 ++- python/openmat/_clib.py | 24 ++ python/openmat/tensor.py | 192 +++++++++- python/tests/test_inplace.py | 222 +++++++++++ scripts/bench_inplace.py | 140 +++++++ src/python/openmat_capi_impl.inc | 206 +++++++++++ tests/CMakeLists.txt | 1 + tests/test_inplace.cpp | 363 ++++++++++++++++++ 15 files changed, 2028 insertions(+), 82 deletions(-) create mode 100644 python/tests/test_inplace.py create mode 100644 scripts/bench_inplace.py create mode 100644 tests/test_inplace.cpp diff --git a/CLAUDE.md b/CLAUDE.md index 58d1556..8f7522c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -4,7 +4,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co ## Build -Requirements: NVIDIA GPU, CUDA Toolkit ≥ 11.2 (for `cudaMallocAsync`), CMake ≥ 3.24 (for `CMAKE_CUDA_ARCHITECTURES=native`), C++17/CUDA 17 compiler, OpenMP (`find_package(OpenMP REQUIRED)` in `CMakeLists.txt`; bundled as `libgomp` with a stock GCC install, nothing extra to install on most systems). Verified building and passing all 14 suites on CUDA 13.0 / GCC 13.3 / CMake 3.28 / GB10 (sm_121) — the 11 correctness suites run in 4.7 s; the three timing/soak suites are separate. +Requirements: NVIDIA GPU, CUDA Toolkit ≥ 11.2 (for `cudaMallocAsync`), CMake ≥ 3.24 (for `CMAKE_CUDA_ARCHITECTURES=native`), C++17/CUDA 17 compiler, OpenMP (`find_package(OpenMP REQUIRED)` in `CMakeLists.txt`; bundled as `libgomp` with a stock GCC install, nothing extra to install on most systems). Verified building and passing all 15 suites on CUDA 13.0 / GCC 13.3 / CMake 3.28 / GB10 (sm_121) — the 12 correctness suites run in 5.2 s; the three timing/soak suites are separate. ```bash # Full clean rebuild (also refreshes compile_commands.json in the repo root) @@ -38,11 +38,11 @@ cd build && ctest # all suites ./build/tests/test_arithmetic --gtest_filter="TensorArithmetic.CPUOperations" ``` -Suites: `test_arithmetic`, `test_fused_ops`, `test_device_transfer`, `test_factory`, `test_reductions`, `test_benchmarks`, `test_reshape`, `test_transpose`, `test_streams`, `test_allocator_stream`, `test_host_pool`, `test_contiguous`, `test_stress`, `test_stream_perf`. +Suites: `test_arithmetic`, `test_fused_ops`, `test_device_transfer`, `test_factory`, `test_reductions`, `test_benchmarks`, `test_reshape`, `test_transpose`, `test_streams`, `test_allocator_stream`, `test_host_pool`, `test_contiguous`, `test_inplace`, `test_stress`, `test_stream_perf`. `test_benchmarks`, `test_stress`, and `test_stream_perf` are timing/soak suites, not correctness suites — they are slow and their numbers are meaningless in a Debug build. `StreamPerf.ParallelFanOut` in particular asserts wall-clock against wall-clock and goes red under load; CI does not gate on those two suites. -**Every test that touches the device starts with `OM_REQUIRE_CUDA();`** ([tests/test_helpers.h](tests/test_helpers.h)) — a `cudaGetDeviceCount` check that `GTEST_SKIP`s instead of failing. That is what makes `ctest` green on a machine with no GPU (110 skipped, 60 host-side tests run) and is what the CPU-only CI job relies on; a new GPU test without it turns that job red. The Python equivalents are the `requires_cuda` marker and the `device` fixture in [python/tests/conftest.py](python/tests/conftest.py). +**Every test that touches the device starts with `OM_REQUIRE_CUDA();`** ([tests/test_helpers.h](tests/test_helpers.h)) — a `cudaGetDeviceCount` check that `GTEST_SKIP`s instead of failing. That is what makes `ctest` green on a machine with no GPU (130 skipped, 88 host-side tests run) and is what the CPU-only CI job relies on; a new GPU test without it turns that job red. The Python equivalents are the `requires_cuda` marker and the `device` fixture in [python/tests/conftest.py](python/tests/conftest.py). ## Debugging asynchronous kernel errors @@ -68,7 +68,7 @@ OPENMAT_DEBUG_SYNC=1 ./build/tests/test_streams # no rebuild needed [.github/workflows/ci.yml](.github/workflows/ci.yml), on every push and PR. Two jobs: - **cpu** — a `nvidia/cuda:*-devel` container on a GitHub-hosted runner: toolkit, no driver, no device. It configures with an explicit `-DCMAKE_CUDA_ARCHITECTURES="75;86;89"` (the `native` default needs a GPU to resolve) and puts `/usr/local/cuda/lib64/stubs` first in `CMAKE_LIBRARY_PATH`, where the link-time-only `libcuda.so` lives. Its job is the full compile+link — the macro machinery and the mandatory explicit instantiations only fail there — plus the host-side tests. -- **gpu** — a self-hosted runner labelled `self-hosted,linux,gpu`, needing nvcc and CMake ≥ 3.24 on PATH. Release build at `native` arch, the correctness suites (`ctest -E "test_benchmarks|test_stream_perf"` — 12 of the 14, `test_stress` included), the Python suite, then `compute-sanitizer --tool memcheck --leak-check full` over `test_stress`, `test_allocator_stream` and `test_streams` (~30 s total). memcheck is the only thing that reports a stream-ownership violation near the call site instead of as an illegal access in an unrelated kernel later on. +- **gpu** — a self-hosted runner labelled `self-hosted,linux,gpu`, needing nvcc and CMake ≥ 3.24 on PATH. Release build at `native` arch, the correctness suites (`ctest -E "test_benchmarks|test_stream_perf"` — 13 of the 15, `test_stress` included), the Python suite, then `compute-sanitizer --tool memcheck --leak-check full` over `test_stress`, `test_allocator_stream` and `test_streams` (~30 s total). memcheck is the only thing that reports a stream-ownership violation near the call site instead of as an illegal access in an unrelated kernel later on. ## Python package @@ -87,7 +87,7 @@ OPENMAT_LIB=build/OpenMat.so python python/test_bindings.py # smoke script cd python && pytest # pytest suite (python/tests/) ``` -Python suites: `test_tensor.py` (the original surface), `test_tensor_api.py` (metadata, indexing, shape/fused ops, buffer protocols), `test_dtypes.py`, `test_streams.py`. [python/tests/conftest.py](python/tests/conftest.py) provides a `device` fixture that runs a test on both backends and a `requires_cuda` marker. +Python suites: `test_tensor.py` (the original surface), `test_tensor_api.py` (metadata, indexing, shape/fused ops, buffer protocols), `test_dtypes.py`, `test_streams.py`, `test_inplace.py`. [python/tests/conftest.py](python/tests/conftest.py) provides a `device` fixture that runs a test on both backends and a `requires_cuda` marker. [python/hatch_build.py](python/hatch_build.py) is the custom hatch hook `pyproject.toml` points at: it copies `build/OpenMat.so` (or `$OPENMAT_LIB`) into `openmat/` so the wheel bundles it, and only warns if no library is present — an sdist should not need a CUDA toolchain. The copied `openmat/OpenMat.so` is gitignored by the root `*.so` rule. @@ -108,7 +108,7 @@ auto c = a.add(b, s); // enqueues on s; caller must s.synchronize() `om::Stream` ([headers/stream.h](headers/stream.h)) is a move-only RAII wrapper. The default constructor calls `cudaStreamCreate` and owns the handle; `Stream(cudaStream_t)` wraps an existing handle without owning it; `Stream::default_stream()` returns a non-owning wrapper around `nullptr`, which is how the synchronous API reuses the async code path with zero duplication. -**When adding an op, implement the `(args, const Stream&)` overload and make the no-stream one a one-line delegate.** Doing it the other way round breaks the single-source-of-truth invariant the whole `tensor.inl` is built on. +**When adding an op, implement the `(args, out, const Stream&)` overload and make everything else a one-line delegate.** Doing it the other way round breaks the single-source-of-truth invariant the whole `tensor.inl` is built on. See "In-place ops and caller-provided destinations" below for the three forms and which one is the real implementation. ### Dispatch: two paths, one of them mostly dead @@ -117,17 +117,46 @@ auto c = a.add(b, s); // enqueues on s; caller must s.synchronize() - `DEFINE_DEVICE_DISPATCH_BINARY_INL(OP)` — defines the free function `_OP(lhs, rhs, dst, DEVICE_TYPE)` that switches at runtime. - `DEFINE_DEVICE_DISPATCH_UNARY_H/INL` — same for tensor⊕scalar ops. -**But `Tensor`'s stream overloads no longer go through them.** Since the stream refactor, [headers/tensor.inl](headers/tensor.inl) branches on `device_type()` itself and calls `add_cpu(...)` / `launch_add(..., s.get())` directly, because the `_dispatch` structs have no `cudaStream_t` parameter. Today only `_fill` still uses the dispatch path (from `Tensor::fill`). Treat the `_dispatch` macros as legacy: adding a new op means wiring `tensor.inl` directly, and adding a dispatch registration only if you also need the stream-less free function. +**But `Tensor` no longer goes through them at all.** Since the stream refactor, [headers/tensor.inl](headers/tensor.inl) branches on `device_type()` itself and calls `add_cpu(...)` / `launch_add(..., s.get())` directly, because the `_dispatch` structs have no `cudaStream_t` parameter. `fill` was the last holdout and is now `fill_(value, s)` for the same reason, so the `_dispatch` machinery has **no callers left**. Treat it as legacy: adding a new op means wiring `tensor.inl` directly, and adding a dispatch registration only if you actually want the stream-less free function. Effective data flow for `a + b`: ``` Tensor::operator+() → Tensor::add(rhs, Stream::default_stream()) [tensor.inl] - → add_cpu(...) or launch_add(..., stream) [ops/cpu/ or ops/kernels/] - → flat CPU loop or rank-specialized CUDA kernel + → allocates the result, then + Tensor::add_out(rhs, out, stream) [tensor.inl — the real body] + → add_cpu(...) or launch_add(..., stream) [ops/cpu/ or ops/kernels/] + → flat CPU loop or contiguous fast path / rank-specialized CUDA kernel ``` +`a.add_(b)` enters the same chain one level down, at `add_out(b, *this, stream)`. + +### In-place ops and caller-provided destinations + +Every op exists in three forms, and only one of them is a real implementation: + +```cpp +auto c = a.add(b, s); // allocates a result, then calls add_out +a.add_out(b, out, s); // ← the body: writes into a destination that exists +a.add_(b, s); // == a.add_out(b, a, s) +``` + +`add_out` returns `Tensor&` (the destination), so calls chain; `add_` returns `*this`. The no-stream spelling of each is a one-line delegate with `Stream::default_stream()`. `operator+=`/`-=`/`*=`/`/=` are delegates to the `_` family. + +The families: `add`/`sub`/`mul`/`div` (tensor and scalar), `apply`, `apply_binary`, `relu`, `sigmoid` and `fill_` have all three forms. `matmul`, `transpose` and `permute` have `_out` but **no** in-place form — their kernels read elements they do not write, so a destination sharing a buffer with an operand would read values already overwritten. `_check_alias_none` rejects that at the call site instead of returning a plausible wrong answer. + +Why the elementwise family *can* alias: the CPU loop, the contiguous GPU fast path and the rank-specialized kernels all read index i and write index i, so `dst == lhs` is exactly as correct as a separate buffer. That rests on each buffer being one flat run, which is true of every tensor the library builds — `_check_alias_elementwise` re-checks `is_contiguous()` anyway, so the day a strided view can alias a *different* region of the same allocation (roadmap P2) it throws rather than computing the wrong thing. + +Two consequences worth knowing: + +- **The contiguous kernels dropped `__restrict__`** ([headers/ops/kernels/contiguous.cuh](headers/ops/kernels/contiguous.cuh)) — in-place calls them with `dst` equal to a source, which is precisely the aliasing the qualifier promises does not happen. Measured, it cost nothing: `add` at 16 M went 230.2-233.4 → 228.6-230.7 GB/s, `relu` 232.9-234.6 → 233.4-235.5. The read-only path comes from the explicit `__ldg` in `device_load`, not from restrict, and each pointer is touched once per thread. Re-measure before putting it back. +- **A destination that is not freshly allocated carries a stream caveat the allocating path does not.** `cudaMallocAsync` memory is stream-ordered, so enqueueing into `out` on a stream other than the one `out` was allocated on is only correct once the caller has ordered the two (an event, or a synchronize). `Tensor` cannot check that, so it is documented, not enforced — unlike the ownership invariant below, which it does enforce by storing `m_Stream`. + +The `_out` core is also where operand validation now lives: shape and device are checked before dispatch, so a GPU `add` between mismatched shapes throws instead of running off the end of a buffer the way it used to. + +`scripts/bench_inplace.py` is the harness; [benchmark_report.md §9](benchmark_report.md#9-every-op-allocated-its-own-result--in-place-and-out-forms-added) has the numbers. The short version: 1.2-1.6× below ~64 K on both backends where allocation is a large fraction of the op, shrinking to 1.04× at 16 M where the kernel dominates. The memory argument does not shrink — `tests/test_inplace.cpp` and `python/tests/test_inplace.py` assert `data_ptr` is unchanged across a 100-step loop, which is the property a correct-but-reallocating implementation would fail while passing every value check. + ### Memory and views **`Tensor`** ([headers/tensor.cuh](headers/tensor.cuh), [headers/tensor.inl](headers/tensor.inl)) — owning N-D tensor: shape, row-major strides, raw `T*`, a `Device`, an `om::Stream m_Stream`, and a `unique_ptr>`. Copy deep-copies; move transfers ownership and nulls the source. There is a private `Tensor(shape, device, Stream)` constructor used by the stream overloads so a result tensor is allocated and freed on the stream that produced it. @@ -203,7 +232,7 @@ Every consumer that includes `tensor.cuh` — and therefore re-instantiates thes **Explicit instantiations** in [src/ops/kernels/fused_op.cu](src/ops/kernels/fused_op.cu) must list every `(T, Op)` pair used from a `.cpp` translation unit. A new functor or a new `Compose` combination without an instantiation is a link error. Calls from `.cu` files instantiate implicitly and hide the problem. -`Tensor` surface: `apply(op[, stream])`, `apply_binary(rhs, op)`, `scale_shift`, `shift_scale`, `relu`, `sigmoid`, `fused_add_mul`, `fused_sub_mul`, `fused_mul_add`, `fused_div_add`. `apply` and `apply_binary` both have a real CPU loop branch — the CUDA-only limitation noted in [docs/roadmap.md](docs/roadmap.md) §4.2 has been fixed. +`Tensor` surface: `apply(op[, stream])`, `apply_binary(rhs, op[, stream])`, `scale_shift`, `shift_scale`, `relu`, `sigmoid`, `fused_add_mul`, `fused_sub_mul`, `fused_mul_add`, `fused_div_add`, plus the `_out` and `_` forms of `apply`, `apply_binary`, `relu` and `sigmoid`. `apply` and `apply_binary` both have a real CPU loop branch — the CUDA-only limitation noted in [docs/roadmap.md](docs/roadmap.md) §4.2 has been fixed, and `apply_binary` gained the stream overload §4.4 was asking for when it was rebuilt on `apply_binary_out`. The fused `fused_*` helpers have no dedicated in-place spelling: `apply_binary_(rhs, BinaryCompose<…>{…})` with the same functor is the way to get one. ## Reductions @@ -244,15 +273,18 @@ Beyond the tensor families the library exports a dtype-independent runtime API: Adding a Python-visible method means three edits: the function in [openmat_capi_impl.inc](src/python/openmat_capi_impl.inc) (once — both dtypes get it), the `ctypes` `restype`/`argtypes` in `_declare_dtype()` in [python/openmat/_clib.py](python/openmat/_clib.py), and the wrapper in [python/openmat/tensor.py](python/openmat/tensor.py). Ops with a `(args, const Stream&)` C++ overload get a `_stream` sibling in the `.inc` and a `stream=None` kwarg in Python. +The in-place and destination-provided families cross the boundary as **int-returning** functions (`om_tensor_float_add_inplace`, `..._add_out`, `..._add_scalar_inplace`, …), not handle-returning ones: nothing is allocated and no ownership changes hands, so there is no pointer to check and no `_wrap` to do. The C symbols spell the C++ trailing underscore as `_inplace` — `add_` would be a legal C identifier but a confusing export. Python's `add_`/`add_out` return `self`/`out` so calls chain, and `__iadd__` and friends return `self` so `a += b` does not silently rebind the name to a new tensor. + The Python package is `Tensor` + `Stream` + a `DType` registry ([python/openmat/_dtypes.py](python/openmat/_dtypes.py)); host tensors expose `__array_interface__` (zero-copy `np.asarray`), CUDA tensors `__cuda_array_interface__`. See [python/README.md](python/README.md) for the user-facing surface. ## Benchmarking -Four harnesses, all requiring a **Release** build (`build-release/`) — Debug numbers are meaningless — and a `bench-env` venv holding NumPy and PyTorch, which the library itself does not need: +Five harnesses, all requiring a **Release** build (`build-release/`) — Debug numbers are meaningless — and a `bench-env` venv holding NumPy and PyTorch, which the library itself does not need: - [scripts/bench_vs.py](scripts/bench_vs.py) — the main cross-framework table (CPU + CUDA + transfers). `--quick` for a smoke run, `--no-cuda` to skip the device. - [scripts/bench_rank_sweep.py](scripts/bench_rank_sweep.py) — the same 16 M buffer reshaped to ranks 1–5, which is what verifies the contiguous fast path is actually engaging. `bench_vs.py` only ever uses rank-1 shapes and would not notice a regression here. - [scripts/bench_omp.py](scripts/bench_omp.py) — run once per `OMP_NUM_THREADS` value; the 1-vs-20 delta is the OpenMP A/B, and a `1.00×` row is how you confirm an op is *not* parallelized (`min`/`max`, `sum`, `apply`/`apply_binary`). +- [scripts/bench_inplace.py](scripts/bench_inplace.py) — the in-place / `out=` families against the allocating forms, per op and over a 32-step chain. Every case builds its own operands: an earlier draft that reused one destination across the three ops at each size reported a CUDA `relu_` "regression" of 0.38× that does not exist (1.05× re-measured). - `OPENMAT_HOST_CACHE_BYTES=0` — restores the pre-`HostPool` allocator, the A/B behind §3. Two traps that have already produced wrong conclusions in the reports: diff --git a/benchmark_report.md b/benchmark_report.md index 4d50aed..ffc45f9 100644 --- a/benchmark_report.md +++ b/benchmark_report.md @@ -494,6 +494,76 @@ the rank-1 layout, GPU `add` at 16 M now runs at **236.4 GB/s against PyTorch's 232.4** — 1.02× faster, where the previous report had it 1.47× slower. `relu` went from 1.9× slower to parity, and `x*s+t` from parity to 1.97× faster. +## 9. Every op allocated its own result — in-place and `out=` forms added + +Until now every `Tensor` method was `const` and returned a freshly allocated +tensor. In a training loop that is one allocation *and* one free per op per +iteration, and a peak footprint equal to the whole expression graph rather than +the working set. PyTorch exposes `add_`/`mul_` and `out=` for exactly this. + +The launchers already took a separate `dst`, so the work was on the public +surface: `add_out(rhs, out)` is now the single implementation of each op, the +allocating form builds `out` and calls it, and `add_(rhs)` calls it with +`*this` as `out`. See CLAUDE.md for the surface and the aliasing rules. + +**One op, measured 2026-09-04, GB10, Release** (min of 7 batches, through the +Python bindings so allocation and synchronization are inside the timing — +`scripts/bench_inplace.py`): + +| device | n | op | allocating | in-place | `out=` | in-place speedup | +|---|---:|---|---:|---:|---:|---:| +| cpu | 1 K | add | 1.79 µs | 1.22 µs | 1.38 µs | 1.47× | +| cpu | 1 K | relu | 1.31 µs | 0.80 µs | 0.99 µs | 1.65× | +| cpu | 64 K | add | 6.34 µs | 5.70 µs | 5.81 µs | 1.11× | +| cpu | 16 M | add | 2029 µs | 1661 µs | 2015 µs | 1.22× | +| cpu | 16 M | mul (scalar) | 1413 µs | 972 µs | 1340 µs | 1.45× | +| cuda | 1 K | add | 7.93 µs | 6.48 µs | 6.65 µs | 1.22× | +| cuda | 64 K | add | 8.21 µs | 6.75 µs | 6.92 µs | 1.22× | +| cuda | 1 M | add | 14.05 µs | 12.60 µs | 12.72 µs | 1.12× | +| cuda | 16 M | add | 862 µs | 827 µs | 823 µs | 1.04× | + +**A 32-step chain** (`w += g`, the shape a training step actually has): + +| device | n | allocating | in-place | speedup | +|---|---:|---:|---:|---:| +| cpu | 1 K | 61.4 µs | 39.4 µs | 1.56× | +| cpu | 64 K | 206 µs | 187 µs | 1.10× | +| cpu | 16 M | 63.1 ms | 54.6 ms | 1.16× | +| cuda | 1 K | 264 µs | 212 µs | 1.25× | +| cuda | 64 K | 273 µs | 221 µs | 1.24× | +| cuda | 1 M | 489 µs | 412 µs | 1.19× | +| cuda | 16 M | 30.8 ms | 26.6 ms | 1.16× | + +Read the shape of it, not the individual rows. The saving is largest where +allocation is a large fraction of the op (1.2–1.6× below ~64 K on both +backends, and it is the same ~1.4 µs per op on CPU and ~1.5 µs on CUDA that §4 +already identified as fixed overhead), and it shrinks toward 1.04× at 16 M +where the kernel itself dominates and `HostPool` / `cudaMallocAsync` are +already recycling the block. That last figure is the honest ceiling for the +*time* argument. + +The memory argument does not shrink: an in-place chain holds one buffer whatever +its length, and the test suite asserts exactly that (`data_ptr` is unchanged +across a 100-step loop, in `tests/test_inplace.cpp` and +`python/tests/test_inplace.py`) — a correct-but-reallocating implementation +would pass every value check and fail those. + +**A trap this script already sprang.** An earlier draft reused one destination +tensor across the three ops at each size and reported CUDA `relu_` at 0.38× — +a 2.6× *regression* that does not exist. Re-measured with per-case operands it +is 1.05×. Every case now builds its own; and the CPU rows at 1 M were dropped +from the table because they swing between 11 µs and 250 µs run to run on a +loaded machine, which is contention, not the op. + +The change also removed `__restrict__` from the three kernels in +`contiguous.cuh`: in-place calls them with `dst` equal to a source, which is +exactly the aliasing the qualifier promises does not happen. It cost nothing — +`add` at 16 M went from 230.2–233.4 GB/s to 228.6–230.7 and `relu` from +232.9–234.6 to 233.4–235.5, i.e. inside run-to-run spread. The read-only path +comes from the explicit `__ldg` in `device_load`, not from restrict-driven +inference, and each pointer is touched once per thread, so there was nothing +for the qualifier to buy. + ## Full results ### CPU, default threading (20 OpenMP threads) @@ -668,6 +738,10 @@ OPENMAT_LIB=build-release/OpenMat.so PYTHONPATH=python \ OPENMAT_LIB=build-release/OpenMat.so PYTHONPATH=python \ bench-env/bin/python scripts/bench_rank_sweep.py +# §9 — in-place / out= against the allocating forms +OPENMAT_LIB=build-release/OpenMat.so PYTHONPATH=python \ + bench-env/bin/python scripts/bench_inplace.py + # §7 — OpenMP scaling, one invocation per thread count for t in 1 20; do OMP_NUM_THREADS=$t OPENMAT_LIB=build-release/OpenMat.so PYTHONPATH=python \ diff --git a/compile_commands.json b/compile_commands.json index a24f119..22746f2 100644 --- a/compile_commands.json +++ b/compile_commands.json @@ -185,6 +185,12 @@ "file": "/home/apalese/OpenMat/tests/test_contiguous.cpp", "output": "tests/CMakeFiles/test_contiguous.dir/test_contiguous.cpp.o" }, +{ + "directory": "/home/apalese/OpenMat/build/tests", + "command": "/usr/bin/c++ -I/home/apalese/OpenMat/headers -I/usr/local/cuda/targets/sbsa-linux/include -isystem /home/apalese/OpenMat/build/_deps/googletest-src/googletest/include -isystem /home/apalese/OpenMat/build/_deps/googletest-src/googletest -g -fopenmp -o CMakeFiles/test_inplace.dir/test_inplace.cpp.o -c /home/apalese/OpenMat/tests/test_inplace.cpp", + "file": "/home/apalese/OpenMat/tests/test_inplace.cpp", + "output": "tests/CMakeFiles/test_inplace.dir/test_inplace.cpp.o" +}, { "directory": "/home/apalese/OpenMat/build/_deps/googletest-build/googlemock", "command": "/usr/bin/c++ -I/home/apalese/OpenMat/build/_deps/googletest-src/googlemock/include -I/home/apalese/OpenMat/build/_deps/googletest-src/googlemock -isystem /home/apalese/OpenMat/build/_deps/googletest-src/googletest/include -isystem /home/apalese/OpenMat/build/_deps/googletest-src/googletest -g -std=c++17 -Wall -Wshadow -Wno-error=dangling-else -DGTEST_HAS_PTHREAD=1 -fexceptions -Wextra -Wno-unused-parameter -Wno-missing-field-initializers -o CMakeFiles/gmock.dir/src/gmock-all.cc.o -c /home/apalese/OpenMat/build/_deps/googletest-src/googlemock/src/gmock-all.cc", diff --git a/docs/roadmap.md b/docs/roadmap.md index c1b0756..b986ae2 100644 --- a/docs/roadmap.md +++ b/docs/roadmap.md @@ -4,7 +4,7 @@ Questo documento elenca le funzionalità del framework, il loro stato attuale e **Legenda:** ✅ implementato · ⚠️ implementato parzialmente o con divergenze rispetto al progetto originale · ❌ da fare -Ultimo allineamento al codice: 4 settembre 2026 (benchmark rieseguiti su GB10, Release; vedi [`benchmark_report.md`](../benchmark_report.md)). +Ultimo allineamento al codice: 4 settembre 2026 (in-place / `out=` aggiunti, §8; benchmark rieseguiti su GB10, Release; vedi [`benchmark_report.md`](../benchmark_report.md)). --- @@ -205,11 +205,9 @@ Il bug latente è risolto. `apply` e `apply_binary` in [`headers/tensor.inl`](.. Entrambi passano per `float` internamente (`static_cast(x) > 0.0f`, `expf(-static_cast(x))`): è ciò che li fa funzionare senza modifiche su `float16_t`. -### 4.4 Overload sullo stream per `apply_binary` ❌ +### 4.4 Overload sullo stream per `apply_binary` ✅ -Asimmetria residua: `apply` ha l'overload `(Op, const Stream&)`, `apply_binary` no. Di conseguenza `scale_shift`, `shift_scale` e i quattro `fused_*_*` sono chiamate sincrone sullo stream di default, mentre `relu` e `sigmoid` no. - -Il lavoro è meccanico e segue `apply`: accettare `const Stream& s`, costruire l'output con il costruttore privato `Tensor(shape, device, Stream)` e passare `s.get()` al launcher — `launch_apply_binary_op` accetta già un `cudaStream_t`. +Risolto come effetto collaterale di §8: `apply_binary` è stato riscritto sopra `apply_binary_out(rhs, op, out, s)`, che uno stream lo prende per costruzione, e le forme allocante e in-place ci delegano. `scale_shift`, `shift_scale` e i quattro `fused_*_*` restano sincroni sullo stream di default perché non hanno un overload proprio — chiamarli su uno stream significa passare per `apply`/`apply_binary` con lo stesso functor. --- @@ -347,6 +345,34 @@ Suite di test: `python/tests/test_tensor.py`, `test_tensor_api.py`, `test_dtypes --- +## 8. Operazioni in-place e destinazione fornita dal chiamante ✅ + +Ogni metodo di `Tensor` era `const` e allocava il risultato. In un ciclo di training questo significa un'allocazione **e** una deallocazione per operazione per iterazione, e un picco di memoria pari a tutto il grafo dell'espressione invece che al working set. PyTorch espone `add_`, `mul_` e `out=` esattamente per questo. + +I launcher accettavano già una `dst` separata, quindi il lavoro era tutto sulla superficie pubblica. Ogni operazione ha ora tre forme, di cui **una sola è un'implementazione vera**: + +```cpp +auto c = a.add(b, s); // alloca il risultato, poi chiama add_out +a.add_out(b, out, s); // ← il corpo: scrive in una destinazione che esiste già +a.add_(b, s); // == a.add_out(b, a, s) +``` + +`add_out` restituisce `Tensor&` (la destinazione) e `add_` restituisce `*this`, quindi le chiamate si concatenano; le forme senza stream sono delegate di una riga con `Stream::default_stream()`. Ci sono anche `operator+=`/`-=`/`*=`/`/=`. + +**Famiglie coperte:** `add`/`sub`/`mul`/`div` (tensore e scalare), `apply`, `apply_binary`, `relu`, `sigmoid`, `fill_`. `matmul`, `transpose` e `permute` hanno solo `_out`: i loro kernel leggono elementi che non scrivono, quindi una destinazione che condivide il buffer con un operando leggerebbe valori già sovrascritti — `_check_alias_none` lo rifiuta al call site invece di restituire un risultato plausibile e sbagliato. + +**Perché la famiglia elementwise *può* fare aliasing:** il loop CPU, il fast path contiguo GPU e i kernel per rank leggono l'indice i e scrivono l'indice i, quindi `dst == lhs` è corretto esattamente quanto un buffer separato. Questo dipende dal fatto che ogni buffer sia una singola sequenza piatta: `_check_alias_elementwise` ricontrolla `is_contiguous()`, così il giorno in cui una view con stride potrà puntare a una regione *diversa* della stessa allocazione (§2.1b) l'operazione solleva un'eccezione invece di calcolare la cosa sbagliata. + +Conseguenza collaterale: i tre kernel di [`contiguous.cuh`](../headers/ops/kernels/contiguous.cuh) hanno perso `__restrict__`, che è precisamente la promessa che l'in-place viola. Misurato, non è costato nulla (`add` a 16M da 230.2-233.4 a 228.6-230.7 GB/s, `relu` da 232.9-234.6 a 233.4-235.5): il percorso read-only viene dall'`__ldg` esplicito in `device_load`, non da inferenza sul restrict. + +Il `_out` è anche il punto in cui vive ora la validazione degli operandi (shape e device controllati prima del dispatch), quindi un `add` GPU fra shape diverse solleva un'eccezione invece di leggere oltre la fine di un buffer come faceva prima. + +**Numeri** ([`benchmark_report.md` §9](../benchmark_report.md#9-every-op-allocated-its-own-result--in-place-and-out-forms-added), harness `scripts/bench_inplace.py`): 1.2-1.6× sotto i ~64 K elementi su entrambi i backend, dove l'allocazione è una frazione grande dell'operazione, che si riduce a 1.04× a 16 M dove domina il kernel e `HostPool`/`cudaMallocAsync` stanno già riciclando il blocco. L'argomento sulla memoria invece non si riduce: `tests/test_inplace.cpp` e `python/tests/test_inplace.py` verificano che `data_ptr` non cambi lungo un ciclo di 100 passi — la proprietà che un'implementazione corretta ma ri-allocante fallirebbe superando ogni controllo sui valori. + +**Non fatto:** `matmul` batch/in-place resta fuori (§5.1), e i quattro `fused_*_*` non hanno una forma in-place dedicata — `apply_binary_(rhs, BinaryCompose<…>{…})` con lo stesso functor è il modo di ottenerla. + +--- + ## Priorità suggerite | Done | Priorità | Item | @@ -367,7 +393,8 @@ Suite di test: `python/tests/test_tensor.py`, `test_tensor_api.py`, `test_dtypes | ✅ | — | 6.6 — `matmul_cpu` `ikj` + tiling L2 + OpenMP: 1.81 → 123 GFLOP/s a 1024³ (§5) | | ✅ | — | 6.7 — `reduce_sum_cpu` a 8 lane: 7.7 → 36.9 GB/s, ora davanti a NumPy (§1.2) | | ✅ | — | 6.8 — `HostPool`/`PinnedHostPool`: round-trip 100 × 64 MB da 4.3 a 51.1 GB/s | -| | **Alta** | 4.4 — Overload stream per `apply_binary` (chiude l'asimmetria di 6.1) | +| ✅ | — | 8 — Operazioni in-place (`add_`, `mul_`, `relu_`, `fill_`) e overload `_out` con destinazione fornita dal chiamante | +| ✅ | — | 4.4 — Overload stream per `apply_binary` (arrivato con §8: `apply_binary_out` prende lo stream per costruzione) | | | **Alta** | 1.1 — Unarie `abs`, `sqrt`, `exp`, `log` via `apply` (costo basso, alto valore) | | | **Alta** | 6.4b — OpenMP su `apply`/`apply_binary`: stesso `if(_total > 65536)` di 6.4, ~2× atteso a 16M su tutta la famiglia fused | | | Media | 6.4c — `min`/`max` con `parallel for reduction(min:)/(max:)` fra thread (2.7× dietro NumPy); **non** la clausola `simd`, già misurata e rimossa | diff --git a/headers/ops/kernels/contiguous.cuh b/headers/ops/kernels/contiguous.cuh index 44a6621..f747395 100644 --- a/headers/ops/kernels/contiguous.cuh +++ b/headers/ops/kernels/contiguous.cuh @@ -47,6 +47,16 @@ // A grid-stride loop was also tried and is not used: capping the grid to a few // waves per SM costs 8-12 %, and at the exact block count the loop body runs // once and buys nothing over a plain bounds check. +// +// None of the three kernels marks its pointers __restrict__, and that is +// deliberate: the in-place family (Tensor::add_, relu_, …) calls them with dst +// equal to one of the sources, which is exactly the aliasing __restrict__ +// promises does not happen. The promise bought nothing here — each pointer is +// read or written once per thread, and the read-only path comes from the +// explicit __ldg in device_load, not from restrict-driven inference. Measured +// on the reference GB10 at 16 M elements, removing it moved `add` from +// 230.2-233.4 GB/s to 228.6-230.7 and `relu` from 232.9-234.6 to 233.4-235.5, +// i.e. nothing outside run-to-run spread. Re-measure before putting it back. // ───────────────────────────────────────────────────────────────────────────── namespace om { @@ -83,8 +93,8 @@ namespace detail { // ill-formed and a member-wise copy is free to lower to two 2-byte accesses, // which is the whole thing this path exists to avoid. template - __global__ void contig_binary_kernel(const T* __restrict__ a, const T* __restrict__ b, - T* __restrict__ dst, size_t n_pack, size_t n, Op op) + __global__ void contig_binary_kernel(const T* a, const T* b, + T* dst, size_t n_pack, size_t n, Op op) { const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; @@ -114,7 +124,7 @@ namespace detail { } template - __global__ void contig_unary_kernel(const T* __restrict__ src, T* __restrict__ dst, + __global__ void contig_unary_kernel(const T* src, T* dst, size_t n_pack, size_t n, Op op) { const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; @@ -141,7 +151,7 @@ namespace detail { } template - __global__ void contig_fill_kernel(T* __restrict__ dst, T value, size_t n_pack, size_t n) + __global__ void contig_fill_kernel(T* dst, T value, size_t n_pack, size_t n) { const size_t i = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; diff --git a/headers/tensor.cuh b/headers/tensor.cuh index b366eee..48995ce 100644 --- a/headers/tensor.cuh +++ b/headers/tensor.cuh @@ -88,10 +88,9 @@ namespace om } - void fill(const value_type& value) - { - _fill(this->view(), value, this->device_type()); - } + // Kept for source compatibility; fill_ is the canonical spelling and + // the one that takes a stream. + void fill(const value_type& value) { this->fill_(value); } Tensor add(const Tensor& rhs) const; @@ -159,6 +158,9 @@ namespace om Tensor scale_shift(value_type scale, value_type shift) const; Tensor shift_scale(value_type shift, value_type scale) const; + template + Tensor apply_binary(const Tensor& rhs, Op op, const Stream& s) const; + template Tensor apply_binary(const Tensor& rhs, Op op) const; @@ -174,6 +176,116 @@ namespace om Tensor fused_div_add(const Tensor& rhs, value_type shift) const; + // ── Destination-provided overloads ────────────────────────────────── + // + // `a.add_out(b, out)` writes a+b into `out`'s existing buffer and + // returns a reference to it, so a loop running the same op every + // iteration allocates once instead of once per iteration. `out` must + // already carry the result's shape and sit on the same device; it may + // be one of the operands (that is what the in-place family below is). + // + // These are the single implementation of each op: the allocating forms + // above build the result and call straight into them, and the in-place + // forms pass *this as the destination. Nothing else branches on + // device_type(). + + Tensor& add_out(const Tensor& rhs, Tensor& out, const Stream& s) const; + Tensor& add_out(const Tensor& rhs, Tensor& out) const; + Tensor& sub_out(const Tensor& rhs, Tensor& out, const Stream& s) const; + Tensor& sub_out(const Tensor& rhs, Tensor& out) const; + Tensor& mul_out(const Tensor& rhs, Tensor& out, const Stream& s) const; + Tensor& mul_out(const Tensor& rhs, Tensor& out) const; + Tensor& div_out(const Tensor& rhs, Tensor& out, const Stream& s) const; + Tensor& div_out(const Tensor& rhs, Tensor& out) const; + + Tensor& add_out(const value_type& scalar, Tensor& out, const Stream& s) const; + Tensor& add_out(const value_type& scalar, Tensor& out) const; + Tensor& sub_out(const value_type& scalar, Tensor& out, const Stream& s) const; + Tensor& sub_out(const value_type& scalar, Tensor& out) const; + Tensor& mul_out(const value_type& scalar, Tensor& out, const Stream& s) const; + Tensor& mul_out(const value_type& scalar, Tensor& out) const; + Tensor& div_out(const value_type& scalar, Tensor& out, const Stream& s) const; + Tensor& div_out(const value_type& scalar, Tensor& out) const; + + // matmul, transpose and permute each read an index they do not write, + // so unlike the elementwise family `out` may not be an operand. + Tensor& matmul_out(const Tensor& rhs, Tensor& out, const Stream& s) const; + Tensor& matmul_out(const Tensor& rhs, Tensor& out) const; + Tensor& transpose_out(Tensor& out, const Stream& s) const; + Tensor& transpose_out(Tensor& out) const; + Tensor& permute_out(const std::vector& axes, Tensor& out, const Stream& s) const; + Tensor& permute_out(const std::vector& axes, Tensor& out) const; + + template + Tensor& apply_out(Op op, Tensor& out, const Stream& s) const; + template + Tensor& apply_out(Op op, Tensor& out) const; + template + Tensor& apply_binary_out(const Tensor& rhs, Op op, + Tensor& out, const Stream& s) const; + template + Tensor& apply_binary_out(const Tensor& rhs, Op op, + Tensor& out) const; + + Tensor& relu_out(Tensor& out, const Stream& s) const; + Tensor& relu_out(Tensor& out) const; + Tensor& sigmoid_out(Tensor& out, const Stream& s) const; + Tensor& sigmoid_out(Tensor& out) const; + + // ── In-place ops ──────────────────────────────────────────────────── + // + // Trailing underscore, as in PyTorch: the result replaces this + // tensor's contents, no allocation and no free happen, and the return + // is *this so calls chain. Every one of them is `X_out(..., *this)`. + // + // They are only defined for ops whose kernels read and write the same + // index — the elementwise families, `apply`, `apply_binary` and + // `fill`. matmul, transpose and permute have no in-place form. + + Tensor& add_(const Tensor& rhs, const Stream& s); + Tensor& add_(const Tensor& rhs); + Tensor& sub_(const Tensor& rhs, const Stream& s); + Tensor& sub_(const Tensor& rhs); + Tensor& mul_(const Tensor& rhs, const Stream& s); + Tensor& mul_(const Tensor& rhs); + Tensor& div_(const Tensor& rhs, const Stream& s); + Tensor& div_(const Tensor& rhs); + + Tensor& add_(const value_type& scalar, const Stream& s); + Tensor& add_(const value_type& scalar); + Tensor& sub_(const value_type& scalar, const Stream& s); + Tensor& sub_(const value_type& scalar); + Tensor& mul_(const value_type& scalar, const Stream& s); + Tensor& mul_(const value_type& scalar); + Tensor& div_(const value_type& scalar, const Stream& s); + Tensor& div_(const value_type& scalar); + + Tensor& operator+=(const Tensor& rhs); + Tensor& operator-=(const Tensor& rhs); + Tensor& operator*=(const Tensor& rhs); + Tensor& operator/=(const Tensor& rhs); + Tensor& operator+=(const value_type& scalar); + Tensor& operator-=(const value_type& scalar); + Tensor& operator*=(const value_type& scalar); + Tensor& operator/=(const value_type& scalar); + + template + Tensor& apply_(Op op, const Stream& s); + template + Tensor& apply_(Op op); + template + Tensor& apply_binary_(const Tensor& rhs, Op op, const Stream& s); + template + Tensor& apply_binary_(const Tensor& rhs, Op op); + + Tensor& relu_(const Stream& s); + Tensor& relu_(); + Tensor& sigmoid_(const Stream& s); + Tensor& sigmoid_(); + + Tensor& fill_(const value_type& value, const Stream& s); + Tensor& fill_(const value_type& value); + Tensor to(const Device& target) const; Tensor cpu() const; Tensor cuda() const; @@ -214,6 +326,34 @@ namespace om // there is no such thing as pinned device memory. Tensor(const std::vector& shape, const Device& dv, Stream stream, bool pinned); + // Throws unless `out` can receive a result of shape `shape` produced on + // this tensor's device. `who` names the caller in the message. + void _check_out(const Tensor& out, + const std::vector& shape, + const char* who) const; + + // Throws unless `rhs` is a legal second operand for an elementwise op. + void _check_operand(const Tensor& rhs, const char* who) const; + + // Every elementwise path — the CPU loop, the contiguous GPU fast path + // and the rank-specialized kernels — reads index i and writes index i, + // so a destination that *is* an operand is exactly as correct as a + // separate one. That equivalence rests on all three buffers being one + // flat run each; it says nothing about a strided view aliasing a + // different region of the same allocation, which is what a real view + // type (roadmap P2) would make possible. Nothing in the library can + // build one today, so this throws when it meets one rather than + // quietly computing the wrong answer. + void _check_alias_elementwise(const Tensor& out, + const Tensor* rhs, + const char* who) const; + + // For ops that read an index they do not write (matmul, transpose, + // permute), where sharing a buffer is simply wrong. + void _check_alias_none(const Tensor& out, + const Tensor* rhs, + const char* who) const; + void _compute_strides(); inline size_t _compute_flat_index(const std::vector& indices) const; diff --git a/headers/tensor.inl b/headers/tensor.inl index 846d676..6d3805c 100644 --- a/headers/tensor.inl +++ b/headers/tensor.inl @@ -1,6 +1,30 @@ #include "tensor.cuh" #include +namespace om { +namespace detail { + // Validates `axes` against `shape` and returns the permuted shape. Shared + // by permute() and permute_out() so the two cannot drift apart on what + // counts as a legal axis list. + inline std::vector permuted_shape(const std::vector& shape, + const std::vector& axes) + { + const size_t r = shape.size(); + if (axes.size() != r) + throw std::invalid_argument("permute: axes length must match tensor rank"); + std::vector seen(r, false); + for (size_t a : axes) { + if (a >= r) throw std::out_of_range("permute: axis value out of range"); + if (seen[a]) throw std::invalid_argument("permute: duplicate axis"); + seen[a] = true; + } + std::vector out(r); + for (size_t d = 0; d < r; ++d) out[d] = shape[axes[d]]; + return out; + } +} // namespace detail +} // namespace om + template om::Tensor::Tensor(const std::vector& shape, const Device& dv) : m_Shape(shape), m_Device(dv), m_Stream(Stream::default_stream()), @@ -301,24 +325,21 @@ om::Tensor om::Tensor::shift_scale(value_type shift, val template template -om::Tensor om::Tensor::apply_binary(const Tensor& rhs, Op op) const +om::Tensor om::Tensor::apply_binary(const Tensor& rhs, Op op, + const Stream& s) const { - Tensor out(this->shape(), this->device()); - if (this->device_type() == DEVICE_TYPE::CPU) { - auto lhs_v = this->view(); - auto rhs_v = rhs.view(); - auto dst_v = out.view(); - if (!lhs_v.match(rhs_v)) - throw std::invalid_argument("apply_binary: tensors must have the same shape"); - size_t n = lhs_v.size(); - for (size_t i = 0; i < n; ++i) - dst_v[i] = op(lhs_v[i], rhs_v[i]); - } else { - launch_apply_binary_op(this->view(), rhs.view(), out.view(), op); - } + Tensor out(this->shape(), this->device(), Stream(s.get())); + this->apply_binary_out(rhs, op, out, s); return out; } +template +template +om::Tensor om::Tensor::apply_binary(const Tensor& rhs, Op op) const +{ + return this->apply_binary(rhs, op, Stream::default_stream()); +} + template om::Tensor om::Tensor::relu() const { @@ -478,15 +499,82 @@ void om::Tensor::copyToDevice(value_type *dest) const throw std::runtime_error("Tensor::copyToDevice: memory already on device"); } -// ── Stream overload implementations ───────────────────────────────────────── -// Each method dispatches CPU ops directly (stream irrelevant on CPU) and GPU -// ops via the stream-aware launch_* functions. The no-stream methods above -// delegate here with Stream::default_stream() (wraps nullptr → synchronous). +// ── Destination-provided implementations ──────────────────────────────────── +// +// This is where every op actually happens. `X_out` takes the destination the +// caller already owns; the allocating overload builds one and calls straight +// in here; the in-place `X_` passes *this. So each op still has exactly one +// place that branches on device_type(), which is the same single-source-of- +// truth rule the stream refactor established for the (args, Stream) form. +// +// CPU ops ignore the stream; GPU ops enqueue on it via the launch_* family. +// The no-stream forms delegate with Stream::default_stream() (wraps nullptr → +// synchronous). +// +// A destination that is not the freshly allocated one carries a caveat the +// allocating path does not have: cudaMallocAsync memory is stream-ordered, so +// enqueueing into `out` on a stream other than the one `out` was allocated on +// is only safe once the caller has ordered the two (an event, or a +// synchronize). Tensor cannot check that, so it is documented, not enforced. + +template +void om::Tensor::_check_out(const Tensor& out, + const std::vector& shape, + const char* who) const +{ + if (out.m_Shape != shape) + throw std::invalid_argument(std::string(who) + + ": destination shape does not match the result shape"); + if (out.device_type() != this->device_type() || out.m_Device.m_Id != m_Device.m_Id) + throw std::invalid_argument(std::string(who) + + ": destination must live on the same device as the operands"); + if (out.m_Data == nullptr) + throw std::invalid_argument(std::string(who) + ": destination has been moved from"); +} template -om::Tensor om::Tensor::add(const Tensor& rhs, const Stream& s) const +void om::Tensor::_check_operand(const Tensor& rhs, const char* who) const { - Tensor out(this->shape(), this->device(), Stream(s.get())); + if (rhs.m_Shape != m_Shape) + throw std::invalid_argument(std::string(who) + ": tensors must have the same shape"); + if (rhs.device_type() != this->device_type() || rhs.m_Device.m_Id != m_Device.m_Id) + throw std::invalid_argument(std::string(who) + ": tensors must live on the same device"); +} + +template +void om::Tensor::_check_alias_elementwise(const Tensor& out, + const Tensor* rhs, + const char* who) const +{ + const bool aliased = (out.m_Data == m_Data) || (rhs && out.m_Data == rhs->m_Data); + if (!aliased) return; + + const bool flat = this->view().is_contiguous() && out.view().is_contiguous() && + (!rhs || rhs->view().is_contiguous()); + if (!flat) + throw std::invalid_argument(std::string(who) + + ": destination may only share a buffer with an operand while every " + "operand is contiguous"); +} + +template +void om::Tensor::_check_alias_none(const Tensor& out, + const Tensor* rhs, + const char* who) const +{ + if (out.m_Data == m_Data || (rhs && out.m_Data == rhs->m_Data)) + throw std::invalid_argument(std::string(who) + + ": destination must not be one of the operands"); +} + +template +om::Tensor& om::Tensor::add_out(const Tensor& rhs, + Tensor& out, + const Stream& s) const +{ + _check_operand(rhs, "add_out"); + _check_out(out, m_Shape, "add_out"); + _check_alias_elementwise(out, &rhs, "add_out"); if (this->device_type() == DEVICE_TYPE::CPU) add_cpu(this->view(), rhs.view(), out.view()); else @@ -495,9 +583,34 @@ om::Tensor om::Tensor::add(const Tensor& rhs } template -om::Tensor om::Tensor::sub(const Tensor& rhs, const Stream& s) const +om::Tensor& om::Tensor::add_out(const Tensor& rhs, + Tensor& out) const +{ return this->add_out(rhs, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::add(const Tensor& rhs, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->add_out(rhs, out, s); + return out; +} + +template +om::Tensor& om::Tensor::add_(const Tensor& rhs, const Stream& s) +{ return this->add_out(rhs, *this, s); } + +template +om::Tensor& om::Tensor::add_(const Tensor& rhs) +{ return this->add_out(rhs, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::sub_out(const Tensor& rhs, + Tensor& out, + const Stream& s) const +{ + _check_operand(rhs, "sub_out"); + _check_out(out, m_Shape, "sub_out"); + _check_alias_elementwise(out, &rhs, "sub_out"); if (this->device_type() == DEVICE_TYPE::CPU) sub_cpu(this->view(), rhs.view(), out.view()); else @@ -506,9 +619,34 @@ om::Tensor om::Tensor::sub(const Tensor& rhs } template -om::Tensor om::Tensor::mul(const Tensor& rhs, const Stream& s) const +om::Tensor& om::Tensor::sub_out(const Tensor& rhs, + Tensor& out) const +{ return this->sub_out(rhs, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::sub(const Tensor& rhs, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->sub_out(rhs, out, s); + return out; +} + +template +om::Tensor& om::Tensor::sub_(const Tensor& rhs, const Stream& s) +{ return this->sub_out(rhs, *this, s); } + +template +om::Tensor& om::Tensor::sub_(const Tensor& rhs) +{ return this->sub_out(rhs, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::mul_out(const Tensor& rhs, + Tensor& out, + const Stream& s) const +{ + _check_operand(rhs, "mul_out"); + _check_out(out, m_Shape, "mul_out"); + _check_alias_elementwise(out, &rhs, "mul_out"); if (this->device_type() == DEVICE_TYPE::CPU) mul_cpu(this->view(), rhs.view(), out.view()); else @@ -517,9 +655,34 @@ om::Tensor om::Tensor::mul(const Tensor& rhs } template -om::Tensor om::Tensor::div(const Tensor& rhs, const Stream& s) const +om::Tensor& om::Tensor::mul_out(const Tensor& rhs, + Tensor& out) const +{ return this->mul_out(rhs, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::mul(const Tensor& rhs, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->mul_out(rhs, out, s); + return out; +} + +template +om::Tensor& om::Tensor::mul_(const Tensor& rhs, const Stream& s) +{ return this->mul_out(rhs, *this, s); } + +template +om::Tensor& om::Tensor::mul_(const Tensor& rhs) +{ return this->mul_out(rhs, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::div_out(const Tensor& rhs, + Tensor& out, + const Stream& s) const +{ + _check_operand(rhs, "div_out"); + _check_out(out, m_Shape, "div_out"); + _check_alias_elementwise(out, &rhs, "div_out"); if (this->device_type() == DEVICE_TYPE::CPU) div_cpu(this->view(), rhs.view(), out.view()); else @@ -528,9 +691,33 @@ om::Tensor om::Tensor::div(const Tensor& rhs } template -om::Tensor om::Tensor::add(const value_type& scalar, const Stream& s) const +om::Tensor& om::Tensor::div_out(const Tensor& rhs, + Tensor& out) const +{ return this->div_out(rhs, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::div(const Tensor& rhs, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->div_out(rhs, out, s); + return out; +} + +template +om::Tensor& om::Tensor::div_(const Tensor& rhs, const Stream& s) +{ return this->div_out(rhs, *this, s); } + +template +om::Tensor& om::Tensor::div_(const Tensor& rhs) +{ return this->div_out(rhs, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::add_out(const value_type& scalar, + Tensor& out, + const Stream& s) const +{ + _check_out(out, m_Shape, "add_out"); + _check_alias_elementwise(out, nullptr, "add_out"); if (this->device_type() == DEVICE_TYPE::CPU) add_k_cpu(this->view(), scalar, out.view()); else @@ -539,9 +726,33 @@ om::Tensor om::Tensor::add(const value_type& scalar, con } template -om::Tensor om::Tensor::sub(const value_type& scalar, const Stream& s) const +om::Tensor& om::Tensor::add_out(const value_type& scalar, + Tensor& out) const +{ return this->add_out(scalar, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::add(const value_type& scalar, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->add_out(scalar, out, s); + return out; +} + +template +om::Tensor& om::Tensor::add_(const value_type& scalar, const Stream& s) +{ return this->add_out(scalar, *this, s); } + +template +om::Tensor& om::Tensor::add_(const value_type& scalar) +{ return this->add_out(scalar, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::sub_out(const value_type& scalar, + Tensor& out, + const Stream& s) const +{ + _check_out(out, m_Shape, "sub_out"); + _check_alias_elementwise(out, nullptr, "sub_out"); if (this->device_type() == DEVICE_TYPE::CPU) sub_k_cpu(this->view(), scalar, out.view()); else @@ -550,9 +761,33 @@ om::Tensor om::Tensor::sub(const value_type& scalar, con } template -om::Tensor om::Tensor::mul(const value_type& scalar, const Stream& s) const +om::Tensor& om::Tensor::sub_out(const value_type& scalar, + Tensor& out) const +{ return this->sub_out(scalar, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::sub(const value_type& scalar, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->sub_out(scalar, out, s); + return out; +} + +template +om::Tensor& om::Tensor::sub_(const value_type& scalar, const Stream& s) +{ return this->sub_out(scalar, *this, s); } + +template +om::Tensor& om::Tensor::sub_(const value_type& scalar) +{ return this->sub_out(scalar, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::mul_out(const value_type& scalar, + Tensor& out, + const Stream& s) const +{ + _check_out(out, m_Shape, "mul_out"); + _check_alias_elementwise(out, nullptr, "mul_out"); if (this->device_type() == DEVICE_TYPE::CPU) mul_k_cpu(this->view(), scalar, out.view()); else @@ -561,9 +796,33 @@ om::Tensor om::Tensor::mul(const value_type& scalar, con } template -om::Tensor om::Tensor::div(const value_type& scalar, const Stream& s) const +om::Tensor& om::Tensor::mul_out(const value_type& scalar, + Tensor& out) const +{ return this->mul_out(scalar, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::mul(const value_type& scalar, const Stream& s) const { Tensor out(this->shape(), this->device(), Stream(s.get())); + this->mul_out(scalar, out, s); + return out; +} + +template +om::Tensor& om::Tensor::mul_(const value_type& scalar, const Stream& s) +{ return this->mul_out(scalar, *this, s); } + +template +om::Tensor& om::Tensor::mul_(const value_type& scalar) +{ return this->mul_out(scalar, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::div_out(const value_type& scalar, + Tensor& out, + const Stream& s) const +{ + _check_out(out, m_Shape, "div_out"); + _check_alias_elementwise(out, nullptr, "div_out"); if (this->device_type() == DEVICE_TYPE::CPU) div_k_cpu(this->view(), scalar, out.view()); else @@ -572,7 +831,67 @@ om::Tensor om::Tensor::div(const value_type& scalar, con } template -om::Tensor om::Tensor::matmul(const Tensor& rhs, const Stream& s) const +om::Tensor& om::Tensor::div_out(const value_type& scalar, + Tensor& out) const +{ return this->div_out(scalar, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::div(const value_type& scalar, const Stream& s) const +{ + Tensor out(this->shape(), this->device(), Stream(s.get())); + this->div_out(scalar, out, s); + return out; +} + +template +om::Tensor& om::Tensor::div_(const value_type& scalar, const Stream& s) +{ return this->div_out(scalar, *this, s); } + +template +om::Tensor& om::Tensor::div_(const value_type& scalar) +{ return this->div_out(scalar, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::operator+=(const Tensor& rhs) +{ return this->add_(rhs); } + +template +om::Tensor& om::Tensor::operator+=(const value_type& scalar) +{ return this->add_(scalar); } + +template +om::Tensor& om::Tensor::operator-=(const Tensor& rhs) +{ return this->sub_(rhs); } + +template +om::Tensor& om::Tensor::operator-=(const value_type& scalar) +{ return this->sub_(scalar); } + +template +om::Tensor& om::Tensor::operator*=(const Tensor& rhs) +{ return this->mul_(rhs); } + +template +om::Tensor& om::Tensor::operator*=(const value_type& scalar) +{ return this->mul_(scalar); } + +template +om::Tensor& om::Tensor::operator/=(const Tensor& rhs) +{ return this->div_(rhs); } + +template +om::Tensor& om::Tensor::operator/=(const value_type& scalar) +{ return this->div_(scalar); } + +// ── matmul / transpose / permute ──────────────────────────────────────────── +// Same destination-provided core, minus the in-place forms: each of these +// kernels reads elements it does not write, so a destination that shares a +// buffer with an operand would read values the kernel has already overwritten. + +template +om::Tensor& om::Tensor::matmul_out(const Tensor& rhs, + Tensor& out, + const Stream& s) const { if (this->rank() != 2 || rhs.rank() != 2) throw std::runtime_error("matmul: both tensors must be 2D matrices"); @@ -580,7 +899,10 @@ om::Tensor om::Tensor::matmul(const Tensor& if (K != K2) throw std::runtime_error("matmul: inner dimensions must match (A.cols=" + std::to_string(K) + " != B.rows=" + std::to_string(K2) + ")"); - Tensor out({M, N}, this->device(), Stream(s.get())); + if (rhs.device_type() != this->device_type() || rhs.device().m_Id != m_Device.m_Id) + throw std::invalid_argument("matmul_out: tensors must live on the same device"); + _check_out(out, {M, N}, "matmul_out"); + _check_alias_none(out, &rhs, "matmul_out"); if (this->device_type() == DEVICE_TYPE::CPU) matmul_cpu(this->view(), rhs.view(), out.view()); else @@ -589,12 +911,30 @@ om::Tensor om::Tensor::matmul(const Tensor& } template -om::Tensor om::Tensor::transpose(const Stream& s) const +om::Tensor& om::Tensor::matmul_out(const Tensor& rhs, + Tensor& out) const +{ return this->matmul_out(rhs, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::matmul(const Tensor& rhs, const Stream& s) const +{ + if (this->rank() != 2 || rhs.rank() != 2) + throw std::runtime_error("matmul: both tensors must be 2D matrices"); + if (m_Shape[1] != rhs.shape()[0]) + throw std::runtime_error("matmul: inner dimensions must match (A.cols=" + + std::to_string(m_Shape[1]) + " != B.rows=" + std::to_string(rhs.shape()[0]) + ")"); + Tensor out({m_Shape[0], rhs.shape()[1]}, this->device(), Stream(s.get())); + this->matmul_out(rhs, out, s); + return out; +} + +template +om::Tensor& om::Tensor::transpose_out(Tensor& out, const Stream& s) const { if (this->rank() != 2) throw std::runtime_error("transpose: tensor must be rank-2 (use permute for higher ranks)"); - size_t M = m_Shape[0], N = m_Shape[1]; - Tensor out({N, M}, this->device(), Stream(s.get())); + _check_out(out, {m_Shape[1], m_Shape[0]}, "transpose_out"); + _check_alias_none(out, nullptr, "transpose_out"); if (this->device_type() == DEVICE_TYPE::CPU) transpose_cpu(this->view(), out.view()); else @@ -603,32 +943,55 @@ om::Tensor om::Tensor::transpose(const Stream& s) const } template -om::Tensor om::Tensor::permute(const std::vector& axes, const Stream& s) const +om::Tensor& om::Tensor::transpose_out(Tensor& out) const +{ return this->transpose_out(out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::transpose(const Stream& s) const { - size_t r = this->rank(); - if (axes.size() != r) - throw std::invalid_argument("permute: axes length must match tensor rank"); - std::vector seen(r, false); - for (size_t a : axes) { - if (a >= r) throw std::out_of_range("permute: axis value out of range"); - if (seen[a]) throw std::invalid_argument("permute: duplicate axis"); - seen[a] = true; - } - std::vector out_shape(r); - for (size_t d = 0; d < r; ++d) out_shape[d] = m_Shape[axes[d]]; - Tensor out(out_shape, this->device(), Stream(s.get())); + if (this->rank() != 2) + throw std::runtime_error("transpose: tensor must be rank-2 (use permute for higher ranks)"); + Tensor out({m_Shape[1], m_Shape[0]}, this->device(), Stream(s.get())); + this->transpose_out(out, s); + return out; +} + +template +om::Tensor& om::Tensor::permute_out(const std::vector& axes, + Tensor& out, + const Stream& s) const +{ + const std::vector out_shape = detail::permuted_shape(m_Shape, axes); + _check_out(out, out_shape, "permute_out"); + _check_alias_none(out, nullptr, "permute_out"); if (this->device_type() == DEVICE_TYPE::CPU) - permute_cpu(this->view(), out.view(), axes.data(), r); + permute_cpu(this->view(), out.view(), axes.data(), this->rank()); else - launch_permute(this->view(), out.view(), axes.data(), r, s.get()); + launch_permute(this->view(), out.view(), axes.data(), this->rank(), s.get()); + return out; +} + +template +om::Tensor& om::Tensor::permute_out(const std::vector& axes, + Tensor& out) const +{ return this->permute_out(axes, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::permute(const std::vector& axes, const Stream& s) const +{ + Tensor out(detail::permuted_shape(m_Shape, axes), this->device(), Stream(s.get())); + this->permute_out(axes, out, s); return out; } +// ── apply / apply_binary / fill ───────────────────────────────────────────── + template template -om::Tensor om::Tensor::apply(Op op, const Stream& s) const +om::Tensor& om::Tensor::apply_out(Op op, Tensor& out, const Stream& s) const { - Tensor out(this->shape(), this->device(), Stream(s.get())); + _check_out(out, m_Shape, "apply_out"); + _check_alias_elementwise(out, nullptr, "apply_out"); if (this->device_type() == DEVICE_TYPE::CPU) { auto src = this->view(); auto dst = out.view(); @@ -642,18 +1005,125 @@ om::Tensor om::Tensor::apply(Op op, const Stream& s) con } template -om::Tensor om::Tensor::relu(const Stream& s) const +template +om::Tensor& om::Tensor::apply_out(Op op, Tensor& out) const +{ return this->apply_out(op, out, Stream::default_stream()); } + +template +template +om::Tensor om::Tensor::apply(Op op, const Stream& s) const +{ + Tensor out(this->shape(), this->device(), Stream(s.get())); + this->apply_out(op, out, s); + return out; +} + +template +template +om::Tensor& om::Tensor::apply_(Op op, const Stream& s) +{ return this->apply_out(op, *this, s); } + +template +template +om::Tensor& om::Tensor::apply_(Op op) +{ return this->apply_out(op, *this, Stream::default_stream()); } + +template +template +om::Tensor& om::Tensor::apply_binary_out(const Tensor& rhs, Op op, + Tensor& out, + const Stream& s) const { - return this->apply(ReLU{}, s); + _check_operand(rhs, "apply_binary_out"); + _check_out(out, m_Shape, "apply_binary_out"); + _check_alias_elementwise(out, &rhs, "apply_binary_out"); + if (this->device_type() == DEVICE_TYPE::CPU) { + auto lhs_v = this->view(); + auto rhs_v = rhs.view(); + auto dst_v = out.view(); + size_t n = lhs_v.size(); + for (size_t i = 0; i < n; ++i) + dst_v[i] = op(lhs_v[i], rhs_v[i]); + } else { + launch_apply_binary_op(this->view(), rhs.view(), out.view(), op, s.get()); + } + return out; } +template +template +om::Tensor& om::Tensor::apply_binary_out(const Tensor& rhs, Op op, + Tensor& out) const +{ return this->apply_binary_out(rhs, op, out, Stream::default_stream()); } + +template +template +om::Tensor& om::Tensor::apply_binary_(const Tensor& rhs, Op op, const Stream& s) +{ return this->apply_binary_out(rhs, op, *this, s); } + +template +template +om::Tensor& om::Tensor::apply_binary_(const Tensor& rhs, Op op) +{ return this->apply_binary_out(rhs, op, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::relu_out(Tensor& out, const Stream& s) const +{ return this->apply_out(ReLU{}, out, s); } + +template +om::Tensor& om::Tensor::relu_out(Tensor& out) const +{ return this->apply_out(ReLU{}, out, Stream::default_stream()); } + +template +om::Tensor om::Tensor::relu(const Stream& s) const +{ return this->apply(ReLU{}, s); } + +template +om::Tensor& om::Tensor::relu_(const Stream& s) +{ return this->apply_out(ReLU{}, *this, s); } + +template +om::Tensor& om::Tensor::relu_() +{ return this->apply_out(ReLU{}, *this, Stream::default_stream()); } + +template +om::Tensor& om::Tensor::sigmoid_out(Tensor& out, const Stream& s) const +{ return this->apply_out(Sigmoid{}, out, s); } + +template +om::Tensor& om::Tensor::sigmoid_out(Tensor& out) const +{ return this->apply_out(Sigmoid{}, out, Stream::default_stream()); } + template om::Tensor om::Tensor::sigmoid(const Stream& s) const +{ return this->apply(Sigmoid{}, s); } + +template +om::Tensor& om::Tensor::sigmoid_(const Stream& s) +{ return this->apply_out(Sigmoid{}, *this, s); } + +template +om::Tensor& om::Tensor::sigmoid_() +{ return this->apply_out(Sigmoid{}, *this, Stream::default_stream()); } + +// fill_ writes into the tensor it is called on by definition, so it is the +// only member of the in-place family with no allocating counterpart. It is +// also the last op that went through the legacy _fill dispatch struct, which +// has no cudaStream_t parameter; going direct is what gives it a stream. +template +om::Tensor& om::Tensor::fill_(const value_type& value, const Stream& s) { - return this->apply(Sigmoid{}, s); + if (this->device_type() == DEVICE_TYPE::CPU) + fill_cpu(this->view(), value); + else + launch_fill(this->view(), value, s.get()); + return *this; } -// ── End stream overload implementations ────────────────────────────────────── +template +om::Tensor& om::Tensor::fill_(const value_type& value) +{ return this->fill_(value, Stream::default_stream()); } + template om::Tensor om::Tensor::transpose() const diff --git a/python/README.md b/python/README.md index 07aec59..3c34430 100644 --- a/python/README.md +++ b/python/README.md @@ -77,7 +77,8 @@ s.synchronize() ``` `stream=` is accepted by `add`, `sub`, `mul`, `div`, `matmul`, `transpose`, -`permute`, `relu`, `sigmoid`, `cpu`, `cuda`, `to` and `Tensor.from_list`. +`permute`, `relu`, `sigmoid`, `cpu`, `cuda`, `to`, `fill`, `Tensor.from_list`, +and every member of the in-place / `_out` families below. Memory a tensor got from a stream's pool must be freed on that same stream, so every result holds a reference on the stream it was produced on. The reference @@ -85,6 +86,48 @@ count lives in the C layer, which is why `Stream.close()` is safe even while tensors from that stream are still alive — the `cudaStream_t` goes away when the last of them does. +## In-place ops and `out=` + +Every op above allocates its result. In a loop that is one allocation and one +free per op per iteration, and a peak footprint equal to the whole expression. +The trailing-underscore family writes back into the tensor it is called on, and +the `_out` family into a destination you already own; both return the tensor +they wrote to, so calls chain. + +```python +w = om.Tensor.zeros([1024], device="cuda") +g = om.Tensor.full([1024], 0.5, device="cuda") + +for _ in range(steps): + w.add_(g) # w's buffer never moves + w.relu_() + +w -= g # += -= *= /= are the same thing +w *= 0.9 + +out = om.Tensor.zeros([1024], device="cuda") +for batch in batches: # one destination for the whole loop + a.mul_out(b, out) +``` + +In-place: `add_`, `sub_`, `mul_`, `div_` (tensor or scalar), `relu_`, +`sigmoid_`, `fill_`, and `+= -= *= /=`. + +Destination-provided: `add_out`, `sub_out`, `mul_out`, `div_out`, `relu_out`, +`sigmoid_out`, `matmul_out`, `transpose_out`, `permute_out`. `out` must already +have the result's shape, dtype and device. + +`matmul`, `transpose` and `permute` have no in-place form and reject a +destination that is one of their operands — their kernels read elements they do +not write, so the answer would be wrong. The elementwise family is free to +alias, which is exactly what `add_` is. + +One caveat when you pass both `out` and `stream`: memory a tensor got from a +stream's pool must be used in an order that stream can see, so enqueueing into a +destination allocated on a *different* stream is only correct once you have +ordered the two yourself. The allocating forms cannot hit this — they allocate +on the stream they run on. + ## API | group | members | @@ -94,6 +137,8 @@ last of them does. | data | `numpy()`, `tolist()`, `flat()`, `item()`, `fill()`, `copy()`, `t[i, j]`, `t[i, j] = v` | | device | `cpu()`, `cuda()`, `to(device)`, `astype()`, `synchronize()` | | arithmetic | `+ - * / @`, reflected and scalar forms, `add`/`sub`/`mul`/`div`/`matmul` | +| in-place | `add_`, `sub_`, `mul_`, `div_`, `relu_`, `sigmoid_`, `fill_`, `+= -= *= /=` | +| destination-provided | `add_out`, `sub_out`, `mul_out`, `div_out`, `relu_out`, `sigmoid_out`, `matmul_out`, `transpose_out`, `permute_out` | | reductions | `sum()`, `mean()`, `min()`, `max()` | | shape | `reshape`, `flatten`, `squeeze`, `unsqueeze`, `transpose`/`T`, `permute` | | fused | `relu`, `sigmoid`, `scale_shift`, `shift_scale`, `fused_add_mul`, `fused_sub_mul`, `fused_mul_add`, `fused_div_add` | diff --git a/python/openmat/_clib.py b/python/openmat/_clib.py index c101f01..e343c1a 100644 --- a/python/openmat/_clib.py +++ b/python/openmat/_clib.py @@ -132,6 +132,30 @@ def d(name, restype, argtypes): d(f"{name}_scalar", _p, [_p, C, _cp, _i]) d(f"{name}_scalar_stream", _p, [_p, C, _p, _cp, _i]) + # in-place and destination-provided ops — all int-returning: nothing is + # allocated and no handle changes hands, so there is no pointer to check + for name in ("add", "sub", "mul", "div"): + d(f"{name}_inplace", _i, [_p, _p, _cp, _i]) + d(f"{name}_inplace_stream", _i, [_p, _p, _p, _cp, _i]) + d(f"{name}_out", _i, [_p, _p, _p, _cp, _i]) + d(f"{name}_out_stream", _i, [_p, _p, _p, _p, _cp, _i]) + d(f"{name}_scalar_inplace", _i, [_p, C, _cp, _i]) + d(f"{name}_scalar_inplace_stream", _i, [_p, C, _p, _cp, _i]) + d(f"{name}_scalar_out", _i, [_p, C, _p, _cp, _i]) + d(f"{name}_scalar_out_stream", _i, [_p, C, _p, _p, _cp, _i]) + d("matmul_out", _i, [_p, _p, _p, _cp, _i]) + d("matmul_out_stream", _i, [_p, _p, _p, _p, _cp, _i]) + for name in ("relu", "sigmoid"): + d(f"{name}_inplace", _i, [_p, _cp, _i]) + d(f"{name}_inplace_stream", _i, [_p, _p, _cp, _i]) + d(f"{name}_out", _i, [_p, _p, _cp, _i]) + d(f"{name}_out_stream", _i, [_p, _p, _p, _cp, _i]) + d("transpose_out", _i, [_p, _p, _cp, _i]) + d("transpose_out_stream", _i, [_p, _p, _p, _cp, _i]) + d("permute_out", _i, [_p, _sp, _sz, _p, _cp, _i]) + d("permute_out_stream", _i, [_p, _sp, _sz, _p, _p, _cp, _i]) + d("fill_stream", _i, [_p, C, _p, _cp, _i]) + # reductions for name in ("sum", "mean", "min", "max"): d(name, C, [_p, _cp, _i]) diff --git a/python/openmat/tensor.py b/python/openmat/tensor.py index af25530..4fbf6c3 100644 --- a/python/openmat/tensor.py +++ b/python/openmat/tensor.py @@ -409,12 +409,20 @@ def copy_to_device(self, dest: "Tensor") -> "Tensor": _check_int(self._fn("to_device")(self._h, dest.data_ptr(), eb, _ERR_LEN), eb) return dest - def fill(self, value: Scalar) -> "Tensor": + def fill(self, value: Scalar, stream: Optional[Stream] = None) -> "Tensor": eb = _errbuf() - _check_int(self._fn("fill")(self._h, self._dt.ctype(self._dt.py_type(value)), - eb, _ERR_LEN), eb) + val = self._dt.ctype(self._dt.py_type(value)) + if stream is None: + _check_int(self._fn("fill")(self._h, val, eb, _ERR_LEN), eb) + else: + _check_int(self._fn("fill_stream")(self._h, val, _handle_of(stream), + eb, _ERR_LEN), eb) return self + def fill_(self, value: Scalar, stream: Optional[Stream] = None) -> "Tensor": + """Alias of fill(), spelled like the rest of the in-place family.""" + return self.fill(value, stream) + def _index_array(self, key): idx = key if isinstance(key, tuple) else (key,) rank = self.rank @@ -523,6 +531,184 @@ def div(self, other, stream: Optional[Stream] = None) -> "Tensor": def matmul(self, other: "Tensor", stream: Optional[Stream] = None) -> "Tensor": return self._binop_tt(other, "matmul", stream) + # ── in-place and destination-provided ops ───────────────────────────── + # + # Every method here writes into a buffer that already exists — this + # tensor's own for the trailing-underscore family, the caller's `out` for + # the `_out` family — and returns the tensor it wrote to, so calls chain. + # No handle is created or destroyed, which is the point: a loop that runs + # the same op every iteration stops allocating one result per iteration. + # + # A destination that is not freshly allocated carries a caveat the + # allocating forms do not: memory from cudaMallocAsync is stream-ordered, + # so passing a `stream` other than the one `out` (or self) was allocated + # on is only safe once the caller has ordered the two. + + def _check_out(self, out: "Tensor"): + self._same_dtype(out) + if out._h is None: + raise ValueError("destination tensor has been closed") + + def _inplace_tt(self, other: "Tensor", name: str, + stream: Optional[Stream] = None) -> "Tensor": + self._same_dtype(other) + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_inplace")(self._h, other._h, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_inplace_stream")(self._h, other._h, + _handle_of(stream), eb, _ERR_LEN) + _check_int(rc, eb) + return self + + def _inplace_ts(self, scalar: Scalar, name: str, + stream: Optional[Stream] = None) -> "Tensor": + if not isinstance(scalar, (int, float)): + raise TypeError( + f"expected a Tensor or a number, got {type(scalar).__name__}") + val = self._dt.ctype(self._dt.py_type(scalar)) + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_scalar_inplace")(self._h, val, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_scalar_inplace_stream")(self._h, val, + _handle_of(stream), eb, _ERR_LEN) + _check_int(rc, eb) + return self + + def _dispatch_inplace(self, other, name, stream=None) -> "Tensor": + if isinstance(other, Tensor): + return self._inplace_tt(other, name, stream) + return self._inplace_ts(other, name, stream) + + def add_(self, other, stream: Optional[Stream] = None) -> "Tensor": + """self += other, without allocating. Returns self.""" + return self._dispatch_inplace(other, "add", stream) + + def sub_(self, other, stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_inplace(other, "sub", stream) + + def mul_(self, other, stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_inplace(other, "mul", stream) + + def div_(self, other, stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_inplace(other, "div", stream) + + def relu_(self, stream: Optional[Stream] = None) -> "Tensor": + return self._inplace_unary("relu", stream) + + def sigmoid_(self, stream: Optional[Stream] = None) -> "Tensor": + return self._inplace_unary("sigmoid", stream) + + def _inplace_unary(self, name: str, stream: Optional[Stream] = None) -> "Tensor": + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_inplace")(self._h, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_inplace_stream")(self._h, _handle_of(stream), + eb, _ERR_LEN) + _check_int(rc, eb) + return self + + def __iadd__(self, other): + return self._dispatch_inplace(other, "add") if _is_operand(other) else NotImplemented + + def __isub__(self, other): + return self._dispatch_inplace(other, "sub") if _is_operand(other) else NotImplemented + + def __imul__(self, other): + return self._dispatch_inplace(other, "mul") if _is_operand(other) else NotImplemented + + def __itruediv__(self, other): + return self._dispatch_inplace(other, "div") if _is_operand(other) else NotImplemented + + def _out_tt(self, other: "Tensor", name: str, out: "Tensor", + stream: Optional[Stream] = None) -> "Tensor": + self._same_dtype(other) + self._check_out(out) + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_out")(self._h, other._h, out._h, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_out_stream")(self._h, other._h, out._h, + _handle_of(stream), eb, _ERR_LEN) + _check_int(rc, eb) + return out + + def _out_ts(self, scalar: Scalar, name: str, out: "Tensor", + stream: Optional[Stream] = None) -> "Tensor": + if not isinstance(scalar, (int, float)): + raise TypeError( + f"expected a Tensor or a number, got {type(scalar).__name__}") + self._check_out(out) + val = self._dt.ctype(self._dt.py_type(scalar)) + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_scalar_out")(self._h, val, out._h, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_scalar_out_stream")(self._h, val, out._h, + _handle_of(stream), eb, _ERR_LEN) + _check_int(rc, eb) + return out + + def _dispatch_out(self, other, name, out, stream=None) -> "Tensor": + if isinstance(other, Tensor): + return self._out_tt(other, name, out, stream) + return self._out_ts(other, name, out, stream) + + def add_out(self, other, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + """Write self + other into `out` (same shape and device). Returns out.""" + return self._dispatch_out(other, "add", out, stream) + + def sub_out(self, other, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_out(other, "sub", out, stream) + + def mul_out(self, other, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_out(other, "mul", out, stream) + + def div_out(self, other, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._dispatch_out(other, "div", out, stream) + + def matmul_out(self, other: "Tensor", out: "Tensor", + stream: Optional[Stream] = None) -> "Tensor": + """Write self @ other into `out`. `out` may not be an operand.""" + return self._out_tt(other, "matmul", out, stream) + + def _out_unary(self, name: str, out: "Tensor", + stream: Optional[Stream] = None) -> "Tensor": + self._check_out(out) + eb = _errbuf() + if stream is None: + rc = self._fn(f"{name}_out")(self._h, out._h, eb, _ERR_LEN) + else: + rc = self._fn(f"{name}_out_stream")(self._h, out._h, _handle_of(stream), + eb, _ERR_LEN) + _check_int(rc, eb) + return out + + def relu_out(self, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._out_unary("relu", out, stream) + + def sigmoid_out(self, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._out_unary("sigmoid", out, stream) + + def transpose_out(self, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + return self._out_unary("transpose", out, stream) + + def permute_out(self, axes, out: "Tensor", stream: Optional[Stream] = None) -> "Tensor": + self._check_out(out) + if len(axes) == 1 and isinstance(axes[0], (list, tuple)): + axes = axes[0] + arr, n = _shape_array(axes) + eb = _errbuf() + if stream is None: + rc = self._fn("permute_out")(self._h, arr, n, out._h, eb, _ERR_LEN) + else: + rc = self._fn("permute_out_stream")(self._h, arr, n, out._h, + _handle_of(stream), eb, _ERR_LEN) + _check_int(rc, eb) + return out + # Operators return NotImplemented for foreign operands so Python can fall # back to the other object's reflected method (and raise the standard # TypeError if there is none). diff --git a/python/tests/test_inplace.py b/python/tests/test_inplace.py new file mode 100644 index 0000000..1d4e51c --- /dev/null +++ b/python/tests/test_inplace.py @@ -0,0 +1,222 @@ +"""In-place ops (add_, mul_, relu_, fill_) and the destination-provided family. + +Every test asserts two things: the values are right, and `data_ptr()` did not +move. The second is the point of the feature — an implementation that quietly +allocated a fresh result and rebound the handle would satisfy every value +assertion here. +""" +import pytest + +import openmat as om + +from .conftest import requires_cuda + + +def flat(t): + return t.cpu().flat() if t.is_cuda else t.flat() + + +# ── in-place ──────────────────────────────────────────────────────────────── + +def test_inplace_tensor_ops(device): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device=device) + b = om.Tensor.full([4], 2.0, device=device) + p = a.data_ptr() + + assert flat(a.add_(b)) == [3.0, 4.0, 5.0, 6.0] + assert flat(a.sub_(b)) == [1.0, 2.0, 3.0, 4.0] + assert flat(a.mul_(b)) == [2.0, 4.0, 6.0, 8.0] + assert flat(a.div_(b)) == [1.0, 2.0, 3.0, 4.0] + + assert a.data_ptr() == p + assert flat(b) == [2.0, 2.0, 2.0, 2.0], "rhs was modified" + + +def test_inplace_scalar_ops(device): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device=device) + p = a.data_ptr() + + a.add_(1.0).mul_(2.0).sub_(2.0).div_(2.0) + + assert flat(a) == [1.0, 2.0, 3.0, 4.0] + assert a.data_ptr() == p + + +def test_inplace_returns_self(device): + a = om.Tensor([1.0, 2.0], device=device) + assert a.add_(1.0) is a + assert a.relu_() is a + assert a.fill_(0.0) is a + + +def test_augmented_assignment_does_not_rebind(device): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device=device) + before = a + p = a.data_ptr() + + a += 1.0 + a *= 2.0 + a -= 2.0 + a /= 2.0 + + assert a is before, "+= rebound the name to a new tensor" + assert a.data_ptr() == p + assert flat(a) == [1.0, 2.0, 3.0, 4.0] + + +def test_inplace_unary_and_fill(device): + a = om.Tensor([-2.0, -1.0, 0.0, 3.0], device=device) + p = a.data_ptr() + + assert flat(a.relu_()) == [0.0, 0.0, 0.0, 3.0] + assert flat(a.fill_(5.0)) == [5.0, 5.0, 5.0, 5.0] + + a.sigmoid_() + assert all(0.0 < v < 1.0 for v in flat(a)) + assert a.data_ptr() == p + + +def test_inplace_matches_out_of_place(device): + a = om.Tensor([[-1.0, 2.0], [3.0, -4.0]], device=device) + b = om.Tensor([[0.5, 1.5], [2.5, 3.5]], device=device) + + want = flat((a * b).relu()) + a.mul_(b).relu_() + assert flat(a) == want + + +def test_inplace_self_operand(device): + a = om.Tensor([1.0, 2.0, 3.0], device=device) + a.add_(a) + assert flat(a) == [2.0, 4.0, 6.0] + + +def test_inplace_int_dtype(device): + a = om.Tensor([1, 2, 3, 4], dtype=om.int32, device=device) + a.mul_(3) + assert flat(a) == [3, 6, 9, 12] + + +def test_inplace_shape_mismatch_raises(device): + a = om.Tensor.zeros([4], device=device) + b = om.Tensor.zeros([5], device=device) + with pytest.raises(RuntimeError): + a.add_(b) + + +def test_inplace_dtype_mismatch_raises(): + a = om.Tensor.zeros([4], dtype=om.float32) + b = om.Tensor.zeros([4], dtype=om.int32) + with pytest.raises(TypeError): + a.add_(b) + + +# ── destination-provided ──────────────────────────────────────────────────── + +def test_out_reuses_one_destination(device): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device=device) + b = om.Tensor.full([4], 2.0, device=device) + out = om.Tensor.zeros([4], device=device) + p = out.data_ptr() + + assert flat(a.add_out(b, out)) == [3.0, 4.0, 5.0, 6.0] + assert flat(a.mul_out(b, out)) == [2.0, 4.0, 6.0, 8.0] + assert flat(a.sub_out(1.0, out)) == [0.0, 1.0, 2.0, 3.0] + assert flat(a.relu_out(out)) == [1.0, 2.0, 3.0, 4.0] + + assert out.data_ptr() == p + assert flat(a) == [1.0, 2.0, 3.0, 4.0], "operand was clobbered" + + +def test_out_returns_the_destination(device): + a = om.Tensor([1.0, 2.0], device=device) + out = om.Tensor.zeros([2], device=device) + assert a.add_out(1.0, out) is out + + +def test_out_matmul_transpose_permute(device): + a = om.Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], device=device) + b = om.Tensor.full([3, 2], 1.0, device=device) + + mm = om.Tensor.zeros([2, 2], device=device) + assert flat(a.matmul_out(b, mm)) == flat(a @ b) + + tr = om.Tensor.zeros([3, 2], device=device) + assert flat(a.transpose_out(tr)) == flat(a.T) + + pm = om.Tensor.zeros([3, 2], device=device) + assert flat(a.permute_out([1, 0], pm)) == flat(a.permute(1, 0)) + + +def test_out_wrong_shape_raises(device): + a = om.Tensor.zeros([4], device=device) + out = om.Tensor.zeros([5], device=device) + with pytest.raises(RuntimeError): + a.add_out(1.0, out) + + +@requires_cuda +def test_out_wrong_device_raises(): + a = om.Tensor.zeros([4], device="cuda") + out = om.Tensor.zeros([4], device="cpu") + with pytest.raises(RuntimeError): + a.add_out(1.0, out) + + +def test_out_aliased_destination_rejected_for_matmul(device): + a = om.Tensor([[1.0, 2.0], [3.0, 4.0]], device=device) + with pytest.raises(RuntimeError): + a.matmul_out(a, a) + with pytest.raises(RuntimeError): + a.transpose_out(a) + + +# ── streams ───────────────────────────────────────────────────────────────── + +@requires_cuda +def test_inplace_on_a_stream(): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device="cuda") + b = om.Tensor.full([4], 2.0, device="cuda") + p = a.data_ptr() + + with om.Stream() as s: + a.mul_(b, stream=s).add_(1.0, stream=s).relu_(stream=s) + + assert a.cpu().flat() == [3.0, 5.0, 7.0, 9.0] + assert a.data_ptr() == p + + +@requires_cuda +def test_out_on_a_stream(): + a = om.Tensor([1.0, 2.0, 3.0, 4.0], device="cuda") + b = om.Tensor.full([4], 2.0, device="cuda") + out = om.Tensor.zeros([4], device="cuda") + p = out.data_ptr() + + with om.Stream() as s: + a.add_out(b, out, stream=s) + + assert out.cpu().flat() == [3.0, 4.0, 5.0, 6.0] + assert out.data_ptr() == p + + +@requires_cuda +def test_fill_on_a_stream(): + a = om.Tensor.zeros([8], device="cuda") + with om.Stream() as s: + a.fill_(3.0, stream=s) + assert a.cpu().flat() == [3.0] * 8 + + +# ── the reason the feature exists ─────────────────────────────────────────── + +def test_loop_keeps_one_buffer(device): + w = om.Tensor.zeros([256], device=device) + g = om.Tensor.full([256], 0.5, device=device) + p = w.data_ptr() + + for _ in range(50): + w.add_(g) + assert w.data_ptr() == p + + assert flat(w)[0] == pytest.approx(25.0) diff --git a/scripts/bench_inplace.py b/scripts/bench_inplace.py new file mode 100644 index 0000000..05f4db5 --- /dev/null +++ b/scripts/bench_inplace.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python3 +""" +In-place / destination-provided ops vs the allocating forms. + +Answers two separate questions that the same table is easy to conflate: + + * what one op costs — `a + b` (allocate a result, compute, free it) against + `dst.add_(b)` and `a.add_out(b, dst)` on a buffer that already exists; + * what a *loop* costs — N iterations of the same op, which is where the + allocating form pays its cost N times and holds every intermediate alive. + +Each case gets its own freshly built operands: reusing one destination across +cases lets an earlier case leave values (or an allocator state) behind that +changes the next one's timing, which is exactly how an earlier draft of this +script produced a 0.38x "regression" that does not exist. + + OPENMAT_LIB=build-release/OpenMat.so PYTHONPATH=python \ + bench-env/bin/python scripts/bench_inplace.py +""" +from __future__ import annotations + +import argparse +import statistics +import time + +import openmat as om + +SIZES = [1 << 10, 1 << 16, 1 << 20, 1 << 24] + + +def measure(fn, sync, *, min_time=0.05, trials=7): + for _ in range(3): + fn() + sync() + + reps = 1 + while reps < 100_000: + sync() + t0 = time.perf_counter() + for _ in range(reps): + fn() + sync() + dt = time.perf_counter() - t0 + if dt >= min_time: + break + reps = int(reps * max(2.0, min(50.0, min_time / max(dt, 1e-9)))) + 1 + + samples = [] + for _ in range(trials): + sync() + t0 = time.perf_counter() + for _ in range(reps): + fn() + sync() + samples.append((time.perf_counter() - t0) / reps) + return min(samples) + + +def ops(device): + """(label, out-of-place, in-place, destination-provided) built per case.""" + def build(n): + a = om.Tensor.full([n], 1.5, device=device) + b = om.Tensor.full([n], 2.0, device=device) + d = om.Tensor.zeros([n], device=device) + return a, b, d + + return [ + ("add (tensor)", build, + lambda a, b, d: a + b, + lambda a, b, d: d.add_(b), + lambda a, b, d: a.add_out(b, d)), + ("mul (scalar)", build, + lambda a, b, d: a * 2.0, + lambda a, b, d: d.mul_(2.0), + lambda a, b, d: a.mul_out(2.0, d)), + ("relu", build, + lambda a, b, d: a.relu(), + lambda a, b, d: d.relu_(), + lambda a, b, d: a.relu_out(d)), + ] + + +def single_op_table(device, sync): + print(f"\n── {device}: one op ──") + print(f"{'n':>10} {'op':14} {'alloc us':>10} {'in-place us':>12} " + f"{'out= us':>10} {'in-place':>9} {'out=':>7}") + for n in SIZES: + for label, build, oop, inplace, outp in ops(device): + a, b, d = build(n) + t_alloc = measure(lambda: oop(a, b, d), sync) + t_in = measure(lambda: inplace(a, b, d), sync) + t_out = measure(lambda: outp(a, b, d), sync) + print(f"{n:>10} {label:14} {t_alloc*1e6:10.2f} {t_in*1e6:12.2f} " + f"{t_out*1e6:10.2f} {t_alloc/t_in:8.2f}x {t_alloc/t_out:6.2f}x") + + +def loop_table(device, sync, steps=32): + """A `steps`-long chain, the shape a training loop actually has.""" + print(f"\n── {device}: a {steps}-step chain (w += g each step) ──") + print(f"{'n':>10} {'alloc us':>10} {'in-place us':>12} {'speedup':>8}") + for n in SIZES: + g = om.Tensor.full([n], 0.5, device=device) + + def allocating(): + w = om.Tensor.zeros([n], device=device) + for _ in range(steps): + w = w + g + return w + + w_ip = om.Tensor.zeros([n], device=device) + + def in_place(): + w_ip.fill_(0.0) + for _ in range(steps): + w_ip.add_(g) + return w_ip + + t_a = measure(allocating, sync, trials=5) + t_i = measure(in_place, sync, trials=5) + print(f"{n:>10} {t_a*1e6:10.2f} {t_i*1e6:12.2f} {t_a/t_i:7.2f}x") + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--no-cuda", action="store_true") + args = ap.parse_args() + + devices = ["cpu"] + if not args.no_cuda and om.cuda_is_available(): + devices.append("cuda") + + for device in devices: + sync = om.synchronize if device == "cuda" else (lambda: None) + single_op_table(device, sync) + loop_table(device, sync) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/python/openmat_capi_impl.inc b/src/python/openmat_capi_impl.inc index ac86b8d..1391e1b 100644 --- a/src/python/openmat_capi_impl.inc +++ b/src/python/openmat_capi_impl.inc @@ -330,6 +330,212 @@ OM_DEFINE_BINOP_TS(div, div) #undef OM_DEFINE_BINOP_TS +// ───────────────────────────────────────────────────────────────────────────── +// In-place arithmetic and destination-provided ops +// +// These return int, not a handle: nothing is allocated and no ownership +// changes — the result lands in a tensor the caller already owns, and every +// handle involved stays valid. That is the point of them. A loop that would +// otherwise create and destroy one tensor per op per iteration keeps a single +// buffer for its whole life. +// +// `_out` takes the destination handle last. It may be one of the operands for +// the elementwise family (that is what the in-place form is); matmul_out, +// transpose_out and permute_out reject an aliased destination on the C++ side. +// ───────────────────────────────────────────────────────────────────────────── + +#define OM_DEFINE_INPLACE_TT(name, method) \ +int OM_FN(name##_inplace)(void* lhs, const void* rhs, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(lhs)->method##_(*static_cast(rhs)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_inplace_stream)(void* lhs, const void* rhs, void* stream, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(lhs)->method##_(*static_cast(rhs), \ + om_deref_stream(stream)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_out)(const void* lhs, const void* rhs, void* out, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(lhs)->method##_out( \ + *static_cast(rhs), *static_cast(out)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_out_stream)(const void* lhs, const void* rhs, void* out, \ + void* stream, char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(lhs)->method##_out( \ + *static_cast(rhs), *static_cast(out), \ + om_deref_stream(stream)); \ + return 0; \ + ) \ +} + +OM_DEFINE_INPLACE_TT(add, add) +OM_DEFINE_INPLACE_TT(sub, sub) +OM_DEFINE_INPLACE_TT(mul, mul) +OM_DEFINE_INPLACE_TT(div, div) + +#undef OM_DEFINE_INPLACE_TT + +// matmul has no in-place form: its kernel reads elements it does not write. +int OM_FN(matmul_out)(const void* lhs, const void* rhs, void* out, + char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(lhs)->matmul_out(*static_cast(rhs), + *static_cast(out)); + return 0; + ) +} + +int OM_FN(matmul_out_stream)(const void* lhs, const void* rhs, void* out, void* stream, + char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(lhs)->matmul_out(*static_cast(rhs), + *static_cast(out), + om_deref_stream(stream)); + return 0; + ) +} + +#define OM_DEFINE_INPLACE_TS(name, method) \ +int OM_FN(name##_scalar_inplace)(void* handle, OM_T scalar, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( static_cast(handle)->method##_(scalar); return 0; ) \ +} \ +int OM_FN(name##_scalar_inplace_stream)(void* handle, OM_T scalar, void* stream, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->method##_(scalar, om_deref_stream(stream)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_scalar_out)(const void* handle, OM_T scalar, void* out, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->method##_out(scalar, \ + *static_cast(out)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_scalar_out_stream)(const void* handle, OM_T scalar, void* out, \ + void* stream, char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->method##_out(scalar, \ + *static_cast(out), \ + om_deref_stream(stream)); \ + return 0; \ + ) \ +} + +OM_DEFINE_INPLACE_TS(add, add) +OM_DEFINE_INPLACE_TS(sub, sub) +OM_DEFINE_INPLACE_TS(mul, mul) +OM_DEFINE_INPLACE_TS(div, div) + +#undef OM_DEFINE_INPLACE_TS + +#define OM_DEFINE_INPLACE_UNARY(name) \ +int OM_FN(name##_inplace)(void* handle, char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( static_cast(handle)->name##_(); return 0; ) \ +} \ +int OM_FN(name##_inplace_stream)(void* handle, void* stream, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->name##_(om_deref_stream(stream)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_out)(const void* handle, void* out, char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->name##_out(*static_cast(out)); \ + return 0; \ + ) \ +} \ +int OM_FN(name##_out_stream)(const void* handle, void* out, void* stream, \ + char* errbuf, int errbuf_len) \ +{ \ + OM_GUARD_INT( \ + static_cast(handle)->name##_out(*static_cast(out), \ + om_deref_stream(stream)); \ + return 0; \ + ) \ +} + +OM_DEFINE_INPLACE_UNARY(relu) +OM_DEFINE_INPLACE_UNARY(sigmoid) + +#undef OM_DEFINE_INPLACE_UNARY + +int OM_FN(transpose_out)(const void* handle, void* out, char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(handle)->transpose_out(*static_cast(out)); + return 0; + ) +} + +int OM_FN(transpose_out_stream)(const void* handle, void* out, void* stream, + char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(handle)->transpose_out(*static_cast(out), + om_deref_stream(stream)); + return 0; + ) +} + +int OM_FN(permute_out)(const void* handle, const size_t* axes, size_t n, void* out, + char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(handle)->permute_out( + std::vector(axes, axes + n), *static_cast(out)); + return 0; + ) +} + +int OM_FN(permute_out_stream)(const void* handle, const size_t* axes, size_t n, void* out, + void* stream, char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(handle)->permute_out( + std::vector(axes, axes + n), *static_cast(out), + om_deref_stream(stream)); + return 0; + ) +} + +// fill() has always existed; this is the same thing with a stream. +int OM_FN(fill_stream)(void* handle, OM_T value, void* stream, + char* errbuf, int errbuf_len) +{ + OM_GUARD_INT( + static_cast(handle)->fill_(value, om_deref_stream(stream)); + return 0; + ) +} + // ───────────────────────────────────────────────────────────────────────────── // Reductions — synchronous, return a host scalar // ───────────────────────────────────────────────────────────────────────────── diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b3baa1c..91590c9 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -27,3 +27,4 @@ add_om_test(test_stress test_stress.cpp) add_om_test(test_stream_perf test_stream_perf.cpp) add_om_test(test_host_pool test_host_pool.cpp) add_om_test(test_contiguous test_contiguous.cpp) +add_om_test(test_inplace test_inplace.cpp) diff --git a/tests/test_inplace.cpp b/tests/test_inplace.cpp new file mode 100644 index 0000000..43c6ed1 --- /dev/null +++ b/tests/test_inplace.cpp @@ -0,0 +1,363 @@ +// In-place ops (add_, mul_, relu_, fill_, …) and the destination-provided +// overloads (add_out, matmul_out, …). +// +// Two things are being asserted throughout, and the second is the point of the +// feature: that the result is right, and that producing it did not move the +// tensor's buffer. data_ptr equality before and after is what proves no +// allocation happened — a correct-but-reallocating implementation would pass +// every value check here. +#include "test_helpers.h" +#include + +namespace { + +Tensor host_seq(const std::vector& shape, float start = 1.0f) +{ + size_t n = std::accumulate(shape.begin(), shape.end(), size_t{1}, std::multiplies<>()); + std::vector v(n); + for (size_t i = 0; i < n; ++i) v[i] = start + static_cast(i); + return Tensor::from_vector(v, shape); +} + +Tensor dev_seq(const std::vector& shape, float start = 1.0f) +{ + return host_seq(shape, start).cuda(); +} + +const void* buf(const Tensor& t) { return t.view().data; } + +} // namespace + +// ── in-place, CPU ─────────────────────────────────────────────────────────── + +TEST(InPlace, CPUTensorOps) { + Tensor a = host_seq({4}); // 1 2 3 4 + Tensor b = Tensor::full({4}, 2.0f); + const void* p = buf(a); + + a.add_(b); + EXPECT_EQ(to_host(a), (std::vector{3, 4, 5, 6})); + a.sub_(b); + EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})); + a.mul_(b); + EXPECT_EQ(to_host(a), (std::vector{2, 4, 6, 8})); + a.div_(b); + EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})); + + EXPECT_EQ(buf(a), p) << "in-place op reallocated the buffer"; +} + +TEST(InPlace, CPUScalarOps) { + Tensor a = host_seq({4}); + const void* p = buf(a); + + a.add_(1.0f).mul_(2.0f).sub_(2.0f).div_(2.0f); + EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})); + EXPECT_EQ(buf(a), p); +} + +TEST(InPlace, CPUCompoundOperators) { + Tensor a = host_seq({4}); + Tensor b = Tensor::full({4}, 3.0f); + const void* p = buf(a); + + a += b; EXPECT_EQ(to_host(a), (std::vector{4, 5, 6, 7})); + a -= b; EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})); + a *= 2.0f; EXPECT_EQ(to_host(a), (std::vector{2, 4, 6, 8})); + a /= 2.0f; EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})); + + EXPECT_EQ(buf(a), p); +} + +TEST(InPlace, CPUUnaryAndFill) { + Tensor a = Tensor::from_vector({-2.0f, -1.0f, 0.0f, 3.0f}, {4}); + const void* p = buf(a); + + a.relu_(); + EXPECT_EQ(to_host(a), (std::vector{0, 0, 0, 3})); + + a.fill_(7.0f); + EXPECT_EQ(to_host(a), (std::vector{7, 7, 7, 7})); + + a.apply_(Mul{0.5f}); + EXPECT_EQ(to_host(a), (std::vector{3.5f, 3.5f, 3.5f, 3.5f})); + + Tensor b = Tensor::full({4}, 2.0f); + a.apply_binary_(b, BinaryMul{}); + EXPECT_EQ(to_host(a), (std::vector{7, 7, 7, 7})); + + EXPECT_EQ(buf(a), p); +} + +TEST(InPlace, CPUSigmoidMatchesOutOfPlace) { + Tensor a = Tensor::from_vector({-1.0f, 0.0f, 1.0f, 2.0f}, {4}); + Tensor expected = a.sigmoid(); + a.sigmoid_(); + const auto got = to_host(a), want = to_host(expected); + for (size_t i = 0; i < got.size(); ++i) EXPECT_FLOAT_EQ(got[i], want[i]); +} + +// The rhs must survive the call unchanged even when it shares nothing but the +// shape — the elementwise loop writes only into the destination. +TEST(InPlace, RhsIsNotModified) { + Tensor a = host_seq({4}); + Tensor b = host_seq({4}, 10.0f); + a.add_(b); + EXPECT_EQ(to_host(b), (std::vector{10, 11, 12, 13})); +} + +// x.add_(x) is the degenerate case: destination and *both* operands alias. +TEST(InPlace, SelfOperandDoubles) { + Tensor a = host_seq({4}); + a.add_(a); + EXPECT_EQ(to_host(a), (std::vector{2, 4, 6, 8})); +} + +// ── in-place, GPU ─────────────────────────────────────────────────────────── + +TEST(InPlace, GPUTensorOps) { + OM_REQUIRE_CUDA(); + Tensor a = dev_seq({4}); + Tensor b = Tensor::full({4}, 2.0f).cuda(); + const void* p = buf(a); + + a.add_(b); + EXPECT_EQ(to_host(a.cpu()), (std::vector{3, 4, 5, 6})); + a.mul_(b); + EXPECT_EQ(to_host(a.cpu()), (std::vector{6, 8, 10, 12})); + a.sub_(b); + EXPECT_EQ(to_host(a.cpu()), (std::vector{4, 6, 8, 10})); + a.div_(b); + EXPECT_EQ(to_host(a.cpu()), (std::vector{2, 3, 4, 5})); + + EXPECT_EQ(buf(a), p); +} + +TEST(InPlace, GPUScalarAndUnary) { + OM_REQUIRE_CUDA(); + Tensor a = Tensor::from_vector({-2.0f, -1.0f, 1.0f, 2.0f}, {4}).cuda(); + const void* p = buf(a); + + a.mul_(2.0f); + EXPECT_EQ(to_host(a.cpu()), (std::vector{-4, -2, 2, 4})); + a.relu_(); + EXPECT_EQ(to_host(a.cpu()), (std::vector{0, 0, 2, 4})); + a.fill_(1.5f); + EXPECT_EQ(to_host(a.cpu()), (std::vector{1.5f, 1.5f, 1.5f, 1.5f})); + a.apply_(Add{0.5f}); + EXPECT_EQ(to_host(a.cpu()), (std::vector{2, 2, 2, 2})); + + EXPECT_EQ(buf(a), p); +} + +TEST(InPlace, GPUSelfOperandDoubles) { + OM_REQUIRE_CUDA(); + Tensor a = dev_seq({4}); + a.add_(a); + EXPECT_EQ(to_host(a.cpu()), (std::vector{2, 4, 6, 8})); +} + +// The packed path (pack_width > 1) and its sub-pack tail are a separate code +// path from the scalar one; run a size that leaves a tail through it aliased. +TEST(InPlace, GPUPackedTailAliased) { + OM_REQUIRE_CUDA(); + for (size_t n : {size_t{1}, size_t{7}, size_t{1021}, size_t{4096}}) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) v[i] = static_cast(i % 5); + Tensor a = Tensor::from_vector(v, {n}).cuda(); + a.add_(a); + std::vector got(n); + a.cpu().copyToHost(got.data()); + for (size_t i = 0; i < n; ++i) + ASSERT_EQ(got[i], static_cast(2 * (i % 5))) << "n=" << n << " i=" << i; + } +} + +// Every rank goes through the same contiguous fast path, but the fallback is +// selected per launcher; check the ranks the launchers special-case. +TEST(InPlace, GPUEveryRank) { + OM_REQUIRE_CUDA(); + const std::vector> shapes = { + {64}, {8, 8}, {4, 4, 4}, {2, 4, 2, 4}, {2, 2, 2, 2, 4} + }; + for (const auto& sh : shapes) { + Tensor a = dev_seq(sh); + Tensor want = a.add(a); + a.add_(a); + EXPECT_EQ(to_host(a.cpu()), to_host(want.cpu())) << "rank " << sh.size(); + } +} + +// ── in-place on a stream ──────────────────────────────────────────────────── + +TEST(InPlace, GPUStreamOverload) { + OM_REQUIRE_CUDA(); + Stream s; + Tensor a = dev_seq({256}); + Tensor b = Tensor::full({256}, 2.0f).cuda(); + const void* p = buf(a); + + a.mul_(b, s).add_(1.0f, s).relu_(s); + s.synchronize(); + + const auto got = to_host(a.cpu()); + for (size_t i = 0; i < got.size(); ++i) + ASSERT_FLOAT_EQ(got[i], 2.0f * static_cast(i + 1) + 1.0f); + EXPECT_EQ(buf(a), p); +} + +// ── destination-provided overloads ────────────────────────────────────────── + +TEST(OutParam, CPUReusesOneDestination) { + Tensor a = host_seq({4}); + Tensor b = Tensor::full({4}, 2.0f); + Tensor out({4}); + const void* p = buf(out); + + a.add_out(b, out); + EXPECT_EQ(to_host(out), (std::vector{3, 4, 5, 6})); + a.mul_out(b, out); + EXPECT_EQ(to_host(out), (std::vector{2, 4, 6, 8})); + a.mul_out(3.0f, out); + EXPECT_EQ(to_host(out), (std::vector{3, 6, 9, 12})); + a.relu_out(out); + EXPECT_EQ(to_host(out), (std::vector{1, 2, 3, 4})); + a.apply_binary_out(b, BinaryAdd{}, out); + EXPECT_EQ(to_host(out), (std::vector{3, 4, 5, 6})); + + EXPECT_EQ(buf(out), p) << "the destination was replaced instead of written"; + EXPECT_EQ(to_host(a), (std::vector{1, 2, 3, 4})) << "operand was clobbered"; +} + +TEST(OutParam, ReturnsTheDestination) { + Tensor a = host_seq({4}); + Tensor out({4}); + Tensor& r = a.add_out(1.0f, out); + EXPECT_EQ(&r, &out); +} + +TEST(OutParam, GPUReusesOneDestination) { + OM_REQUIRE_CUDA(); + Tensor a = dev_seq({4}); + Tensor b = Tensor::full({4}, 2.0f).cuda(); + Tensor out({4}, Device("cuda:0")); + const void* p = buf(out); + + a.add_out(b, out); + EXPECT_EQ(to_host(out.cpu()), (std::vector{3, 4, 5, 6})); + a.div_out(b, out); + EXPECT_EQ(to_host(out.cpu()), (std::vector{0.5f, 1.0f, 1.5f, 2.0f})); + EXPECT_EQ(buf(out), p); +} + +TEST(OutParam, MatmulTransposePermute) { + Tensor a = host_seq({2, 3}); // [[1,2,3],[4,5,6]] + Tensor b = Tensor::full({3, 2}, 1.0f); + + Tensor mm({2, 2}); + a.matmul_out(b, mm); + EXPECT_EQ(to_host(mm), to_host(a.matmul(b))); + + Tensor tr({3, 2}); + a.transpose_out(tr); + EXPECT_EQ(to_host(tr), to_host(a.transpose())); + + Tensor pm({3, 2}); + a.permute_out({1, 0}, pm); + EXPECT_EQ(to_host(pm), to_host(a.permute({1, 0}))); +} + +TEST(OutParam, GPUMatmulTransposePermute) { + OM_REQUIRE_CUDA(); + Tensor a = dev_seq({2, 3}); + Tensor b = Tensor::full({3, 2}, 1.0f).cuda(); + Device gpu("cuda:0"); + + Tensor mm({2, 2}, gpu); + a.matmul_out(b, mm); + EXPECT_EQ(to_host(mm.cpu()), to_host(a.matmul(b).cpu())); + + Tensor tr({3, 2}, gpu); + a.transpose_out(tr); + EXPECT_EQ(to_host(tr.cpu()), to_host(a.transpose().cpu())); + + Tensor pm({3, 2}, gpu); + a.permute_out({1, 0}, pm); + EXPECT_EQ(to_host(pm.cpu()), to_host(a.permute({1, 0}).cpu())); +} + +// ── validation ────────────────────────────────────────────────────────────── + +TEST(OutParam, WrongDestinationShapeThrows) { + Tensor a = host_seq({4}); + Tensor out({5}); + EXPECT_THROW(a.add_out(1.0f, out), std::invalid_argument); + EXPECT_THROW(a.relu_out(out), std::invalid_argument); +} + +TEST(OutParam, WrongDestinationDeviceThrows) { + OM_REQUIRE_CUDA(); + Tensor a = dev_seq({4}); + Tensor out({4}); // host destination, device operand + EXPECT_THROW(a.add_out(1.0f, out), std::invalid_argument); +} + +TEST(OutParam, MismatchedOperandThrows) { + Tensor a = host_seq({4}); + Tensor b = host_seq({5}); + Tensor out({4}); + EXPECT_THROW(a.add_out(b, out), std::invalid_argument); + EXPECT_THROW(a.add_(b), std::invalid_argument); +} + +TEST(OutParam, MixedDeviceOperandsThrow) { + OM_REQUIRE_CUDA(); + Tensor host = host_seq({4}); + Tensor dev = dev_seq({4}); + EXPECT_THROW(host.add_(dev), std::invalid_argument); + EXPECT_THROW(dev.add_(host), std::invalid_argument); +} + +// matmul, transpose and permute read elements they do not write, so unlike the +// elementwise family they cannot service an aliased destination — and say so +// instead of returning a plausible wrong answer. +TEST(OutParam, AliasedDestinationRejectedForNonElementwise) { + Tensor a = host_seq({2, 2}); + Tensor b = Tensor::full({2, 2}, 1.0f); + EXPECT_THROW(a.matmul_out(b, a), std::invalid_argument); + EXPECT_THROW(a.matmul_out(b, b), std::invalid_argument); + EXPECT_THROW(a.transpose_out(a), std::invalid_argument); + EXPECT_THROW(a.permute_out({1, 0}, a), std::invalid_argument); +} + +// ── the reason the feature exists ─────────────────────────────────────────── + +// A loop of out-of-place ops allocates a fresh result every iteration; the +// same loop written in-place keeps one buffer for its whole life. +TEST(InPlace, LoopKeepsOneBuffer) { + Tensor w = Tensor::zeros({1024}); + Tensor g = Tensor::full({1024}, 0.5f); + const void* p = buf(w); + + for (int step = 0; step < 100; ++step) { + w.add_(g); + ASSERT_EQ(buf(w), p) << "buffer moved at step " << step; + } + const auto got = to_host(w); + EXPECT_FLOAT_EQ(got.front(), 50.0f); + EXPECT_FLOAT_EQ(got.back(), 50.0f); +} + +TEST(OutParam, LoopKeepsOneDestination) { + OM_REQUIRE_CUDA(); + Tensor a = Tensor::full({1024}, 1.0f).cuda(); + Tensor b = Tensor::full({1024}, 2.0f).cuda(); + Tensor out({1024}, Device("cuda:0")); + const void* p = buf(out); + + for (int step = 0; step < 100; ++step) { + a.mul_out(b, out); + ASSERT_EQ(buf(out), p) << "destination moved at step " << step; + } + EXPECT_FLOAT_EQ(to_host(out.cpu()).front(), 2.0f); +}