[perf] investigate performance issue in `LazySegtree` with `ModInt`

[harry0000]

Summary

When using LazySegtree in the AtCoder Library Practice Contest, the program consumes a large amount of memory and results in TLE (Time Limit Exceeded). This issue is particularly noticeable when running under Scala Native.

Observed Behavior

Environment:

AtCoder 2023/1 Language Update - Google スプレッドシート

• Scala 3.3.0 (Scala Native 0.4.14)
• Scala (Dotty 3.3.0)
• openjdk-20-jdk-headless

Contest tasks

• K - Range Affine Range Sum

Submissions

K - Range Affine Range Sum

• https://atcoder.jp/contests/practice2/submissions/46781423
  • Dotty 3.3.0 / 4972 ms / 294588 KiB
• https://atcoder.jp/contests/practice2/submissions/70007496
  • Scala Native 0.4.14 / 5620 ms / 1396188 KiB

case	result	time	memory
max_random_00	TLE	5620 ms	1393420 KiB
max_random_01	TLE	5617 ms	1375172 KiB
max_random_02	TLE	5614 ms	1396188 KiB
random_00	TLE	5616 ms	1319488 KiB
random_01	TLE	5613 ms	1302620 KiB
random_02	AC	3710 ms	842776 KiB

(reference) L - Lazy Segment Tree

• https://atcoder.jp/contests/practice2/submissions/46865651
  • Dotty 3.3.0 / 1590 ms / 149388 KiB
• https://atcoder.jp/contests/practice2/submissions/70007637
  • Scala Native 0.4.14 / 935 ms / 286028 KiB

Hypotheses

• Memory leak in Scala Native generated code

Steps to Reproduce

Tasks

• Reproduce and confirm the issue
• Run profiling (JMH, async-profiler, etc.)
• Identify bottleneck
• Propose optimization approach
• Benchmark optimized version
• Create pull request

Related

• Specialization for Scala 3 scala/scala3#18044
• Runtime / Garbage Collector Settings — Scala Native 0.5.9 documentation
• 新ジャッジ運用開始のお知らせ - AtCoder
• AtCoder言語アップデート2024-2025 - Google スプレッドシート

Notes

Previously, the AtCoder administrators uploaded the test cases for the AtCoder Library Practice Contest to Dropbox for me, but the public access has been temporarily suspended, so the test data is currently unavailable.

• https://atcoder.jp/posts/1376
• https://atcoder.jp/posts/20

[harry0000]

Reading from standard input alone takes 4,246 ms and consumes 705,316 KiB of memory, leaving less than 800 ms available for processing the Segtree.

https://atcoder.jp/contests/practice2/submissions/46770349

  def main(_args: Array[String]): Unit = {
    val Array(n, q) = io.StdIn.readLine().split(' ').map(_.toInt)
    val a = io.StdIn.readLine().split(' ').map(ai => S(ai.toLong))

    // val segtree = LazySegtree(a)
    rep(q, {
      val query = io.StdIn.readLine().split(' ').map(_.toInt)

      if (query(0) == 0) {
        val Array(_, l, r, b, c) = query
        // segtree.applyRange(l, r, Pair(b.toLong, c.toLong))
      } else {
        val Array(_, l, r) = query
        println()
        // println(segtree.prod(l, r).a)
      }
    })
  }

[harry0000]

It seems that in Dotty, rewriting a foreach loop into an inline-expanded while loop results in a slight performance improvement.

	time	memory	submission
before	4872 ms	324172 KiB	70112691
after	4576 ms	308416 KiB	70185985

-  private val _1_to_log = 1 to log
-  private val _1_to_log_rev = _1_to_log.reverse
+  private inline def _1_to_log[R](f: Int => Unit): Unit = {
+    var i = 1
+    while (i <= log) {
+      f(i)
+      i += 1
+    }
+  }
+  private inline def _log_to_1[R](f: Int => Unit): Unit = {
+    var i = log
+    while (i >= 1) {
+      f(i)
+      i -= 1
+    }
+  }

Scala Native also seems to be effective.

• before: 70113170
• after: 70186795

[harry0000]

Reading from standard input alone takes 4,246 ms and consumes 705,316 KiB of memory, leaving less than 800 ms available for processing the Segtree.

https://atcoder.jp/contests/practice2/submissions/46770349

We can use FastScanner¹. It takes only 301ms to read from standard input.

• https://atcoder.jp/contests/practice2/submissions/70282759
• 301 ms / 12524 KiB

FastScanner

import java.io.InputStream

final class FastScanner(private val in: InputStream) {
  private val buffer = new Array[Byte](1 << 16)
  private var ptr = 0
  private var len = 0

  private def readByte(): Int = {
    if (ptr >= len) {
      len = in.read(buffer)
      ptr = 0
      if (len <= 0) { return -1 }
    }
    val res = buffer(ptr) & 0xff
    ptr += 1
    res
  }

  private def skipWhitespace(): Int = {
    var b = readByte()
    while (b != -1 && b <= 32) {
      b = readByte()
    }
    b
  }

  def next(): String = {
    val sb = new java.lang.StringBuilder
    var b = skipWhitespace()
    while (b != -1 && b > 32) {
      sb.append(b.toChar)
      b = readByte()
    }
    sb.toString
  }

  def nextInt(): Int = {
    var b = skipWhitespace()
    var sign = 1
    if (b == '-') {
      sign = -1
      b = readByte()
    }
    var num = 0
    while (b != -1 && b > 32) {
      num = num * 10 + (b - '0')
      b = readByte()
    }
    sign * num
  }

  def nextLong(): Long = {
    var b = skipWhitespace()
    var sign = 1L
    if (b == '-') {
      sign = -1L
      b = readByte()
    }
    var num = 0L
    while (b != -1 && b > 32) {
      num = num * 10 + (b - '0')
      b = readByte()
    }
    sign * num
  }

  def nextLine(): String = {
    val sb = new java.lang.StringBuilder
    var b = readByte()
    while (b != -1 && b != '\n') {
      if (b != '\r') {
        sb.append(b.toChar)
      }
      b = readByte()
    }
    sb.toString
  }
}

However, even this still causes TLE and MLE in Scala Native.

• https://atcoder.jp/contests/practice2/submissions/70283494
• 5616 ms / 1594076 KiB

Footnotes

1. generated by Codex ↩

[harry0000]

However, even this still causes TLE and MLE in Scala Native.

Ultimately, it seems we can achieve AC by abandoning the use of StaticModInt and avoiding the generation of large numbers of objects.

  val Mod = 998_244_353L

  final case class S(var sum: Int, var size: Int)
  object S {
    def apply(value: Int): S = new S(value, 1)
    def e(): S = S(0, 0)
  }

  given Monoid[S] with {
    final override def e(): S = S.e()
    final override def combine(a: S, b: S): S = S(
      ((a.sum.toLong + b.sum) % Mod).toInt,
      a.size + b.size
    )
  }

  final case class F(b: Int, c: Int)
  object F {
    def id(): F = F(1, 0)
  }
  
  given MapMonoid[S, F] with {
    final override def id(): F = F.id()
    final override def mapping(f: F, s: S): S = {
      val F(b, c) = f
      s.sum = ((s.sum.toLong * b + s.size.toLong * c) % Mod).toInt
      s
    }
    final override def composition(l: F, r: F): F = {
      val nb = (l.b.toLong * r.b) % Mod
      val nc = (l.b.toLong * r.c + l.c) % Mod
      F(nb.toInt, nc.toInt)
    }
  }

Scala 3.3.0 (Scala Native 0.4.14)

• https://atcoder.jp/contests/practice2/submissions/70283774
• 3860 ms / 568536 KiB

Scala (Dotty 3.3.0)

• https://atcoder.jp/contests/practice2/submissions/70283762
• 3697 ms / 195124 KiB