SIGN IN SIGN UP

Fix StorageBuffer size in PrepackNode (#2811)

Summary:
Pull Request resolved: https://github.com/pytorch/executorch/pull/2811

Fixing a bug I wrote, but in my defense the bug was in the code I copied. 😛

This change decreases the size of our staging `StorageBuffer`s which were previously bigger than necessary.

There's a few numbers describing size at play:
```
numel -> CPU buffer
nbytes -> numel * cpu_dtype
gpu_numel -> GPU texture
gpu_nbytes -> gpu_numel * gpu_dtype
```

Firstly, `StorageBuffer`'s ctor takes `numel` not `nbytes`. Secondly, we should use the CPU size not the GPU size which may be aligned up to a multiple of 4.
ghstack-source-id: 220995380
exported-using-ghexport

Reviewed By: copyrightly, SS-JIA

Differential Revision: D55619076

fbshipit-source-id: 9a63ee6111ce9b16b1aecba051c29a662e6d12dd
J
Jorge Pineda committed
b1efd43cdf7f3ac6b578e18610ff40d2514300fa
Parent: 31d5c61
Committed by Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> on 4/2/2024, 9:15:10 PM