DTT performance benchmarks

This example benchmarks the DTT performance for the different algorithms available in Shamrock

import random
import time

import matplotlib.colors as colors
import matplotlib.pyplot as plt
import numpy as np

import shamrock

# If we use the shamrock executable to run this script instead of the python interpreter,
# we should not initialize the system as the shamrock executable needs to handle specific MPI logic
if not shamrock.sys.is_initialized():
    shamrock.change_loglevel(1)
    shamrock.sys.init("0:0")

Main benchmark functions

bounding_box = shamrock.math.AABB_f64_3((0.0, 0.0, 0.0), (1.0, 1.0, 1.0))


def benchmark_dtt_core(N, theta_crit, compression_level, nb_repeat=10):
    times = []
    random.seed(111)
    max_mem_delta = 0
    for i in range(nb_repeat):
        positions = shamrock.algs.mock_buffer_f64_3(
            random.randint(0, 1000000), N, bounding_box.lower, bounding_box.upper
        )
        tree = shamrock.tree.CLBVH_u64_f64_3()
        tree.rebuild_from_positions(positions, bounding_box, compression_level)
        shamrock.backends.reset_mem_info_max()
        mem_info_before = shamrock.backends.get_mem_perf_info()
        times.append(shamrock.tree.benchmark_clbvh_dual_tree_traversal(tree, theta_crit) * 1000)
        mem_info_after = shamrock.backends.get_mem_perf_info()

        mem_delta = (
            mem_info_after.max_allocated_byte_device - mem_info_before.max_allocated_byte_device
        )
        max_mem_delta = max(max_mem_delta, mem_delta)
    return times, max_mem_delta


def benchmark_dtt(N, theta_crit, compression_level, nb_repeat=10):
    times, max_mem_delta = benchmark_dtt_core(N, theta_crit, compression_level, nb_repeat)
    return min(times), max(times), sum(times) / nb_repeat, max_mem_delta

Run the performance test for all parameters

def run_performance_sweep(compression_level, threshold_run):

    # Define parameter ranges
    # logspace as array
    particle_counts = np.logspace(2, 7, 10).astype(int).tolist()
    theta_crits = [0.1, 0.3, 0.5, 0.7, 0.9]

    # Initialize results matrix
    results_mean = np.zeros((len(theta_crits), len(particle_counts)))
    results_min = np.zeros((len(theta_crits), len(particle_counts)))
    results_max = np.zeros((len(theta_crits), len(particle_counts)))
    results_max_mem_delta = np.zeros((len(theta_crits), len(particle_counts)))

    print(f"Particle counts: {particle_counts}")
    print(f"Theta_crit values: {theta_crits}")
    print(f"Compression level: {compression_level}")

    total_runs = len(particle_counts) * len(theta_crits)
    current_run = 0

    for i, theta_crit in enumerate(theta_crits):
        exceed_mem = False
        for j, N in enumerate(particle_counts):
            current_run += 1

            if exceed_mem:
                print(
                    f"[{current_run:2d}/{total_runs}] Skipping N={N:5d}, theta_crit={theta_crit:.1f}"
                )
                results_mean[i, j] = np.nan
                results_min[i, j] = np.nan
                results_max[i, j] = np.nan
                continue

            print(
                f"[{current_run:2d}/{total_runs}] Running N={N:5d}, theta_crit={theta_crit:.1f}...",
                end=" ",
            )

            start_time = time.time()
            min_time, max_time, mean_time, max_mem_delta = benchmark_dtt(
                N, theta_crit, compression_level
            )
            elapsed = time.time() - start_time

            results_mean[i, j] = mean_time
            results_min[i, j] = min_time
            results_max[i, j] = max_time
            results_max_mem_delta[i, j] = max_mem_delta

            print(f"mean={mean_time:.3f}ms (took {elapsed:.1f}s)")

            if max_mem_delta > threshold_run:
                exceed_mem = True

    return (
        particle_counts,
        theta_crits,
        results_mean,
        results_min,
        results_max,
        results_max_mem_delta,
    )

Create checkerboard plot with execution times and relative performance to reference algorithm

def create_checkerboard_plot(
    particle_counts,
    theta_crits,
    results_data,
    compression_level,
    algname,
    max_axis_value,
    reference_data,
    results_max_mem_delta,
):
    """Create checkerboard plot with execution times"""

    fig, ax = plt.subplots(figsize=(12, 8))

    # Calculate relative performance compared to reference algorithm
    # results_data / reference_data gives the ratio (>1 means slower, <1 means faster)
    relative_performance = results_data / reference_data

    # Create the heatmap with relative performance values
    # Create a masked array to handle NaN values (skipped benchmarks) as white
    masked_relative = np.ma.masked_invalid(relative_performance)

    # Use a diverging colormap: red for better performance (<1), green for worse (>1)
    # RdYlGn_r (reversed) has green for high values (worse) and red for low values (better)
    cmap = plt.cm.RdYlGn_r.copy()  # Green for >1 (slower), Red for <1 (faster)
    cmap.set_bad(color="white")  # Set NaN values to white

    # Set the color scale limits for relative performance
    vmin = 0.5
    vmax = 1.5

    im = ax.imshow(
        masked_relative, cmap=cmap, aspect="auto", interpolation="nearest", vmin=vmin, vmax=vmax
    )

    # Set ticks and labels
    ax.set_xticks(range(len(particle_counts)))
    ax.set_yticks(range(len(theta_crits)))
    ax.set_xticklabels([f"{N//1000}k" if N >= 1000 else str(N) for N in particle_counts])
    ax.set_yticklabels([f"{theta:.1f}" for theta in theta_crits])

    # Add labels
    ax.set_xlabel("Particle Count")
    ax.set_ylabel("Theta Critical")
    ax.set_title(
        f"Dual Tree Traversal Performance\n(Colors: Relative to Reference, Text: Absolute Time in ms)\ncompression level = {compression_level} algorithm = {algname}",
        pad=20,
    )

    # Add text annotations showing the values
    for i in range(len(theta_crits)):
        for j in range(len(particle_counts)):
            value = results_data[i, j]

            if np.isnan(value):
                # For skipped benchmarks, show "SKIPPED" in black on white background
                # ax.text(j, i, 'SKIPPED', ha='center', va='center',
                #       color='black', fontweight='bold', fontsize=8)
                pass
            else:
                perf = relative_performance[i, j]
                mem_delta = results_max_mem_delta[i, j] / 1e6
                text_color = "black"
                ax.text(
                    j,
                    i,
                    f"{value:.2f}ms\n{perf:.2f}\n{mem_delta:.2f}MB",
                    ha="center",
                    va="center",
                    color=text_color,
                    fontweight="bold",
                    fontsize=10,
                )

    # Add colorbar for relative performance
    cbar = plt.colorbar(im, ax=ax, shrink=0.8)
    cbar.set_label("Relative performance (time / reference time)")
    cbar.ax.tick_params(labelsize=10)

    # Add custom tick labels for better interpretation
    tick_positions = [0.1, 0.2, 0.5, 1.0, 2.0, 3.0]
    cbar.set_ticks([pos for pos in tick_positions if vmin <= pos <= vmax])

    # Improve layout
    plt.tight_layout()

    # Add grid for better readability
    ax.set_xticks(np.arange(len(particle_counts)) - 0.5, minor=True)
    ax.set_yticks(np.arange(len(theta_crits)) - 0.5, minor=True)
    ax.grid(which="minor", color="black", linestyle="-", linewidth=1, alpha=0.3)

    return fig, ax

List current implementation

current_impl = shamrock.tree.get_current_impl_clbvh_dual_tree_traversal_impl()

print(current_impl)

impl_param(impl_name="scan_multipass", params="")

List all implementations available

all_default_impls = shamrock.tree.get_default_impl_list_clbvh_dual_tree_traversal()

print(all_default_impls)

[impl_param(impl_name="reference", params=""), impl_param(impl_name="parallel_select", params=""), impl_param(impl_name="scan_multipass", params="")]

Run the performance benchmarks for all implementations

results = {}

for default_impl in all_default_impls:
    shamrock.tree.set_impl_clbvh_dual_tree_traversal(default_impl.impl_name, default_impl.params)

    print(f"Running DTT performance benchmarks for {default_impl.impl_name}...")

    compression_level = 4

    threshold_run = 5e6
    # Run the performance sweep
    particle_counts, theta_crits, results_mean, results_min, results_max, results_max_mem_delta = (
        run_performance_sweep(compression_level, threshold_run)
    )

    results[default_impl.impl_name + " " + default_impl.params] = {
        "particle_counts": particle_counts,
        "theta_crits": theta_crits,
        "results_mean": results_mean,
        "results_min": results_min,
        "results_max": results_max,
        "results_max_mem_delta": results_max_mem_delta,
        "name": default_impl.impl_name + " " + default_impl.params,
    }

Info: setting dtt implementation to impl : reference                                 [tree][rank=0]
Running DTT performance benchmarks for reference...
Particle counts: [100, 359, 1291, 4641, 16681, 59948, 215443, 774263, 2782559, 10000000]
Theta_crit values: [0.1, 0.3, 0.5, 0.7, 0.9]
Compression level: 4
[ 1/50] Running N=  100, theta_crit=0.1... mean=1.755ms (took 0.0s)
[ 2/50] Running N=  359, theta_crit=0.1... mean=1.067ms (took 0.0s)
[ 3/50] Running N= 1291, theta_crit=0.1... mean=1.307ms (took 0.0s)
[ 4/50] Running N= 4641, theta_crit=0.1... mean=3.113ms (took 0.1s)
[ 5/50] Running N=16681, theta_crit=0.1... mean=23.228ms (took 0.4s)
[ 6/50] Skipping N=59948, theta_crit=0.1
[ 7/50] Skipping N=215443, theta_crit=0.1
[ 8/50] Skipping N=774263, theta_crit=0.1
[ 9/50] Skipping N=2782559, theta_crit=0.1
[10/50] Skipping N=10000000, theta_crit=0.1
[11/50] Running N=  100, theta_crit=0.3... mean=0.834ms (took 0.0s)
[12/50] Running N=  359, theta_crit=0.3... mean=1.090ms (took 0.0s)
[13/50] Running N= 1291, theta_crit=0.3... mean=1.178ms (took 0.0s)
[14/50] Running N= 4641, theta_crit=0.3... mean=2.764ms (took 0.1s)
[15/50] Running N=16681, theta_crit=0.3... mean=21.239ms (took 0.3s)
[16/50] Skipping N=59948, theta_crit=0.3
[17/50] Skipping N=215443, theta_crit=0.3
[18/50] Skipping N=774263, theta_crit=0.3
[19/50] Skipping N=2782559, theta_crit=0.3
[20/50] Skipping N=10000000, theta_crit=0.3
[21/50] Running N=  100, theta_crit=0.5... mean=0.868ms (took 0.0s)
[22/50] Running N=  359, theta_crit=0.5... mean=1.037ms (took 0.0s)
[23/50] Running N= 1291, theta_crit=0.5... mean=1.228ms (took 0.0s)
[24/50] Running N= 4641, theta_crit=0.5... mean=2.661ms (took 0.1s)
[25/50] Running N=16681, theta_crit=0.5... mean=12.086ms (took 0.2s)
[26/50] Running N=59948, theta_crit=0.5... mean=73.430ms (took 1.1s)
[27/50] Skipping N=215443, theta_crit=0.5
[28/50] Skipping N=774263, theta_crit=0.5
[29/50] Skipping N=2782559, theta_crit=0.5
[30/50] Skipping N=10000000, theta_crit=0.5
[31/50] Running N=  100, theta_crit=0.7... mean=0.995ms (took 0.0s)
[32/50] Running N=  359, theta_crit=0.7... mean=0.888ms (took 0.0s)
[33/50] Running N= 1291, theta_crit=0.7... mean=1.064ms (took 0.0s)
[34/50] Running N= 4641, theta_crit=0.7... mean=2.019ms (took 0.1s)
[35/50] Running N=16681, theta_crit=0.7... mean=7.198ms (took 0.2s)
[36/50] Running N=59948, theta_crit=0.7... mean=33.120ms (took 0.7s)
[37/50] Skipping N=215443, theta_crit=0.7
[38/50] Skipping N=774263, theta_crit=0.7
[39/50] Skipping N=2782559, theta_crit=0.7
[40/50] Skipping N=10000000, theta_crit=0.7
[41/50] Running N=  100, theta_crit=0.9... mean=1.002ms (took 0.0s)
[42/50] Running N=  359, theta_crit=0.9... mean=1.205ms (took 0.0s)
[43/50] Running N= 1291, theta_crit=0.9... mean=1.053ms (took 0.0s)
[44/50] Running N= 4641, theta_crit=0.9... mean=1.762ms (took 0.1s)
[45/50] Running N=16681, theta_crit=0.9... mean=4.814ms (took 0.2s)
[46/50] Running N=59948, theta_crit=0.9... mean=19.865ms (took 0.5s)
[47/50] Skipping N=215443, theta_crit=0.9
[48/50] Skipping N=774263, theta_crit=0.9
[49/50] Skipping N=2782559, theta_crit=0.9
[50/50] Skipping N=10000000, theta_crit=0.9
Info: setting dtt implementation to impl : parallel_select                           [tree][rank=0]
Running DTT performance benchmarks for parallel_select...
Particle counts: [100, 359, 1291, 4641, 16681, 59948, 215443, 774263, 2782559, 10000000]
Theta_crit values: [0.1, 0.3, 0.5, 0.7, 0.9]
Compression level: 4
[ 1/50] Running N=  100, theta_crit=0.1... mean=1.313ms (took 0.0s)
[ 2/50] Running N=  359, theta_crit=0.1... mean=1.179ms (took 0.0s)
[ 3/50] Running N= 1291, theta_crit=0.1... mean=1.574ms (took 0.0s)
[ 4/50] Running N= 4641, theta_crit=0.1... mean=8.599ms (took 0.1s)
[ 5/50] Running N=16681, theta_crit=0.1... mean=92.707ms (took 1.1s)
[ 6/50] Skipping N=59948, theta_crit=0.1
[ 7/50] Skipping N=215443, theta_crit=0.1
[ 8/50] Skipping N=774263, theta_crit=0.1
[ 9/50] Skipping N=2782559, theta_crit=0.1
[10/50] Skipping N=10000000, theta_crit=0.1
[11/50] Running N=  100, theta_crit=0.3... mean=0.943ms (took 0.0s)
[12/50] Running N=  359, theta_crit=0.3... mean=1.084ms (took 0.0s)
[13/50] Running N= 1291, theta_crit=0.3... mean=1.654ms (took 0.0s)
[14/50] Running N= 4641, theta_crit=0.3... mean=9.200ms (took 0.1s)
[15/50] Running N=16681, theta_crit=0.3... mean=93.288ms (took 1.1s)
[16/50] Skipping N=59948, theta_crit=0.3
[17/50] Skipping N=215443, theta_crit=0.3
[18/50] Skipping N=774263, theta_crit=0.3
[19/50] Skipping N=2782559, theta_crit=0.3
[20/50] Skipping N=10000000, theta_crit=0.3
[21/50] Running N=  100, theta_crit=0.5... mean=0.947ms (took 0.0s)
[22/50] Running N=  359, theta_crit=0.5... mean=1.097ms (took 0.0s)
[23/50] Running N= 1291, theta_crit=0.5... mean=1.645ms (took 0.0s)
[24/50] Running N= 4641, theta_crit=0.5... mean=8.865ms (took 0.1s)
[25/50] Running N=16681, theta_crit=0.5... mean=68.973ms (took 0.8s)
[26/50] Running N=59948, theta_crit=0.5... mean=515.644ms (took 5.5s)
[27/50] Skipping N=215443, theta_crit=0.5
[28/50] Skipping N=774263, theta_crit=0.5
[29/50] Skipping N=2782559, theta_crit=0.5
[30/50] Skipping N=10000000, theta_crit=0.5
[31/50] Running N=  100, theta_crit=0.7... mean=0.966ms (took 0.0s)
[32/50] Running N=  359, theta_crit=0.7... mean=1.039ms (took 0.0s)
[33/50] Running N= 1291, theta_crit=0.7... mean=1.647ms (took 0.0s)
[34/50] Running N= 4641, theta_crit=0.7... mean=8.317ms (took 0.1s)
[35/50] Running N=16681, theta_crit=0.7... mean=44.213ms (took 0.6s)
[36/50] Running N=59948, theta_crit=0.7... mean=272.458ms (took 3.1s)
[37/50] Skipping N=215443, theta_crit=0.7
[38/50] Skipping N=774263, theta_crit=0.7
[39/50] Skipping N=2782559, theta_crit=0.7
[40/50] Skipping N=10000000, theta_crit=0.7
[41/50] Running N=  100, theta_crit=0.9... mean=0.951ms (took 0.0s)
[42/50] Running N=  359, theta_crit=0.9... mean=0.985ms (took 0.0s)
[43/50] Running N= 1291, theta_crit=0.9... mean=1.532ms (took 0.0s)
[44/50] Running N= 4641, theta_crit=0.9... mean=5.570ms (took 0.1s)
[45/50] Running N=16681, theta_crit=0.9... mean=30.386ms (took 0.4s)
[46/50] Running N=59948, theta_crit=0.9... mean=166.187ms (took 2.0s)
[47/50] Skipping N=215443, theta_crit=0.9
[48/50] Skipping N=774263, theta_crit=0.9
[49/50] Skipping N=2782559, theta_crit=0.9
[50/50] Skipping N=10000000, theta_crit=0.9
Info: setting dtt implementation to impl : scan_multipass                            [tree][rank=0]
Running DTT performance benchmarks for scan_multipass...
Particle counts: [100, 359, 1291, 4641, 16681, 59948, 215443, 774263, 2782559, 10000000]
Theta_crit values: [0.1, 0.3, 0.5, 0.7, 0.9]
Compression level: 4
[ 1/50] Running N=  100, theta_crit=0.1... mean=3.548ms (took 0.1s)
[ 2/50] Running N=  359, theta_crit=0.1... mean=5.249ms (took 0.1s)
[ 3/50] Running N= 1291, theta_crit=0.1... mean=7.061ms (took 0.1s)
[ 4/50] Running N= 4641, theta_crit=0.1... mean=10.032ms (took 0.1s)
[ 5/50] Skipping N=16681, theta_crit=0.1
[ 6/50] Skipping N=59948, theta_crit=0.1
[ 7/50] Skipping N=215443, theta_crit=0.1
[ 8/50] Skipping N=774263, theta_crit=0.1
[ 9/50] Skipping N=2782559, theta_crit=0.1
[10/50] Skipping N=10000000, theta_crit=0.1
[11/50] Running N=  100, theta_crit=0.3... mean=3.792ms (took 0.1s)
[12/50] Running N=  359, theta_crit=0.3... mean=5.151ms (took 0.1s)
[13/50] Running N= 1291, theta_crit=0.3... mean=7.014ms (took 0.1s)
[14/50] Running N= 4641, theta_crit=0.3... mean=10.191ms (took 0.1s)
[15/50] Skipping N=16681, theta_crit=0.3
[16/50] Skipping N=59948, theta_crit=0.3
[17/50] Skipping N=215443, theta_crit=0.3
[18/50] Skipping N=774263, theta_crit=0.3
[19/50] Skipping N=2782559, theta_crit=0.3
[20/50] Skipping N=10000000, theta_crit=0.3
[21/50] Running N=  100, theta_crit=0.5... mean=3.523ms (took 0.1s)
[22/50] Running N=  359, theta_crit=0.5... mean=5.194ms (took 0.1s)
[23/50] Running N= 1291, theta_crit=0.5... mean=7.422ms (took 0.1s)
[24/50] Running N= 4641, theta_crit=0.5... mean=10.060ms (took 0.1s)
[25/50] Skipping N=16681, theta_crit=0.5
[26/50] Skipping N=59948, theta_crit=0.5
[27/50] Skipping N=215443, theta_crit=0.5
[28/50] Skipping N=774263, theta_crit=0.5
[29/50] Skipping N=2782559, theta_crit=0.5
[30/50] Skipping N=10000000, theta_crit=0.5
[31/50] Running N=  100, theta_crit=0.7... mean=3.562ms (took 0.1s)
[32/50] Running N=  359, theta_crit=0.7... mean=5.142ms (took 0.1s)
[33/50] Running N= 1291, theta_crit=0.7... mean=7.078ms (took 0.1s)
[34/50] Running N= 4641, theta_crit=0.7... mean=9.638ms (took 0.1s)
[35/50] Running N=16681, theta_crit=0.7... mean=15.659ms (took 0.3s)
[36/50] Skipping N=59948, theta_crit=0.7
[37/50] Skipping N=215443, theta_crit=0.7
[38/50] Skipping N=774263, theta_crit=0.7
[39/50] Skipping N=2782559, theta_crit=0.7
[40/50] Skipping N=10000000, theta_crit=0.7
[41/50] Running N=  100, theta_crit=0.9... mean=3.928ms (took 0.1s)
[42/50] Running N=  359, theta_crit=0.9... mean=5.695ms (took 0.1s)
[43/50] Running N= 1291, theta_crit=0.9... mean=7.337ms (took 0.1s)
[44/50] Running N= 4641, theta_crit=0.9... mean=9.673ms (took 0.1s)
[45/50] Running N=16681, theta_crit=0.9... mean=13.692ms (took 0.3s)
[46/50] Skipping N=59948, theta_crit=0.9
[47/50] Skipping N=215443, theta_crit=0.9
[48/50] Skipping N=774263, theta_crit=0.9
[49/50] Skipping N=2782559, theta_crit=0.9
[50/50] Skipping N=10000000, theta_crit=0.9

Plot the performance benchmarks for all implementations

dump_folder = "_to_trash"

import os

# Create the dump directory if it does not exist
if shamrock.sys.world_rank() == 0:
    os.makedirs(dump_folder, exist_ok=True)

ref_key = "reference "
largest_refalg_value = np.nanmax(results[ref_key]["results_min"])

i = 0
# iterate over the results
for k, v in results.items():

    # Get the results for this algorithm
    particle_counts = v["particle_counts"]
    theta_crits = v["theta_crits"]
    results_min = v["results_min"]
    results_max_mem_delta = v["results_max_mem_delta"]

    # Get reference algorithm results for comparison
    reference_min = results[ref_key]["results_min"]

    # Create and display the plot
    fig, ax = create_checkerboard_plot(
        particle_counts,
        theta_crits,
        results_min,
        compression_level,
        v["name"],
        largest_refalg_value,
        reference_min,
        results_max_mem_delta,
    )

    plt.savefig(f"{dump_folder}/benchmark-dtt-performance-{i}.pdf")
    i += 1

plt.show()

• 
• 
• 

Estimated memory usage: 119 MB