koen-dejonghe/NaiveMemoryManagement.scala

## NaiveMemoryManagement.scala
package jtorch.cpu

import java.util.concurrent.atomic.AtomicLong

import com.typesafe.scalalogging.LazyLogging

import scala.concurrent.{Await, Future}
import scala.concurrent.duration._
import scala.language.postfixOps

object NaiveMemoryManagement extends App with LazyLogging {

  logger.info("*** starting sequential **************************")
  sequential()
  logger.info("*** starting parallel ****************************")
  parallel()

  def sequential(): Unit = {
    val t3 = MyTensor.zeros(100, 100) // this one will only get garbage collected at the end of the program

    for (i <- 1 to 100) {
      MyTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
      Thread.sleep(1)
    }

    logger.info("DONE")
    logger.info(t3.cPtr.toString)
    logger.info(t3.payload.toString)
    logger.info(TH.THFloatTensor_desc(t3.payload).getStr) // this should still work
    logger.info(TH.THFloatTensor_get2d(t3.payload, 10, 10).toString)
  }

  def parallel(): Unit = {

    import scala.concurrent.ExecutionContext.Implicits.global

    val t3 = MyTensor.zeros(100, 100) // this one will only get garbage collected at the end of the program

    val futures = Future.sequence {
      (1 to 100).map { _ =>
        Future {
          MyTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
          Thread.sleep(1)
        }
      }
    }

    Await.result(futures, 10 seconds)

    logger.info("DONE")
    logger.info(t3.cPtr.toString)
    logger.info(t3.payload.toString)
    logger.info(TH.THFloatTensor_desc(t3.payload).getStr) // this should still work
    logger.info(TH.THFloatTensor_get2d(t3.payload, 10, 10).toString)
  }

}

case class MyTensor private (payload: SWIGTYPE_p_THFloatTensor,
                             cPtr: Long,
                             size: Long) extends LazyLogging {
  override def finalize(): Unit = {
    THJNI.THFloatTensor_free(cPtr)
    val memSize = MyTensor.memoryWaterMark.addAndGet(-size)
    logger.info(s"freeing $cPtr (mem = $memSize)")
  }
}

object MyTensor extends LazyLogging {

  val threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB

  val memoryWaterMark = new AtomicLong(0)

  def memCheck(size: Long): Unit =
    if (memoryWaterMark.addAndGet(size) > threshold) {
      System.gc()
    }

  def zeros(d1: Long, d2: Long): MyTensor = {
    val tensor = makeTensorOfZeros(d1, d2)
    logger.info(s"creating ${tensor.cPtr}")
    memCheck(tensor.size)
    tensor
  }

  // boiler plate to create a Torch tensor of floats
  def makeTensorOfZeros(d1: Long, d2: Long): MyTensor = {
    val size: SWIGTYPE_p_THLongStorage = TH.THLongStorage_newWithSize2(d1, d2)
    val cPtr = THJNI.THFloatTensor_newWithSize2d(d1, d2)
    val t = new SWIGTYPE_p_THFloatTensor(cPtr, false)
    TH.THFloatTensor_zeros(t, size)
    MyTensor(t, cPtr, d1 * d2 * 4) // float = 4 bytes
  }

}
	package jtorch.cpu

	import java.util.concurrent.atomic.AtomicLong

	import com.typesafe.scalalogging.LazyLogging

	import scala.concurrent.{Await, Future}
	import scala.concurrent.duration._
	import scala.language.postfixOps

	object NaiveMemoryManagement extends App with LazyLogging {

	logger.info("* starting sequential ************************")
	sequential()
	logger.info("* starting parallel **************************")
	parallel()

	def sequential(): Unit = {
	val t3 = MyTensor.zeros(100, 100) // this one will only get garbage collected at the end of the program

	for (i <- 1 to 100) {
	MyTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
	Thread.sleep(1)
	}

	logger.info("DONE")
	logger.info(t3.cPtr.toString)
	logger.info(t3.payload.toString)
	logger.info(TH.THFloatTensor_desc(t3.payload).getStr) // this should still work
	logger.info(TH.THFloatTensor_get2d(t3.payload, 10, 10).toString)
	}

	def parallel(): Unit = {

	import scala.concurrent.ExecutionContext.Implicits.global

	val t3 = MyTensor.zeros(100, 100) // this one will only get garbage collected at the end of the program

	val futures = Future.sequence {
	(1 to 100).map { _ =>
	Future {
	MyTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
	Thread.sleep(1)
	}
	}
	}

	Await.result(futures, 10 seconds)

	logger.info("DONE")
	logger.info(t3.cPtr.toString)
	logger.info(t3.payload.toString)
	logger.info(TH.THFloatTensor_desc(t3.payload).getStr) // this should still work
	logger.info(TH.THFloatTensor_get2d(t3.payload, 10, 10).toString)
	}

	}

	case class MyTensor private (payload: SWIGTYPE_p_THFloatTensor,
	cPtr: Long,
	size: Long) extends LazyLogging {
	override def finalize(): Unit = {
	THJNI.THFloatTensor_free(cPtr)
	val memSize = MyTensor.memoryWaterMark.addAndGet(-size)
	logger.info(s"freeing $cPtr (mem = $memSize)")
	}
	}

	object MyTensor extends LazyLogging {

	val threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB

	val memoryWaterMark = new AtomicLong(0)

	def memCheck(size: Long): Unit =
	if (memoryWaterMark.addAndGet(size) > threshold) {
	System.gc()
	}

	def zeros(d1: Long, d2: Long): MyTensor = {
	val tensor = makeTensorOfZeros(d1, d2)
	logger.info(s"creating ${tensor.cPtr}")
	memCheck(tensor.size)
	tensor
	}

	// boiler plate to create a Torch tensor of floats
	def makeTensorOfZeros(d1: Long, d2: Long): MyTensor = {
	val size: SWIGTYPE_p_THLongStorage = TH.THLongStorage_newWithSize2(d1, d2)
	val cPtr = THJNI.THFloatTensor_newWithSize2d(d1, d2)
	val t = new SWIGTYPE_p_THFloatTensor(cPtr, false)
	TH.THFloatTensor_zeros(t, size)
	MyTensor(t, cPtr, d1 * d2 * 4) // float = 4 bytes
	}

	}