[fix] better handle numpy arrays when their layouts are not C-order

digikar99 · Apr 7, 2024 · 4394ebf · 4394ebf
1 parent e40a87b
commit 4394ebf
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 14 deletions.
diff --git a/src/gil-gc.lisp b/src/gil-gc.lisp
@@ -102,8 +102,11 @@
       ("PyArray_Data"       nil)
       ("PyArray_GetItem"    nil)
       ("PyArray_SetItem"    nil)
+      ("PyArray_Is_C_Contiguous" nil)
+      ("PyArray_Is_F_Contiguous" nil)
       ("memcpy"             nil)
-      ("PyArray_element_type_from_array" nil))
+      ("PyArray_element_type_from_array" nil)
+      ("PY4CL_PyArray_FromArray"  :new))
   :test #'equal)
 
 (defmacro pyforeign-funcall (name-and-options &rest args)

diff --git a/src/lispifiers.lisp b/src/lispifiers.lisp
@@ -86,23 +86,38 @@
                            (pyforeign-funcall "PyArray_element_type_from_array"
                                               :pointer o :pointer)))))
          (from-vec  (pyforeign-funcall "PyArray_Data" :pointer o :pointer))
-         ;; FIXME: Strides
          (array     (make-array dims :element-type element-type))
          (num-bytes (* (array-element-type-num-bytes array)
                        (reduce #'* dims :initial-value 1))))
-    (if (type= t element-type)
-        (let* ((numpy-array-data (pyforeign-funcall "PyArray_Data"
-                                                    :pointer o :pointer)))
-          (loop :for idx :below (array-total-size array)
-                :do (setf (row-major-aref array idx)
-                          (lispify
-                           (pyforeign-funcall "PyArray_GetItem"
-                                              :pointer o
-                                              :pointer (inc-pointer numpy-array-data
-                                                                    (cl:* idx 8))
-                                              :pointer)))))
+    (when (zerop (foreign-funcall "PyArray_Is_C_Contiguous"
+                                  :pointer o :int))
+      (let ((new-pyarray
+              (pyforeign-funcall
+               "PY4CL_PyArray_FromArray"
+               :pointer o
+               :pointer (pyforeign-funcall
+                         "PyArray_Descr_from_element_type_code"
+                         :string (array-element-typecode array)
+                         :pointer)
+               :int (mem-ref (foreign-symbol-pointer
+                              "PyArray_C_Contiguous")
+                             :int)
+               :pointer)))
+        (setq from-vec (pyforeign-funcall "PyArray_Data"
+                                          :pointer new-pyarray
+                                          :pointer))))
+    (if (type= element-type t)
+        (loop :for idx :below (array-total-size array)
+              :do (setf (row-major-aref array idx)
+                        (lispify
+                         (pyforeign-funcall "PyArray_GetItem"
+                                            :pointer o
+                                            :pointer (inc-pointer from-vec
+                                                                  (cl:* idx 8))
+                                            :pointer))))
         (with-pointer-to-vector-data (to-vec (array-storage array))
-          (pyforeign-funcall "memcpy" :pointer to-vec :pointer from-vec :int num-bytes)))
+          (pyforeign-funcall "memcpy" :pointer to-vec :pointer from-vec
+                                      :int num-bytes)))
     array))
 
   ;; TODO: Test these aka find reference in documentation for why this works

diff --git a/src/py4cl-utils.c b/src/py4cl-utils.c
@@ -95,3 +95,17 @@ void* PyArray_GetItem(PyArrayObject* arr, void* itemptr){
 int PyArray_SetItem(PyArrayObject* arr, void* itemptr, PyObject* obj){
   return PyArray_SETITEM(arr, itemptr, obj);
 }
+
+const int PyArray_C_Contiguous = NPY_ARRAY_C_CONTIGUOUS;
+const int PyArray_F_Contiguous = NPY_ARRAY_F_CONTIGUOUS;
+
+int PyArray_Is_C_Contiguous(PyArrayObject* arr){
+  return PyArray_IS_C_CONTIGUOUS(arr);
+}
+int PyArray_Is_F_Contiguous(PyArrayObject* arr){
+  return PyArray_IS_F_CONTIGUOUS(arr);
+}
+
+PyObject* PY4CL_PyArray_FromArray(PyArrayObject *op, PyArray_Descr *newtype, int requirements){
+  return PyArray_FromArray(op, newtype, requirements);
+}