Allocate CPU output buffers outside torch's allocator

CLAUDE.md

@@ -12,7 +12,16 @@ generated locally with FFmpeg; end-to-end cloud-storage tests live in the
 consuming project.
 
 `tch` comes from crates.io; its version maps 1:1 to a libtorch version
-(0.24.0 ↔ torch 2.11), so bump them together.
+(0.24.0 ↔ torch 2.11, 0.26.0 ↔ torch 2.13). `Cargo.toml` accepts a range
+so one source tree builds against either torch; `Cargo.lock` holds the
+default. To build against another torch, run
+`cargo update -p tch -p torch-sys --precise <version>` (keep `tch` and
+`torch-sys` on the same version).
+
+`build.rs` compiles `csrc/tensor_from_blob.cpp` against the libtorch headers
+(located the same way torch-sys locates them); `src/alloc.rs` uses it to give
+decoded-output buffers to torch via `from_blob` with a deleter instead of
+`at::empty`. Keep the two in sync if the C API changes.
 
 ## Git
 

Cargo.toml

@@ -27,7 +27,13 @@ log = { version = "0.4.27", features = ["kv"] }
 pyo3 = "0.25.0"
 pyo3-log = "0.12.4"
 rsmpeg = { version = "0.17.0", features = ["ffmpeg7_1"] }
-tch = { version = "0.26.0", features = ["python-extension"] }
+# Each tch release targets one libtorch (0.24 <-> torch 2.11, 0.26 <-> torch 2.13).
+# The source builds against both; Cargo.lock picks the default, and a
+# build for another torch selects its match with
+# `cargo update -p tch -p torch-sys --precise <version>`.
+tch = { version = ">=0.24.0, <0.27", features = ["python-extension"] }
+# Same version as tch; only used for the opaque `C_tensor` handle in src/alloc.rs.
+torch-sys = ">=0.24.0, <0.27"
 tempfile = "3.20.0"
 thiserror = "2.0.15"
 libloading = "0.8"
@@ -38,6 +44,10 @@ tokio = { version = "1.47.1", features = ["time"] }
 extension-module = ["pyo3/extension-module"]
 default = ["extension-module"]
 
+[build-dependencies]
+# Compiles csrc/tensor_from_blob.cpp against the libtorch headers.
+cc = "1.2"
+
 [dev-dependencies]
 mel_spec = "0.3.3"
 ndarray = "0.16.1"

README.md

@@ -416,16 +416,26 @@ provider-specific support — the URL is handed to FFmpeg's http protocol.
 - **`s3://` works with AWS but not with your S3-compatible store** — check
   whether the store requires path-style addressing
   (`AVTENSOR_S3_FORCE_PATH_STYLE=1`) and that `AWS_ENDPOINT_URL_S3` is set.
+- **Resident memory grows when decoded tensors are freed on a different
+  thread than the one that decoded them** — torch's CPU allocator is not
+  always plain malloc (the aarch64 Linux wheels embed mimalloc, which keeps
+  blocks freed cross-thread mapped and resident until the allocating thread
+  exits). avtensor therefore allocates its output buffers with the system
+  allocator and hands them to torch via `from_blob` with a deleter, so a free
+  is an immediate `munmap` on any thread. `AVTENSOR_OUTPUT_ALLOCATOR=torch`
+  switches back to `torch.empty`-style allocation.
 
 ## Contributing
 
 Issues and pull requests are welcome. Project layout:
 
 ```
 src/
+  alloc.rs       output-buffer allocation (system allocator + torch from_blob)
   decoder/       demuxing, decoding, filter graphs (mod.rs), cloud AVIO reader (io.rs)
   ffi/           PyO3 bindings: request/response types, decode_asset
   util/          gcs/s3 URI handling, memory, test media generation
+csrc/            C++ shim for the one libtorch call tch-rs does not wrap (built by build.rs)
 avtensor.pyi     Python type stubs, shipped with the wheel
 ```
 

build.rs

@@ -0,0 +1,125 @@
+//! Compiles the C++ shim in `csrc/` that exposes `at::from_blob` with a
+//! deleter to Rust (used by `src/alloc.rs`). tch-rs only wraps the non-owning
+//! overload, so this is the one piece of libtorch we have to reach through
+//! C++. The include directories and C++ ABI are resolved the same way
+//! torch-sys resolves them, so both shims link against the same libtorch.
+
+use std::env;
+use std::path::PathBuf;
+use std::process::Command;
+
+const PYTHON_PROBE: &str = r"
+import torch
+from torch.utils import cpp_extension
+print('CXX11_ABI:', int(torch._C._GLIBCXX_USE_CXX11_ABI))
+for path in cpp_extension.include_paths():
+    print('INCLUDE:', path)
+";
+
+struct LibtorchConfig {
+    include_dirs: Vec<PathBuf>,
+    cxx11_abi: String,
+}
+
+fn env_var_rerun(name: &str) -> Option<String> {
+    println!("cargo:rerun-if-env-changed={name}");
+    env::var(name).ok()
+}
+
+/// Mirrors torch-sys: `python` inside a virtualenv, `python3` otherwise,
+/// unless maturin/PyO3 point at a specific interpreter.
+fn python_interpreter() -> PathBuf {
+    if let Some(p) = env_var_rerun("PYO3_PYTHON") {
+        return PathBuf::from(p);
+    }
+    if env::var_os("VIRTUAL_ENV").is_some() {
+        PathBuf::from("python")
+    } else {
+        PathBuf::from("python3")
+    }
+}
+
+fn config_from_python() -> Result<LibtorchConfig, String> {
+    let python = python_interpreter();
+    let output = Command::new(&python)
+        .arg("-c")
+        .arg(PYTHON_PROBE)
+        .output()
+        .map_err(|e| format!("running {python:?}: {e}"))?;
+    if !output.status.success() {
+        return Err(format!(
+            "{python:?} failed: {}",
+            String::from_utf8_lossy(&output.stderr)
+        ));
+    }
+    let stdout = String::from_utf8_lossy(&output.stdout);
+    let mut include_dirs = vec![];
+    let mut cxx11_abi = None;
+    for line in stdout.lines() {
+        if let Some(path) = line.strip_prefix("INCLUDE: ") {
+            include_dirs.push(PathBuf::from(path));
+        } else if let Some(abi) = line.strip_prefix("CXX11_ABI: ") {
+            cxx11_abi = Some(abi.trim().to_owned());
+        }
+    }
+    match cxx11_abi {
+        Some(cxx11_abi) if !include_dirs.is_empty() => Ok(LibtorchConfig {
+            include_dirs,
+            cxx11_abi,
+        }),
+        _ => Err(format!("unexpected probe output: {stdout}")),
+    }
+}
+
+fn config_from_libtorch_dir() -> Result<LibtorchConfig, String> {
+    let root = env_var_rerun("LIBTORCH")
+        .map(PathBuf::from)
+        .or_else(|| {
+            // torch-sys exports the lib dir it linked against (`links = "tch"`).
+            env::var("DEP_TCH_LIBTORCH_LIB")
+                .ok()
+                .map(|lib| PathBuf::from(lib).join(".."))
+        })
+        .ok_or("neither LIBTORCH_USE_PYTORCH nor LIBTORCH is set")?;
+    let include_dirs = match env_var_rerun("LIBTORCH_INCLUDE") {
+        Some(include) => vec![
+            PathBuf::from(&include).join("include"),
+            PathBuf::from(&include).join("include/torch/csrc/api/include"),
+        ],
+        None => vec![
+            root.join("include"),
+            root.join("include/torch/csrc/api/include"),
+        ],
+    };
+    Ok(LibtorchConfig {
+        include_dirs,
+        cxx11_abi: env_var_rerun("LIBTORCH_CXX11_ABI").unwrap_or_else(|| "1".to_owned()),
+    })
+}
+
+fn main() {
+    println!("cargo:rerun-if-changed=csrc/tensor_from_blob.cpp");
+
+    let config = if env_var_rerun("LIBTORCH_USE_PYTORCH").is_some() {
+        config_from_python()
+    } else {
+        config_from_libtorch_dir()
+    }
+    .unwrap_or_else(|e| panic!("locating libtorch headers for csrc/tensor_from_blob.cpp: {e}"));
+
+    let mut build = cc::Build::new();
+    build
+        .cpp(true)
+        .pic(true)
+        .warnings(false)
+        .includes(&config.include_dirs)
+        .file("csrc/tensor_from_blob.cpp");
+    if env::var("CARGO_CFG_TARGET_OS").as_deref() == Ok("windows") {
+        build.flag("/std:c++17");
+    } else {
+        build
+            .flag("-std=c++17")
+            .flag(format!("-D_GLIBCXX_USE_CXX11_ABI={}", config.cxx11_abi));
+    }
+    build.compile("avtensor_blob");
+}

csrc/tensor_from_blob.cpp

@@ -0,0 +1,46 @@
+// Exposes `at::from_blob` with a deleter to Rust. tch-rs wraps only the
+// non-owning overload (`at_tensor_of_blob`); this one lets a buffer allocated
+// outside torch's CPU allocator be owned by a tensor and freed by `deleter`
+// when the last reference goes away. See src/alloc.rs for why.
+#include <ATen/ATen.h>
+
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+#include <exception>
+#include <utility>
+
+extern "C" {
+
+typedef void (*avtensor_blob_deleter)(void *data, void *ctx);
+
+// Wraps `data` in a contiguous CPU tensor of `sizes` and element type
+// `scalar_type` (a c10::ScalarType value). Ownership of `data` moves to the
+// tensor's storage: `deleter(data, ctx)` runs when the last tensor sharing
+// that storage is released, on whichever thread releases it.
+//
+// Returns a heap-allocated at::Tensor, the object tch-rs calls `C_tensor`
+// and frees with `at_free`. On failure returns nullptr and, if `err` is not
+// null, stores a malloc'd message in *err that the caller frees.
+void *avtensor_tensor_from_blob(void *data, const int64_t *sizes, size_t ndim,
+                                int32_t scalar_type,
+                                avtensor_blob_deleter deleter, void *ctx,
+                                char **err) {
+  try {
+    auto options = at::TensorOptions()
+                       .dtype(static_cast<at::ScalarType>(scalar_type))
+                       .device(at::kCPU);
+    at::Tensor tensor = at::from_blob(
+        data, at::IntArrayRef(sizes, ndim),
+        [deleter, ctx](void *p) { deleter(p, ctx); }, options);
+    return new at::Tensor(std::move(tensor));
+  } catch (const std::exception &e) {
+    if (err) *err = strdup(e.what());
+    return nullptr;
+  } catch (...) {
+    if (err) *err = strdup("unknown exception in at::from_blob");
+    return nullptr;
+  }
+}
+
+}  // extern "C"

pyproject.toml

@@ -36,7 +36,7 @@ auditwheel = "skip"
 
 # Include some directives to instruct uv to recompile on changes to Rust files.
 [tool.uv]
-cache-keys = [{file = "pyproject.toml"}, {file = "Cargo.toml"}, {file = "**/*.rs"}]
+cache-keys = [{file = "pyproject.toml"}, {file = "Cargo.toml"}, {file = "**/*.rs"}, {file = "csrc/*.cpp"}]
 
 [tool.ruff]
 line-length = 88