Donald Tolley 98f87edd23 Tensor .cuda() very slow with specific array sizes (#138964) ... ### **Pull Request: Optimized Non-Contiguous Tensor Copy for CPU to GPU in PyTorch** #### **Summary** This PR addresses the performance issue identified in [#111570](https://github.com/pytorch/pytorch/issues/111570), where non-contiguous tensors took significantly longer to transfer from CPU to GPU. Through detailed tracing of the call flow, we identified that PyTorch was creating temporary contiguous buffers for non-contiguous tensor transfers, which introduced unnecessary overhead. #### **Tracing the Issue** To pinpoint the cause of the slowdown, we followed the call flow from Python’s `tensor.cuda()` method through PyTorch’s backend, ultimately identifying `copy_kernel_cuda` as the key function responsible for CPU-to-GPU tensor transfers. Here’s a summary of the tracing process: 1. **Python Call: `tensor.cuda()`** - Starting from Python, the `cuda()` method initiates the tensor transfer to the GPU. 2. **`TensorBody.h: cuda()`** - The `cuda()` method calls `to()`, specifying the target device as CUDA. 3. **`Tensor.cpp: TensorBase::to()`** - The `to()` function prepares device and data type options before invoking `_ops::to_dtype_layout::call()`. 4. **Operator Call: `_ops::to_dtype_layout::call()`** - This operator dispatches the request to the backend-specific function responsible for managing the transfer. 5. **`Copy.cpp: copy_()`** - The `copy_()` function performs preliminary checks (e.g., zero-tensor immutability) and proceeds to call `copy_impl()`. 6. **`Copy.cpp: copy_impl()`** - This function sets up a tensor iterator and dispatches the copy operation to the appropriate backend through `copy_stub`. 7. **Dispatch to CUDA: `copy_stub`** - The dispatch mechanism routes the call to the CUDA-specific function, `copy_kernel_cuda`. 8. **`Copy.cu: copy_kernel_cuda()`** - Here, we identified that PyTorch was creating temporary contiguous buffers for 1D and 2D non-contiguous tensors, which slowed down the copy process. This behavior is managed by the `copy_requires_temporaries()` function. #### **Solution** To address this, we modified `copy_kernel_cuda` to handle non-contiguous 1D and 2D tensors directly by using `cudaMemcpy2DAsync`, which allows efficient, stride-aware memory transfers without temporary buffers. Here’s why this approach improves performance: - **Efficiency of `cudaMemcpy2DAsync`**: This CUDA function is optimized for pitched (stride-based) memory transfers, allowing it to handle non-contiguous data layouts effectively by specifying memory strides for source and destination tensors. - **Reduction of Overhead**: By directly copying non-contiguous tensors without intermediate buffers, we eliminate extra memory allocation and achieve faster CPU-to-GPU transfers. - **Asynchronous Execution**: `cudaMemcpy2DAsync` enables asynchronous transfer on the CUDA stream, further improving performance by taking advantage of CUDA's optimized memory handling for non-contiguous layouts. #### **Performance Results** In my testing, I created tensors of size `327680 x 2000` and used slices for transfer performance measurements. The tests show that the average time for transferring a non-contiguous slice (e.g., rows 10,000 to 50,000) from CPU to GPU now closely matches the contiguous case. This improvement indicates that the updated implementation effectively addresses the performance discrepancy. Below are the measured times and validation checks: ```plaintext Average time for contiguous slice (rows 10,000-50,000): 66 ms Average time for non-contiguous slice (rows 10,000-50,000): 66 ms Validation of contiguous and non-contiguous tensor copies: ✅ PASS: Tensor shapes match. ✅ PASS: Tensor contiguity matches. ✅ PASS: Tensor contents match. ✅ PASS: Tensor data types match. ✅ Success: Both contiguous and non-contiguous tensors were copied correctly to the GPU. ``` #### **Conclusion** This PR resolves the identified performance issue by eliminating the need for temporary buffers in non-contiguous 1D and 2D tensor transfers, ensuring faster and more efficient copies from CPU to GPU. Future optimizations could further enhance performance for higher-dimensional non-contiguous tensors. Pull Request resolved: https://github.com/pytorch/pytorch/pull/138964 Approved by: https://github.com/jeffdaily Co-authored-by: Natalia Gimelshein <ngimel@gmail.com> Co-authored-by: Jeff Daily <jeff.daily@amd.com>		2025-01-31 17:05:02 +00:00
..
conda
src	Tensor .cuda() very slow with specific array sizes (#138964)	2025-01-31 17:05:02 +00:00
tools
CMakeLists.txt