Profiling

Quick profiling in your terminal

Note

This is only meant to be used for quick profiling or programmatically accessing the profiling results. For more detailed and GUI friendly profiling proceed to the next section.

Simply replace the use of Base.@time or Base.@timed with Reactant.Profiler.@time or Reactant.Profiler.@timed. We will automatically compile the function if it is not already a Reactant compiled function (with sync=true).

using Reactant

x = Reactant.to_rarray(randn(Float32, 100, 2))
W = Reactant.to_rarray(randn(Float32, 10, 100))
b = Reactant.to_rarray(randn(Float32, 10))

linear(x, W, b) = (W * x) .+ b

Reactant.@time linear(x, W, b)

┌ Debug: Profiling directory: /home/runner/.julia/scratchspaces/3c362404-f566-11ee-1572-e11a4b42c853/reactant_profiling/jl_nFsh1H
└ @ Reactant.Profiler ~/work/Reactant.jl/Reactant.jl/src/Profiler.jl:824
┌ Debug: Starting XProf gRPC server...
└ @ Reactant.Profiler ~/work/Reactant.jl/Reactant.jl/src/Profiler.jl:796
┌ Debug: Initializing XProf stubs for worker service at 0.0.0.0:33271
└ @ Reactant.Profiler ~/work/Reactant.jl/Reactant.jl/src/Profiler.jl:595
┌ Debug: Starting XProf gRPC server on port 33271
└ @ Reactant.Profiler ~/work/Reactant.jl/Reactant.jl/src/Profiler.jl:611
  runtime: 0.00028551s
  compile time: 0.66745776s

Reactant.@timed nrepeat=100 linear(x, W, b)

AggregateProfilingResult(
    runtime = 0.00003461s, 
    compile_time = 0.11153518s,     metrics = Metrics(
        flops = 0.0,         bandwidth_utils = [0.0, 0.0, 0.0],
        uncapped_flops = 0.0, 
        raw_time = 0.00000000s,         raw_flops = 0.0,         bf16_flops = 0.0,         raw_bytes_accessed = [0.0, 0.0, 0.0],
        occurrences = 0,
        raw_flops_rate = NaN,         bf16_flops_rate = NaN,     )
)

Note that the information returned depends on the backend. Specifically CUDA and TPU backends provide more detailed information regarding memory usage and allocation (something like the following will be displayed on GPUs):

AggregateProfilingResult(
    runtime = 0.00003829s, 
    compile_time = 2.18053260s,  # time spent compiling by Reactant
    GPU_0_bfc = MemoryProfileSummary(
        peak_bytes_usage_lifetime = 64.010 MiB,  # peak memory usage over the entire program (lifetime of memory allocator)
        peak_stats = MemoryAggregationStats(
            stack_reserved_bytes = 0 bytes,  # memory usage by stack reservation
            heap_allocated_bytes = 30.750 KiB,  # memory usage by heap allocation
            free_memory_bytes = 23.518 GiB,  # free memory available for allocation or reservation
            fragmentation = 0.514931,  # fragmentation of memory within [0, 1]
            peak_bytes_in_use = 30.750 KiB # The peak memory usage over the entire program
        )
        peak_stats_time = 0.04975365s, 
        memory_capacity = 23.518 GiB # memory capacity of the allocator
    )
    flops = FlopsSummary(
        Flops = 2.8369974648038653e-9,  # [flops / (peak flops * program time)], capped at 1.0
        UncappedFlops = 2.8369974648038653e-9, 
        RawFlops = 4060.0,  # Total FLOPs performed
        BF16Flops = 4060.0,  # Total FLOPs Normalized to the bf16 (default) devices peak bandwidth
        RawTime = 0.00040298422s,  # Raw time in seconds
        RawFlopsRate = 1.0074836180930361e7,  # Raw FLOPs rate in FLOPs/seconds
        BF16FlopsRate = 1.0074836180930361e7,  # BF16 FLOPs rate in FLOPs/seconds
    )
)

Additionally for GPUs and TPUs, we can use the Reactant.@profile macro to profile the function and get information regarding each of the kernels executed.

Reactant.@profile linear(x, W, b)

┌ Debug: Profiling directory: /home/runner/.julia/scratchspaces/3c362404-f566-11ee-1572-e11a4b42c853/reactant_profiling/jl_1LLn1p
└ @ Reactant.Profiler ~/work/Reactant.jl/Reactant.jl/src/Profiler.jl:824
E1103 23:22:00.066688    4424 unified_profile_processor_factory.cc:65] No UnifiedProfileProcessor registered for tool: kernel_stats
W1103 23:22:00.066700    4424 xplane_to_tools_data_with_profile_processor.cc:229] Unified processor not found for tool: kernel_stats, falling back to legacy.

╔================================================================================╗
║ SUMMARY                                                                        ║
╚================================================================================╝

AggregateProfilingResult(
    runtime = 0.00006997s,
    compile_time = 0.10655576s,  # time spent compiling by Reactant
    metrics = Metrics(
        flops = 0.0,  # [flops / (peak flops * program time)], capped at 1.0
        bandwidth_utils = [0.0, 0.0, 0.0],
        uncapped_flops = 0.0,
        raw_time = 0.00000000s,  # Raw time in seconds
        raw_flops = 0.0,  # Total FLOPs performed
        bf16_flops = 0.0,  # Total FLOPs Normalized to the bf16 (default) devices peak bandwidth
        raw_bytes_accessed = [0.0, 0.0, 0.0],
        occurrences = 0,
        raw_flops_rate = NaN,  # Raw FLOPs rate in FLOPs/seconds
        bf16_flops_rate = NaN,  # BF16 FLOPs rate in FLOPs/seconds
    )
)

On GPUs this would look something like the following:

╔================================================================================╗
║ KERNEL STATISTICS                                                              ║
╚================================================================================╝

┌───────────────────┬─────────────┬────────────────┬──────────────┬──────────────┬──────────────┬──────────────┬───────────┬──────────┬────────────┬─────────────┐
│       Kernel Name │ Occurrences │ Total Duration │ Avg Duration │ Min Duration │ Max Duration │ Static Shmem │ Block Dim │ Grid Dim │ TensorCore │ Occupancy % │
├───────────────────┼─────────────┼────────────────┼──────────────┼──────────────┼──────────────┼──────────────┼───────────┼──────────┼────────────┼─────────────┤
│ gemm_fusion_dot_1 │           1 │    0.00000250s │  0.00000250s │  0.00000250s │  0.00000250s │    2.000 KiB │    64,1,1 │    1,1,1 │          ✗ │      100.0% │
│   loop_add_fusion │           1 │    0.00000131s │  0.00000131s │  0.00000131s │  0.00000131s │      0 bytes │    20,1,1 │    1,1,1 │          ✗ │       31.2% │
└───────────────────┴─────────────┴────────────────┴──────────────┴──────────────┴──────────────┴──────────────┴───────────┴──────────┴────────────┴─────────────┘

╔================================================================================╗
║ FRAMEWORK OP STATISTICS                                                        ║
╚================================================================================╝

┌───────────────────┬─────────┬─────────────┬─────────────┬─────────────────┬───────────────┬──────────┬───────────┬──────────────┬──────────┐
│         Operation │    Type │ Host/Device │ Occurrences │ Total Self-Time │ Avg Self-Time │ Device % │ Memory BW │    FLOP Rate │ Bound By │
├───────────────────┼─────────┼─────────────┼─────────────┼─────────────────┼───────────────┼──────────┼───────────┼──────────────┼──────────┤
│ gemm_fusion_dot.1 │ Unknown │      Device │           1 │     0.00000250s │   0.00000250s │   65.55% │ 1.82 GB/s │  1.6 GFLOP/s │      HBM │
│             +/add │     add │      Device │           1 │     0.00000131s │   0.00000131s │   34.45% │ 0.14 GB/s │ 0.05 GFLOP/s │      HBM │
└───────────────────┴─────────┴─────────────┴─────────────┴─────────────────┴───────────────┴──────────┴───────────┴──────────────┴──────────┘

╔================================================================================╗
║ SUMMARY                                                                        ║
╚================================================================================╝

AggregateProfilingResult(
    runtime = 0.00005622s, 
    compile_time = 2.32802137s,  # time spent compiling by Reactant
    GPU_0_bfc = MemoryProfileSummary(
        peak_bytes_usage_lifetime = 64.010 MiB,  # peak memory usage over the entire program (lifetime of memory allocator)
        peak_stats = MemoryAggregationStats(
            stack_reserved_bytes = 0 bytes,  # memory usage by stack reservation
            heap_allocated_bytes = 81.750 KiB,  # memory usage by heap allocation
            free_memory_bytes = 23.518 GiB,  # free memory available for allocation or reservation
            fragmentation = 0.514564,  # fragmentation of memory within [0, 1]
            peak_bytes_in_use = 81.750 KiB # The peak memory usage over the entire program
        )
        peak_stats_time = 0.00608052s, 
        memory_capacity = 23.518 GiB # memory capacity of the allocator
    )
    flops = FlopsSummary(
        Flops = 2.033375207640664e-8,  # [flops / (peak flops * program time)], capped at 1.0
        UncappedFlops = 2.033375207640664e-8, 
        RawFlops = 4060.0,  # Total FLOPs performed
        BF16Flops = 4060.0,  # Total FLOPs Normalized to the bf16 (default) devices peak bandwidth
        RawTime = 0.00005622s,  # Raw time in seconds
        RawFlopsRate = 7.220987105380169e7,  # Raw FLOPs rate in FLOPs/seconds
        BF16FlopsRate = 7.220987105380169e7,  # BF16 FLOPs rate in FLOPs/seconds
    )
)

Capturing traces

When running Reactant, it is possible to capture traces using the XLA profiler. These traces can provide information about where the XLA specific parts of program spend time during compilation or execution. Note that tracing and compilation happen on the CPU even though the final execution is aimed to run on another device such as GPU or TPU. Therefore, including tracing and compilation in a trace will create annotations on the CPU.

Let's setup a simple function which we can then profile

using Reactant

x = Reactant.to_rarray(randn(Float32, 100, 2))
W = Reactant.to_rarray(randn(Float32, 10, 100))
b = Reactant.to_rarray(randn(Float32, 10))

linear(x, W, b) = (W * x) .+ b

linear (generic function with 1 method)

The profiler can be accessed using the Reactant.with_profiler function.

Reactant.with_profiler("./") do
    mylinear = Reactant.@compile linear(x, W, b)
    mylinear(x, W, b)
end

10×2 ConcretePJRTArray{Float32,2}:
   9.39044    11.3488
   5.05005     0.199757
  27.2152     16.8379
  -3.7231      0.507502
 -17.6287      0.484669
  -4.84174    21.5336
   4.35265     3.5596
   6.32858     8.83566
   0.350231  -10.997
   4.69716    23.7649

Running this function should create a folder called plugins in the folder provided to Reactant.with_profiler which will contain the trace files. The traces can then be visualized in different ways.

Note

For more insights about the current state of Reactant, it is possible to fetch device information about allocations using the Reactant.XLA.allocatorstats function.

Profiling compilation time

To see where compilation time goes, use Reactant.Profiler.@timed_compile. It accepts the same options as Reactant.@compile, but returns the compiled function together with a Reactant.Profiler.CompileTimings breaking down the time spent tracing, running MLIR passes and compiling with XLA.

compiled_linear, timings = Reactant.Profiler.@timed_compile linear(x, W, b)

For more control, call Reactant.Profiler.profile_compile_timings directly. It takes the same arguments and keyword arguments as Reactant.Profiler.with_profiler and returns (result, timings). If no trace output directory is given, a temporary one is created and removed afterwards.

compiled_linear, timings = Reactant.Profiler.profile_compile_timings() do
    Reactant.@compile linear(x, W, b)
end

Perfetto UI

The first and easiest way to visualize a captured trace is to use the online perfetto.dev tool. Reactant.with_profiler has a keyword parameter called create_perfetto_link which will create a usable perfetto URL for the generated trace. The function will block execution until the URL has been clicked and the trace is visualized. The URL only works once.

Reactant.with_profiler("./"; create_perfetto_link=true) do
    mylinear = Reactant.@compile linear(x, W, b)
    mylinear(x, W, b)
end

Note

It is recommended to use the Chrome browser to open the perfetto URL.

XProf

XProf is a complete web UI to analyze the log files captured by Reactant. It can be installed in the following manner:

pip install xprof # or xprof-nightly

Launching xprof is then as simple as:

xprof --logdir=./

which will then make the xprof interface available on port :8791 by default.

Tensorboard

Another option to visualize the generated trace files is to use the tensorboard profiler plugin. The tensorboard viewer can offer more details than the timeline view such as visualization for compute graphs.

First install tensorboard and its profiler plugin:

pip install tensorboard tensorboard-plugin-profile

And then run the following in the folder where the plugins folder was generated:

tensorboard --logdir ./

Adding Custom Annotations

By default, the traces contain only information captured from within XLA. The ReactantCore.@annotate macro can be used to annotate traces. It is a no-op when only ReactantCore is loaded, allowing libraries to add annotations without depending on Reactant. Loading Reactant activates the annotations in host-side profiles.

using Reactant

Reactant.Profiler.@annotate "my_annotation" begin
    # Do things...
end

The added annotations will be captured in the traces and can be seen in the different viewers along with the default XLA annotations. When the profiler is not activated, then the custom annotations have no effect and can therefore always be activated.