|
| 1 | +####################################################################### |
| 2 | +# Copyright (c) 2019-present, Blosc Development Team <blosc@blosc.org> |
| 3 | +# All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under a BSD-style license (found in the |
| 6 | +# LICENSE file in the root directory of this source tree) |
| 7 | +####################################################################### |
| 8 | + |
| 9 | +import math |
| 10 | +import os |
| 11 | +import pprint |
| 12 | + |
| 13 | +import numpy as np |
| 14 | +from time import time |
| 15 | +import blosc2 |
| 16 | + |
| 17 | +dtype = np.float32 |
| 18 | +shape = (20_000, 20_000) |
| 19 | +nelem = math.prod(shape) |
| 20 | +large_shape = (50,) + shape |
| 21 | +large_nelem = math.prod(large_shape) |
| 22 | + |
| 23 | +# cparams = blosc2.CParams(codec=blosc2.Codec.BLOSCLZ, clevel=9) |
| 24 | +cparams = blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=9) |
| 25 | +# cparams = blosc2.CParams(codec=blosc2.Codec.ZSTD, clevel=1) |
| 26 | + |
| 27 | +# a = np.ones(shape, dtype=np.float64) |
| 28 | +# a = np.arange(nelem, dtype=np.float64).reshape(shape) |
| 29 | +a = np.linspace(0, 1, nelem, dtype=dtype).reshape(shape) |
| 30 | + |
| 31 | + |
| 32 | +def compute_example(a, large_a=None, intensity=None, cparams=None): |
| 33 | + t0 = time() |
| 34 | + nelem_compute = nelem if type(a) is np.ndarray else large_nelem |
| 35 | + if intensity == "low": |
| 36 | + if type(a) is not np.ndarray: |
| 37 | + a = large_a |
| 38 | + b = a + 1 |
| 39 | + intensity = 1 / 8 # 1 FLOPs / (2 * 4 bytes) |
| 40 | + elif intensity == "medium": |
| 41 | + if type(a) is not np.ndarray: |
| 42 | + a = large_a |
| 43 | + b = np.sqrt(a + 2 * a + (a / 2)) ** 1.2 |
| 44 | + intensity = 22 / 8 # 22 FLOPs / (2 * 4 bytes) |
| 45 | + elif intensity == "high": |
| 46 | + if type(a) is not np.ndarray: |
| 47 | + a = large_a |
| 48 | + b = np.exp(np.sqrt(np.sin(a) ** 2 + (np.cos(a) + np.arctan(a)) ** 3)) |
| 49 | + intensity = 133 / (2 * 4) # 133 FLOPs / (2 * 4 bytes) |
| 50 | + elif intensity.startswith("matmul"): |
| 51 | + if type(a) is not np.ndarray: |
| 52 | + if intensity == "matmul2": |
| 53 | + c = a |
| 54 | + elif intensity == "matmul1": |
| 55 | + n = shape[0] // 2 |
| 56 | + c = a.slice((slice(n, n + n), slice(n, n + n))) |
| 57 | + elif intensity == "matmul0": |
| 58 | + n = shape[0] // 5 |
| 59 | + c = a.slice((slice(n, n + n), slice(n, n + n))) |
| 60 | + else: |
| 61 | + raise ValueError("Invalid intensity") |
| 62 | + b = blosc2.matmul(c, c, cparams=cparams, urlpath="result_array.b2nd", mode="w") |
| 63 | + # print(f"b.chunks: {b.chunks}, b.blocks: {b.blocks}") |
| 64 | + else: |
| 65 | + if intensity == "matmul2": |
| 66 | + c = a |
| 67 | + elif intensity == "matmul1": |
| 68 | + n = shape[0] // 2 |
| 69 | + c = a[n:n+n, n:n+n] |
| 70 | + elif intensity == "matmul0": |
| 71 | + n = shape[0] // 5 |
| 72 | + c = a[n:n+n, n:n+n] |
| 73 | + else: |
| 74 | + raise ValueError("Invalid intensity") |
| 75 | + b = np.matmul(c, c) |
| 76 | + nelem_compute = b.size |
| 77 | + # print(type(c), type(b), nelem) |
| 78 | + intensity = int((2 * c.shape[0]) / 3) |
| 79 | + else: |
| 80 | + raise ValueError("Invalid intensity") |
| 81 | + print(f"Intensity = {intensity}", end=", ") |
| 82 | + if hasattr(b, "compute"): |
| 83 | + b = b.compute(cparams=cparams, urlpath="result_array.b2nd", mode="w") |
| 84 | + # print(f"cratio result = {b.cratio:.2f}") |
| 85 | + t = (time() - t0) |
| 86 | + print(f"Time = {t:.2f}s", end=", ") |
| 87 | + gflops = intensity * nelem_compute / t / 10**9 |
| 88 | + print(f"GFLOPS = {gflops:.2f}", end=", ") |
| 89 | + print(f"Mem BW = {nelem_compute * b.dtype.itemsize * 2 / (t * 10**9):.2f} GB/s") |
| 90 | + return {"GFLOPS": gflops, "Intensity": intensity, "Time": t} |
| 91 | + |
| 92 | +results = {"numpy": {}, "blosc2": {}, "blosc2-nocomp": {}} |
| 93 | + |
| 94 | +print("Shape:", shape) |
| 95 | +print("*** Numpy array") |
| 96 | +intensities = ["low", "medium", "high", "matmul0", "matmul1", "matmul2"] |
| 97 | +for intensity in intensities: |
| 98 | + results["numpy"][intensity] = compute_example(a, None, intensity) |
| 99 | + |
| 100 | +print("*** Blosc2 array (compressed, on-disk)") |
| 101 | +b2a = blosc2.asarray(a, cparams=cparams, urlpath="a_array.b2nd", mode="w") |
| 102 | +large_a = blosc2.linspace(0, 1, large_nelem, dtype=dtype, shape=large_shape, |
| 103 | + cparams=cparams, urlpath="large_array.b2nd", mode="w") |
| 104 | +#print(f"large_a.cratio = {large_a.cratio:.2f}, b2a.cratio = {b2a.cratio:.2f}") |
| 105 | +for intensity in intensities: |
| 106 | + results["blosc2"][intensity] = compute_example(b2a, large_a, intensity, cparams) |
| 107 | + |
| 108 | +print("*** Blosc2 array (no compressed)") |
| 109 | +no_compr_cparams = blosc2.CParams(clevel=0) |
| 110 | +b2a = blosc2.asarray(a, cparams=no_compr_cparams, urlpath="a_array.b2nd", mode="w") |
| 111 | +large_a = blosc2.copy(large_a, cparams=no_compr_cparams, urlpath="large_array2.b2nd", mode="w") |
| 112 | +#print(f"large_a.cratio = {large_a.cratio:.2f}, b2a.cratio = {b2a.cratio:.2f}") |
| 113 | +for intensity in intensities: |
| 114 | + results["blosc2-nocomp"][intensity] = compute_example(b2a, large_a, intensity, no_compr_cparams) |
| 115 | + |
| 116 | +pprint.pprint(results) |
| 117 | + |
| 118 | +os.unlink("a_array.b2nd") |
| 119 | +os.unlink("large_array.b2nd") |
| 120 | +os.unlink("large_array2.b2nd") |
| 121 | +os.unlink("result_array.b2nd") |
0 commit comments