时间戳

首页 > 解决方案 > JVM：如何管理由 JNI 创建的堆外内存

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

问题描述

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

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

“内存管理器”会跟踪消耗的堆外内存量，当达到阈值时，它会调用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()
    }
  }
}

张量本身被包装在一个类中，带有一个 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()但没有其他任何东西）或者你能想到一个更好的解决方案吗？

谢谢你。

标签： javascalajvmswigtorch

解决方案

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

所以，大致如果我们有一个这样的类：

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
  }
}

我们可以有一个实现所谓的“贷款模式”的方法，它给我们一个新的区域，然后处理释放

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

这种方法有一点限制（如果您尝试将 a 分配给Leakable传入的闭包之外的变量run，则不会提升其范围），但即使调用System.gc被禁用也会更快并保证工作。

推荐阅读