Support for GPU/CUDA

[haixuanTao]

I tried this package on the exemple Pytorch and it seems that there's currently no GPU/CUDA support.

Is it a feature that is going to be added in the future?

Thanks in advance for this package :)

[kali]

Thanks for your interest

This is correct. There is no CUDA or GPU support, and it is not on the roadmap. We are not putting too much energy on the PC platform, tract main focus being arm32 and arm64 devices. If something happens on the GPU, it won't come from us.

[haixuanTao]

Make sense.

Best option for GPU ONNX on rust at the moment seems to be https://github.com/nbigaouette/onnxruntime-rs

I'll close the issue.

[dvc94ch]

That might make sense for training the nn on a gpu. Support for gpu in tract is probably more along the lines of doing inference on a mobile gpu. Might make sense to reopen the issue as it's not resolved and other people might be interested in the status of this.

[dsseng]

+1 for reopening. Yes, I beleive support for, e.g., Vulkan, would boost the library capabilities. I have found this project and liked it because Rust is definitely more maintainable than C. With hardware acceleration tract will be a really great tool for mobile as well (mobile devices have GL/CL/Vulkan to accelerate inference, and those are a lot more power-efficient). I would like to work on some HW acceleration approach if such is wanted.

One of ideas is to create a loadable C-ABI libraries for accel methods. It would require creating an API and binary-stable interface to C-style libraries (built with Rust, C/C++ or many other languages), which are discovered and loaded by core using libloading. They would need to have attributes such as AVX2 (arch features necessary to load, btw, SIMD could also help accelerating inference). API should include generic buffer management functions (create, free, pass to dev/to CPU, get location (shm, unified memory, dedicated memory, RDMA etc), import/export platform buffer) and operators' implementations.

I propose implementing either wgpu or direct Vulkan backend. Former is preferred, as it should be able to activate GPU nearly on anything: Windows, Linux (incl. Android) and macOS + MVK with Vulkan backend, Metal for macOS and iOS, fallback OpenGL on Linux (and Android) with older GPUs, DirectX fallback for Windows when Vulkan is not available, WebGPU with WebGL fallback with WebAssembly.

[haixuanTao]

Agreed. As GPU optimisation is completely different. That's great that you split the method into decluttering and optimization.

[kali]

Inspiration comes from LLVM really, and it does show in the wildly inconsistent naming. I think of tract-core as a MIR, while into_optimized translate into a LIR space. into_optimized actually calls the "codegen" method on operators.

[dsseng]

Ah, then I misread that. Exactly, if model simplification takes place earlier, we should split a new backend after there

[haixuanTao]

As your mentioning LLVM, have you checked MLIR ? And do you believe that it could be a better Intermediate representation in the future. I feel like ONNX and tensorflow-IR are both often not enough low-level.

Just curious, as MLIR is still pretty young and support pretty low.

[dsseng]

Hmm, it should be definitely looked at, but for now it's too young, as you mentioned. Using LLVM and MLIR should be some long-term goals

[kali]

FYI, tract-linalg already provides runtime switchable implementations for some high cost operations (matrix multiplication, sigmoid and tanh, mostly). It is accelerated already on SIMD (with specific code for intel v2, and some armv7 and arm8 configurations).

One possible path to accelerating tract when a GPU is available is to implement the matrix multiplication on GPU. I think there is a MVP here with local changes only (in tract-linalg). We could then move on to lowering more operators in tract-linalg, discuss buffer locality and stuff, that would require some awareness from tract-core and tract-data.

One other option is to use tract-core as infrastructure and semantic pivot, but instead of optimizing the network in tract, substitute some operators with external ones from another crate that provides acceleration. This path is used by some teams who have specific hardware with capabilities they want to use without exposing themselves, as it allows things to happen outside of tract. This is not a requirement, such a crate could also be part of tract. Also, a source of friction for external consumer of tract-core is tract-core stability, on which I have not make commitments so far as I considered it internal, even if the API is de-facto pretty stable.

The reason why I closed this as an issue is not that I don't want it to happen: I don't have resources for it. tract is mostly a one-man-show so far, so I must pick my battles. If somebody else picks this one up and start working on it, I can help with guidance and discussing some changes to make it easier, but I can't do the heavy lifting. Until somebody does, this conversation is just daydreaming while I believe issues have to be actionable.

I'm gonna move this topic to the discussion section at this stage.

[dsseng]

Acceleration architecture is like one of TensorFlow with PluggableDevice (lacks docs, but looks great).

[kali]

Let's start here: https://github.com/sonos/tract/blob/main/doc/intro.md :)

[kali]

About the buffer management cost... the rationale with "starting" with the matrix multiplier is that it is the most expensive thing to run, and the one GPU are the most suited for. So we may already have a gain by running these operations on GPU even factoring in the transfert cost, while leaving the cheaper stuff to run on the CPU. This would be a milestone / POC that do not require flipping all of tract codebase on its head.

[dsseng]

Thanks for docs, sorry for not finding it before. Will read and think more about structure and plugging GPU in there

[dsseng]

After discussing I can try to make a PoC with GPU and compare running time of examples, such as ALBERT (it is rather slow on my machine, 3.5 s). It's always better to bench rather than think about cost of such things.

[dsseng]

Trying to create a stateful extension for Ops within linalg crate to implement GPU ops in there first. I guess it would be zero-copy with integrated solutions like my Ryzen laptop or mobile/embedded SoC-s.

[dsseng]

But I do not believe that it is going to be zero-copy by default, might need some hacks around it.

[dsseng]

Are only 4x4, 4x1 and 3x2 matrices being multiplied? Weird. If larger ones are somehow splitted, how can that be disabled to handle full size on the GPU? The crate does not have enough internal docs to be understood and abstracted easily, I might need a direct consulting to get something done. Do you have some chat room like Matrix for that?

[kali]

Lol. No. tract can multiply matrices of any size. The strategy is pretty classic, looping over tiles computed by a micro kernels. See https://tech-blog.sonos.com/posts/the-anatomy-of-efficient-matrix-multipliers/ for an intro on the topic, or the various papers by GotoBLAS and BLIS teams.

You probably don't want to plug inside linalg to POC this. If you do so, you need to implement the MatMul trait and tweak the arch and feature detection in linalg/src/lib.rs to plug yourself in. But it will be some work, because you will also need to implement the fusable operations. It may be better to plug also at the tract-core level, just like @haixuanTao is planning, but just substituing the MatMulUnary operator by a GPU based one.

[dsseng]

Hmm, will try. What about chat, would you be comfortable with that? It would be more productive because of ease to explore the internal structure.

[dsseng]

I ran into a lot of trouble while trying to plug it into linalg ops. E.g. ownership and Clone issues (you cannot easily clone shader object handles and so on to ensure safety for GPU stuff). So I need to explore more and, I guess, plug new impl instead of core or draw some extra split lines between and within modules (but that would be great, since we just upload an input buffer object and GPU does everything giving us a copyable output tensor in the end). Definitely would be quite a large PR.

[haixuanTao]

So, I'm ok to have a chat if need be, anytime you want. I am absolutely not familiar with Tract however. And so @kali should participate if you have Tract questions.

From what I understood the plan is to use wonnx or part of the wonnx implementation for allocating the data and the computation but use Tract-Core instead of the current wonnx-parsing and optimisation code.

I feel like that is not what I think you have in mind. I won't be able to help you out in replacing linalg ops if that's what you are planning to do.

[dsseng]

Well, I don't feel the need in chat now, @ kali explained everything pretty well. I will see the code, but not sure about directly integrating Wonnx into Tract, maybe I will copy some parts of the code when contributing, but not everything directly (I guess architectures are different enough).

[dsseng]

Part of core features is platform-independent (Graph, Nodes, Inlets & Outlets etc). So, I guess, optim and ops should be split out into backend-cpu, and backend-wgpu crate should be created with analogic functions (then something like backend-coral can be made for specialized accelerators or other cases). They should be available as build flags and if both are selected, then TRACT_BACKEND env var will specify priority of backends.

[haixuanTao]

So, you might have to use a CPU for some ops for sure. In the current wonnx design, we do not implement this, but it's a clear limitation. You might not be able to compute atanh for example.

There are also memory and thread constraints on GPU that might force you to fallback on CPU as well.

Yes, you can't read RAM from the GPU, you have to create a buffer between the two.

[dsseng]

atanh(x) is listed within WGSL reference for computation precision. We'll see what needs to be copied in between and check for any optimizations that might improve the perf (needs benches across platforms, btw, I posess Ryzen 5 4500U laptop with unified memory, so driver should just give the pointer to app I guess). If targeting mobiles I guess copying is not a problem at all because of UMA. Discrete desktop GPUs might require benchmark.

[haixuanTao]

Ok, so I didn't thought about integrated GPU and you actually can do a 0 copy apparently: https://www.reddit.com/r/rust/comments/s24uqt/unified_memory_on_m1_macs/?utm_medium=android_app&utm_source=share

I think they recently added atanh. My point was that some ops might not be easily implementable as there might not be an instruction within WGSL.

[dvc94ch]

advanced gpu features like subgroups aren't exposed in wgpu. but for now I guess getting it working is good enough, optimization can happen later. piet-gpu was so unsatisfied with wgpu that @raphlinus wrote his own gpu abstraction layer.

[dsseng]

Well if it goes fine we can then either use wgpu-hal or write direct Vulkan, Metal and DX12 backends. Definitely not for the start

[dsseng]

This returns proper model description, and then some backend methods should assemble this graph on GPU (allocate buffers, set BO properties to them (input and outputs, and maybe something that is computed on CPU should be accessible to it) and then execute on GPU. Am I right, or is there something else existing codebase can help me with and I skipped?

        tract_tensorflow::tensorflow()
        // load the model
        .model_for_path("mobilenet_v2_1.4_224_frozen.pb")?
        // specify input type and shape
        .with_input_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec!(1, 224, 224, 3)))?
        // optimize the model
        .into_typed()?
        .into_decluttered()?

[dsseng]

As far as I understood, into_runnable does create some sort of plan which is then CPU-bound. I didn't really understand how does the plan thing work, but it might be a universal mechanism I should not skip. @kali could you please elaborate to keep me from doing something useless.

[dsseng]

At least Plan needs to be forked to contain buffer pointers instead of buffers in RAM, but smth like order could be reused, but unsure whether that is worth it.

[kali]

The Plan is basically about finding the right order in which to call the operators, and some management of variables dependencies (we want to free intermediate tensors when we are done with them). This is established when you build up the Plan (so when into_runnable() is called). The plan is then an immutable shareable object.

In some cases, we also need a Session, to manage state (for stateful networks). But in the simple stateless case, the user can call Plan::run (it will create a throw-away session).

Now if you just substitute just some operators for more efficient ones before calling into_optimized().into_runnable(), they will implement the same EvalOp::eval method, so into_optimized() and the plan will keep working the same way on the rest of the network. If you the entire model replacement way, then once you've translated the TypedModel, what to do with it depends entirely on what you've implemented on your side. You may need to keep the TypedModel around or not butI don't think optimizing it or using the plan would make much sense at this stage.

[dsseng]

I do not plan to do EvalOp, since that will be less performant for discrete memory case. I want buffers to be allocated on gpu. So, I guess, I should take Plan, but not State.

[dsseng]

And SimpleState will be forked to GpuState in new crate gpu. Looks quite extensible, could be later used to play with MLIR or so. Thanks for getting those ideas to me!

[dsseng]

Am I going the right way? I'm still on the way of forking existing classes to accept GPUTensor instead of Tensor and use GPUOp instead of Op. The key difference is that GPUTensor should hold (doesn't now) pointer to VRAM, not actual data. I guess I'm stuck with doing something entirely wrong.
Diff of trials attached, it is not expected to be buildable:
gpu-7.txt

[dsseng]

Btw is that a bug? When SimplePlan marks two ops as parallel, it keeps an empty element after:

        [
            80,
            91,
        ],
        [],
        [
            93,
        ],

[kali]

I don't think it is bug. There is no "marking ops as parallel" in SimplePlan, btw. I'm assuming this comes from the flush_list: it's the id of previous ops which results can be results once you're the same-index op.

or instance, if #0 is a Source, #1 a Source, #2 is Add(#0,#1), then the flush list should be [ [], [], [0,1] ] : when #2 is done, you can free outputs for #0 and #1.

[dsseng]

Okay. Is everything right with the idea of reading SimplePlan and then doing entirely GPU-specific stuff with own GPU ops?

[kali]

I think so. You may not even need the SimplePlan entirely. The eval_order function may be enough.

[dsseng]

Great. Now thinking about buffers with generic code. Need to allocate them for each step, pass over to next and drop those not necessary for further computations or output.

[dsseng]

Looks like #697 will now be the place to discuss the implementation as initial GPUTensor code is already there.