search

LLNL / acrotensor

A C++ library for computing large scale tensor contractions.
MIT License
36 stars 3 forks source link

GPU-only Tensors #3

Open Linux-cpp-lisp opened 3 years ago

Linux-cpp-lisp commented 3 years ago

Hi all,

I'm considering trying to integrate acrotensor with PyTorch to use it in a model. It looks like it should be fairly trivial to create an acro::Tensor that references the GPU memory of some_pytorch_tensor.contiguous(); the only issue is Tensor.cpp:L192:

if (ddata != nullptr) {
      ACROBATIC_ASSERT(hdata != nullptr, "Acrotensor does not currently support GPU only tensors.");

It seems to me reading the rest of Tensor.cpp that this isn't actually true — it seems like as long as you never try to move the Tensor back to host it should be fine — but I was hoping to confirm this.

If anyone has previously tried to integrate this with PyTorch or has any thoughts on that, I'd also be grateful to hear them.

Thanks for your work on this library!

acfisher commented 5 months ago

I obviously haven't looked at this in quite a while and had no idea folks were looking at it. I believe that should work as long as you don't call anything that needs the data on the host. Did you end up giving it a try?

Linux-cpp-lisp commented 5 months ago

Hi @acfisher ,

It's also been a long time since I looked at this, but I did implement PyTorch bindings at one point, which I think worked fine with regard to this issue. The code is here if you are interested: https://github.com/Linux-cpp-lisp/acrotensor/tree/pytorch.