Zarr Compressor Benchmark

python/benchmarks/zarr-compressors.py

@@ -0,0 +1,187 @@
+# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import argparse
+import statistics
+from dataclasses import dataclass
+from enum import Enum
+from time import perf_counter as clock
+
+import cupy
+import numcodecs.blosc
+import numpy
+from dask.utils import format_bytes, parse_bytes
+from numcodecs.abc import Codec
+
+import kvikio
+import kvikio.defaults
+import kvikio.zarr
+
+
+class Device(Enum):
+ CPU = 1
+ GPU = 2
+
+
+@dataclass
+class Compressor:
+ device: Device
+ codec: Codec
+
+
+compressors = {
+ "lz4-default": Compressor(device=Device.CPU, codec=numcodecs.LZ4()),
+ "lz4-blosc": Compressor(
+ device=Device.CPU, codec=numcodecs.blosc.Blosc(cname="lz4")
+ ),
+ "lz4-nvcomp": Compressor(device=Device.GPU, codec=kvikio.zarr.LZ4()),
+ "snappy-nvcomp": Compressor(device=Device.GPU, codec=kvikio.zarr.Snappy()),
+ "cascaded-nvcomp": Compressor(device=Device.GPU, codec=kvikio.zarr.Cascaded()),
+ "gdeflate-nvcomp": Compressor(device=Device.GPU, codec=kvikio.zarr.Gdeflate()),
+ "bitcomp-nvcomp": Compressor(device=Device.GPU, codec=kvikio.zarr.Bitcomp()),
+}
+
+
+def create_src_data(args, compressor: Compressor):
+ if compressor.device == Device.CPU:
+ return numpy.random.random(args.nelem).astype(args.dtype)
+ if compressor.device == Device.GPU:
+ return cupy.random.random(args.nelem).astype(args.dtype)
+ assert False, "Unknown device type"
+
+
+def run(args, compressor: Compressor):
+ src = create_src_data(args, compressor)
+ dst = numpy.empty_like(src) # Notice, if src is a cupy array dst is as well
+ t0 = clock()
+ a = compressor.codec.encode(src)
+ encode_time = clock() - t0
+
+ t0 = clock()
+ compressor.codec.decode(a, out=dst)
+ decode_time = clock() - t0
+ return encode_time, decode_time
+
+
+def main(args):
+ cupy.cuda.set_allocator(None) # Disable CuPy's default memory pool
+ cupy.arange(10) # Make sure CUDA is initialized
+
+ try:
+ import pynvml.smi
+
+ nvsmi = pynvml.smi.nvidia_smi.getInstance()
+ except ImportError:
+ gpu_name = "Unknown (install pynvml)"
+ mem_total = gpu_name
+ bar1_total = gpu_name
+ else:
+ info = nvsmi.DeviceQuery()["gpu"][0]
+ gpu_name = f"{info['product_name']} (dev #0)"
+ mem_total = format_bytes(
+ parse_bytes(
+ str(info["fb_memory_usage"]["total"]) + info["fb_memory_usage"]["unit"]
+ )
+ )
+ bar1_total = format_bytes(
+ parse_bytes(
+ str(info["bar1_memory_usage"]["total"])
+ + info["bar1_memory_usage"]["unit"]
+ )
+ )
+
+ print("Encode/decode benchmark")
+ print("----------------------------------")
+ print(f"GPU | {gpu_name}")
+ print(f"GPU Memory Total | {mem_total}")
+ print(f"BAR1 Memory Total | {bar1_total}")
+ print("----------------------------------")
+ print(f"nbytes | {format_bytes(args.nbytes)} ({args.nbytes})")
+ print(f"4K aligned | {args.nbytes % 4096 == 0}")
+ print(f"dtype | {args.dtype}")
+ print(f"nruns | {args.nruns}")
+ print("==================================")
+
+ encode_output = ""
+ decode_output = ""
+ # Run each benchmark using the requested APIs
+ for comp_name, comp in ((n, compressors[n]) for n in args.compressors):
+ rs = []
+ ws = []
+ for _ in range(args.n_warmup_runs):
+ encode, decode = run(args, comp)
+ for _ in range(args.nruns):
+ encode, decode = run(args, comp)
+ rs.append(args.nbytes / encode)
+ ws.append(args.nbytes / decode)
+
+ def pprint_api_res(name, samples):
+ mean = statistics.mean(samples) if len(samples) > 1 else samples[0]
+ ret = f"{comp_name} {name}".ljust(24)
+ ret += f"| {format_bytes(mean).rjust(10)}/s".ljust(14)
+ if len(samples) > 1:
+ stdev = statistics.stdev(samples) / mean * 100
+ ret += " ± %5.2f %%" % stdev
+ ret += " ("
+ for sample in samples:
+ ret += f"{format_bytes(sample)}/s, "
+ ret = ret[:-2] + ")" # Replace trailing comma
+ return ret
+
+ encode_output += pprint_api_res("", rs) + "\n"
+ decode_output += pprint_api_res("", ws) + "\n"
+ print("Encode:")
+ print(encode_output)
+ print("Decode:")
+ print(decode_output)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(description="Roundtrip benchmark")
+ parser.add_argument(
+ "-n",
+ "--nbytes",
+ metavar="BYTES",
+ default="10 MiB",
+ type=parse_bytes,
+ help="Message size, which must be a multiple of 8 (default: %(default)s).",
+ )
+ parser.add_argument(
+ "--dtype",
+ default="float32",
+ type=numpy.dtype,
+ help="NumPy datatype to use (default: '%(default)s')",
+ )
+ parser.add_argument(
+ "--nruns",
+ metavar="RUNS",
+ default=1,
+ type=int,
+ help="Number of runs per API (default: %(default)s).",
+ )
+ parser.add_argument(
+ "--n-warmup-runs",
+ default=0,
+ type=int,
+ help="Number of warmup runs (default: %(default)s).",
+ )
+ parser.add_argument(
+ "--compressors",
+ metavar="COMP_LIST",
+ default="all",
+ nargs="+",
+ choices=tuple(compressors.keys()),
+ help="List of compressors to use {%(choices)s} (default: all)",
+ )
+
+ args = parser.parse_args()
+ if "all" in args.compressors:
+ args.compressors = tuple(compressors.keys())
+
+ # Check if size is divisible by size of datatype
+ assert args.nbytes % args.dtype.itemsize == 0
+
+ # Compute/convert to number of elements
+ args.nelem = args.nbytes // args.dtype.itemsize
+
+ main(args)