spark-rapids icon indicating copy to clipboard operation
spark-rapids copied to clipboard

Avoid creating objects for NVTX ranges

Open jlowe opened this issue 1 year ago • 1 comments

Currently we use withResource(new NvtxRange.... when adding NVTX profile ranges to the code. This has a couple of minor drawbacks:

  • a bit verbose
  • creates an object, adding unnecessary GC pressure

We could provide a Scala interface like this instead which is a bit more concise and avoids creating an object for each range:

object Nvtx {
  def range[T](name: String, color: Color)(block: => T): T = {
    NvtxRange.push(name, color)
    try {
      block
    } finally {
      NvtxRange.pop()
    }
}

which would then allow us to transform uses like this:

withResource(new NvtxRange("compute", NvtxColor.GREEN)) { _ =>
  ...
}

to this simpler, more efficient form:

Nvtx.range("compute", NvtxColor.GREEN) {
  ...
}

This would require exposing the NVTX push and pop methods currently hidden under ai.rapids.cudf.NvtxRange.

jlowe avatar Mar 15 '24 15:03 jlowe

Could we optionally include a lit of metrics so that it matches NvtxWithMetrics and then we just totally delete NvtxWithMetrics.

revans2 avatar Mar 15 '24 16:03 revans2