PyTorch PR fixes opcheck preserving pinned CPU memory for CUDA ops
Custom CUDA ops writing to pinned memory now pass opcheck without illegal access errors.
PyTorch's custom operator testing framework, opcheck, copies inputs before running generated checks to prevent mutations of user arguments. The copy path used clone_input, which preserved tensor strides but allocated ordinary CPU storage for pinned CPU tensors. This broke CUDA custom ops that intentionally write to pinned host memory—they would pass eager execution but fail opcheck with an illegal memory access because the pinned tensor was replaced with pageable memory.
The fix ensures clone_input and its fallback clone path preserve pinned CPU storage. For AOT opcheck, the per-run input copy is moved outside the traced function, using an AOT-specific differentiable clone helper that re-pins CPU tensors and keeps Dynamo tensor attributes. This avoids tracing unsupported pin_memory operations while preserving existing non-leaf gradient-check behavior. The alternative—building pinned copies inside the traced AOT function—was rejected because functionalization rewrites the copy to an out-of-place copy that drops pinned status, and aten._pin_memory is not supported by FakeTensor.
- Opcheck now preserves pinned CPU memory when copying inputs, fixing illegal memory access for CUDA custom ops that write to pinned host memory.
- AOT opcheck uses a new differentiable clone helper outside the traced function to re-pin CPU tensors without tracing unsupported pin_memory.
- The fix passes multiple tests including test_opcheck_preserves_pinned_memory_by_default and test_aot_autograd_check_degenerate_cases on both CPU and CUDA.
Why It Matters
Enables reliable testing of custom CUDA ops that rely on pinned memory, preventing silent failures in production pipelines.