Optimize colorize using matmul and inplace operations

[philippjfr]

As the title says, this attempts to optimize the colorize part of the shade operation by avoiding temporary copies and performing a single matmul operation rather than multiple dot operations. In my testing this is about a 10% speedup. My guess is that this could result in even better performance for systems with MKL support.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.34%. Comparing base (f44670c) to head (36a703c).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@ Coverage Diff @@ ## main #1437 +/- ## ======================================= Coverage 88.33% 88.34% ======================================= Files 96 96 Lines 18901 18908 +7 ======================================= + Hits 16696 16704 +8  + Misses 2205 2204 -1 

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hoxbro]

Wasn't this already in place? Is it to avoid calculation along the color_mask axis?

[hoxbro]

Have you tried to play around with Einstein notation or np.tensordot for this? AI generated benchmark:

import numpy as np import timeit def benchmark(H, W, C): # Random test data color_data = np.random.rand(H, W, C) RGB = np.random.rand(C, 3) # Method 1: reshape + matmul def method_matmul(): cd2 = color_data.reshape(-1, C) return (cd2 @ RGB).reshape(H, W, 3) # Method 2: einsum def method_einsum(): return np.einsum('hwc,cj->hwj', color_data, RGB) # Method 3: einsum with optimize=True def method_einsum_opt(): return np.einsum('hwc,cj->hwj', color_data, RGB, optimize=True) # Method 4: tensordot def method_tensordot(): return np.tensordot(color_data, RGB, axes=([2],[0])) # shape: (H, W, 3) # Verify correctness out1 = method_matmul() out2 = method_einsum() out3 = method_einsum_opt() out4 = method_tensordot() assert np.allclose(out1, out2) assert np.allclose(out1, out3) assert np.allclose(out1, out4) # Benchmark time_matmul = timeit.timeit(method_matmul, number=10) time_einsum = timeit.timeit(method_einsum, number=10) time_einsum_opt = timeit.timeit(method_einsum_opt, number=10) time_tensordot = timeit.timeit(method_tensordot, number=10) print(f"H={H}, W={W}, C={C}") print(f"reshape+matmul: {time_matmul:.4f} s") print(f"einsum: {time_einsum:.4f} s") print(f"einsum (optimize=True): {time_einsum_opt:.4f} s") print(f"tensordot: {time_tensordot:.4f} s") print("-" * 50) # Test different shapes benchmark(256, 256, 64) benchmark(512, 512, 64) benchmark(256, 256, 256) benchmark(128, 128, 1024)

H=256, W=256, C=64 reshape+matmul: 0.0284 s einsum: 0.1792 s einsum (optimize=True): 0.0188 s tensordot: 0.0280 s -------------------------------------------------- H=512, W=512, C=64 reshape+matmul: 0.1330 s einsum: 0.5891 s einsum (optimize=True): 0.0341 s tensordot: 0.1112 s -------------------------------------------------- H=256, W=256, C=256 reshape+matmul: 0.1042 s einsum: 0.6018 s einsum (optimize=True): 0.0419 s tensordot: 0.1038 s -------------------------------------------------- H=128, W=128, C=1024 reshape+matmul: 0.0713 s einsum: 0.6043 s einsum (optimize=True): 0.0497 s tensordot: 0.0712 s -------------------------------------------------- 

[philippjfr]

Thanks, not 100% sure what to take from that though I will state that in most cases C << 100.

[hoxbro]

Why I said it was AI-generated. I assume you had a more fully fledged way to actually measure your performance. This was more to show that there was a tiny bit of performance that could be gained here.

[hoxbro]

Can't we use an existing mask for this? I would assume nan_to_mask would do this calculation behind the scenes.

Also note that this will convert inf values to floats.

[codspeed-hq]

CodSpeed Instrumentation Performance Report

Merging #1437 will degrade performances by 17.88%

_{Comparing optimize_colorize (36a703c) with main (a4d57be)}

Summary

❌ 1 regression
✅ 42 untouched

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Benchmarks breakdown

	Benchmark	BASE	HEAD	Change
❌	test_layout[forceatlas2_layout]	68.7 ms	83.6 ms	-17.88%

[philippjfr]

My profiling code

import time import numpy as np import pandas as pd import datashader as ds N = int(10e6) C = 20 def gen_data(N=int(10e6), C=20): xy = np.random.randn(int(N), 2) c = np.random.choice([chr(65+i) for i in range(C)], size=N) df = pd.DataFrame(xy, columns=['x', 'y']) df['c'] = pd.Series(c).astype('category') return df def profile(df, size=1000): W = H = size cvs = ds.Canvas(plot_width=W, plot_height=H) agg = cvs.points(df, x='x', y='y', agg=ds.count_cat('c')) pre = time.monotonic() ds.transfer_functions.shade(agg) return time.monotonic()-pre # Warmup df = gen_data(C=1) profile(df, size=10) results = [] for c in (1, 5, 10, 20): df = gen_data(C=c) for s in range(1000, 6000, 1000): timing = profile(df, size=s) results.append((c, s, timing))

[hoxbro]

Current status:

current = 6b0982b, before = fb2c16e, main = 184ef3c

Benchmark code

import time import numpy as np import pandas as pd import datashader as ds import datashader.transfer_functions as tf N = int(10e6) C = 20 class Profile: def __init__(self, output_file): import cProfile self.profiler = cProfile.Profile() self.output_file = output_file def __enter__(self): self.profiler.enable() return self.profiler def __exit__(self, exc_type, exc_val, exc_tb): self.profiler.disable() self.profiler.dump_stats(self.output_file) class LineProfileContext: def __init__(self, output_file): from line_profiler import LineProfiler self.profiler = LineProfiler() self.output_file = output_file self.functions_to_profile = [] def add_function(self, func): """Add a function to be profiled line-by-line""" self.profiler.add_function(func) self.functions_to_profile.append(func) return func def __enter__(self): self.profiler.enable_by_count() return self def __exit__(self, exc_type, exc_val, exc_tb): self.profiler.disable_by_count() self.profiler.dump_stats(self.output_file) self.profiler.print_stats() def gen_data(N=N, C=C): np.random.seed(1) xy = np.random.randn(int(N), 2) c = np.random.choice([chr(65 + i) for i in range(C)], size=N) df = pd.DataFrame(xy, columns=["x", "y"]) df["c"] = pd.Series(c).astype("category") return df def profile(df, size=1000): W = H = size cvs = ds.Canvas(plot_width=W, plot_height=H) agg = cvs.points(df, x="x", y="y", agg=ds.count_cat("c")) tf.shade(agg) # warm up pre = time.monotonic() # with LineProfileContext("line_profile.lprof") as line_profiler: # line_profiler.add_function(tf._colorize) # tf.shade(agg) # with Profile(output_file="optional.perf"): # ds.transfer_functions.shade(agg) return time.monotonic() - pre # Warmup df = gen_data(C=20) profile(df, size=5000) results = [] for c in (1, 5, 10, 20): df = gen_data(C=c) for s in range(1000, 6000, 1000): timing = profile(df, size=s) results.append((c, s, timing)) print(f"{c=}, {s=}, {timing=}")

Plotting

current = [ # 6b0982b dict(c=1, s=1000, timing=0.07571580299918423), dict(c=1, s=2000, timing=0.295644159999938), dict(c=1, s=3000, timing=0.6464670440000191), dict(c=1, s=4000, timing=1.1230143669999961), dict(c=1, s=5000, timing=1.7188509200004773), dict(c=5, s=1000, timing=0.11476161499922455), dict(c=5, s=2000, timing=0.44561883799906354), dict(c=5, s=3000, timing=0.9790756620004686), dict(c=5, s=4000, timing=1.7118233849996614), dict(c=5, s=5000, timing=2.6856131889999233), dict(c=10, s=1000, timing=0.14612587800002075), dict(c=10, s=2000, timing=0.5675542549997772), dict(c=10, s=3000, timing=1.2379318599996623), dict(c=10, s=4000, timing=2.2251677369986282), dict(c=10, s=5000, timing=3.397321854999973), dict(c=20, s=1000, timing=0.1993868179997662), dict(c=20, s=2000, timing=0.8214870430001611), dict(c=20, s=3000, timing=1.7614306820014463), dict(c=20, s=4000, timing=3.0943053329992836), dict(c=20, s=5000, timing=4.7508491489988955), ] before = [ # fb2c16e dict(c=1, s=1000, timing=0.07645769699956873), dict(c=1, s=2000, timing=0.3170905290007795), dict(c=1, s=3000, timing=0.7142776969994884), dict(c=1, s=4000, timing=1.2551025209995714), dict(c=1, s=5000, timing=1.9599227520011482), dict(c=5, s=1000, timing=0.16230177999932494), dict(c=5, s=2000, timing=0.5520949959991412), dict(c=5, s=3000, timing=1.2177506650004943), dict(c=5, s=4000, timing=2.171157504999428), dict(c=5, s=5000, timing=3.3560801679996075), dict(c=10, s=1000, timing=0.2009295749994635), dict(c=10, s=2000, timing=0.7160231019988714), dict(c=10, s=3000, timing=1.6094946339999296), dict(c=10, s=4000, timing=2.7828460880009516), dict(c=10, s=5000, timing=4.274540911001168), dict(c=20, s=1000, timing=0.2542700350004452), dict(c=20, s=2000, timing=0.9284682460001932), dict(c=20, s=3000, timing=2.0608999519990903), dict(c=20, s=4000, timing=3.6744658019997587), dict(c=20, s=5000, timing=5.747611536000477), ] main = [ # 184ef3c dict(c=1, s=1000, timing=0.0718935530003364), dict(c=1, s=2000, timing=0.31208833799973945), dict(c=1, s=3000, timing=0.7055044320004527), dict(c=1, s=4000, timing=1.2214937410008133), dict(c=1, s=5000, timing=1.899291293000715), dict(c=5, s=1000, timing=0.1668667740013916), dict(c=5, s=2000, timing=0.6655240790005337), dict(c=5, s=3000, timing=1.5014597809986299), dict(c=5, s=4000, timing=2.5980365989998973), dict(c=5, s=5000, timing=4.086923677999948), dict(c=10, s=1000, timing=0.2160664650000399), dict(c=10, s=2000, timing=0.8515692499986471), dict(c=10, s=3000, timing=1.879354708999017), dict(c=10, s=4000, timing=3.3537094929997693), dict(c=10, s=5000, timing=5.179373872999349), dict(c=20, s=1000, timing=0.3006982959996094), dict(c=20, s=2000, timing=1.1935125410000182), dict(c=20, s=3000, timing=2.72798394000165), dict(c=20, s=4000, timing=4.852396742000565), dict(c=20, s=5000, timing=7.331704080999771), ] import hvplot.pandas import pandas as pd fn = ( lambda x: pd.DataFrame(eval(x)) .hvplot.bar(x="s", y="timing", by="c", title=x) .opts(show_grid=True) ) (fn("current") + fn("before") + fn("main")).cols(1)