search

pytorch / torcharrow

High performance model preprocessing library on PyTorch
https://pytorch.org/torcharrow/beta/index.html
BSD 3-Clause "New" or "Revised" License
642 stars 78 forks source link

Use rootMemoryPool to construct execContext (#480) #485

Closed wenleix closed 2 years ago

wenleix commented 2 years ago

Summary: Pull Request resolved: https://github.com/pytorch/torcharrow/pull/480

It originally use memory::getDefaultScopedMemoryPool() since ExecCtx constructor expects unique_ptr<MemoryPool>.

Since https://github.com/facebookincubator/velox/commit/15eeddb8e7d50de39afdbde4049f9a1310edf5e5, ExecCtx now expects MemoryPool*, so we can use &velox::memory::getProcessDefaultMemoryManager().getRoot() returned by TorchArrowGlobalStatic::rootMemoryPool.

This makes TorchArrow always use this process default root memory pool.

Reviewed By: dracifer, vancexu

Differential Revision: D39010989

fbshipit-source-id: 3a383b01163e9866500a6fc8b091080b70684c6e