What makes light and works through oxidization? Flash Powder.
A very minimal rust wrapper for libtorch, using the Torch Stable API only. This is mostly my project to gain a better understanding of how (lib/py)torch works under the hood. I do not recommend using this.
The stable ABI doesn't expose all functionality of libtorch, but a surprising amount of functionality is available, especially if the goal is just to do inference. The example_vgg crate holds an implementation of vgg11.
This was developed for doing postprocessing and inference with an U-Net in my overlay_segmenter.
It follows the rust semantics as closely as possible. This means;
- No unsafe in the public interface, safe behaviour as you'd expect.
- No interior mutability, all methods are const correct.
- Modifying one tensor will not modify another, unless it has a mutable borrow on the other.
- Rust style lifetimes on tensors, either tied together with an explicit lifetime, or owning.
There are three structures fundamental to achieving this:
Tensor; Owning tensor, this owns the data. (thinkVec<u8>)Ten<'_>; Const borrow of Tensor, this has a parent, its lifetime cannot exceed the parent. (think&[u8])TenMut<'_>; Mutable borrow of Tensor, this has a mutable parent, its lifetime cannot exceed the parent. (think&mut [u8])
Under the hood, each of these is a StableTensor and its own tensor handle on the LibTorch side.
This doesn't map perfectly to Torch's operations, for example the .to() method in libtorch sometimes returns a copy, but not always.
So there's some arbitrary choices here, like .to() in this crate always makes a copy.
The main high-level and safe interface lives in the flash_powder crate.
The crate is fairly well documented, here's an overview of existing functionality to get an idea of the semantics as well as the location in crate, most examples are copied from the unit tests.
All functions or methods that can fail return a Result, which when it fails holds an anyhow::Error with the message that was returned by the stable API.
Creating a tensor can be done with the conversion module through TryInto<Tensor>:
use flash_powder as fp;
use flash_powder::Tensor;
// Convert a scalar like so:
let d: Tensor = 5i64.try_into()?;
assert_eq!(d.dim(), 0);
assert_eq!(d.i64_ref(&[])?, &5);
// Or create a 2D Tensor with some floats;
let d: Tensor = [[5.0f32, 3.0, 5.0], [1.0, 2.0, 0.0]].try_into()?;
assert_eq!(d.sizes(), &[2, 3]);
// Change it or move it to a device
let d_as_u8 = d.to(&fp::DType::U8.into())?;
let u8_on_gpu = as_u8.to(&fp::Device::CUDA.into())?;Or create them with any of the factory trait methods:
let a = Tensor::empty(&[5, 5], &Default::default()); // Defaults to cpu, f32
let t = Tensor::randn(&[3, 3], &fp::Device::CPU.into())?; // We can give it a device to create it on
let e = Tensor::zeros(&[6, 6], &fp::DType::U8.into())?; // Or specify a type (or mix these options).The properties of a tensor, like dtype(), device() and sizes() are all provided by the TensorProperties trait from the properties module.
Data access to the data contained in the Tensor is provided through the data module, through the DataRef and DataMut traits.
This exposes (typed) slices created from the tensor's data using zerocopy.
as_<T>(): Access to the value stored in a scalar tensor:&T<T>_ref(indices: &[usize]): Index into the storage to return a reference to a value at the provided index position:&T<T>s_ref(): Access to the entire slice of values:&[T]as_<T>_mut(),<T>_mut(indices: &[usize]),<T>s_mut(): Mutable flavours of these.
The Tensor and Ten implement the CoreMethods trait from the core_methods module.
This provides functionality like flatten, mul, permute and other operations.
let t = Tensor::from(&[0.2015f32, -0.4255, 2.6087])?;
let factor: Tensor = 100.0.try_into()?;
let r = t.mul(&factor)?;
assert_eq!(r.sizes(), &[3]);
assert_eq!(
r.f32s_ref()?,
&[20.149999618530273f32, -42.54999923706055, 260.8699951171875]
);The TenMut also implements the CoreMethodsMut from the same module, this provides some in-place modification and mutable view slicing.
Indexing is provided by index, mutably with i_mut, this is limited to operations that always produce a view in PyTorch, so indexing with indices is not supported as that returns a copy in some cases.
let d = Tensor::from(&[
[1.0f32, 2.0, 3.0, 4.0],
[5.0, 6.0, 7.0, 8.0],
[9.0, 10.0, 11.0, 12.0],
[13.0, 14.0, 15.0, 16.0],
])?;
let z = d.i((1..3, 0..1))?; // Equivalent to PyTorch; z = d[1:3, 0:1]
assert_eq!(z.sizes(), &[2, 1]);
let z = d.i((-3isize..3, -3isize..3))?; // z = d[-3:3, -3:3]
assert_eq!(z.sizes(), &[2, 2]); // #PYTHON list(z.shape)The functional module provides the basic building blocks I needed for vgg & U-net;
adaptive_avg_pool2d, conv2d, conv_transpose2d, interpolate, linear, max_pool2d, relu and upsample.
This is definitely not fully featured, but shows how to dispatch kernels, they are defined in native_functions.yaml, which comes with a README that explains the flags (permalinks to v2.12, be sure to change that to latest).
Kernel dispatches can of course be done out-of-crate.
Finally, the nn aims to be the equivalent of torch.nn.
The nn::Module is the Rust trait equivalent to torch.nn.Module and exposes methods like forward, to, state_dict, load_state_dict, the nn::layer module provides the layers for the functions from functional that all implement Module, as well as Sequential to be able to chain type-erased layers together.
The nn module also provides some helper functionality around the StateDict, which dovetails with the flash_powder_safetensors crate to be able to load tensors from disk easily.
Very minimal (handwritten) bindings for the LibTorch Stable ABI. This system works through a small set of C functions that provide a limited subset of the functionality from libtorch.
The crate is structured after the stable, aoti_torch and headeronly directories.
There's some support tooling in the contrib submodule, but it's mostly there for testing and superseded by the flash_powder crate.
The functionality in this crate is a subset of the upstream functionality, it does not follow Rust lifetimes or safety guarantees.
Helper utilities for working with safetensors are available in the flash_powder_safetensors crate.
Loading a safetensors file into a nn::Module, from the example_vgg code:
// Load safetensors data from disk, deserialize it and wrap it in a reader.
let data = std::fs::read(&weights).expect("Unable to read file");
let tensors = flash_powder_safetensors::safetensors::SafeTensors::deserialize(&data)?;
let reader = flash_powder_safetensors::SafetensorReader::from_safetensors(&tensors);
// Or by memory mapping the file instead of reading to memory, this allows directly moving tensors from disk to gpu.
let mapped_file = flash_powder_safetensors::MappedFile::map(weights)?; // Very thin wrapper around mmap2.
let tensors = mapped_file.to_safetensors()?;
let reader = flash_powder_safetensors::SafetensorReader::from_safetensors(&tensors);
// And load it into the module.
vgg.load_state_dict(&reader, &Default::default())?;It also provides helpers to serialize/deserialize and reader/write to and from an fp::nn::module::StateDict.
Helper utilities for working with the Rust image crate available in the flash_powder_image crate.
The most useful functionality of this crate is that it facilitates reading and writing images directly to and from Tensors:
use flash_powder_image::prelude::*;
// Read an image from disk;
let img = Tensor::read_image("super_cool_image.png")?; // [C, H, W], DType::U8, [0, 255]
// Save an image with:
img.save_image("/tmp/it_was_really_cool.png")?; // Expects [C, H, W], DType::U8 in [0, 255], or float in [0, 1.0].It also handles [B, C, H, W], which creates a row of images, and [V, B, C, H, W] to vertically stack batched image rows.
Additional functionality is provided that's commonly used when handling images:
// You can floatify an image to scale it from integer [0,255] to F32 [0.0, 1.0], you can pass a DType to the ToOptions struct
// to immediately select another data data and/or device.
let img = Tensor::read_image(&path)?.image_floatify(&ToOptions::default())?;
// It can also scale a tensor to be within the [0.0, 1.0] domain for easy visualisation:
combined.image_scale_to_domain()?.save_image("/tmp/visible_tensor.png")?;
// Or through an image::ImageReader:
let img = image::ImageReader::open("/tmp/fp_rgba_f32.png")?.decode()?;
let img_as_tensor = img.to_tensor()?;
// And back
let dynamic_image = img_as_tensor.to_dynamic_image()?;
// And image resize using tensors, using LibTorch's interpolate method.
let d_larger = img_as_tensor.image_resize([100, 100], functional::InterpolateAlgorithm::Nearest)?;Run this to add the dependency to a cargo project in both build-dependencies and dependencies;
cargo add --git https://github.com/iwanders/flash_powder.git flash_powder -F cuda
cargo add --git https://github.com/iwanders/flash_powder.git flash_powder -F cuda --build
Update manually with
cargo update
The minimum PyTorch/libTorch version is 2.13, which was the version under development when I reached out about this ffi
use case with this comment,
a followup issue around lack of error retrieval functionality was created.
The proposed changes were incorporated in this PR, and improved in a followup.
The lack of allocator/deleter for StableIValue was addressed in this PR.
The v2_13 feature was removed in fdb282, the stable ivalue creation workaround in 60b505fd1, with these reinstated it could run on older versions.
Prior to v2.14, there's a tiny memory leak in the conversion between StableIValue's to String, fixed in this PR.
This does mean that kernels that take a string argument will leak their value.
At the time of writing that is only a concern for the div.Tensor_mode kernel, which is used by the div_mode method to do integer division.
Mostly jotting these down for myself, but this may be relevant for anyone coming across this:
PyTorch heavily uses overloading, for example squeeze has the default, dim, dimname and dims overloads.
And mean has six of them. Currently, this is handled a bit ad-hoc.
Maybe this should have a stricter convention...
In a lot of places I used usize, while LibTorch uses i64, in PyTorch you can often use negative indices.
An example where I introduced an 'overload' is TensorProperties::size(&self, dim: usize) and its signed TensorProperties::isize(&self, dim: isize) counterpart.
This is all a bit tricky, because using usize is nice as you can pass lengths from rust directly, but losing the negative values is problematic.
What we could do instead is introduce a Dimension trait that accepts any integer and handles it appropriately?
On indexing, we currently have .i and .i_mut to return Ten and TenMut respectively, but index.Tensor returns a copy, which doesn't fit into that architecture, so we can't index with tensors right now through the nice interface and you manually have to call CoreMethods::index_tensor.
Not all kernels can be dispatched at the moment, Scalar support is missing, see the comment and thread here.
This is most notably a problem with add.Tensor, which is currently worked around with by calling into _foreach_addcmul.Tensor instead, which conveniently takes scalars as a Tensor, subtraction is made with the same kernel.
This feels pretty fragile, and is probably less performant than a normal addition, once Scalar support lands this can be cleaned up.
Once that lands, a lot of methods should probably be changed to take either a Tensor or a Scalar, such that dividing by a single number doesn't require creating a scalar tensor first.
Sometimes chains of borrows are problematic, like let nonzero = counts.eq(&zero)?.squeeze()? will result in a lifetime error as squeeze borrows, but the result of counts.eq goes out of scope, work around this by separating the statement or making it owning with .to_owned()?;.
Should see if this can be made better.
- example_vgg Implements torchvisions' VGG network and shows it produces identical outputs, also leverages the
flash_powder_imageandflash_powder_safetensorscrates. - example_pytorch_extension A prototype pure-rust PyTorch extension library that can be loaded from python and provides kernels.
I want to ensure that the tensors & function arguments follow conventions from the Python side, so there's a heavy emphasis on testing all functions against their Python equivalents. The Python code to test against is interwoven with the Rust code with some helper tooling. Tests should run cleanly in valgrind.
For tests, the equivalent Python PyTorch execution is considered the ground truth and the Rust should produce the
same values.
To be able to easily create reference values in the tests there's a helper tool in ./util/python_truth.py that can
execute python code in rusts' comment blocks and update values in the rust tests accordingly.
This ensures that the equivalent python code is next to the rust code in the unit tests and also facilitates automatic
generation of reference values without manual copy pasting which may introduce errors.
The scope of a particular Python execution is limited to within a (test) function scope;
The following:
/*
#|PYTHON
d = torch.tensor(list(range(1,17)), dtype=torch.float).reshape([1,4,4])
w = torch.tensor([[[1.0, 2.0],[3.0, 4.0]]]).unsqueeze(0)
r = torch.nn.functional.conv2d(d, w)
*/defines what is considered a Python block, this runs the statements in this block in python and stores their values for use in the next block(s) (either Rust or Python).
The values are then used with a rust comment like: // #PYTHON <STATEMENT>, where <STATEMENT> is a single Python statement that will be executed.
This comment is placed after the statement it is applied to, it can apply to both constants and function calls like assert_eq!.
With function calls, the last argument is replaced with the ground truth.
assert_eq!(d.sizes(), &[1, 4, 4]); // #PYTHON list(d.shape)
const GROUND_TRUTH: &[usize] = &[1usize, 4, 4]; // #PYTHON list(d.shape)
assert_eq!(
d.f32_ref()?,
&[
1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
16.0
]
); // #PYTHON list(d.view(-1).tolist())Functionality is limited to integers, floats and 1d arrays, in both reference and (implicit) array form.
By default, the binary processess the entire rust file, it can be constrained to a single test with --test-case test_flash_power_conv2d or --test-case test_flash_power_conv*.
It automatically calls rustfmt to ensure files are always formatted after modification.
# Extract the python code;
./util/python_truth.py extract ./flash_powder/src/native_functions.rs --test-case test_flash_power_conv2d
# Execute the python code;
./util/python_truth.py execute ./flash_powder/src/native_functions.rs
# Execute & substitute the results, write output to /tmp/foo.rs
./util/python_truth.py substitute ./flash_powder/src/native_functions.rs -o /tmp/foo.rs
# Execute & substitute into the input file.
./util/python_truth.py update ./flash_powder/src/native_functions.rs
When developing, something like this is usually helpful:
./util/python_truth.py update ./flash_powder/src/functional.rs && cargo t -- --nocapture
There's some helper tooling in ./util/valgrind to create suppression files against a C++ binary.
These ensure that we ignore some uninitialised values that valgrind finds in the bowels of LibTorch.
Run with these suppressions using valgrind through the runner;
./util/valgrind/valgrind.sh target/debug/deps/torch_stable-5f3b6c1dd8420412