Shamrock DeviceBuffer usage

This example shows how to use the shamrock.backends.DeviceBuffer_f64 class for GPU-accelerated data storage and manipulation.

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")

-> modified loglevel to 0 enabled log types :
log status :
 - Loglevel: 1, enabled log types :
[xxx] Info: xxx ( logger::info )
[xxx] : xxx ( logger::normal )
[xxx] Warning: xxx ( logger::warn )
[xxx] Error: xxx ( logger::err )

Create a new DeviceBuffer for f64 (double precision floats)

buffer = shamrock.backends.DeviceBuffer_f64()

print(f"Initial buffer size: {buffer.get_size()}")

Initial buffer size: 0

Resize the buffer to hold 10 elements

buffer.resize(10)
print(f"Buffer size after resize: {buffer.get_size()}")

Buffer size after resize: 10

Set individual values using set_val_at_idx

for i in range(10):
    buffer.set_val_at_idx(i, float(i * 2.5))

print("Values set individually:")
for i in range(buffer.get_size()):
    print(f"buffer[{i}] = {buffer.get_val_at_idx(i)}")

Values set individually:
buffer[0] = 0.0
buffer[1] = 2.5
buffer[2] = 5.0
buffer[3] = 7.5
buffer[4] = 10.0
buffer[5] = 12.5
buffer[6] = 15.0
buffer[7] = 17.5
buffer[8] = 20.0
buffer[9] = 22.5

Copy data from a Python list/numpy array

data = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0]
buffer.copy_from_stdvec(data)

print(f"\nAfter copying from list: {data}")
for i in range(buffer.get_size()):
    print(f"buffer[{i}] = {buffer.get_val_at_idx(i)}")

After copying from list: [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0]
buffer[0] = 1.1
buffer[1] = 2.2
buffer[2] = 3.3
buffer[3] = 4.4
buffer[4] = 5.5
buffer[5] = 6.6
buffer[6] = 7.7
buffer[7] = 8.8
buffer[8] = 9.9
buffer[9] = 10.0

Copy data back to a Python list

result = buffer.copy_to_stdvec()
print(f"\nCopied back to Python list: {result}")

Copied back to Python list: [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0]

Working with numpy arrays

np_data = np.linspace(0.0, 1.0, 20)
print(f"Original numpy array (size {len(np_data)}): {np_data}")

# Resize buffer to match numpy array size
buffer.resize(len(np_data))
buffer.copy_from_stdvec(np_data.tolist())

print(f"Buffer size after numpy copy: {buffer.get_size()}")

Original numpy array (size 20): [0.         0.05263158 0.10526316 0.15789474 0.21052632 0.26315789
 0.31578947 0.36842105 0.42105263 0.47368421 0.52631579 0.57894737
 0.63157895 0.68421053 0.73684211 0.78947368 0.84210526 0.89473684
 0.94736842 1.        ]
Buffer size after numpy copy: 20

Perform some operations: multiply each element by 2

for i in range(buffer.get_size()):
    current_val = buffer.get_val_at_idx(i)
    buffer.set_val_at_idx(i, current_val * 2.0)

# Copy back and convert to numpy array
modified_data = np.array(buffer.copy_to_stdvec())
print(f"Modified data (multiplied by 2): {modified_data}")

Modified data (multiplied by 2): [0.         0.10526316 0.21052632 0.31578947 0.42105263 0.52631579
 0.63157895 0.73684211 0.84210526 0.94736842 1.05263158 1.15789474
 1.26315789 1.36842105 1.47368421 1.57894737 1.68421053 1.78947368
 1.89473684 2.        ]

Demonstrate with larger dataset

rng = np.random.default_rng(42)  # Use modern random generator with seed for reproducibility
large_data = rng.random(1000)
buffer.resize(len(large_data))
buffer.copy_from_stdvec(large_data.tolist())

print(f"Processed large dataset of size: {buffer.get_size()}")
print(f"First 10 elements: {buffer.copy_to_stdvec()[:10]}")
print(f"Last 10 elements: {buffer.copy_to_stdvec()[-10:]}")

Processed large dataset of size: 1000
First 10 elements: [0.7739560485559633, 0.4388784397520523, 0.8585979199113825, 0.6973680290593639, 0.09417734788764953, 0.9756223516367559, 0.761139701990353, 0.7860643052769538, 0.12811363267554587, 0.45038593789556713]
Last 10 elements: [0.6227087449882082, 0.7508385847448596, 0.7936353838623774, 0.2119546110805426, 0.9237833160931802, 0.4379896029203726, 0.6390136492641646, 0.002308772586742003, 0.993368498411178, 0.2808963859641551]

Estimated memory usage: 162 MB