spark-rapids
spark-rapids copied to clipboard
Avoid creating objects for NVTX ranges
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.
Could we optionally include a lit of metrics so that it matches NvtxWithMetrics and then we just totally delete NvtxWithMetrics.