Add an option to getWriteableTensorData to avoid copy CUDA tensor to CPU

[mrshenli]

Stack from ghstack:

• Enable GPU-to-GPU comm in TensorPipeAgent #44418 [WIP] Use streams from pool on RPC callees
• Add an option to getWriteableTensorData to avoid copy CUDA tensor to CPU #46524 Add an option to getWriteableTensorData to avoid copy CUDA tensor to CPU
• Enable TP_USE_CUDA and TP_ENABLE_CUDA_IPC #46523 Enable TP_USE_CUDA and TP_ENABLE_CUDA_IPC

Differential Revision: D24392794

[facebook-github-bot]

💊 CI failures summary and remediations

As of commit 39396f4 (more details on the Dr. CI page):

---

• 1/1 failures possibly* introduced in this PR
  • 1/1 non-CircleCI failure(s)---

Extra GitHub checks: 1 failed

• Failed: GitHub Actions - clang-format

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 3 times.

[dr-ci]

💊 CI failures summary and remediations

As of commit ad3f662 (more details on the Dr. CI page):

---

✅ None of the CI failures appear to be your fault 💚

---

• 2/2 broken upstream at merge base 9300a27 on Oct 16 from 10:29am to 4:01pm PDT (10 commits; 9300a27 - e6e83bf)

---

🚧 2 fixed upstream failures:

These were probably caused by upstream breakages that were already fixed.

Please rebase on the viable/strict branch (expand for instructions)

Since your merge base is older than viable/strict, run these commands:

git fetch https://github.com/pytorch/pytorch viable/strict
git rebase FETCH_HEAD

Check out the recency history of this "viable master" tracking branch.

• pytorch_linux_xenial_py3_clang7_onnx_ort_test2 on Oct 16 from 10:29am to 4:01pm PDT (10 commits; 9300a27 - e6e83bf)
  • 🔁 rerun
• pytorch_linux_xenial_py3_clang7_onnx_ort_test1 on Oct 16 from 10:29am to 4:01pm PDT (10 commits; 9300a27 - e6e83bf)
  • 🔁 rerun

---

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 14 times.

[lw]

Looks good to me, but perhaps someone from the JIT team should also take a look?

[mrshenli]

Hey @wanchaol, I added this because we need to access WritableTensorData in RPC and would like to avoid the cpu-to-gpu copy. Please let me know if this looks OK to you. If not, we can also try expose WritableTensorData struct from JIT or replicate the logic in RPC.

[wanchaol]

This looks good to me.

[facebook-github-bot]

@mrshenli merged this pull request in 5003fd1.