Memory not release after jit.trace/freeze
Labels: oncall: jit

### 🐛 Describe the bug

Can not use `del` and `gc` to release memory after trace/freeze.
To track memory malloc/release
```
diff --git a/c10/core/impl/alloc_cpu.cpp b/c10/core/impl/alloc_cpu.cpp
index 6ca9ea10967..c4fd33ae701 100644
--- a/c10/core/impl/alloc_cpu.cpp
+++ b/c10/core/impl/alloc_cpu.cpp
@@ -6,6 +6,8 @@
 #include <c10/util/irange.h>
 #include <c10/util/numa.h>
 
+#include <iostream>
+
 // TODO: rename flags to C10
 C10_DEFINE_bool(
     caffe2_cpu_allocator_do_zero_fill,
@@ -94,7 +96,7 @@ void* alloc_cpu(size_t nbytes) {
   } else if (FLAGS_caffe2_cpu_allocator_do_junk_fill) {
     memset_junk(data, nbytes);
   }
-
+  std::cout << "malloc data in c++" << data << " size " << float(nbytes) / 1024 / 1024 << "MB" << std::endl;
   return data;
 }
 
@@ -103,6 +105,7 @@ void free_cpu(void* data) {
   _aligned_free(data);
 #else
   // NOLINTNEXTLINE(cppcoreguidelines-no-malloc)
+  std::cout << "free data " << data << std::endl;
   free(data);
 #endif
 }
```

```python
import time
import psutil, os
import torch
import gc

class M(torch.nn.Module):
    def __init__(self):
          super(M, self).__init__()
          self.w1 = torch.rand(int(1e7), 10)
          print("malloc input", hex(self.w1.data_ptr()),  "size", 1e7 * 100 * 4 / 1024 / 1024, "MB")

    def forward(self, x):
        x = self.w1 + x
        return x


def run_leak():
    process = psutil.Process(os.getpid())
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    a = M().eval()
    input = torch.zeros(int(1e7), 1)
    print("malloc input", hex(input.data_ptr()),  "size", 1e7 * 4 / 1024 / 1024, "MB")
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    print("trace==============")
    a_trace = torch.jit.trace(a, input)

    del(input)
    print("===============================================delete input")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    gc.collect()
    print("gc=================================================")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")

    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    del(a_trace)
    print("===============================================delete a_trace")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    gc.collect()
    print("gc=================================================")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")

    del(a)
    print("===============================================delete a")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")
    gc.collect()
    print("gc=================================================")
    time.sleep(2)
    print("crurent mem usage:", process.memory_info().rss / 1024/1024, "MB")

if __name__ == '__main__':
    run_leak()
    print("process exit===========================")
    exit()

```
Output:
```
crurent mem usage: 282.78515625 MB
malloc data in c++0x7f4c2f8da040 size 381.47MB
malloc w1 0x7f4c2f8da040 size 3814.697265625 MB
malloc data in c++0x7f4c2d2b4040 size 38.147MB
malloc input 0x7f4c2d2b4040 size 38.14697265625 MB
crurent mem usage: 701.67578125 MB
trace==============
...
...
===============================================delete input
crurent mem usage: 713.7421875 MB
gc=================================================
crurent mem usage: 713.7421875 MB
crurent mem usage: 713.7421875 MB
free data 0x7f4c2d2b4040
===============================================delete a_trace
crurent mem usage: 675.59375 MB
gc=================================================
crurent mem usage: 675.59375 MB
===============================================delete a
crurent mem usage: 675.59375 MB
gc=================================================
crurent mem usage: 675.59375 MB
process exit===========================
free data 0x7f4c2f8da040

```

` 0x7f4c2f8da040` is released just before process exit, the `delete` and `gc.collect` do not work.

The release back trace, I can see the last decrese of ref count is from `torch::jit::CompilationUnit` which own the `graph` and `w1` is a node of `graph`. But since I have delete all related variable in python, I do not know whether this `PyObject` is reachable to release the memory`.
```
#26 0x00007fffdf9e885e in torch::jit::Node::~Node (this=0x555559e66a50, __in_chrg=<optimized out>)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/torch/csrc/jit/ir/ir.h:820
#27 0x00007fffdf9e2628 in torch::jit::Graph::~Graph (this=0x555559e61920, __in_chrg=<optimized out>)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/torch/csrc/jit/ir/ir.cpp:2003
#28 0x00007fffdf9a1092 in std::_Sp_counted_ptr<torch::jit::Graph*, (__gnu_cxx::_Lock_policy)2>::_M_dispose (
    this=0x555559e65440) at /usr/include/c++/11/bits/shared_ptr_base.h:348
#29 0x00007fffd9f88ab6 in std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release (this=0x555559e65440)
    at /usr/include/c++/11/bits/shared_ptr_base.h:168
#30 0x00007fffd9f867b5 in std::__shared_count<(__gnu_cxx::_Lock_policy)2>::~__shared_count (this=0x555559e68798, 
    __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr_base.h:702
#31 0x00007fffddf4a334 in std::__shared_ptr<torch::jit::Graph, (__gnu_cxx::_Lock_policy)2>::~__shared_ptr (
    this=0x555559e68790, __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr_base.h:1149
#32 0x00007fffddf4a354 in std::shared_ptr<torch::jit::Graph>::~shared_ptr (this=0x555559e68790, 
    __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr.h:122
#33 0x00007fffdf82aa1e in torch::jit::GraphFunction::~GraphFunction (this=0x555559e68710, __in_chrg=<optimized out>)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/torch/csrc/jit/api/function_impl.h:11
#34 0x00007fffdf82aa5a in torch::jit::GraphFunction::~GraphFunction (this=0x555559e68710, __in_chrg=<optimized out>)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/torch/csrc/jit/api/function_impl.h:11
#35 0x00007fffef24c7ce in std::default_delete<torch::jit::Function>::operator() (this=0x555559e66458, 
    __ptr=0x555559e68710) at /usr/include/c++/11/bits/unique_ptr.h:85
#36 0x00007fffef23c448 in std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> >::~unique_ptr (this=0x555559e66458, __in_chrg=<optimized out>) at /usr/include/c++/11/bits/unique_ptr.h:361
#37 0x00007fffef2ba6f3 in std::_Destroy<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> > > (__pointer=0x555559e66458) at /usr/include/c++/11/bits/stl_construct.h:140
#38 0x00007fffef29f424 in std::_Destroy_aux<false>::__destroy<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> >*> (__first=0x555559e66458, __last=0x555559e66460)
--Type <RET> for more, q to quit, c to continue without paging--
    at /usr/include/c++/11/bits/stl_construct.h:152
#39 0x00007fffef28555c in std::_Destroy<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> >*> (__first=0x555559e66450, __last=0x555559e66460) at /usr/include/c++/11/bits/stl_construct.h:185
#40 0x00007fffef26148f in std::_Destroy<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> >*, std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> > > (__first=0x555559e66450, 
    __last=0x555559e66460) at /usr/include/c++/11/bits/alloc_traits.h:746
#41 0x00007fffef2dad91 in std::vector<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> >, std::allocator<std::unique_ptr<torch::jit::Function, std::default_delete<torch::jit::Function> > > >::~vector (
    this=0x555558bede60, __in_chrg=<optimized out>) at /usr/include/c++/11/bits/stl_vector.h:680
#42 0x00007fffef2d1d7a in torch::jit::CompilationUnit::~CompilationUnit (this=0x555558bede60, 
    __in_chrg=<optimized out>)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/torch/csrc/jit/api/compilation_unit.h:48
#43 0x00007fffef2eed48 in __gnu_cxx::new_allocator<torch::jit::CompilationUnit>::destroy<torch::jit::CompilationUnit> (
    this=0x555558bede60, __p=0x555558bede60) at /usr/include/c++/11/ext/new_allocator.h:162
#44 0x00007fffef2eebf3 in std::allocator_traits<std::allocator<torch::jit::CompilationUnit> >::destroy<torch::jit::CompilationUnit> (__a=..., __p=0x555558bede60) at /usr/include/c++/11/bits/alloc_traits.h:531
#45 0x00007fffef2ee407 in std::_Sp_counted_ptr_inplace<torch::jit::CompilationUnit, std::allocator<torch::jit::CompilationUnit>, (__gnu_cxx::_Lock_policy)2>::_M_dispose (this=0x555558bede50)
    at /usr/include/c++/11/bits/shared_ptr_base.h:528
#46 0x00007fffee809326 in std::_Sp_counted_base<(__gnu_cxx::_Lock_policy)2>::_M_release (this=0x555558bede50)
    at /usr/include/c++/11/bits/shared_ptr_base.h:168
#47 0x00007fffee803465 in std::__shared_count<(__gnu_cxx::_Lock_policy)2>::~__shared_count (this=0x7ffb98f11310, 
    __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr_base.h:702
#48 0x00007fffeed6dc84 in std::__shared_ptr<torch::jit::CompilationUnit, (__gnu_cxx::_Lock_policy)2>::~__shared_ptr (
    this=0x7ffb98f11308, __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr_base.h:1149
#49 0x00007fffeed6dcce in std::shared_ptr<torch::jit::CompilationUnit>::~shared_ptr (this=0x7ffb98f11308, 
    __in_chrg=<optimized out>) at /usr/include/c++/11/bits/shared_ptr.h:122
#50 0x00007fffef259523 in pybind11::class_<torch::jit::CompilationUnit, std::shared_ptr<torch::jit::CompilationUnit> >::dealloc (v_h=...)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/third_party/pybind11/include/pybind11/pybind11.h:1863
#51 0x00007fffee7fd46a in pybind11::detail::clear_instance (self=0x7ffb98f112f0)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/third_party/pybind11/include/pybind11/detail/class.h:424
#52 0x00007fffee7fd581 in pybind11::detail::pybind11_object_dealloc (self=0x7ffb98f112f0)
    at /home/haozhe/rebase/frameworks.ai.pytorch.private-cpu/third_party/pybind11/include/pybind11/detail/class.h:448
#53 0x0000555555663fac in _Py_Dealloc (op=<optimized out>)
    at /tmp/build/80754af9/python-split_1634043551344/work/Objects/object.c:2215
#54 _Py_DECREF () at /tmp/build/80754af9/python-split_1634043551344/work/Include/object.h:478
```



### Versions

Collecting environment information...
PyTorch version: 2.1.0a0+gitd9f822b
Is debug build: True
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 11.1.0-1ubuntu1~20.04) 11.1.0
Clang version: 9.0.1-12 
CMake version: version 3.22.1
Libc version: glibc-2.35

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.0-60-generic-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   52 bits physical, 57 bits virtual
CPU(s):                          128
On-line CPU(s) list:             0-127
Thread(s) per core:              2
Core(s) per socket:              32
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           106
Model name:                      Intel(R) Xeon(R) Platinum 8358 CPU @ 2.60GHz
Stepping:                        6
CPU MHz:                         2600.000
CPU max MHz:                     3400.0000
CPU min MHz:                     800.0000
BogoMIPS:                        5200.00
L1d cache:                       3 MiB
L1i cache:                       2 MiB
L2 cache:                        80 MiB
L3 cache:                        96 MiB
NUMA node0 CPU(s):               0-31,64-95
NUMA node1 CPU(s):               32-63,96-127
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 invpcid_single ssbd mba ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm rdt_a avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb intel_pt avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local split_lock_detect wbnoinvd dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg tme avx512_vpopcntdq la57 rdpid fsrm md_clear pconfig flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.21.2
[pip3] torch==2.1.0a0+gitd9f822b
[pip3] torchvision==0.15.0a0+135a0f9
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] mkl-include               2022.1.0                 pypi_0    pypi
[conda] mkl-static                2022.1.0                 pypi_0    pypi
[conda] numpy                     1.21.2           py38hd8d4704_0  
[conda] numpy-base                1.21.2           py38h2b8c604_0  
[conda] torch                     1.11.0                   pypi_0    pypi
[conda] torchvision               0.15.0a0+135a0f9           dev_0    <develop>

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel