<jupyter_start><jupyter_text>metalcompute and NumPyThis notebook shows how to use NumPy (np) arrays with metalcompute (on macOS)np arrays can be passed directly as arguments to metalcompute kernels for inputmetalcompute buffers can also be wrapped using np for easier access and usageFirst we import both numpy (np) and metalcompute (mc). (This assumes both are installed already in the python environment using pip install numpy & pip install metalcompute)<jupyter_code>import numpy as np
import metalcompute as mc<jupyter_output><empty_output><jupyter_text>Next we create a metalcompute device. On an M1 family mac, this will be the built in GPU.<jupyter_code>dev = mc.Device()<jupyter_output><empty_output><jupyter_text>This is a simple kernel function called "add", to add together (element-wise) two arrays of float32 values, and write the result to a third array.<jupyter_code>kernel = dev.kernel("""
// We always need these two lines in a metal kernel
#include <metal_stdlib>
using namespace metal;

// This is the add function
kernel void add(
    // These are the two input arrays, const as we will not write to them
    const device float* a [[ buffer(0) ]],
    const device float *b [[ buffer(1) ]],
    // This is the output array
    device float *c [[ buffer(2) ]],
    // This is the index of the current kernel instance
    uint id [[ thread_position_in_grid ]]) 
{
    // This is the add operation: c = a + b (for each element)
    c[id] = a[id] + b[id]; 
}
""")

add_fn = kernel.function("add")<jupyter_output><empty_output><jupyter_text>We will test this with 256MB buffers<jupyter_code>count = 1024*1024*64
size = count * 4 # 256MB<jupyter_output><empty_output><jupyter_text>Next we create two np test arrays, a_np and b_np, and calculate the sum of those as a reference<jupyter_code>a_np = np.arange(count,dtype='f') # f32 array
b_np = (count - a_np).astype('f') # Cast to f32 array
c_np = a_np + b_np # Calculate reference result<jupyter_output><empty_output><jupyter_text>Now we create metalcompute buffers with copies of the np data, and space for the result<jupyter_code>a = dev.buffer(a_np) # Create mc buffer as copy of a_np
b = dev.buffer(b_np) # Create mc buffer as copy of b_np
c = dev.buffer(size) # Space for the result<jupyter_output><empty_output><jupyter_text>Now we do the add calculation using metalcompute<jupyter_code>handle = add_fn(count,a,b,c)<jupyter_output><empty_output><jupyter_text>Next we wait for the compute to finish (by deleting the handle) and check the result against the numpy version.Note how the metalcompute buffer can be wrapped into a numpy array using np.frombuffer<jupyter_code>del handle # Will block until the compute has finished
assert((c_np == np.frombuffer(c,dtype='f')).all())<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.0882564841498559, 'fraction_numerical': 0.009726224783861672, 'mean_word_length': 5.212527964205816, 'pattern_counts': {'":': 0, '<': 34, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '4833', 'license': 'MIT', 'n_tokens_mistral': 882, 'n_tokens_neox': 856, 'n_words': 331, 'path': 'examples/metalcompute_numpy.ipynb', 'repo_name': 'ProteusMRIgHIFU/py-metal-compute', 'size': '2776'}