CUDA second update

Author: t81dev

AGENTS.md

@@ -57,4 +57,4 @@ This file helps AI agents discover and understand how to work with this reposito
 - Added `scripts/ternary_quantization_benchmark.py` plus `BENCHMARKS.md` so contributors can reproduce a Fashion-MNIST FP32/PTQ/QAT benchmark and log accuracy/latency/storage for each mode; README now links the benchmark doc.
 - Rewrote `pyproject.toml` with valid TOML sections so editable installs (and `pip install -e '.[torch]'`) can parse the metadata cleanly before building the extension.
 - Restructured `README.md` into a onboarding-focused front door and added companion docs (`docs/use-cases.md`, `docs/hardware.md`, `docs/api-overview.md`, `docs/python-install.md`, `docs/torch.md`, `examples/README.md`) so heavy reference material lives outside the visitor-facing overview.
-- Added optional CUDA/ROCm toggles plus a GPU dispatcher sketch (`include/t81/linalg/gemm_gpu.hpp`, `src/linalg/{gemm_cuda.cu,gemm_dispatch.cpp,gemm_rocm.cpp}`) so future teams can wire the new `where`/`clamp`/`lerp`/`addcmul` helpers into GPU kernels, and expanded `python/bindings.cpp` + `tests/python/test_gpu_ops.py` to expose and exercise the new APIs.
+- Added optional CUDA/ROCm toggles plus a GPU dispatcher sketch (`include/t81/linalg/gemm_gpu.hpp`, `src/linalg/{gemm_cuda.cu,gemm_dispatch.cpp,gemm_rocm.cpp}`) so future teams can wire the new `where`/`clamp`/`lerp`/`addcmul` helpers into GPU kernels, introduced `t81::TensorMetadata` + Python helpers (`python/bindings.cpp`) that extract metadata from NumPy/Torch tensors, and expanded `tests/python/test_gpu_ops.py` to cover the metadata-backed bindings on both CPU and GPU paths.

README.md

@@ -97,6 +97,10 @@ target_link_libraries(... t81::t81lib)
 
 `pip install .[torch]` unlocks the `t81lib`/`t81` namespace, NumPy quantization helpers, and the `t81.torch`/`t81.nn` layers that mix ternary weights with FP32/BF16 biases. Jump deeper via [docs/python-api.md](docs/python-api.md), [docs/python-cookbook.md](docs/python-cookbook.md), and [docs/torch.md](docs/torch.md).
 
+## GPU backends & tensor metadata
+
+Enable CUDA/ROCm through the optional `-DUSE_CUDA=ON` and `-DUSE_ROCM=ON` flags during CMake configuration so the Python bindings link against the new GPU kernels (`python/CMakeLists.txt`). Once enabled, `t81lib.where`, `t81lib.clamp`, `t81lib.lerp`, and `t81lib.addcmul` accept either NumPy buffers or PyTorch tensors and route the work directly to CUDA/HIP kernels via the lightweight [`t81::TensorMetadata`](include/t81/tensor_metadata.hpp) ABI. The metadata struct carries device/dtype/shape/stride info plus raw `data_ptr`, letting the dispatcher avoid host copies and keep outputs on-device. When torch is installed, `t81lib` automatically wraps GPU tensors; when only NumPy is available it falls back to CPU buffers. Consult [docs/torch.md](docs/torch.md) and `python/bindings.cpp` for the extraction helpers and lifetime semantics.
+
 ## CLI helpers
 
 `t81-convert`, `t81-gguf`, and `t81-qat` automate quantize→export→train flows with progress reporting and validation hooks. Browse [docs/references/cli-usage.md](docs/references/cli-usage.md), [docs/diagrams/cli-workflows-mermaid.md](docs/diagrams/cli-workflows-mermaid.md), and [examples/cli-examples.md](examples/cli-examples.md) for recipes.

docs/ROADMAP.md

@@ -96,6 +96,16 @@ The codebase is thoughtfully commented, consistently formatted by `clang-format`
   2. Create a `scripts/` benchmark script that trains the model, applies `t81` quantization, and logs accuracy, model size, and latency.
   3. Document the benchmark/results in a new `BENCHMARKS.md` (link from `README.md`) so the community can reproduce and compare.
 
+### Recommendation 5: Harden the GPU Tensor Metadata Path
+
+* **Why**: GPU acceleration for ternary ops unlocks real inference speedups only if the dispatcher can safely read raw device pointers and keep outputs on-device, and the Python bindings need a stable contract (TensorMetadata) to wrap Torch/NumPy tensors the same way.
+* **Benefits**: Enables CUDA/ROCm kernels (where/clamp/lerp/addcmul) to run without host copies, makes GPU/CPU fallbacks explicit, and gives PyTorch users a predictable ABI.
+* **Effort**: Medium.
+* **Implementation**:
+  1. Continue documenting the `t81::TensorMetadata` ABI (`device_type`, `dtype`, `sizes`, `strides`, `data_ptr`, `storage_offset`, `requires_sync`) so future kernels and Python helpers rely on the same struct.
+  2. Expand the dispatcher to record broadcasting/non-contiguous layouts (or copy them to contiguous buffers when needed) so `backend_available == CUDA/ROCm` paths no longer require perfect contiguity.
+  3. Run the CUDA 12.4 / ROCm 6+ CI jobs, validate the latency/accuracy guardrails in `tests/python/test_gpu_ops.py`, and flag any ABI mismatches between `torch.Tensor` and `TensorMetadata` so the GPU helpers can be certified before release.
+
 ---
 
 The above roadmap tightly ties into the existing structure (`docs/`, `scripts/`, CLI helpers) while making the project more approachable, ensuring quality, highlighting the Python story, and showcasing real-world value through benchmarks.

python/bindings.cpp

@@ -163,7 +163,6 @@ namespace {
             try {
                 return py::module_::import("torch");
             } catch (const py::error_already_set &) {
-                py::gil_scoped_release release;
                 return py::object();
             }
         }();
@@ -213,7 +212,7 @@ namespace {
         if (info.size < 0 || info.itemsize <= 0) {
             throw py::value_error("buffer has invalid dimensions or element size");
         }
-        if (info.format != "f" && info.itemsize != static_cast<py::ssize_t>(sizeof(float))) {
+        if (info.itemsize != static_cast<py::ssize_t>(sizeof(float))) {
             throw py::value_error("buffer must hold float32 elements");
         }
         if (!buffer_is_c_contiguous(info)) {
@@ -233,7 +232,7 @@ namespace {
 
     PyTensorHandle metadata_from_buffer(const py::object &obj) {
         py::buffer buffer(obj);
-        const auto info = buffer.request(true);
+        const auto info = buffer.request(false);
         TensorMetadata metadata = metadata_from_buffer_info(info);
         PyTensorHandle handle;
         handle.meta = std::move(metadata);
@@ -641,8 +640,6 @@ PYBIND11_MODULE(t81lib, module) {
         t81::linalg::detail::backend_available(t81::linalg::Backend::CUDA);
     module.attr("HAS_ROCM_BACKEND") =
         t81::linalg::detail::backend_available(t81::linalg::Backend::ROCm);
-    using FloatArray =
-        py::array_t<float, py::array::c_style | py::array::forcecast>;
 
     module.def(
         "gemm_ternary",
@@ -687,28 +684,16 @@ PYBIND11_MODULE(t81lib, module) {
 
     module.def(
         "where",
-        [](FloatArray condition, FloatArray x, FloatArray y) {
-            const auto cond_info = condition.request();
-            const auto x_info = x.request();
-            const auto y_info = y.request();
-            const std::size_t total = buffer_element_count(cond_info);
-            if (buffer_element_count(x_info) != total ||
-                buffer_element_count(y_info) != total) {
-                throw py::value_error(
-                    "where requires inputs with the same number of elements");
-            }
-            py::array_t<float> output(cond_info.shape);
-            const auto out_info = output.request();
-            const auto cond_span = std::span<const float>{
-                static_cast<const float *>(cond_info.ptr), total};
-            const auto x_span = std::span<const float>{
-                static_cast<const float *>(x_info.ptr), total};
-            const auto y_span = std::span<const float>{
-                static_cast<const float *>(y_info.ptr), total};
-            auto out_span = std::span<float>{
-                static_cast<float *>(out_info.ptr), total};
-            t81::linalg::detail::where(cond_span, x_span, y_span, out_span);
-            return output;
+        [](py::object condition, py::object x, py::object y) {
+            const auto cond_handle = extract_tensor_handle(condition);
+            const auto x_handle = extract_tensor_handle(x);
+            const auto y_handle = extract_tensor_handle(y);
+            auto out_handle = create_output_like(x_handle);
+            t81::linalg::detail::where(cond_handle.meta,
+                                       x_handle.meta,
+                                       y_handle.meta,
+                                       out_handle.meta);
+            return out_handle.owner;
         },
         py::arg("condition"),
         py::arg("x"),
@@ -717,20 +702,15 @@ PYBIND11_MODULE(t81lib, module) {
 
     module.def(
         "clamp",
-        [](FloatArray values, float min_value, float max_value) {
+        [](py::object values, float min_value, float max_value) {
             if (min_value > max_value) {
                 throw py::value_error("min_value must be <= max_value");
             }
-            const auto info = values.request();
-            const std::size_t total = buffer_element_count(info);
-            py::array_t<float> output(info.shape);
-            const auto out_info = output.request();
-            std::span<const float> input_span{
-                static_cast<const float *>(info.ptr), total};
-            std::span<float> out_span{
-                static_cast<float *>(out_info.ptr), total};
-            t81::linalg::detail::clamp(input_span, min_value, max_value, out_span);
-            return output;
+            const auto input_handle = extract_tensor_handle(values);
+            auto out_handle = create_output_like(input_handle);
+            t81::linalg::detail::clamp(
+                input_handle.meta, out_handle.meta, min_value, max_value);
+            return out_handle.owner;
         },
         py::arg("x"),
         py::arg("min"),
@@ -739,28 +719,16 @@ PYBIND11_MODULE(t81lib, module) {
 
     module.def(
         "lerp",
-        [](FloatArray start, FloatArray end, FloatArray weight) {
-            const auto start_info = start.request();
-            const auto end_info = end.request();
-            const auto weight_info = weight.request();
-            const std::size_t total = buffer_element_count(start_info);
-            if (buffer_element_count(end_info) != total ||
-                buffer_element_count(weight_info) != total) {
-                throw py::value_error(
-                    "lerp requires start, end, and weight to align");
-            }
-            py::array_t<float> output(start_info.shape);
-            const auto out_info = output.request();
-            std::span<const float> start_span{
-                static_cast<const float *>(start_info.ptr), total};
-            std::span<const float> end_span{
-                static_cast<const float *>(end_info.ptr), total};
-            std::span<const float> weight_span{
-                static_cast<const float *>(weight_info.ptr), total};
-            std::span<float> out_span{
-                static_cast<float *>(out_info.ptr), total};
-            t81::linalg::detail::lerp(start_span, end_span, weight_span, out_span);
-            return output;
+        [](py::object start, py::object end, py::object weight) {
+            const auto start_handle = extract_tensor_handle(start);
+            const auto end_handle = extract_tensor_handle(end);
+            const auto weight_handle = extract_tensor_handle(weight);
+            auto out_handle = create_output_like(start_handle);
+            t81::linalg::detail::lerp(start_handle.meta,
+                                      end_handle.meta,
+                                      weight_handle.meta,
+                                      out_handle.meta);
+            return out_handle.owner;
         },
         py::arg("start"),
         py::arg("end"),
@@ -769,32 +737,20 @@ PYBIND11_MODULE(t81lib, module) {
 
     module.def(
         "addcmul",
-        [](FloatArray input,
-           FloatArray tensor1,
-           FloatArray tensor2,
+        [](py::object input,
+           py::object tensor1,
+           py::object tensor2,
            float value) {
-            const auto input_info = input.request();
-            const auto tensor1_info = tensor1.request();
-            const auto tensor2_info = tensor2.request();
-            const std::size_t total = buffer_element_count(input_info);
-            if (buffer_element_count(tensor1_info) != total ||
-                buffer_element_count(tensor2_info) != total) {
-                throw py::value_error(
-                    "addcmul requires input and tensors to share the same size");
-            }
-            py::array_t<float> output(input_info.shape);
-            const auto out_info = output.request();
-            std::span<const float> input_span{
-                static_cast<const float *>(input_info.ptr), total};
-            std::span<const float> tensor1_span{
-                static_cast<const float *>(tensor1_info.ptr), total};
-            std::span<const float> tensor2_span{
-                static_cast<const float *>(tensor2_info.ptr), total};
-            std::span<float> out_span{
-                static_cast<float *>(out_info.ptr), total};
-            t81::linalg::detail::addcmul(
-                input_span, tensor1_span, tensor2_span, value, out_span);
-            return output;
+            const auto input_handle = extract_tensor_handle(input);
+            const auto tensor1_handle = extract_tensor_handle(tensor1);
+            const auto tensor2_handle = extract_tensor_handle(tensor2);
+            auto out_handle = create_output_like(input_handle);
+            t81::linalg::detail::addcmul(input_handle.meta,
+                                         tensor1_handle.meta,
+                                         tensor2_handle.meta,
+                                         value,
+                                         out_handle.meta);
+            return out_handle.owner;
         },
         py::arg("input"),
         py::arg("tensor1"),

src/linalg/gemm_dispatch.cpp

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <cstdint>
 #include <stdexcept>
 
 namespace t81::linalg::detail {

tests/python/test_gpu_ops.py

@@ -6,6 +6,11 @@
 
 import pytest
 
+try:
+    import torch
+except ImportError:
+    torch = None
+
 try:
     import t81lib
 except ImportError as exc:
@@ -92,3 +97,28 @@ def test_clamp_rocm_latency_and_accuracy():
     for index, value in enumerate(result):
         expected = min(1.0, max(-1.0, values[index]))
         assert math.isclose(value, expected, abs_tol=1e-6)
+
+
+@pytest.mark.skipif(torch is None, reason="torch not installed")
+def test_where_torch_cpu_tensor():
+    cond = torch.tensor([1.0, 0.0, -1.0], dtype=torch.float32)
+    x = torch.tensor([1.1, 2.2, 3.3], dtype=torch.float32)
+    y = torch.tensor([-1.0, -2.0, -3.0], dtype=torch.float32)
+    result = t81lib.where(cond, x, y)
+    assert isinstance(result, torch.Tensor)
+    assert torch.allclose(result, torch.tensor([1.1, -2.0, -3.0], dtype=torch.float32))
+
+
+@pytest.mark.skipif(
+    torch is None or not torch.cuda.is_available(),
+    reason="requires CUDA-capable torch tensor",
+)
+def test_where_torch_cuda_tensor():
+    cond = torch.tensor([1.0, 0.0], dtype=torch.float32, device="cuda")
+    x = torch.tensor([4.0, 4.0], dtype=torch.float32, device="cuda")
+    y = torch.tensor([0.0, -4.0], dtype=torch.float32, device="cuda")
+    result = t81lib.where(cond, x, y)
+    assert isinstance(result, torch.Tensor)
+    assert result.device.type == "cuda"
+    expected = torch.tensor([4.0, -4.0], dtype=torch.float32, device="cuda")
+    assert torch.allclose(result, expected)