autograph v0.1.0

[charles-r-earp]

This is the first release of autograph rebuilt on SPIR-V compute shaders that can be compiled from Rust source with rust-gpu!

Compute Shaders

All computations are implemented in either Rust or GLSL (to be replaced by Rust), and this API is publicly exposed so that external crates can develop their own routines. Shader code targeting SPIR-V is portable and is compiled at runtime for devices supporting Vulkan, Metal, and DX12 API's.

Datasets

The library includes MNIST and Iris datasets to make it easy to get started and these are used in examples.

Machine Learning

High level traits like Train, Test, and Infer are provided to create a common interface for different algorithms.

KMeans

An implementation of the KMeans classifier, demonstrated in the examples.

Neural Networks

Networks can be constructed as a structure of Layers, including:

• Convolutions
• ReLU
• MaxPool
• Dense

Each of these layers implement Layer and Forward traits, which can be derived to reduce boiler plate.

#[derive(Layer, Forward, Clone, Debug, Serialize, Deserialize)]
struct Lenet5 {
    #[autograph(layer)]
    conv1: Conv,
    #[autograph(layer)]
    relu1: Relu,
    #[autograph(layer)]
    pool1: MaxPool,
    #[autograph(layer)]
    conv2: Conv,
    #[autograph(layer)]
    relu2: Relu,
    #[autograph(layer)]
    pool2: MaxPool,
    #[autograph(layer)]
    dense1: Dense,
    #[autograph(layer)]
    relu3: Relu,
    #[autograph(layer)]
    dense2: Dense,
    #[autograph(layer)]
    relu4: Relu,
    #[autograph(layer)]
    dense3: Dense,
}

Similarly, backward ops can be defined using the Autograd and Backward traits, where Autograd can be derived in much the same way that Layer is.

#[derive(Autograd)]
struct DenseBackward {
    // Use vertex / optional_vertex for Variables and Parameters
    #[autograph(vertex)]
    input: Variable2,
    #[autograph(vertex)]
    weight: Parameter2,
    #[autograph(optional_vertex)]
    bias: Option<Parameter1>,
}

The intent is that users can write their own custom, modular layers and functions which can be defined from the high level down to custom shader code, all implemented in Rust.

Status

The crate is fairly minimal, missing implementations for some data types, not supporting bf16 for convolutions and pooling layers, with many functions like matrix multiplication internal and not publicly exposed. Things that are potential work items:

• Fully support bf16 in Neural Networks, with a nicer means to convert from f32 to bf16 and back for Variables and Parameters.
• Render the backward "graph" using petgraph for visualization and debugging purposes.
• Profiling tools for evaluating key functions / shaders and for improving the engine itself.
• Port GLSL to Rust, rust-gpu barriers are not working yet and need to reduce the need for code duplication particularly for bf16.
• Improve performance, particularly the GEMM implementation.
• Implement more operations and algorithms:
  • MeanPool is implemented but backward is not yet working.
  • Binary ops like addition are easy but not yet implemented due to uncertainty over API (in regards to Residual layers etc with more than 2 inputs).
  • SGD with momentum not yet implemented, implement other optimizers.
• Model parallelism supported but not tested or optimized. Data parallelism is intended to override Layer::update() to perform an all reduce (ie mean) over the the gradients for each parameter duplicated on several devices prior to the optimization step.

Contributors

Thank you to those that have contributed to the project!

• @AlbertoGP
• @nkconnor

---

This discussion was created from the release autograph v0.1.0.

[TrueDoctor]

Would it make sense to use wgpu instead of gfx-hal as gfx-rs no longer seems to be actively developed?

[charles-r-earp]

So actually initially I had MKL / CUDA and later HIP backends. And I tried to get OpenCL working but performance wasn't to par. I then tried wgpu / compute shaders, and I liked it so much that I decided to scrap everything else and rewrite targeting that, particularly after discovering rust-gpu.
However, while wgpu would probably be fine, it doesn't expose much in the way of error handling, so not only is there not a lot of control, there's often no way to catch errors / avoid panics. It seems like some of these issues are being addressed but that was the motivation for using the lower level gfx-hal api.
Unfortunately this didn't work correctly on Mac, so I tried again with wgpu and ran into similar issues, and then I settled on vulkano, which offers a higher level api but also exposes more, targeting Vulkan as well as MoltenVK for Apple platforms. This is more or less fully functional, and has about the same performance, but I'm currently trying to hit my performance goal before cleaning it up and merging.

[TrueDoctor]

Thanks, that is very insightful. I should probably give you some context on what I'm trying to do:
I'm working on a node based graphics editor called graphite.

At the core lies a node based visual programming language using nodes, with nodes essentially just being rust functions.
For effects like for example a blur, the idea ist to write one "kernel" that describes the output operation for one pixel and feed that into a node that slides a window across the soruce picture. All of this is fully hackable so if for example the user decided they didn't like the color bleeding of a regular blur they can just edit the "kernel" node, modify the rust code (probably via an embedded editor) to only take into account the average brightnes of the surrounding pixels.
For this to work we need to be able to turn a node graph desrcription into rust-gpu code, I worked on that yesterday and it already looks quite promising:

# [cfg (target_arch = "spirv")]pub mod gpu {
# [repr (C )]pub struct PushConsts {
    n : u32 , node : u32 ,
    }
    use super :: * ;
    use spirv_std :: glam :: UVec3 ;
    # [allow (unused )]# [spirv (compute (threads (64)))]pub fn "add"(# [spirv (global_invocation_id )]global_id : UVec3 , # [spirv (storage_buffer , descriptor_set = 0, binding = 0)]a : & [u32 ], # [spirv (storage_buffer , descriptor_set = 0, binding = 1)]y : & mut [u32 ], # [spirv (push_constant )]push_consts : & PushConsts , ){
    fn node_graph (input : u32 )-> u32 {
        let n0 = graphene_core :: value :: ValueNode :: new (input );
            let n1 = graphene_core :: value :: ValueNode :: new (1u32);
            let n2 = graphene_core :: ops :: AddNode :: new ((& n0 , & n1 ));
            n2 . eval ()
        }
        let gid = global_id . x as usize ;
        if global_id . x < push_consts . n {
        y [gid ]= node_graph (a [gid ]);

        }

    }

}

(This has to be the most cumbersome way to increment a value by one possible) But it is only generated from a few nodes that could be configured in a user interface, so that is quite neat.

I originally started to dabble around with vulkano + rust-gpu but could for the life of me not get the push constants to work because they weren't picked up for some reason by vulkano even though they were present in the assembly.
While searching the internet for usage examples, I think I clicked on the list of dependants for vulkan on crates.io and that is how I found autograph. Even though I don't do ML I instantly fell in love with the simple and streamlined dev experiance.
Gpu computing always has this high barrier of entry just for the setup to get your code runnig which has always bothered me when I tried to use the gpu to accelerate computations. Part of my goal is to make this significantly easier going forward so a broader audience can use gpu computing in their projects.

Because the rust ui framework landscape is not very mature quite yet, we opted to for the moment make our editor web based to get platform agnosticism without much as much dev overhead, but that mean that our code has to be able to be run in the web browser + also since gfx-hal seems to no longer be actively developed we would like to opt for wgpu going forward.
Two days ago I read through some of the code in engine.rs and tried to estimate how hard it would be to integrate wgpu.
I came to the conclusion that it would probably be less efford for me to write a minimal gpu runtime witth a similar interface to autograph, mainly because that forces me to understand everything that is going on under the hood + might be more light weight as a result (we still have to distribute the final binary to users and want to avoid page load times as much as possible). The actual on the fly compilation would probably be done by a dedicated compilation server for now because it is not feasible yet to execute rustc on wasm.

So what is my point in writing all of this: I really love the idea of autograph but we have slightly different reqirements and I'm wondering if it makes sense to adapt autograph to our needs or first experiment with writing a minimal runtime and then potentially integrate it in the future.

Also as a side node, would it be beneficial for your users to have a visual interface to design their nn? If so we could potentially collaborate on that. It is quite tricky to find proper abstractions that retain all the power provided by the gpu e.g. local(shared?, workgroup?) memory.
Anyway, I'm working on this over here: https://github.com/GraphiteEditor/NodeGraphExperiments/tree/async
And thanks again for your very comprehensive reply, I didn't even really consider error reporting/handling as a factor before

[charles-r-earp]

Awesome! This looks very cool.

As far as the interface I was inspired in part by ocl.

When I ported to wgpu and then vulkano, I was hoping to do more validation in with .compute_pass(), potentially make some calls safe. In fact, I even thought about doing inline gpu code like accel. Anyway, essentially the only validation I'm currently doing is checking that the push constant range (ie the size of the struct) in the source matches the struct (or series of arguments) provided. Likewise, I also check to make sure that an immutable buffer argument is not modified in the shader. Initially I was doing this via a requirement to mark an argument as "NonWritable", ie this should be equivalent to a &[_] in rust. I actually did a PR to rust-gpu to add this but for whatever reason they don't use that. So I have some code in shader.rs that scans entry functions for every use of an argument (and it's derivatives) and also determines whether it's mutable or not.
wgpu does this now with naga, but naga wasn't and still may not support certain functions, particularly atomics for my use case. But that would probably be a good start and you can use that with or without wgpu.
I did run into the same issue with vulkano, push constants worked on glsl but not with rust-gpu. It looks like push constants can be declared differently, glsl has it's own way, and vulkano expects that and isn't conformant to the spec. Maybe that could be patched.
If you compile the source from rust then validation is less of an issue, but since autograph takes spirv then it has to leave it up to the user to ensure that the arguments provided are the right types, in bounds, etc.

In terms of basic desgin, autograph uses precompiled rust-gpu shaders to spirv. These are then included in the binary and loaded at runtime. This is easy to extend externally, as users can emulate this pattern, and don't need to provide sources at compile time, or at device creation. It even supports JIT ie runtime compiled shaders, which would honestly be extremely useful and is commonly used in HPC and ML libs, but seems like it would be difficult with rust-gpu.

If you want to have user provided gpu code in rust, then you either need to package the nightly compiler for that rust-gpu version within your executable / package, or the user has to install it themselves. Maybe you can ask to install it during install idk. But that's kind of a problem in terms of user generated code, which is why autograph uses precompiled sources to allow it to run on stable. Maybe if this is a power user feature, requiring the extra install might be less of an issue.
I'm also not sure if this is fast enough for hot reloading. Ideally, you could partial compile a function, then link to kind of a template (ie the entry function, push constants, etc), point being that you want to retain compilation artifacts, deps, so that it's just recompilation each time. A naive approach using cargo / spirv-builder would probably take too much time, and take too much space, as each one has to compile it's own set of deps. This is something I might experiment with but it looks challenging.

One idea I had was targeting naga instead, and then you can potentially link functions together in naga at runtime, override constants, etc, eliminating the expensive compile.

So if I'm understanding you correctly, it looks like each node is a function (with some gui parameters), that can be coded in rust, and then those are fused together into a single gpu shader that is executed? Or at least, a few ops are fused and then run in some sequence?

 fn node_graph (input : u32 )-> u32 {
    let n0 = graphene_core :: value :: ValueNode :: new (input );
        let n1 = graphene_core :: value :: ValueNode :: new (1u32);
        let n2 = graphene_core :: ops :: AddNode :: new ((& n0 , & n1 ));
        n2 . eval ()
    }

Is this part meant to be generated or provided by the user? This is from the node graph? And then each node is a struct implementing a trait?

It seems like it would be nicer for the user if this was just a function, or even just an closure:

// Add 
|n0, n1| n0 + n1

This means that you essentially have duck typing which might be a lot nicer if you're doing arithmetic, an it would fit into one line in a text box for a very simple ui. The argument names can correspond to the inputs of that node. Note that you can do inline structs, traits, and fn's. But obviously if the user wants to import external crates, that's tricky. I assume this is probably going to be simple enough that won't be necessary.

If you look at autograph_derive I have an autobind proc macro that adds the descriptor_set and bindings, which is particularly useful when working with macros and generated code, reordering of arguments, etc.

autograph is dynamic ie imperative, similar to pytorch, so it doesn't know what operations will occur before they are called. This leads to say input x, alloc y, compute y, drop x, alloc z, compute z, drop y, etc. And a naive system would have to wait for the gpu to finish after each alloc / dealloc. Further, it sub allocates large buffers, to again reduce the cost of allocation, and it can do this concurrently. AFAIK wgpu allocates an entire buffer for each buffer, and destroys it when you drop it, which isn't great for a lot of small buffers.
If you know all the ops you will do ahead of time, then this isn't as useful as you could allocate exactly what you need, you don't have to worry about creating a lot of temp buffers for intermediate values. You could potentially even reuse command buffers, but I only see this being useful if you were doing a lot of images / or video, since they would need to be recreated if you change parameters / nodes. If you can fit all of your inputs into one buffer, you can do a bindless system, with dynamic offsets, or use push constants for offsets. This is in theory much faster, but might not be worth it.

[TrueDoctor]

Sorry for the delay, I was really busy last week but do have some free time again now.
The inline code in accel looks really cool, I think I can pretty much just reuse the derive macro for compiling rust-gpu code.
The input sanitization/checking is something I have not put that much thought into until now because in my use case I have other mechanisms that already do the type checking for me, but it would be great the have a more robust, ideally usecase agnostic, system to deal with that. We might be able to get them to merge the mutability checks after all sounds quite useful, I'll keep that in mind.
Using naga to link the shaders ad hoc sound like a great idea, we could basically compile the individual nodes via a proc macro and include them in the binary and just link them at runtime as a tier1 assembler. This could potentially be further optimized in the future by having a comilation server (akin to the rust playground) that can compile a non realtime optimized version of the shader.

So if I'm understanding you correctly, it looks like each node is a function (with some gui parameters), that can be coded in rust, and then those are fused together into a single gpu shader that is executed? Or at least, a few ops are fused and then run in some sequence?

Each node is a rust function (with a bit of extra magic like potentially a state).

 fn node_graph (input : u32 )-> u32 {
    let n0 = graphene_core :: value :: ValueNode :: new (input );
        let n1 = graphene_core :: value :: ValueNode :: new (1u32);
        let n2 = graphene_core :: ops :: AddNode :: new ((& n0 , & n1 ));
        n2 . eval ()
    }

> Is this part meant to be generated or provided by the user? This is from the node graph? And then each node is a struct implementing a trait?

This would be generated from the visual representation of the node graph, this is obviously a pretty complicated example but it should be possible to reach turing completenes by just visually chaining nodes because many of our users won't be programmers but rather artists and forcing them to write actual code might be a bit intimidating.
But closures are of course still supported, you can just add an fn node in the editor which basically allows you to enter an arbitrary closure which is then turned into a node. I just wanted to provide an example that showcased the chaining of different nodes.

Further, it sub allocates large buffers, to again reduce the cost of allocation, and it can do this concurrently. AFAIK wgpu allocates an entire buffer for each buffer, and destroys it when you drop it, which isn't great for a lot of small buffers.

Yeah I stumbled upon that and wasn't sure how to best deal with that. I could just alloc one big buffer from wgpu and use the rust alloc api + any of the already implemented rust allocators (e.g. wee alloc) to manage the memory. I'd really like to avioid implementing my own allocator because that one would probably be much buggier and less efficient than tested solutions.

On another note, do you potentially have time for a call at some point?
I feel like that might be a more efficient way get me up to speed with the topic. You could just join our Discord and shoot me a message there, or suggest something else if your interrested. (my timezone is GMT + 2)

[TrueDoctor]

We could also try to get @kvark or @cwfitzgerald on the same call I feel like I need any input I can get if it is supposed to come out halfway decent in the end :D