Multi device bugfix

[asglover]

This is a PR to harden multi device behavior.

• It adds testing to make sure all tensors are on the same device
• it adds usage of "device guards". This is how PyTorch signals to apis which do not explicitly take device, which device to use.
• It also adjust how BLAS handles work. BLAS handles must differ across devices. PyTorch keeps a pool of BLAS handles. This PR switches OEQ to using PyTorch's BLAS handles. Just a note that this current design does not set the BLAS precision mode which means it inherits from pytorch settings. I think that's a reasonable design because people can control precision with the idiomatic pytorch methods. If you think this calculation should never be in TF32, then I can add a manual setting command.
• It adds stream tests to that should fail before the changes from the PR come in on multi device. And should be fixed later.

[asglover]

This is a draft PR for now because I'm making these changes based on pytorch documentation and recommendations for multi device. I need to get access to multi device machines in order to test properly. The lack of AMD machine availability for testing is very annoying. I don't have any access through work (currently). And AMD didn't respond to my application for developer grants for this project.

[asglover]

@vbharadwaj-bk I tested this on a multi device host and verified the problems the branch is trying to address exist in main and verified that using the device guard works as intended. There are a few aspects to this change.

• having the streams come from the correct stream for the device.
  - This happens at the C++ level by using the torch APIs defined for the purpose.
• using the pytorch cuBLAS handles.
  • BLAS handles are one per device, and the easiest way to handle this is to reuse the handles that pytorch hands out to other pytorch ops for calling blas
  • from rocBLAS to hipBLAS because torch uses hipBLAS
• It makes the workspace be a "buffer" associated with a module, this will make the buffer follow the module around as you send it from device to device. It also selects "persist" = false, so that the the workspace is not serialized during export. Which is nice because they are just zeros.
• I added an all on the same device check to the C++ level. It's different than the existing check tensor, I can unify the mechanisms or leave it separate.
• Adds a check that CPU pointers are on the right device.

I have not tested this on a AMD machine. I cannot get access. I need help testing this aspect of the change.

[vbharadwaj-bk]

This rolls in a few different changes. Some comments:

1. Prefer the stream_ordering_test.py bio e in a separate PR (and would like a few more details about the failure mode you are trying to catch here).

2. The multi device changes are good.

3. [No action needed] Regarding the BLAS handle - there's a tricky problem with the existing codebase. cuBLAS is not ABI-stable across major versions, so if the user has a CUDAToolkit that's not the same as PyTorch's version (which could well be the case, as PyTorch installs just the NVIDIA libraries it needs from a Python distribution channel), then we have a mismatch and may not be able to find certain symbols at link time. So far, it has worked because we load torch first and the symbols it loads are consistent with the ones we have compiled against.

Switching to PyTorch's BLAS handle is a step in the right direction; there's no good way to solve this problem without distributing multiple package versions.

Can run some of the testing in a bit.

[asglover]

Thanks for the review.

1. I'll take a look at this and come back with some changes.

2. Thanks!

3. This is an interesting point. I'll think more about this. It would be nice to be able to somehow use the same cuBLAS library as PyTorch when compiling our package. If we precompile, we'd have to do it for multiple versions. I wonder if by luck, the subsection of the API that we use happens to be stable? hmmm, thanks for calling this out.

Whatever the solution, it would be nice to have a strategy that we can port for other accelerators.