NVIDIA Tools Extension API: Python ? C/C++?? ??? ??????? ?? ?? ??

PyData? CUDA? ???? ?? ?? ??? ?? ??? ??? ???? ??? ???? ??? ????, CPU ? GPU ??? ?? ????? ? ?? ??? ?????.   Python ??????? ??? ????? ??? ??? cProfile? ?? ?? ?????? PySpy/Viztracer? ?? C ????? ?? ??? ??? ? ?? ?????? ????.   Python ????? ? ?? ?? GPU?? ???? ??? ?????? ? ????.   ???? ???/?????? ???? ????? Python? GPU?? ???? ?? ???? ? ??? ?? ??? ?????.   ??? ??? ??? ?????. NVTX(NVIDIA Tools Extension)? Nsight Systems? ?? CPU ? GPU ??? ????? ??? ?????.

NVTX? Python, C, C++? ??? ?? ??? ? ??? ?? ?? ????????.   Nsight Systems? ???? ?? Python ??, Pandas/NumPy? ?? PyData ?????, ??? PyData ?????? ?? C/C++ ???? ??? ? ????!   Nsight Systems? ?? ????? ????(GPU)?? ???? ?? ?? ????? ???? CUDA? ?? ??? ?? ?????.   ?, cudaMalloc, cudaMemcpy ?? ?? CUDA ??? “?” ? ????.

?? ??

NVTX? ?? ?? ???? ?? ?? ?? ??? ???? ? ??? ? ????.   Python ???? @nvtx.annotate() ?????? ????? nvtx.annotate(..):? ?? ???? ???? ???? ??? ??? ??? ? ????.   ?? ??,

import time
import nvtx


@nvtx.annotate(“f()”, color="purple")
def f():
    for i in range(5):
        with nvtx.annotate("loop", color="red"):
            time.sleep(i)

f()

??? ??? ???? ????? ??? ?? ? ????. ??? ?? ??? ?? ??? ???? ????? NVIDIA Nsight Systems? ?? ?? ??????? ?? ???? ???. ???? ??, f, for ??? ??? ?????.   ?? ?? Nsight Systems CLI? ???? ????? ??? ? ????.

> nsys profile -t nvtx,osrt --force-overwrite=true --stats=true --output=quickstart python nvtx-quickstart.py

-t nvtx,osrt ??? nsys? ???? ? ??? ?????.   ? ???? nvtx ?? ? OSRT(OS ???) ??(??/?? ?)???.

nsys? python ????? ?? ??? qdrep ??? sqlite ???????? ? ?? ??? ?????.   Nsight Systems UI(?? OS?? ?? ??)? qdrep ??? ???? ?????? ????? ???? ? ????.

???? ?? ?

?? ??? ?? ?? ???? ???? ????. ?? ??? ???? ????. ??? ?? ????? ??? ?? ? ?? ???????. ?, ?? ?? f()?? ??? ??? ? ???? ??? ?????? ???? ?? ?? ? ?? ?????? ?? ? ???? ??? ??? ? ????. nsys? ??? ?? ???? ????.

NVTX Push-Pop ?? ??(???)

py-spy? ????? NVTX? ?? ????? ???? ?? ?????? ? ????. 

@nvtx.annotate("mean", color="blue")
def compute_mean(x):
    return x.mean()

@nvtx.annotate("mock generate data", color="red")
def read_data(size):
    return np.random.random((size, size))

def big_computation(size):
    time.sleep(0.5)  # mock warmup
    data = read_data(size)
    return compute_mean(data)

@nvtx.annotate("main")
def main():
    ctx = multiprocessing.get_context("spawn")
    sizes = [4096, 1024]
    procs = [
        ctx.Process(name="calculate", target=big_computation, args=[sizes[i]])
        for i in range(2)
    ]

    for proc in procs:
        proc.start()

    for proc in procs:
        proc.join()

?? ?? Nsight Systems CLI? ???? ????? ??? ? ????.

> MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 nsys profile -t nvtx,osrt,cuda --stats=true --force-overwrite=true --output=multiprocessing python nvtx-multiprocess.py

? ????? ? ?? ????? ???? ?? ?? ???? ???? ??? ?????.   ? ?? ??????? ?? ???? 4096×4096 ??? ????, ? ?? ??????? ?? ???? 1024×1024 ??? ?????.   ????? ? ? ?? ?? ??? ??(??? ?? ??? ???) 2?? ??? ??(?? ??? ??? ?? ??? ??? ??? ??? ?? ???) 3????.   ??? ? ?? ?? ??? ??? ? ? ???, ?? ??? ??? ???? ????? ???? ??? ? ????.

• ????? ????? ??? ? ??? ??? ??? ??? ???? ?? ? ? ????!
• ??? ?? ??? 0? ??? ?? 4096×4096 ???? ?? ??? ????? ?? ? ?? ??? ?????.

GPU ? ??(Cuda ??)

“NVTX ?? ??”?? Nsight Systems? ??? ?? ???? ?? ????, ?? NVIDIA GPU? ??? ? ? ????.   ?, GPU ???? ???? ?(cudaMalloc, ??? ??, ???)? ?? ?? ??? ???? ? ?? ?? ??? ???? ??? ?? ???, ?? ????, ?? ????? ???? Python? GPU? ??? ? ????.   ?? ?? ??? ?????.

with nvtx.annotate("data generation", color="purple"):
     imgs = [cupy.random.random((1024, 1024)) for i in range(100)] 

GPU? ???? ?? NumPy ?? ?????? CuPy? ???? NumPy ???? ????? ???? ?????. ?? ??? GPU? ???? ???? ? ??? ??? ?????. ?? ?? cupy ??(cupy_random, gen_sequence, generate_seed)? ??? ? ? ?? ?? ??? ?? ??? ???? ??? ???? ??? ??? ? ? ????. ?? ????? ? ?? ?? ?? ??? ???? ? ??? ???? ? ??? ??? ?????.

def squared_diff(x, y):
    return (x - y) * (x - y)

with nvtx.annotate("compute", color="yellow"):
    squares = [squared_diff(arr_x, arr_y) for arr_x, arr_y in list(zip(imgs, imgs[1:]))]

??? ?? ? ??? ??

? ??? ?? ?? GPU? ??? ??? ??? ??? ????? ??? ? ????. ????, ? ?? ? ??????? ??? ? ????. ?? ????! cudaMalloc? 0? ?? ?? ???? ?? GPU ??? ???? ??? ???? ??? ?? ? ????. ?? RAPIDS RMM ? ???? ???? ??? ?? GPU ??? ??? ???? ?? ?? ??? ?? ??? ??? ? ????.

cupy.cuda.set_allocator(rmm.rmm_cupy_allocator)
rmm.reinitialize(pool_allocator=True, initial_pool_size=2 ** 32)  # 4GBs

with nvtx.annotate("data generation w pool", color="purple"):
    imgs = [cupy.random.random((1024, 1024)) for i in range(100)]

? ?? ??? squared_difference? ?? cupy.fuse? ???? ??? ??? ???? ? ????.   ?? ???? GPU ??????.

with nvtx.annotate("compute fast", color="orange"):
    squares = [
        cupy.fuse(squared_diff(arr_x, arr_y))
        for arr_x, arr_y in list(zip(imgs, imgs[1:]))
    ]

??? ? ????? ??? ?? NVTX

?? ??? ???? ????? ??? nsys? ??? ?? ??? ?????. ????? ??? ???? ????? ??? ??? ??? ???? ??? ???? ??? ??? ?? ?? ??? ?????. ????? ???? NVTX? ??? ????? ???? ???? ?? ? ?? ?????? ?? ???? ??? ? ????. ?? ?? GPU Dataframe ?????? RAPIDS cuDF? ???? ??? ?? ??? ??? ??? ???? ???? ???. ?? ?????? RMM ? ???? ???? GPU ??? ??? ???? ??? ?? ? ????.

rmm.reinitialize(pool_allocator=True, initial_pool_size=2 ** 32)  # 4GBs

gcdf1 = cdf1.groupby("id").mean().reset_index()
gcdf2 = cdf2.groupby("id").mean().reset_index()

gcdf2.join(gcdf1, on=["id"], lsuffix="_l")

?? ??? ??? ?? ?? ??? ??? ?? ???? ?? ??? ??? ? ? ????. ?? groupby-aggregation(groupby().mean(), groupby().count(), …)? ??? ? ??? ????? ????? agg ???? ??? ? ????. agg? ???? groupby ??? ? ?? ?????.

with nvtx.annotate("agg-all", color="blue"):
    cdf.groupby("name").agg(["count", "mean", "sum"])

with nvtx.annotate("split-agg", color="red"):
    cdf.groupby("name").count(), cdf.groupby("name").mean(), cdf.groupby("name").count()

?? ???? agg-all ? split-agg ???? ? ?? ?? ? ? ????. cuDF ???? ??? ??? ? ? ??? agg([“count”, “mean”, “sum”])? ?? group-agg ?? 1?? count(), mean(), sum()? ?? group-agg ?? 3?? ???? ? ? ????.

?? ??

?? ??? ??? ??? ??? ????? ???? ?? ???? ????? ???? ??? ??? ??? ?? ???? ???? ???. ?? ??? ???? ?? ??? ??? ? ? ??? ??? ????. ?? ???? ??? ?? ???? ????.

????? ??? ? ?? ??? ?? ???? ??? ?? ????. ?? ??? ?? ?? ???? ??? 0, ??? 1 ??? ???? ???? ?? ? ????. ?? ?? C ????? ????? ?? Python ????? ???? ????.

?? ? ?? ??? ???, ????? ??? ?????. C, C++, Python ???? ??? GitHub ?????? ?? ???? ?? ???? ??? ?????. ??? ??? ?? ?? PR? ???? ??? ???????. ?? ??? ??? ? ?? Slack ??? ???? ????.

??

pip ?? conda? ???? NVTX? ?? ????? ? ????.

ython -m pip install nvtx

??

tconda install -c conda-forge nvtx

? ??? ???(? ?? ? ??)? ??? ? GitHub ????? ??? ? ????.

? ???? ??? SDK? ???? ?? ???, ?? ???, ?? ??, ??, ?? ??, ???? NVIDIA ??? ???? ??? ??? ??? ??? ? ????. ?? ??? ???? NVIDIA? ?? ????? ???? ? ??? ??? ??? ?????? ???? ??? ??? ???.