koen-dejonghe/NaiveMemoryManagement.scala

## NaiveMemoryManagement.scala
package org.apache.mxnet

import java.util.concurrent.atomic.AtomicLong

import com.typesafe.scalalogging.LazyLogging

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

import org.apache.mxnet.Base._


object NaiveMemoryManagement extends App with LazyLogging {

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

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

    for (i <- 1 to 100) {
      RawTensor.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.payload.toString)
  }

  def parallel(): Unit = {

    import scala.concurrent.ExecutionContext.Implicits.global

    val t3 = RawTensor.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 {
          RawTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
          Thread.sleep(10)
        }
      }
    }

    Await.result(futures, 10 seconds)

    logger.info("DONE")
    logger.info(t3.payload.toString)
  }

}

case class RawTensor private (payload: Long, size: Long) extends LazyLogging {
  override def finalize(): Unit = {
    _LIB.mxNDArrayFree(payload)
    val memSize = RawTensor.hiMemMark.addAndGet(-size)
    logger.info(s"freeing $payload (mem = $memSize)")
  }
}

object RawTensor extends LazyLogging {
  val ctx: Context = Context.defaultCtx

  val threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB
  val hiMemMark = new AtomicLong(0)

  def memCheck(size: Long): Unit =
    if (hiMemMark.addAndGet(size) > threshold) {
      logger.info("invoking gc")
      System.gc()
    }

  def zeros(d1: Int, d2: Int): RawTensor = {
    val size = d1 * d2 * 4
    memCheck(size)

    val ref: NDArrayHandleRef = new NDArrayHandleRef

    _LIB.mxNDArrayCreateEx(shape = Array(d1, d2),
      ndim = 2,
      devType = ctx.deviceTypeid,
      devId = ctx.deviceId,
      delayAlloc = 0,
      dtype = Base.MX_REAL_TYPE.id,
      out = ref)

    logger.info(s"creating ${ref.value}")

    RawTensor(ref.value, size)
  }
}
	package org.apache.mxnet

	import java.util.concurrent.atomic.AtomicLong

	import com.typesafe.scalalogging.LazyLogging

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

	import org.apache.mxnet.Base._


	object NaiveMemoryManagement extends App with LazyLogging {

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

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

	for (i <- 1 to 100) {
	RawTensor.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.payload.toString)
	}

	def parallel(): Unit = {

	import scala.concurrent.ExecutionContext.Implicits.global

	val t3 = RawTensor.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 {
	RawTensor.zeros(3000, 3000) // these will get GC'ed as soon as as System.gc() is called
	Thread.sleep(10)
	}
	}
	}

	Await.result(futures, 10 seconds)

	logger.info("DONE")
	logger.info(t3.payload.toString)
	}

	}

	case class RawTensor private (payload: Long, size: Long) extends LazyLogging {
	override def finalize(): Unit = {
	_LIB.mxNDArrayFree(payload)
	val memSize = RawTensor.hiMemMark.addAndGet(-size)
	logger.info(s"freeing $payload (mem = $memSize)")
	}
	}

	object RawTensor extends LazyLogging {
	val ctx: Context = Context.defaultCtx

	val threshold: Long = 2L * 1024L * 1024L * 1024L // 2 GB
	val hiMemMark = new AtomicLong(0)

	def memCheck(size: Long): Unit =
	if (hiMemMark.addAndGet(size) > threshold) {
	logger.info("invoking gc")
	System.gc()
	}

	def zeros(d1: Int, d2: Int): RawTensor = {
	val size = d1 * d2 * 4
	memCheck(size)

	val ref: NDArrayHandleRef = new NDArrayHandleRef

	_LIB.mxNDArrayCreateEx(shape = Array(d1, d2),
	ndim = 2,
	devType = ctx.deviceTypeid,
	devId = ctx.deviceId,
	delayAlloc = 0,
	dtype = Base.MX_REAL_TYPE.id,
	out = ref)

	logger.info(s"creating ${ref.value}")

	RawTensor(ref.value, size)
	}
	}