baldand / py-metal-compute

A python library to run metal compute kernels on macOS
MIT License
68 stars 11 forks source link

Buffers from arrays length of [1] treated as it would be as size instead of an array #21

Closed spichardo closed 1 year ago

spichardo commented 2 years ago

Hi,

I'm noticing some inconsistent behavior with arrays that have a size of [1], in some cases, it throws a error: Could not make buffer. Passing buffers of length [1] is useful when passing single value parameters to a kernel For example,

A=dev.buffer(np.array([0]).astype(np.float32))

throws the error, but

A=dev.buffer(np.array([100.343]).astype(np.float32))

"works", but actually it rather creates an array of 100 bytes instead of an array of 4 bytes (with float value of 100.0) , when pulled back it shows this

np.frombuffer(A,dtype='f').shape
25

The issue is in py-metal-compute/src/metalcompute.c https://github.com/baldand/py-metal-compute/blob/a50481ec5c583afc29753bf3c17871560e9b88b0/src/metalcompute.c#L561

The problem is that an array of [1] will pass the test of being accepted as PyNumber_Long, while probably the best would be to be treated with PyObject_GetBuffer. Maybe adding a check with PyObject_CheckBuffer or ask directly if this is an NDarray.

As temporary fix, packing the parameters as an small array is a simple easy approach, but still, the inconsistency may trigger some weird problems as some code may pass an array of length [1]

Happy to test and submit a PR to see if PyObject_CheckBuffer can help distinguish a long value from a buffer such as ndarrays.