Not similar in the way of "Decorate any function and now it's a thread on the GPU", but Candle been pretty neat for experimenting with ML on Rust, and easy to move things between CPU and GPU, more of a library than a DSL though: https://github.com/huggingface/candle

It seems somewhat similar to rust-gpu https://github.com/Rust-GPU/rust-gpu

I'm not totally sure what it is, but I believe there is something for running Rust code on the GPU easily

You could use wgpu to replicate this demo.

https://wgpu.rs

I hate doing self-promotion, but this is basically exactly what CubeCL does. CubeCL is a bit more limited because as a proc macro we can't see any real type info, but it's the closest thing I'm aware of. Other solutions need a bunch of boilerplate and custom (nightly-only) compiler backends.