Skip to content

Commit f0929db

Browse files
FrancescAltedFrancescAlted
committed
Add bench on tensordot out-of-core
1 parent c3a931a commit f0929db

1 file changed

Lines changed: 392 additions & 0 deletions

File tree

‎bench/ndarray/tensordot_pure_persistent.ipynb‎

Lines changed: 392 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,392 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "4805cb5f-cff6-46f0-97a7-caf6b46cf30c",
6+
"metadata": {
7+
"ExecuteTime": {
8+
"end_time": "2025-10-13T05:29:01.209170Z",
9+
"start_time": "2025-10-13T05:29:01.205387Z"
10+
}
11+
},
12+
"source": [
13+
"### Tensordot performance comparison between Blosc2 and Dask+Zarr with persistent storage"
14+
]
15+
},
16+
{
17+
"cell_type": "code",
18+
"execution_count": null,
19+
"id": "b95648d5a1f442e7",
20+
"metadata": {
21+
"ExecuteTime": {
22+
"end_time": "2025-10-13T05:29:02.508649Z",
23+
"start_time": "2025-10-13T05:29:01.216017Z"
24+
}
25+
},
26+
"outputs": [],
27+
"source": [
28+
"%load_ext memprofiler\n",
29+
"from time import time\n",
30+
"import numpy as np\n",
31+
"import blosc2\n",
32+
"import dask\n",
33+
"import dask.array as da\n",
34+
"import zarr\n",
35+
"from numcodecs import Blosc\n",
36+
"import h5py\n",
37+
"import hdf5plugin\n",
38+
"# It looks like b2h5py does not significantly accelerates this workload\n",
39+
"# import b2h5py.auto\n",
40+
"# assert(b2h5py.is_fast_slicing_enabled())"
41+
]
42+
},
43+
{
44+
"cell_type": "code",
45+
"execution_count": null,
46+
"id": "27d7d27956970325",
47+
"metadata": {
48+
"ExecuteTime": {
49+
"end_time": "2025-10-13T05:29:03.107498Z",
50+
"start_time": "2025-10-13T05:29:03.105334Z"
51+
}
52+
},
53+
"outputs": [],
54+
"source": [
55+
"# --- Experiment Setup ---\n",
56+
"N = 600\n",
57+
"shape_a = (N,) * 3\n",
58+
"shape_b = (N,) * 3\n",
59+
"shape_out = (N,) * 2\n",
60+
"chunks = (150,) * 3\n",
61+
"chunks_out = (150,) * 2\n",
62+
"dtype = np.float64\n",
63+
"cparams = blosc2.CParams(codec=blosc2.Codec.LZ4, clevel=1)\n",
64+
"compressor = Blosc(cname='lz4', clevel=1, shuffle=Blosc.SHUFFLE)\n",
65+
"h5compressor = hdf5plugin.Blosc2(cname='lz4', clevel=1, filters=hdf5plugin.Blosc2.SHUFFLE)\n",
66+
"create = True\n",
67+
"scheduler = \"single-threaded\" if blosc2.nthreads == 1 else \"threads\""
68+
]
69+
},
70+
{
71+
"cell_type": "code",
72+
"execution_count": null,
73+
"id": "e8d44803821da66c",
74+
"metadata": {
75+
"ExecuteTime": {
76+
"end_time": "2025-10-13T05:29:03.111527Z",
77+
"start_time": "2025-10-13T05:29:03.109952Z"
78+
}
79+
},
80+
"outputs": [],
81+
"source": [
82+
"# --- Numpy array creation ---\n",
83+
"if create:\n",
84+
" t0 = time()\n",
85+
" matrix_numpy = np.linspace(0, 1, N**3).reshape(shape_a)\n",
86+
" print(f\"N={N}, Numpy array creation = {time() - t0:.2f} s\")"
87+
]
88+
},
89+
{
90+
"cell_type": "code",
91+
"execution_count": null,
92+
"id": "bcc8a4eb914d7b9",
93+
"metadata": {
94+
"ExecuteTime": {
95+
"end_time": "2025-10-13T05:29:03.115097Z",
96+
"start_time": "2025-10-13T05:29:03.113517Z"
97+
}
98+
},
99+
"outputs": [],
100+
"source": [
101+
"# --- Blosc2 array creation ---\n",
102+
"if create:\n",
103+
" t0 = time()\n",
104+
" matrix_a_blosc2 = blosc2.asarray(matrix_numpy, cparams=cparams, chunks=chunks, urlpath=\"a.b2nd\", mode=\"w\")\n",
105+
" matrix_b_blosc2 = blosc2.asarray(matrix_numpy, cparams=cparams, chunks=chunks, urlpath=\"b.b2nd\", mode=\"w\")\n",
106+
" print(f\"N={N}, Array creation = {time() - t0:.2f} s\")"
107+
]
108+
},
109+
{
110+
"cell_type": "code",
111+
"execution_count": null,
112+
"id": "7ef51b03b68daf87",
113+
"metadata": {
114+
"ExecuteTime": {
115+
"end_time": "2025-10-13T05:29:03.121131Z",
116+
"start_time": "2025-10-13T05:29:03.117815Z"
117+
}
118+
},
119+
"outputs": [],
120+
"source": [
121+
"# Re-open the arrays\n",
122+
"t0 = time()\n",
123+
"matrix_a_blosc2 = blosc2.open(\"a.b2nd\", mode=\"r\")\n",
124+
"matrix_b_blosc2 = blosc2.open(\"b.b2nd\", mode=\"r\")\n",
125+
"print(f\"N={N}, Blosc2 array opening = {time() - t0:.2f} s\")"
126+
]
127+
},
128+
{
129+
"cell_type": "markdown",
130+
"id": "cd22e0f7-93ea-4559-bc63-cc6ae70b40c4",
131+
"metadata": {
132+
"ExecuteTime": {
133+
"end_time": "2025-10-13T05:29:23.021598Z",
134+
"start_time": "2025-10-13T05:29:13.886484Z"
135+
}
136+
},
137+
"source": [
138+
"# Tensordot computation with Blosc2"
139+
]
140+
},
141+
{
142+
"cell_type": "code",
143+
"execution_count": null,
144+
"id": "f6656fa5-5a6e-4d9c-9e86-bd422da1ae35",
145+
"metadata": {
146+
"ExecuteTime": {
147+
"end_time": "2025-10-13T05:29:07.116802Z",
148+
"start_time": "2025-10-13T05:29:03.126994Z"
149+
}
150+
},
151+
"outputs": [],
152+
"source": [
153+
"%%mprof_run 1.Blosc2::1.from_blosc2_to_blosc2\n",
154+
"# --- Tensordot computation ---\n",
155+
"for axis in ((0, 1), (1, 2), (2, 0)):\n",
156+
" t0 = time()\n",
157+
" lexpr = blosc2.lazyexpr(\"tensordot(matrix_a_blosc2, matrix_b_blosc2, axes=(axis, axis))\")\n",
158+
" out_blosc2 = lexpr.compute(urlpath=\"out.b2nd\", mode=\"w\", chunks=chunks_out)\n",
159+
" print(f\"axes={axis}, Blosc2 Performance = {time() - t0:.2f} s\")"
160+
]
161+
},
162+
{
163+
"cell_type": "code",
164+
"execution_count": null,
165+
"id": "8b2d0173c2233e8a",
166+
"metadata": {
167+
"ExecuteTime": {
168+
"end_time": "2025-10-13T05:33:48.548609Z",
169+
"start_time": "2025-10-13T05:33:48.539641Z"
170+
}
171+
},
172+
"outputs": [],
173+
"source": [
174+
"# --- HDF5 array creation ---\n",
175+
"if create:\n",
176+
" t0 = time()\n",
177+
" f = h5py.File(\"a_b_out.h5\", \"w\")\n",
178+
" f.create_dataset(\"a\", data=matrix_numpy, dtype=dtype, chunks=chunks, **h5compressor)\n",
179+
" f.create_dataset(\"b\", data=matrix_numpy, dtype=dtype, chunks=chunks, **h5compressor)\n",
180+
" f.create_dataset(\"out\", shape=shape_out, dtype=dtype, chunks=chunks_out, **h5compressor)\n",
181+
" print(f\"N={N}, HDF5 array creation = {time() - t0:.2f} s\")\n",
182+
" f.close()\n",
183+
"\n",
184+
"# Re-open the HDF5 arrays\n",
185+
"t0 = time()\n",
186+
"f = h5py.File(\"a_b_out.h5\", \"a\")\n",
187+
"matrix_a_hdf5 = f[\"a\"]\n",
188+
"matrix_b_hdf5 = f[\"b\"]\n",
189+
"out_hdf5 = f[\"out\"]"
190+
]
191+
},
192+
{
193+
"cell_type": "code",
194+
"execution_count": null,
195+
"id": "1f2d7065a801cb23",
196+
"metadata": {
197+
"ExecuteTime": {
198+
"end_time": "2025-10-13T05:29:13.857438Z",
199+
"start_time": "2025-10-13T05:29:07.134420Z"
200+
}
201+
},
202+
"outputs": [],
203+
"source": [
204+
"%%mprof_run 2.Blosc2::2.from_hdf5_to_hdf5\n",
205+
"# --- Tensordot computation with HDF5 ---\n",
206+
"for axis in ((0, 1), (1, 2), (2, 0)):\n",
207+
" t0 = time()\n",
208+
" blosc2.evaluate(\"tensordot(matrix_a_hdf5, matrix_b_hdf5, axes=(axis, axis))\", out=out_hdf5)\n",
209+
" print(f\"axes={axis}, HDF5 Performance = {time() - t0:.2f} s\")"
210+
]
211+
},
212+
{
213+
"cell_type": "code",
214+
"execution_count": null,
215+
"id": "2ef837e4e109515c",
216+
"metadata": {
217+
"ExecuteTime": {
218+
"end_time": "2025-10-13T05:29:13.870072Z",
219+
"start_time": "2025-10-13T05:29:13.867910Z"
220+
}
221+
},
222+
"outputs": [],
223+
"source": [
224+
"# --- Zarr array creation ---\n",
225+
"if create:\n",
226+
" t0 = time()\n",
227+
" matrix_a_zarr = zarr.open_array(\"a.zarr\", mode=\"w\", shape=shape_a, chunks=chunks,\n",
228+
" dtype=dtype, compressor=compressor, zarr_format=2)\n",
229+
" matrix_a_zarr[:] = matrix_numpy\n",
230+
"\n",
231+
" matrix_b_zarr = zarr.open_array(\"b.zarr\", mode=\"w\", shape=shape_b, chunks=chunks,\n",
232+
" dtype=dtype, compressor=compressor, zarr_format=2)\n",
233+
" matrix_b_zarr[:] = matrix_numpy\n",
234+
" print(f\"N={N}, Zarr array creation = {time() - t0:.2f} s\")"
235+
]
236+
},
237+
{
238+
"cell_type": "code",
239+
"execution_count": null,
240+
"id": "1185f8c3d421ef0d",
241+
"metadata": {
242+
"ExecuteTime": {
243+
"end_time": "2025-10-13T05:29:13.880901Z",
244+
"start_time": "2025-10-13T05:29:13.874433Z"
245+
}
246+
},
247+
"outputs": [],
248+
"source": [
249+
"# --- Re-open the Zarr arrays ---\n",
250+
"t0 = time()\n",
251+
"matrix_a_zarr = zarr.open(\"a.zarr\", mode=\"r\")\n",
252+
"matrix_b_zarr = zarr.open(\"b.zarr\", mode=\"r\")\n",
253+
"print(f\"N={N}, Zarr array opening = {time() - t0:.2f} s\")"
254+
]
255+
},
256+
{
257+
"cell_type": "code",
258+
"execution_count": null,
259+
"id": "c58bca30-70b3-4fc5-9514-7a0909f0cd86",
260+
"metadata": {
261+
"ExecuteTime": {
262+
"end_time": "2025-10-13T05:29:23.021598Z",
263+
"start_time": "2025-10-13T05:29:13.886484Z"
264+
}
265+
},
266+
"outputs": [],
267+
"source": [
268+
"%%mprof_run 2.Blosc2::1.from_zarr_to_zarr\n",
269+
"# --- Tensordot computation with Blosc2\n",
270+
"zout2 = zarr.open_array(\"out2.zarr\", mode=\"w\", shape=shape_out, chunks=chunks_out,\n",
271+
" dtype=dtype, compressor=compressor, zarr_format=2)\n",
272+
"for axis in ((0, 1), (1, 2), (2, 0)):\n",
273+
" t0 = time()\n",
274+
" blosc2.evaluate(\"tensordot(matrix_a_zarr, matrix_b_zarr, axes=(axis, axis))\", out=zout2)\n",
275+
" print(f\"axes={axis}, Blosc2 Performance = {time() - t0:.2f} s\")"
276+
]
277+
},
278+
{
279+
"cell_type": "markdown",
280+
"id": "f6257b5d-be65-415b-a9f5-e32a4c2d07c5",
281+
"metadata": {
282+
"ExecuteTime": {
283+
"end_time": "2025-10-13T05:33:18.928446Z",
284+
"start_time": "2025-10-13T05:33:07.317979Z"
285+
}
286+
},
287+
"source": [
288+
"# --- Tensordot computation with Dask"
289+
]
290+
},
291+
{
292+
"cell_type": "code",
293+
"execution_count": null,
294+
"id": "6097a8dd1f4673be",
295+
"metadata": {
296+
"ExecuteTime": {
297+
"end_time": "2025-10-13T05:34:08.678218Z",
298+
"start_time": "2025-10-13T05:33:52.684622Z"
299+
}
300+
},
301+
"outputs": [],
302+
"source": [
303+
"%%mprof_run 3.Dask::2.from_hdf5_to_hdf5\n",
304+
"# --- Tensordot computation with Dask (to_zarr) ---\n",
305+
"matrix_a_dask = da.from_array(matrix_a_hdf5, chunks=chunks)\n",
306+
"matrix_b_dask = da.from_array(matrix_b_hdf5, chunks=chunks)\n",
307+
"with dask.config.set(scheduler=scheduler, num_workers=blosc2.nthreads):\n",
308+
" for axis in ((0, 1), (1, 2), (2, 0)):\n",
309+
" t0 = time()\n",
310+
" dexpr = da.tensordot(matrix_a_dask, matrix_b_dask, axes=(axis, axis))\n",
311+
" da.to_hdf5('a_b_out.h5', '/out', dexpr, chunks=chunks_out)\n",
312+
" print(f\"axes={axis}, Dask Performance = {time() - t0:.2f} s\")\n",
313+
"f.close()"
314+
]
315+
},
316+
{
317+
"cell_type": "code",
318+
"execution_count": null,
319+
"id": "d3b54cac-36d6-491f-bd11-d5b86d58697a",
320+
"metadata": {
321+
"ExecuteTime": {
322+
"end_time": "2025-10-13T05:33:18.928446Z",
323+
"start_time": "2025-10-13T05:33:07.317979Z"
324+
}
325+
},
326+
"outputs": [],
327+
"source": [
328+
"%%mprof_run 3.Dask::1.from_zarr_to_zarr\n",
329+
"# --- Tensordot computation with Dask (to_zarr) ---\n",
330+
"matrix_a_dask = da.from_zarr(matrix_a_zarr, chunks=chunks)\n",
331+
"matrix_b_dask = da.from_zarr(matrix_b_zarr, chunks=chunks)\n",
332+
"zout = zarr.open_array(\"out.zarr\", mode=\"w\", shape=shape_out, chunks=chunks_out,\n",
333+
" dtype=dtype, compressor=compressor, zarr_format=2)\n",
334+
"with dask.config.set(scheduler=scheduler, num_workers=blosc2.nthreads):\n",
335+
" for axis in ((0, 1), (1, 2), (2, 0)):\n",
336+
" t0 = time()\n",
337+
" dexpr = da.tensordot(matrix_a_dask, matrix_b_dask, axes=(axis, axis))\n",
338+
" da.to_zarr(dexpr, zout, chunks=chunks_out)\n",
339+
" print(f\"axes={axis}, Dask Performance = {time() - t0:.2f} s\")"
340+
]
341+
},
342+
{
343+
"cell_type": "code",
344+
"execution_count": null,
345+
"id": "7447c635f3a870b7",
346+
"metadata": {
347+
"ExecuteTime": {
348+
"end_time": "2025-10-13T05:34:12.483993Z",
349+
"start_time": "2025-10-13T05:34:12.439333Z"
350+
}
351+
},
352+
"outputs": [],
353+
"source": [
354+
"%mprof_plot .* -t \"tensordot ({N}, {N}, {N}) -- Number of threads: {blosc2.nthreads}\""
355+
]
356+
},
357+
{
358+
"cell_type": "code",
359+
"execution_count": null,
360+
"id": "ca55545c401fff05",
361+
"metadata": {
362+
"ExecuteTime": {
363+
"end_time": "2025-10-13T05:29:50.560064Z",
364+
"start_time": "2025-10-13T05:29:50.558637Z"
365+
}
366+
},
367+
"outputs": [],
368+
"source": []
369+
}
370+
],
371+
"metadata": {
372+
"kernelspec": {
373+
"display_name": "Python 3 (ipykernel)",
374+
"language": "python",
375+
"name": "python3"
376+
},
377+
"language_info": {
378+
"codemirror_mode": {
379+
"name": "ipython",
380+
"version": 3
381+
},
382+
"file_extension": ".py",
383+
"mimetype": "text/x-python",
384+
"name": "python",
385+
"nbconvert_exporter": "python",
386+
"pygments_lexer": "ipython3",
387+
"version": "3.13.5"
388+
}
389+
},
390+
"nbformat": 4,
391+
"nbformat_minor": 5
392+
}

0 commit comments

Comments
 (0)