Expand API: BLAS, reductions, statistics, index ops, bitwise; type & FFI fixes

.github/workflows/ci.yml

@@ -17,6 +17,11 @@ jobs:
     steps:
       - uses: actions/checkout@v4
 
+      - uses: cachix/cachix-action@v16
+        with:
+          name: arrayfire-haskell
+          authToken: '${{ secrets.CACHIX_AUTH_TOKEN }}'
+
       # macos-latest is Apple Silicon, but ArrayFire only ships an x86_64 macOS
       # binary, so the flake's darwin output is x86_64-darwin. Build it under
       # Rosetta 2: ensure Rosetta is present and let Nix build x86_64-darwin

.gitignore

@@ -7,3 +7,4 @@ result/
 cabal.project.local
 tags
 /.stack-work/
+/.ghc.environment*

arrayfire.cabal

@@ -1,6 +1,6 @@
 cabal-version:       3.0
 name:                arrayfire
-version:             0.7.1.0
+version:             0.8.0.0
 synopsis:            Haskell bindings to the ArrayFire general-purpose GPU library
 homepage:            https://github.com/arrayfire/arrayfire-haskell
 license:             BSD-3-Clause
@@ -177,6 +177,7 @@ test-suite test
     ArrayFire.ImageSpec
     ArrayFire.IndexSpec
     ArrayFire.LAPACKSpec
+    ArrayFire.NumericalSpec
     ArrayFire.RandomSpec
     ArrayFire.SignalSpec
     ArrayFire.SparseSpec

include/algorithm.h

@@ -34,3 +34,12 @@ af_err af_sort_by_key(af_array *out_keys, af_array *out_values, const af_array k
 af_err af_set_unique(af_array *out, const af_array in, const bool is_sorted);
 af_err af_set_union(af_array *out, const af_array first, const af_array second, const bool is_unique);
 af_err af_set_intersect(af_array *out, const af_array first, const af_array second, const bool is_unique);
+af_err af_sum_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_sum_by_key_nan(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim, const double nanval);
+af_err af_product_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_product_by_key_nan(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim, const double nanval);
+af_err af_min_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_max_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_all_true_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_any_true_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);
+af_err af_count_by_key(af_array *keys_out, af_array *vals_out, const af_array keys, const af_array vals, const int dim);

include/blas.h

@@ -5,3 +5,4 @@ af_err af_dot(af_array *out, const af_array lhs, const af_array rhs, const af_ma
 af_err af_dot_all(double *real, double *imag, const af_array lhs, const af_array rhs, const af_mat_prop optLhs, const af_mat_prop optRhs);
 af_err af_transpose(af_array *out, af_array in, const bool conjugate);
 af_err af_transpose_inplace(af_array in, const bool conjugate);
+af_err af_gemm(af_array *out, const af_mat_prop optLhs, const af_mat_prop optRhs, const void *alpha, const af_array lhs, const af_array rhs, const void *beta);

include/statistics.h

@@ -15,3 +15,4 @@ af_err af_stdev_all(double *real, double *imag, const af_array in);
 af_err af_median_all(double *realVal, double *imagVal, const af_array in);
 af_err af_corrcoef(double *realVal, double *imagVal, const af_array X, const af_array Y);
 af_err af_topk(af_array *values, af_array *indices, const af_array in, const int k, const int dim, const af_topk_function order);
+af_err af_meanvar(af_array *mean, af_array *var, const af_array in, const af_array weights, const af_var_bias bias, const dim_t dim);

src/ArrayFire/Algorithm.hs

@@ -26,6 +26,9 @@
 --------------------------------------------------------------------------------
 module ArrayFire.Algorithm where
 
+import Data.Word (Word32)
+import Foreign.C.Types (CBool)
+
 import ArrayFire.FFI
 import ArrayFire.Internal.Algorithm
 import ArrayFire.Internal.Types
@@ -152,13 +155,13 @@ max x (fromIntegral -> n) = x `op1` (\p a -> af_max p a n)
 -- [1 1 1 1]
 --         0
 allTrue
-  :: forall a. AFType a
+  :: AFType a
   => Array a
   -- ^ Array input
   -> Int
   -- ^ Dimension along which to see if all elements are True
-  -> Array a
-  -- ^ Will contain the maximum of all values in the input array along dim
+  -> Array CBool
+  -- ^ Will contain 1 where all elements along dim are true, 0 otherwise
 allTrue x (fromIntegral -> n) =
   x `op1` (\p a -> af_all_true p a n)
 
@@ -169,13 +172,13 @@ allTrue x (fromIntegral -> n) =
 -- [1 1 1 1]
 --         0
 anyTrue
-  :: forall a . AFType a
+  :: AFType a
   => Array a
   -- ^ Array input
   -> Int
-  -- ^ Dimension along which to see if all elements are True
-  -> Array a
-  -- ^ Returns if all elements are true
+  -- ^ Dimension along which to see if any elements are True
+  -> Array CBool
+  -- ^ Will contain 1 where any element along dim is true, 0 otherwise
 anyTrue x (fromIntegral -> n) =
   (x `op1` (\p a -> af_any_true p a n))
 
@@ -193,7 +196,7 @@ count
   -- ^ Dimension along which to count
   -> Array Int
   -- ^ Count of all elements along dimension
-count x (fromIntegral -> n) = x `op1d` (\p a -> af_count p a n)
+count x (fromIntegral -> n) = x `op1` (\p a -> af_count p a n)
 
 -- | Sum all elements in an 'Array' along all dimensions
 --
@@ -323,7 +326,7 @@ imin
   -- ^ Input array
   -> Int
   -- ^ The dimension along which the minimum value is extracted
-  -> (Array a, Array a)
+  -> (Array a, Array Word32)
   -- ^ will contain the minimum of all values along dim, will also contain the location of minimum of all values in in along dim
 imin a (fromIntegral -> n) = op2p a (\x y z -> af_imin x y z n)
 
@@ -343,7 +346,7 @@ imax
   -- ^ Input array
   -> Int
   -- ^ The dimension along which the minimum value is extracted
-  -> (Array a, Array a)
+  -> (Array a, Array Word32)
   -- ^ will contain the maximum of all values in in along dim, will also contain the location of maximum of all values in in along dim
 imax a (fromIntegral -> n) = op2p a (\x y z -> af_imax x y z n)
 
@@ -471,8 +474,8 @@ where'
   :: AFType a
   => Array a
   -- ^ Is the input array.
-  -> Array a
-  -- ^ will contain indices where input array is non-zero
+  -> Array Word32
+  -- ^ Indices where input array is non-zero
 where' = (`op1` af_where)
 
 -- | First order numerical difference along specified dimension.
@@ -565,7 +568,7 @@ sortIndex
   -- ^ Dimension along `sortIndex` is performed
   -> Bool
   -- ^ Return results in ascending order
-  -> (Array a, Array a)
+  -> (Array a, Array Word32)
   -- ^ Contains the sorted, contains indices for original input
 sortIndex a (fromIntegral -> n) (fromIntegral . fromEnum -> b) =
   a `op2p` (\p1 p2 p3 -> af_sort_index p1 p2 p3 n b)
@@ -657,3 +660,137 @@ setIntersect
   -- ^ Intersection of first and second array
 setIntersect a1 a2 (fromIntegral . fromEnum -> b) =
   op2 a1 a2 (\x y z -> af_set_intersect x y z b)
+
+-- | Sum values in 'Array' grouped by keys along a dimension.
+--
+-- Each contiguous run of equal keys in @keys@ produces one output element.
+-- Returns @(keys_out, vals_out)@.
+--
+-- >>> sumByKey (vector @Int 5 [1,1,2,2,2]) (vector @Double 5 [10,20,1,2,3]) 0
+-- (ArrayFire Array
+-- [2 1 1 1]
+--    1   2,
+-- ArrayFire Array
+-- [2 1 1 1]
+--    30.0000   6.0000)
+sumByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array (contiguous equal keys form a group)
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension along which to reduce
+  -> (Array Int, Array a)
+  -- ^ (reduced keys, reduced values)
+sumByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_sum_by_key ko vo k v dim)
+
+-- | 'sumByKey' replacing NaN values with a substitute before summing.
+sumByKeyNaN
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> Double
+  -- ^ Substitute for NaN values
+  -> (Array Int, Array a)
+  -- ^ (reduced keys, reduced values)
+sumByKeyNaN keys vals (fromIntegral -> dim) nanval =
+  op2p2kv keys vals (\ko vo k v -> af_sum_by_key_nan ko vo k v dim nanval)
+
+-- | Product of values in 'Array' grouped by keys along a dimension.
+productByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+productByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_product_by_key ko vo k v dim)
+
+-- | 'productByKey' replacing NaN values with a substitute before multiplying.
+productByKeyNaN
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> Double
+  -- ^ Substitute for NaN values
+  -> (Array Int, Array a)
+productByKeyNaN keys vals (fromIntegral -> dim) nanval =
+  op2p2kv keys vals (\ko vo k v -> af_product_by_key_nan ko vo k v dim nanval)
+
+-- | Minimum of values in 'Array' grouped by keys along a dimension.
+minByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+minByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_min_by_key ko vo k v dim)
+
+-- | Maximum of values in 'Array' grouped by keys along a dimension.
+maxByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+maxByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_max_by_key ko vo k v dim)
+
+-- | True if all values are true within each key group.
+allTrueByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array (treated as boolean)
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+allTrueByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_all_true_by_key ko vo k v dim)
+
+-- | True if any value is true within each key group.
+anyTrueByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array (treated as boolean)
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+anyTrueByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_any_true_by_key ko vo k v dim)
+
+-- | Count non-zero values within each key group.
+countByKey
+  :: AFType a
+  => Array Int
+  -- ^ Keys array
+  -> Array a
+  -- ^ Values array
+  -> Int
+  -- ^ Dimension
+  -> (Array Int, Array a)
+countByKey keys vals (fromIntegral -> dim) =
+  op2p2kv keys vals (\ko vo k v -> af_count_by_key ko vo k v dim)