add reorderVector function for tensor libraries (e.g. hTensor) to implement tensor transpose

3 files changed, 86 insertions, 1 deletions

diff --git a/packages/base/src/Internal/C/vector-aux.c b/packages/base/src/Internal/C/vector-aux.c
index 1cef27d..dcd6c0b 100644
--- a/packages/base/src/Internal/C/vector-aux.c
+++ b/packages/base/src/Internal/C/vector-aux.c
@@ -1533,3 +1533,54 @@ int chooseQ(KIVEC(cond),KQVEC(lt),KQVEC(eq),KQVEC(gt),QVEC(r)) {
     CHOOSE_IMP
 }
 
+//////////////////// reorder /////////////////////////
+
+#define REORDER_IMP                                                                     \
+    REQUIRES(kn == stridesn && stridesn == dimsn ,BAD_SIZE);                            \
+    int i,j,l;                                                                          \
+    for (i=1,j=0,l=0;l<kn;++l) {                                                        \
+        kp[l] = 0;                                                                      \
+        i *= dimsp[l];                                                                  \
+        j += (dimsp[l]-1) * stridesp[l];                                                \
+    }                                                                                   \
+    REQUIRES(i <= vn && j < rn ,BAD_SIZE);                                              \
+    for (i=0,j=0;;i++) {                                                                \
+        rp[i] = vp[j];                                                                  \
+        for(l=kn-1;;l--) {                                                              \
+            ++kp[l];                                                                    \
+            if (kp[l] < dimsp[l]) {                                                     \
+                j += stridesp[l];                                                       \
+                break;                                                                  \
+            } else {                                                                    \
+                if (l == 0) {                                                           \
+                    return 0;                                                           \
+                }                                                                       \
+                kp[l] = 0;                                                              \
+                j -= (dimsp[l]-1) * stridesp[l];                                        \
+            }                                                                           \
+        }                                                                               \
+    }
+
+int reorderF(IVEC(k), KIVEC(strides),KIVEC(dims),KFVEC(v),FVEC(r)) {
+    REORDER_IMP
+}
+
+int reorderD(IVEC(k), KIVEC(strides),KIVEC(dims),KDVEC(v),DVEC(r)) {
+    REORDER_IMP
+}
+
+int reorderI(IVEC(k), KIVEC(strides),KIVEC(dims),KIVEC(v),IVEC(r)) {
+    REORDER_IMP
+}
+
+int reorderL(IVEC(k), KIVEC(strides),KIVEC(dims),KLVEC(v),LVEC(r)) {
+    REORDER_IMP
+}
+
+int reorderC(IVEC(k), KIVEC(strides),KIVEC(dims),KCVEC(v),CVEC(r)) {
+    REORDER_IMP
+}
+
+int reorderQ(IVEC(k), KIVEC(strides),KIVEC(dims),KQVEC(v),QVEC(r)) {
+    REORDER_IMP
+}
diff --git a/packages/base/src/Internal/Matrix.hs b/packages/base/src/Internal/Matrix.hs
index 0135288..fbf8295 100644
--- a/packages/base/src/Internal/Matrix.hs
+++ b/packages/base/src/Internal/Matrix.hs
@@ -285,6 +285,7 @@ class (Storable a) => Element a where
     remapM   :: Matrix CInt -> Matrix CInt -> Matrix a -> Matrix a
     rowOp    :: Int -> a -> Int -> Int -> Int -> Int -> Matrix a -> IO ()
     gemm     :: Vector a -> Matrix a -> Matrix a -> Matrix a -> IO ()
+    reorderV :: Vector CInt-> Vector CInt-> Vector a -> Vector a -- see reorderVector for documentation
 
 
 instance Element Float where
@@ -298,6 +299,7 @@ instance Element Float where
     remapM     = remapF
     rowOp      = rowOpAux c_rowOpF
     gemm       = gemmg c_gemmF
+    reorderV   = reorderAux c_reorderF
 
 instance Element Double where
     constantD  = constantAux cconstantR
@@ -310,6 +312,7 @@ instance Element Double where
     remapM     = remapD
     rowOp      = rowOpAux c_rowOpD
     gemm       = gemmg c_gemmD
+    reorderV   = reorderAux c_reorderD
 
 instance Element (Complex Float) where
     constantD  = constantAux cconstantQ
@@ -322,6 +325,7 @@ instance Element (Complex Float) where
     remapM     = remapQ
     rowOp      = rowOpAux c_rowOpQ
     gemm       = gemmg c_gemmQ
+    reorderV   = reorderAux c_reorderQ
 
 instance Element (Complex Double) where
     constantD  = constantAux cconstantC
@@ -334,6 +338,7 @@ instance Element (Complex Double) where
     remapM     = remapC
     rowOp      = rowOpAux c_rowOpC
     gemm       = gemmg c_gemmC
+    reorderV   = reorderAux c_reorderC
 
 instance Element (CInt) where
     constantD  = constantAux cconstantI
@@ -346,6 +351,7 @@ instance Element (CInt) where
     remapM     = remapI
     rowOp      = rowOpAux c_rowOpI
     gemm       = gemmg c_gemmI
+    reorderV   = reorderAux c_reorderI
 
 instance Element Z where
     constantD  = constantAux cconstantL
@@ -358,6 +364,7 @@ instance Element Z where
     remapM     = remapL
     rowOp      = rowOpAux c_rowOpL
     gemm       = gemmg c_gemmL
+    reorderV   = reorderAux c_reorderL
 
 -------------------------------------------------------------------
 
@@ -580,6 +587,33 @@ foreign import ccall unsafe "gemm_mod_int64_t" c_gemmML :: Z -> Tgemm Z
 
 --------------------------------------------------------------------------------
 
+reorderAux f s d v = unsafePerformIO $ do
+    k <- createVector (dim s)
+    r <- createVector (dim v)
+    (k # s # d # v #! r) f #| "reorderV"
+    return r
+
+type Reorder x = CV CInt (CV CInt (CV CInt (CV x (CV x (IO CInt)))))
+
+foreign import ccall unsafe "reorderD" c_reorderD :: Reorder Double
+foreign import ccall unsafe "reorderF" c_reorderF :: Reorder Float
+foreign import ccall unsafe "reorderI" c_reorderI :: Reorder CInt
+foreign import ccall unsafe "reorderC" c_reorderC :: Reorder (Complex Double)
+foreign import ccall unsafe "reorderQ" c_reorderQ :: Reorder (Complex Float)
+foreign import ccall unsafe "reorderL" c_reorderL :: Reorder Z
+
+-- | Transpose an array with dimensions @dims@ by making a copy using @strides@. For example, for an array with 3 indices,
+--   @(reorderVector strides dims v) ! ((i * dims ! 1 + j) * dims ! 2 + k) == v ! (i * strides ! 0 + j * strides ! 1 + k * strides ! 2)@
+--   This function is intended to be used internally by tensor libraries..
+reorderVector :: Element a
+                    => Vector CInt -- ^ @strides@: array strides
+                    -> Vector CInt -- ^ @dims@: array dimensions of new array @v@
+                    -> Vector a    -- ^ @v@: flattened input array
+                    -> Vector a    -- ^ @v'@: flattened output array
+reorderVector = reorderV
+
+--------------------------------------------------------------------------------
+
 foreign import ccall unsafe "saveMatrix" c_saveMatrix
     :: CString -> CString -> Double ::> Ok
 
diff --git a/packages/base/src/Numeric/LinearAlgebra/Devel.hs b/packages/base/src/Numeric/LinearAlgebra/Devel.hs
index 941b597..e974fc4 100644
--- a/packages/base/src/Numeric/LinearAlgebra/Devel.hs
+++ b/packages/base/src/Numeric/LinearAlgebra/Devel.hs
@@ -55,7 +55,7 @@ module Numeric.LinearAlgebra.Devel(
     GMatrix(..),
 
     -- * Misc
-    toByteString, fromByteString, showInternal
+    toByteString, fromByteString, showInternal, reorderVector
 
 ) where