Tensor.grad id 不固定

[wyg1997]

Code to reproduce bug

import oneflow as flow

a = flow.ones(2, 3).requires_grad_()
b = flow.ones(2, 3).requires_grad_()

c = a + b
c.sum().backward()

print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))

输出为：

140517641953136
140517641952048
140517641953136
140517641952048
140517641953136
140517641952048

[wyg1997]

如果把 a.grad 在 python 中增加引用计数后，再打印的话 id 就固定了，估计问题在 C++ Tensor 和 py_ptr 的绑定处理中

import oneflow as flow

a = flow.ones(2, 3).requires_grad_()
b = flow.ones(2, 3).requires_grad_()

c = a + b
c.sum().backward()

a_grad = a.grad
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))
print(id(a.grad))

[wyg1997]

原因

此问题是目前 C++ Tensor 对象和 PyTensorObject 的引用关系有缺陷造成的：

class Tensor {
  ...;
  PyObject* py_obj_;
};

struct PyTensorObject {
  ...;
  std::shared_ptr<Tensor> data;
};

其指针绑定和解绑的时机为：

• CppTensor 被返回到 Python 端时，PyTensorObject 对象绑定 CppTensor，同时 CppTensor 也绑定 PyTensorObject 指针，由于前者是 shared_ptr 绑定，后者是裸指针绑定，所以没有循环引用问题 PyTensorObject 可以在引用计数为 0 时释放。
• PyTensorObject 存活的情况下，再次试图把 CppTensor 返回到 Python 端时，会直接返回绑定的 py_obj_ 指针，这样在 python 端 id() 得到的结果是一样的。

  a_grad0 = a.grad
  a_grad1 = a.grad
  print(id(a_grad0) == id(a_grad1))  # True
  

• PyTensorObject 引用计算为 0，尝试释放时，会双向的解除绑定。也正是因为这个，PyTensorObject 在一次次创建时，内存位置是不固定的，就导致了本 issue 发现的问题。

解决方案

参考了 PyTorch 的一些设计，针对 OneFlow 设计上的缺陷想到以下一种方案：

之前的方案从始至终都是 PyTensorObject 持有 CppTensor，这样就导致每一次 PyTensorObject 重新被唤起时都是一个新的对象，那么它的内存地址确实是不固定的。

a_grad = a.grad
print(id(a_grad))
del a_grad  # PyTensorObject deallocate
a_grad = a.grad  # PyTensorObject maybe allocate at a new address

如果希望每次分配的 PyTensorObject 对象是同一个地址的话，只能让它不要被析构，当 CppTensor 返回到 Python 端时，永远返回的是同一个 PyTensorObject 对象。

所以 CppTensor 和 PyTensorObject 应该是一个双向的引用关系，互相持有着对方的生命周期，但直接互持就是经典的循环引用问题了，所以可以通过以下规则来保证单向引用且生命周期被正确管理：

// CppTensor 上加一个 owns_pyobj_ flag，表示持有权
class Tensor {
  ...;
  PyObject* py_obj_;
  bool owns_pyobj_;
};

• CppTensor 第一次被返回到 Python 端时，创建PyTensorObject 对象，此时持有权交给 PyTensorObject 对象，Tensor.owns_pyobj_ 置为 false，Tensor.py_obj_ 绑定上新创建的 PyTensorObject 对象（不增加引用计数），PyTensorObject.data 绑定 CppTensor。
• CppTensor 第二次被返回到 Python 端时，由于 Tensor.owns_pyobj_ 为 false，表示有存活的 PyTensorObject 对象，直接返回 Tensor.py_obj_ 指针并增加引用计数。
• PyTensorObject 尝试析构时，Tensor.owns_pyobj_ 置为 true，表示持有权交还给 CppTensor，PyTensorObject.data 置空，解除对 CppTensor 的引用，PyTensorObject 的引用计数+1，不对 PyTensorObject 做内存上的析构
• CppTensor 再次被返回到 Python 端时，Tensor.owns_pyobj_ 置为 false，PyTensorObject.data 绑定 CppTensor，返回 Tensor.py_obj_ 但不增加引用计数（因为上一步尝试析构时已经增加过了）
• CppTensor 析构时，释放 PyTensorObject 的内存。

[strint]

PyTensorObject 尝试析构时，Tensor.owns_pyobj_ 置为 true，表示持有权交还给 CppTensor，PyTensorObject.data 置空，解除对 CppTensor 的引用，PyTensorObject 的引用计数+1，不对 PyTensorObject 做内存上的析构

看起来就是缓存了一个 PyTensorObject 到 CppTensor 上？

[wyg1997]

PyTensorObject 尝试析构时，Tensor.owns_pyobj_ 置为 true，表示持有权交还给 CppTensor，PyTensorObject.data 置空，解除对 CppTensor 的引用，PyTensorObject 的引用计数+1，不对 PyTensorObject 做内存上的析构

看起来就是缓存了一个 PyTensorObject 到 CppTensor 上？

是的，之前这一步是做一个双向的清除，导致了解绑，但实际上应该在 python 和 cpp 都析构时才做解绑