extra/math/blas/matrices/matrices.factor

   1 USING: accessors alien alien.c-types arrays byte-arrays combinators
   2 combinators.short-circuit fry kernel locals macros
   3 math math.blas.cblas math.blas.vectors math.blas.vectors.private
   4 math.complex math.functions math.order functors words
   5 sequences sequences.merged sequences.private shuffle
   6 specialized-arrays.direct.float specialized-arrays.direct.double
   7 specialized-arrays.float specialized-arrays.double ;
   8 IN: math.blas.matrices
   9
  10 TUPLE: blas-matrix-base underlying ld rows cols transpose ;
  11
  12 : Mtransposed? ( matrix -- ? )
  13     transpose>> ; inline
  14 : Mwidth ( matrix -- width )
  15     dup Mtransposed? [ rows>> ] [ cols>> ] if ; inline
  16 : Mheight ( matrix -- height )
  17     dup Mtransposed? [ cols>> ] [ rows>> ] if ; inline
  18
  19 GENERIC: n*M.V+n*V! ( alpha A x beta y -- y=alpha*A.x+b*y )
  20 GENERIC: n*V(*)V+M! ( alpha x y A -- A=alpha*x(*)y+A )
  21 GENERIC: n*V(*)Vconj+M! ( alpha x y A -- A=alpha*x(*)yconj+A )
  22 GENERIC: n*M.M+n*M! ( alpha A B beta C -- C=alpha*A.B+beta*C )
  23
  24 <PRIVATE
  25
  26 : (blas-transpose) ( matrix -- integer )
  27     transpose>> [ CblasTrans ] [ CblasNoTrans ] if ;
  28
  29 GENERIC: (blas-matrix-like) ( data ld rows cols transpose exemplar -- matrix )
  30
  31 : (validate-gemv) ( A x y -- )
  32     {
  33         [ drop [ Mwidth  ] [ length>> ] bi* = ]
  34         [ nip  [ Mheight ] [ length>> ] bi* = ]
  35     } 3&&
  36     [ "Mismatched matrix and vectors in matrix-vector multiplication" throw ]
  37     unless ;
  38
  39 :: (prepare-gemv)
  40     ( alpha A x beta y >c-arg -- order A-trans m n alpha A-data A-ld x-data x-inc beta y-data y-inc
  41                                  y )
  42     A x y (validate-gemv)
  43     CblasColMajor
  44     A (blas-transpose)
  45     A rows>>
  46     A cols>>
  47     alpha >c-arg call
  48     A underlying>>
  49     A ld>>
  50     x underlying>>
  51     x inc>>
  52     beta >c-arg call
  53     y underlying>>
  54     y inc>>
  55     y ; inline
  56
  57 : (validate-ger) ( x y A -- )
  58     {
  59         [ nip  [ length>> ] [ Mheight ] bi* = ]
  60         [ nipd [ length>> ] [ Mwidth  ] bi* = ]
  61     } 3&&
  62     [ "Mismatched vertices and matrix in vector outer product" throw ]
  63     unless ;
  64
  65 :: (prepare-ger)
  66     ( alpha x y A >c-arg -- order m n alpha x-data x-inc y-data y-inc A-data A-ld
  67                             A )
  68     x y A (validate-ger)
  69     CblasColMajor
  70     A rows>>
  71     A cols>>
  72     alpha >c-arg call
  73     x underlying>>
  74     x inc>>
  75     y underlying>>
  76     y inc>>
  77     A underlying>>
  78     A ld>>
  79     A f >>transpose ; inline
  80
  81 : (validate-gemm) ( A B C -- )
  82     {
  83         [ drop [ Mwidth  ] [ Mheight ] bi* = ]
  84         [ nip  [ Mheight ] bi@ = ]
  85         [ nipd [ Mwidth  ] bi@ = ]
  86     } 3&&
  87     [ "Mismatched matrices in matrix multiplication" throw ]
  88     unless ;
  89
  90 :: (prepare-gemm)
  91     ( alpha A B beta C >c-arg -- order A-trans B-trans m n k alpha A-data A-ld B-data B-ld beta C-data C-ld
  92                                  C )
  93     A B C (validate-gemm)
  94     CblasColMajor
  95     A (blas-transpose)
  96     B (blas-transpose)
  97     C rows>>
  98     C cols>>
  99     A Mwidth
 100     alpha >c-arg call
 101     A underlying>>
 102     A ld>>
 103     B underlying>>
 104     B ld>>
 105     beta >c-arg call
 106     C underlying>>
 107     C ld>>
 108     C f >>transpose ; inline
 109
 110 : (>matrix) ( arrays >c-array -- c-array ld rows cols transpose )
 111     '[ <merged> @ ] [ length dup ] [ first length ] tri f ; inline
 112
 113 PRIVATE>
 114
 115 ! XXX should do a dense clone
 116 M: blas-matrix-base clone
 117     [
 118         [ {
 119             [ underlying>> ]
 120             [ ld>> ]
 121             [ cols>> ]
 122             [ element-type heap-size ]
 123         } cleave * * memory>byte-array ]
 124         [ {
 125             [ ld>> ]
 126             [ rows>> ]
 127             [ cols>> ]
 128             [ transpose>> ]
 129         } cleave ]
 130         bi
 131     ] keep (blas-matrix-like) ;
 132
 133 ! XXX try rounding stride to next 128 bit bound for better vectorizin'
 134 : <empty-matrix> ( rows cols exemplar -- matrix )
 135     [ element-type [ * ] dip <c-array> ]
 136     [ 2drop ]
 137     [ f swap (blas-matrix-like) ] 3tri ;
 138
 139 : n*M.V+n*V ( alpha A x beta y -- alpha*A.x+b*y )
 140     clone n*M.V+n*V! ;
 141 : n*V(*)V+M ( alpha x y A -- alpha*x(*)y+A )
 142     clone n*V(*)V+M! ;
 143 : n*V(*)Vconj+M ( alpha x y A -- alpha*x(*)yconj+A )
 144     clone n*V(*)Vconj+M! ;
 145 : n*M.M+n*M ( alpha A B beta C -- alpha*A.B+beta*C )
 146     clone n*M.M+n*M! ;
 147
 148 : n*M.V ( alpha A x -- alpha*A.x )
 149     1.0 2over [ Mheight ] dip <empty-vector>
 150     n*M.V+n*V! ; inline
 151
 152 : M.V ( A x -- A.x )
 153     1.0 -rot n*M.V ; inline
 154
 155 : n*V(*)V ( alpha x y -- alpha*x(*)y )
 156     2dup [ length>> ] bi@ pick <empty-matrix>
 157     n*V(*)V+M! ;
 158 : n*V(*)Vconj ( alpha x y -- alpha*x(*)yconj )
 159     2dup [ length>> ] bi@ pick <empty-matrix>
 160     n*V(*)Vconj+M! ;
 161
 162 : V(*) ( x y -- x(*)y )
 163     1.0 -rot n*V(*)V ; inline
 164 : V(*)conj ( x y -- x(*)yconj )
 165     1.0 -rot n*V(*)Vconj ; inline
 166
 167 : n*M.M ( alpha A B -- alpha*A.B )
 168     2dup [ Mheight ] [ Mwidth ] bi* pick <empty-matrix>
 169     1.0 swap n*M.M+n*M! ;
 170
 171 : M. ( A B -- A.B )
 172     1.0 -rot n*M.M ; inline
 173
 174 :: (Msub) ( matrix row col height width -- data ld rows cols )
 175     matrix ld>> col * row + matrix element-type heap-size *
 176     matrix underlying>> <displaced-alien>
 177     matrix ld>>
 178     height
 179     width ;
 180
 181 :: Msub ( matrix row col height width -- sub )
 182     matrix dup transpose>>
 183     [ col row width height ]
 184     [ row col height width ] if (Msub)
 185     matrix transpose>> matrix (blas-matrix-like) ;
 186
 187 TUPLE: blas-matrix-rowcol-sequence
 188     parent inc rowcol-length rowcol-jump length ;
 189 C: <blas-matrix-rowcol-sequence> blas-matrix-rowcol-sequence
 190
 191 INSTANCE: blas-matrix-rowcol-sequence sequence
 192
 193 M: blas-matrix-rowcol-sequence length
 194     length>> ;
 195 M: blas-matrix-rowcol-sequence nth-unsafe
 196     {
 197         [
 198             [ rowcol-jump>> ]
 199             [ parent>> element-type heap-size ]
 200             [ parent>> underlying>> ] tri
 201             [ * * ] dip <displaced-alien>
 202         ]
 203         [ rowcol-length>> ]
 204         [ inc>> ]
 205         [ parent>> ]
 206     } cleave (blas-vector-like) ;
 207
 208 : (Mcols) ( A -- columns )
 209     { [ ] [ drop 1 ] [ rows>> ] [ ld>> ] [ cols>> ] }
 210     cleave <blas-matrix-rowcol-sequence> ;
 211 : (Mrows) ( A -- rows )
 212     { [ ] [ ld>> ] [ cols>> ] [ drop 1 ] [ rows>> ] }
 213     cleave <blas-matrix-rowcol-sequence> ;
 214
 215 : Mrows ( A -- rows )
 216     dup transpose>> [ (Mcols) ] [ (Mrows) ] if ;
 217 : Mcols ( A -- cols )
 218     dup transpose>> [ (Mrows) ] [ (Mcols) ] if ;
 219
 220 : n*M! ( n A -- A=n*A )
 221     [ (Mcols) [ n*V! drop ] with each ] keep ;
 222
 223 : n*M ( n A -- n*A )
 224     clone n*M! ; inline
 225
 226 : M*n ( A n -- A*n )
 227     swap n*M ; inline
 228 : M/n ( A n -- A/n )
 229     recip swap n*M ; inline
 230
 231 : Mtranspose ( matrix -- matrix^T )
 232     [ {
 233         [ underlying>> ]
 234         [ ld>> ] [ rows>> ]
 235         [ cols>> ]
 236         [ transpose>> not ]
 237     } cleave ] keep (blas-matrix-like) ;
 238
 239 M: blas-matrix-base equal?
 240     {
 241         [ [ Mwidth ] bi@ = ]
 242         [ [ Mcols ] bi@ [ = ] 2all? ]
 243     } 2&& ;
 244
 245 <<
 246
 247 FUNCTOR: (define-blas-matrix) ( TYPE T U C -- )
 248
 249 VECTOR      IS ${TYPE}-blas-vector
 250 <VECTOR>    IS <${TYPE}-blas-vector>
 251 >ARRAY      IS >${TYPE}-array
 252 TYPE>ARG    IS ${TYPE}>arg
 253 XGEMV       IS cblas_${T}gemv
 254 XGEMM       IS cblas_${T}gemm
 255 XGERU       IS cblas_${T}ger${U}
 256 XGERC       IS cblas_${T}ger${C}
 257
 258 MATRIX      DEFINES ${TYPE}-blas-matrix
 259 <MATRIX>    DEFINES <${TYPE}-blas-matrix>
 260 >MATRIX     DEFINES >${TYPE}-blas-matrix
 261
 262 WHERE
 263
 264 TUPLE: MATRIX < blas-matrix-base ;
 265 : <MATRIX> ( underlying ld rows cols transpose -- matrix )
 266     MATRIX boa ; inline
 267
 268 M: MATRIX element-type
 269     drop TYPE ;
 270 M: MATRIX (blas-matrix-like)
 271     drop <MATRIX> execute ;
 272 M: VECTOR (blas-matrix-like)
 273     drop <MATRIX> execute ;
 274 M: MATRIX (blas-vector-like)
 275     drop <VECTOR> execute ;
 276
 277 : >MATRIX ( arrays -- matrix )
 278     [ >ARRAY execute underlying>> ] (>matrix)
 279     <MATRIX> execute ;
 280
 281 M: VECTOR n*M.V+n*V!
 282     [ TYPE>ARG execute ] (prepare-gemv)
 283     [ XGEMV execute ] dip ;
 284 M: MATRIX n*M.M+n*M!
 285     [ TYPE>ARG execute ] (prepare-gemm)
 286     [ XGEMM execute ] dip ;
 287 M: MATRIX n*V(*)V+M!
 288     [ TYPE>ARG execute ] (prepare-ger)
 289     [ XGERU execute ] dip ;
 290 M: MATRIX n*V(*)Vconj+M!
 291     [ TYPE>ARG execute ] (prepare-ger)
 292     [ XGERC execute ] dip ;
 293
 294 ;FUNCTOR
 295
 296
 297 : define-real-blas-matrix ( TYPE T -- )
 298     "" "" (define-blas-matrix) ;
 299 : define-complex-blas-matrix ( TYPE T -- )
 300     "u" "c" (define-blas-matrix) ;
 301
 302 "float"          "s" define-real-blas-matrix
 303 "double"         "d" define-real-blas-matrix
 304 "float-complex"  "c" define-complex-blas-matrix
 305 "double-complex" "z" define-complex-blas-matrix
 306
 307 >>