JVM：如何管理 JNI 创建的堆外内存

Question

我正在围绕 Torch libraries. I'm using Swig 构建一个 Scala 包装器来构建粘合层。它允许我在堆外创建张量，我只能通过显式调用库的静态方法来释放它。但是，我想以命令的方式使用张量，而不必担心释放内存，就像 Java 中的任何普通对象一样。

我能想到的唯一方法是按以下方式（误）使用 JVM 的垃圾收集器：

A 'memory manager' 跟踪消耗的堆外内存量，当达到阈值时，它调用 System.gc()。

object MemoryManager {    
  val Threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB
  val FloatSize = 4
  private val hiMemMark = new AtomicLong(0)

  def dec(size: Long): Long = hiMemMark.addAndGet(-size * FloatSize)
  def inc(size: Long): Long = hiMemMark.addAndGet(size * FloatSize)

  def memCheck(size: Long): Unit = {
    val level = inc(size)
    if (level > Threshold) {
      System.gc()
    }
  }
}

张量本身被包装在一个 class 中，使用 finalize 方法释放堆外内存，如下所示：

class Tensor private (val payload: SWIGTYPE_p_THFloatTensor) {
  def numel: Int = TH.THFloatTensor_numel(payload)

  override def finalize(): Unit = {
    val memSize = MemoryManager.dec(numel)
    TH.THFloatTensor_free(payload)
  }    
}

张量的创建是通过通知内存管理器的工厂方法完成的。例如，要创建零张量：

object Tensor {
  def zeros(shape: List[Int]): Tensor = {
      MemoryManager.memCheck(shape.product)
      val storage = ... // boilerplate
      val t = TH.THFloatTensor_new
      TH.THFloatTensor_zeros(t, storage)
      new Tensor(t)
  }
}

我意识到这是一种幼稚的方法，但我可以摆脱这种做法吗？它似乎工作正常，当运行并行时也是如此（这会产生大量对 System.gc() 的多余调用，但除此之外什么都没有）或者你能想到更好的解决方案吗？

谢谢。

Answer 1

有一个更具确定性的选项 - 显式管理的内存区域

所以，如果我们有一个像这样的 class：

class Region private () {
  private val registered = ArrayBuffer.empty[() => Unit]
  def register(finalizer: () => Unit): Unit = registered += finalizer
  def releaseAll(): Unit = {
    registered.foreach(f => f()) // todo - will leak if f() throws
  }
}

我们可以有一个方法实现所谓的 "Loan pattern" 给我们一个新的区域然后处理释放

object Region {
  def run[A](f: Region => A): A = {
    val r = new Region
    try f(r) finally r.releaseAll()
  }
}

那么需要手动释放的东西可以描述为采用隐式 Region:

class Leakable(i: Int)(implicit r: Region) {
  // Class body is constructor body, so you can register finalizers
  r.register(() => println(s"Deallocated foo $i"))

  def foo() = println(s"Foo: $i")
}

您可以以完全没有样板的方式使用：

Region.run { implicit r =>
  val a = new Leakable(1)
  val b = new Leakable(2)
  b.foo()
  a.foo()
}

此代码产生以下输出：

Foo: 2
Foo: 1
Deallocated foo 1
Deallocated foo 2

这种方法有一点限制（如果你试图将一个Leakable赋给一个在run中传递的闭包之外的变量，它的作用域将不会被提升），但会更快并保证即使调用 System.gc 被禁用也能正常工作。

JVM：如何管理 JNI 创建的堆外内存

JVM: How to manage off-heap memory created by JNI

java

scala

jvm

swig

torch