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