Added benchmark code

arrayfire/array_object.py

@@ -746,27 +746,48 @@ def __getitem__(self, key: IndexKey, /) -> Array:
         ----------
         self : Array
             Array instance.
-        key : int | slice | tuple[int | slice, ...] | Array
+        key : int | slice | tuple[int | slice | Array, ...] | Array
             Index key.
 
         Returns
         -------
         out : Array
             An array containing the accessed value(s). The returned array must have the same data type as self.
         """
-        # TODO
-        # API Specification - key: Union[int, slice, ellipsis, tuple[Union[int, slice, ellipsis], ...], array].
-        # consider using af.span to replace ellipsis during refactoring
         out = Array()
         ndims = self.ndim
 
-        if isinstance(key, Array) and key == afbool.c_api_value:
+        indexing = key
+
+        if isinstance(key, int | float | slice): # when indexing with one dimension, treat it as indexing a flat array
+            ndims = 1
+        elif isinstance(key, Array): # when indexing with one array, treat it as indexing a flat array
             ndims = 1
-            if wrapper.count_all(key.arr) == 0:  # HACK was count() method before
-                return out
+            if key.is_bool:
+                indexing = wrapper.where(key.arr)
+            else:
+                indexing = key.arr
+        elif isinstance(key, tuple):
+            key_list = []
+            for elem in key:
+                if isinstance(elem, Array):
+                    if elem.is_bool:
+                        key_list.append(wrapper.where(elem.arr))
+                    else:
+                        key_list.append(elem.arr)
+                else:
+                    key_list.append(elem)
+            indexing = tuple(key_list)
+
+        out._arr = wrapper.index_gen(self._arr, ndims, wrapper.get_indices(indexing))  # type: ignore[arg-type]
+
+        if isinstance(key, Array) and key.is_bool:
+            wrapper.release_array(indexing)
+        elif isinstance(key, tuple):
+            for i in range(len(key)):
+                if isinstance(key[i], Array) and key[i].is_bool:
+                    wrapper.release_array(indexing[i])
 
-        # HACK known issue
-        out._arr = wrapper.index_gen(self._arr, ndims, wrapper.get_indices(key))  # type: ignore[arg-type]
         return out
 
     def __index__(self) -> int:
@@ -781,8 +802,19 @@ def __len__(self) -> int:
         return self.shape[0] if self.shape else 0
 
     def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> None:
-        ndims = self.ndim
+        """
+        Assigns self[key] = value
+
+        Parameters
+        ----------
+        self : Array
+            Array instance.
+        key : int | slice | tuple[int | slice | Array, ...] | Array
+            Index key.
+        value: int | float | complex | bool | Array
 
+        """
+        ndims = self.ndim
         is_array_with_bool = isinstance(key, Array) and type(key) is afbool
 
         if is_array_with_bool:
@@ -803,12 +835,41 @@ def __setitem__(self, key: IndexKey, value: int | float | bool | Array, /) -> No
             other_arr = value.arr
             del_other = False
 
-        indices = wrapper.get_indices(key)  # type: ignore[arg-type]  # FIXME
-        out = wrapper.assign_gen(self._arr, other_arr, ndims, indices)
+        indexing = key
+        if isinstance(key, int | float | slice): # when indexing with one dimension, treat it as indexing a flat array
+            ndims = 1
+        elif isinstance(key, Array): # when indexing with one array, treat it as indexing a flat array
+            ndims = 1
+            if key.is_bool:
+                indexing = wrapper.where(key.arr)
+            else:
+                indexing = key.arr
+        elif isinstance(key, tuple):
+            key_list = []
+            for elem in key:
+                if isinstance(elem, Array):
+                    if elem.is_bool:
+                        locs = wrapper.where(elem.arr)
+                        key_list.append(locs)
+                    else:
+                        key_list.append(elem.arr)
+                else:
+                    key_list.append(elem)
+            indexing = tuple(key_list)
+
+        out = wrapper.assign_gen(self._arr, other_arr, ndims, wrapper.get_indices(indexing))
+
+        if isinstance(key, Array) and key.is_bool:
+            wrapper.release_array(indexing)
+        elif isinstance(key, tuple):
+            for i in range(len(key)):
+                if isinstance(key[i], Array) and key[i].is_bool:
+                    wrapper.release_array(indexing[i])
 
         wrapper.release_array(self._arr)
         if del_other:
             wrapper.release_array(other_arr)
+
         self._arr = out
 
     def __str__(self) -> str:
@@ -1144,7 +1205,10 @@ def _get_processed_index(key: IndexKey, shape: tuple[int, ...]) -> tuple[int, ..
     if isinstance(key, tuple):
         return tuple(_index_to_afindex(key[i], shape[i]) for i in range(len(key)))
 
-    return (_index_to_afindex(key, shape[0]),) + shape[1:]
+    size = 1
+    for dim_size in shape:
+        size *= dim_size
+    return (_index_to_afindex(key, size),)
 
 
 def _index_to_afindex(key: int | float | complex | bool | slice | wrapper.ParallelRange | Array, axis: int) -> int:
@@ -1168,6 +1232,10 @@ def _index_to_afindex(key: int | float | complex | bool | slice | wrapper.Parall
 
 
 def _slice_to_length(key: slice, axis: int) -> int:
+    start = key.start
+    stop = key.stop
+    step = key.step
+
     if key.start is None:
         start = 0
     elif key.start < 0:

arrayfire/library/array_functions.py

@@ -954,6 +954,7 @@ def replace(lhs: Array, rhs: Array | int | float, conditional: Array, /) -> None
     wrapper.replace_scalar(lhs.arr, conditional.arr, rhs)
 
 
+@afarray_as_array
 def select(lhs: Array | int | float, rhs: Array | int | float, conditional: Array, /) -> Array:
     """
     Conditionally selects elements from one of two sources (ArrayFire arrays or scalars) based on a condition array.

arrayfire/library/linear_algebra.py

@@ -92,6 +92,7 @@ def gemm(
     rhs_opts: MatProp = MatProp.NONE,
     alpha: int | float = 1.0,
     beta: int | float = 0.0,
+    accum: Array = None
 ) -> Array:
     """
     Performs BLAS general matrix multiplication (GEMM) on two Array instances.
@@ -125,6 +126,10 @@ def gemm(
     beta : int | float, optional
         Scalar multiplier for the existing matrix C in the accumulation. Default is 0.0.
 
+    accum: Array, optional
+        A 2-dimensional, real or complex array representing the matrix C in the accumulation.
+        Default is None (no accumulation).
+
     Returns
     -------
     Array
@@ -135,7 +140,10 @@ def gemm(
     - The data types of `lhs` and `rhs` must be compatible.
     - Batch operations are not supported in this version.
     """
-    return cast(Array, wrapper.gemm(lhs.arr, rhs.arr, lhs_opts, rhs_opts, alpha, beta))
+    accumulator = None
+    if isinstance(accum, Array):
+        accumulator = accum.arr
+    return cast(Array, wrapper.gemm(lhs.arr, rhs.arr, lhs_opts, rhs_opts, alpha, beta, accumulator))
 
 
 @afarray_as_array

arrayfire/library/mathematical_functions.py

@@ -261,12 +261,11 @@ def atan2(x1: int | float | Array, x2: int | float | Array, /) -> Array:
     return process_c_function(x1, x2, wrapper.atan2)
 
 
-@afarray_as_array
 def cplx(x1: int | float | Array, /, x2: int | float | Array | None = None) -> Array:
     if x2 is None:
         if not isinstance(x1, Array):
             raise TypeError("x1 can not be int or tuple when x2 is None.")
-        return cast(Array, wrapper.cplx(x1.arr))
+        return Array.from_afarray(Array, wrapper.cplx(x1.arr))
     else:
         return process_c_function(x1, x2, wrapper.cplx2)
 

benchmarks/src/README.md

@@ -0,0 +1,36 @@
+Benchmarks
+===========
+
+## Setting up environment
+
+```sh
+    python -m pip install -r requirements.txt
+```
+
+## Benchmark parameters
+
+The benchmark packages, rounds, array sizes, and numeric type may be specified on the constants at the top of [pytest_benchmark/common.py](pytest_benchmark/common.py).
+
+Alternatively, they may be specified individually at the top of each test file.
+
+
+## Running
+
+These are the steps to run the benchmarks, and produce the graphs 
+
+Run the benchmarks and store the results in `results.json`
+```sh
+    pytest .\pytest_benchmark --benchmark-json=results.json
+```
+
+To create graphs and store the timing results after creating the `results.json`, run:
+```sh
+    python graphs.py
+```
+
+To modify the tests being shown, modify the `TESTS` list at the top of the `graphs.py` file.
+To modify the labels shown, modify `PKG_LABELS`
+To modify the hardware display, modify `HARDWARE` 
+
+## Notes
+When running with `dpnp`, set the environment variable `DPNP_RAISE_EXCEPION_ON_NUMPY_FALLBACK` to 0.