basis/math/vectors/simd/simd.factor

   1 USING: accessors alien.c-types arrays byte-arrays classes combinators
   2 cpu.architecture effects fry functors generalizations generic
   3 generic.parser kernel lexer literals macros math math.functions
   4 math.vectors math.vectors.private math.vectors.simd.intrinsics namespaces parser
   5 prettyprint.custom quotations sequences sequences.private vocabs
   6 vocabs.loader words ;
   7 QUALIFIED-WITH: alien.c-types c
   8 IN: math.vectors.simd
   9
  10 ERROR: bad-simd-length got expected ;
  11
  12 <<
  13 <PRIVATE
  14 ! Primitive SIMD constructors
  15
  16 GENERIC: new-underlying ( underlying seq -- seq' )
  17
  18 : make-underlying ( seq quot -- seq' )
  19     dip new-underlying ; inline
  20 : change-underlying ( seq quot -- seq' )
  21     '[ underlying>> @ ] keep new-underlying ; inline
  22 PRIVATE>
  23 >>
  24
  25 <PRIVATE
  26
  27 ! Helper for boolean vector literals
  28
  29 : vector-true-value ( class -- value )
  30     { c:float c:double } member? [ -1 bits>double ] [ -1 ] if ; foldable
  31
  32 : vector-false-value ( type -- value )
  33     { c:float c:double } member? [ 0.0 ] [ 0 ] if ; foldable
  34
  35 : boolean>element ( bool/elt type -- elt )
  36     swap {
  37         { t [ vector-true-value  ] }
  38         { f [ vector-false-value ] }
  39         [ nip ]
  40     } case ; inline
  41
  42 PRIVATE>
  43
  44 ! SIMD base type
  45
  46 TUPLE: simd-128
  47     { underlying byte-array read-only initial: $[ 16 <byte-array> ] } ;
  48
  49 GENERIC: simd-element-type ( obj -- c-type )
  50 GENERIC: simd-rep ( simd -- rep )
  51
  52 <<
  53 <PRIVATE
  54
  55 DEFER: simd-construct-op
  56
  57 ! Unboxers for SIMD operations
  58 : if-both-vectors ( a b rep t f -- )
  59     [ 2over [ simd-128? ] both? ] 2dip if ; inline
  60
  61 : if-both-vectors-match ( a b rep t f -- )
  62     [ 3dup [ drop [ simd-128? ] both? ] [ '[ simd-rep _ eq? ] both? ] 3bi and ]
  63     2dip if ; inline
  64
  65 : simd-construct-op ( exemplar quot: ( rep -- v ) -- v )
  66     [ dup simd-rep ] dip curry make-underlying ; inline
  67
  68 : simd-unbox ( a -- a (a) )
  69     [ ] [ underlying>> ] bi ; inline
  70
  71 : v->v-op ( a rep quot: ( (a) rep -- (c) ) fallback-quot -- c )
  72     drop [ simd-unbox ] 2dip 2curry make-underlying ; inline
  73
  74 : vn->v-op ( a n rep quot: ( (a) n rep -- (c) ) fallback-quot -- c )
  75     drop [ simd-unbox ] 3dip 3curry make-underlying ; inline
  76
  77 : v->n-op ( a rep quot: ( (a) rep -- n ) fallback-quot -- n )
  78     drop [ underlying>> ] 2dip call ; inline
  79
  80 : (vv->v-op) ( a b rep quot: ( (a) (b) rep -- (c) ) -- c )
  81     [ [ simd-unbox ] [ underlying>> ] bi* ] 2dip 3curry make-underlying ; inline
  82
  83 : (vv->n-op) ( a b rep quot: ( (a) (b) rep -- n ) -- n )
  84     [ [ underlying>> ] bi@ ] 2dip 3curry call ; inline
  85
  86 : vv->v-op ( a b rep quot: ( (a) (b) rep -- (c) ) fallback-quot -- c )
  87     [ '[ _ (vv->v-op) ] ] [ '[ drop @ ] ] bi* if-both-vectors-match ; inline
  88
  89 : vv'->v-op ( a b rep quot: ( (a) (b) rep -- (c) ) fallback-quot -- c )
  90     [ '[ _ (vv->v-op) ] ] [ '[ drop @ ] ] bi* if-both-vectors ; inline
  91
  92 : vv->n-op ( a b rep quot: ( (a) (b) rep -- n ) fallback-quot -- n )
  93     [ '[ _ (vv->n-op) ] ] [ '[ drop @ ] ] bi* if-both-vectors-match ; inline
  94
  95 PRIVATE>
  96 >>
  97
  98 DEFER: simd-with
  99 DEFER: simd-cast
 100
 101 <<
 102 <PRIVATE
 103
 104 ! SIMD concrete type functor
 105
 106 FUNCTOR: define-simd-128 ( T -- )
 107
 108 A      DEFINES-CLASS ${T}
 109 A-rep  IS            ${T}-rep
 110 >A     DEFINES       >${T}
 111 A-boa  DEFINES       ${T}-boa
 112 A-with DEFINES       ${T}-with
 113 A-cast DEFINES       ${T}-cast
 114 A{     DEFINES       ${T}{
 115
 116 ELT   [ A-rep rep-component-type ]
 117 N     [ A-rep rep-length ]
 118
 119 SET-NTH [ ELT dup c:c-setter c:array-accessor ]
 120
 121 BOA-EFFECT [ N "n" <repetition> >array { "v" } <effect> ]
 122
 123 WHERE
 124
 125 TUPLE: A < simd-128 ;
 126
 127 M: A new-underlying    drop \ A boa ; inline
 128 M: A simd-rep          drop A-rep ; inline
 129 M: A simd-element-type drop ELT ; inline
 130
 131 M: A set-nth-unsafe
 132     [ ELT boolean>element ] 2dip
 133     underlying>> SET-NTH call ; inline
 134
 135 : >A ( seq -- simd ) \ A new clone-like ; inline
 136
 137 M: A like drop dup \ A instance? [ >A ] unless ; inline
 138
 139 : A-with ( n -- v ) \ A new simd-with ; inline
 140 : A-cast ( v -- v' ) \ A new simd-cast ; inline
 141
 142 ! SIMD vectors as sequences
 143
 144 M: A hashcode* underlying>> hashcode* ; inline
 145 M: A clone [ clone ] change-underlying ; inline
 146 M: A length drop N ; inline
 147 M: A nth-unsafe
 148     swap {
 149         {  0 [  0 \ A-rep (simd-select) ] }
 150         {  1 [  1 \ A-rep (simd-select) ] }
 151         {  2 [  2 \ A-rep (simd-select) ] }
 152         {  3 [  3 \ A-rep (simd-select) ] }
 153         {  4 [  4 \ A-rep (simd-select) ] }
 154         {  5 [  5 \ A-rep (simd-select) ] }
 155         {  6 [  6 \ A-rep (simd-select) ] }
 156         {  7 [  7 \ A-rep (simd-select) ] }
 157         {  8 [  8 \ A-rep (simd-select) ] }
 158         {  9 [  9 \ A-rep (simd-select) ] }
 159         { 10 [ 10 \ A-rep (simd-select) ] }
 160         { 11 [ 11 \ A-rep (simd-select) ] }
 161         { 12 [ 12 \ A-rep (simd-select) ] }
 162         { 13 [ 13 \ A-rep (simd-select) ] }
 163         { 14 [ 14 \ A-rep (simd-select) ] }
 164         { 15 [ 15 \ A-rep (simd-select) ] }
 165     } case ; inline
 166 M: A c:byte-length drop 16 ; inline
 167
 168 M: A new-sequence
 169     2dup length =
 170     [ nip [ 16 (byte-array) ] make-underlying ]
 171     [ length bad-simd-length ] if ; inline
 172
 173 M: A equal?
 174     \ A [ drop v= vall? ] [ 3drop f ] if-both-vectors-match ; inline
 175
 176 ! SIMD primitive operations
 177
 178 M: A v+                \ A-rep [ (simd-v+)                ] [ call-next-method ] vv->v-op ; inline
 179 M: A v-                \ A-rep [ (simd-v-)                ] [ call-next-method ] vv->v-op ; inline
 180 M: A vneg              \ A-rep [ (simd-vneg)              ] [ call-next-method ] v->v-op  ; inline
 181 M: A v+-               \ A-rep [ (simd-v+-)               ] [ call-next-method ] vv->v-op ; inline
 182 M: A vs+               \ A-rep [ (simd-vs+)               ] [ call-next-method ] vv->v-op ; inline
 183 M: A vs-               \ A-rep [ (simd-vs-)               ] [ call-next-method ] vv->v-op ; inline
 184 M: A vs*               \ A-rep [ (simd-vs*)               ] [ call-next-method ] vv->v-op ; inline
 185 M: A v*                \ A-rep [ (simd-v*)                ] [ call-next-method ] vv->v-op ; inline
 186 M: A v/                \ A-rep [ (simd-v/)                ] [ call-next-method ] vv->v-op ; inline
 187 M: A vmin              \ A-rep [ (simd-vmin)              ] [ call-next-method ] vv->v-op ; inline
 188 M: A vmax              \ A-rep [ (simd-vmax)              ] [ call-next-method ] vv->v-op ; inline
 189 M: A v.                \ A-rep [ (simd-v.)                ] [ call-next-method ] vv->n-op ; inline
 190 M: A vsqrt             \ A-rep [ (simd-vsqrt)             ] [ call-next-method ] v->v-op  ; inline
 191 M: A sum               \ A-rep [ (simd-sum)               ] [ call-next-method ] v->n-op  ; inline
 192 M: A vabs              \ A-rep [ (simd-vabs)              ] [ call-next-method ] v->v-op  ; inline
 193 M: A vbitand           \ A-rep [ (simd-vbitand)           ] [ call-next-method ] vv->v-op ; inline
 194 M: A vbitandn          \ A-rep [ (simd-vbitandn)          ] [ call-next-method ] vv->v-op ; inline
 195 M: A vbitor            \ A-rep [ (simd-vbitor)            ] [ call-next-method ] vv->v-op ; inline
 196 M: A vbitxor           \ A-rep [ (simd-vbitxor)           ] [ call-next-method ] vv->v-op ; inline
 197 M: A vbitnot           \ A-rep [ (simd-vbitnot)           ] [ call-next-method ] v->v-op  ; inline
 198 M: A vand              \ A-rep [ (simd-vand)              ] [ call-next-method ] vv->v-op ; inline
 199 M: A vandn             \ A-rep [ (simd-vandn)             ] [ call-next-method ] vv->v-op ; inline
 200 M: A vor               \ A-rep [ (simd-vor)               ] [ call-next-method ] vv->v-op ; inline
 201 M: A vxor              \ A-rep [ (simd-vxor)              ] [ call-next-method ] vv->v-op ; inline
 202 M: A vnot              \ A-rep [ (simd-vnot)              ] [ call-next-method ] v->v-op  ; inline
 203 M: A vlshift           \ A-rep [ (simd-vlshift)           ] [ call-next-method ] vn->v-op ; inline
 204 M: A vrshift           \ A-rep [ (simd-vrshift)           ] [ call-next-method ] vn->v-op ; inline
 205 M: A hlshift           \ A-rep [ (simd-hlshift)           ] [ call-next-method ] vn->v-op ; inline
 206 M: A hrshift           \ A-rep [ (simd-hrshift)           ] [ call-next-method ] vn->v-op ; inline
 207 M: A vshuffle-elements \ A-rep [ (simd-vshuffle-elements) ] [ call-next-method ] vn->v-op ; inline
 208 M: A vshuffle-bytes    \ A-rep [ (simd-vshuffle-bytes)    ] [ call-next-method ] vv->v-op ; inline
 209 M: A (vmerge-head)     \ A-rep [ (simd-vmerge-head)       ] [ call-next-method ] vv->v-op ; inline
 210 M: A (vmerge-tail)     \ A-rep [ (simd-vmerge-tail)       ] [ call-next-method ] vv->v-op ; inline
 211 M: A v<=               \ A-rep [ (simd-v<=)               ] [ call-next-method ] vv->v-op ; inline
 212 M: A v<                \ A-rep [ (simd-v<)                ] [ call-next-method ] vv->v-op ; inline
 213 M: A v=                \ A-rep [ (simd-v=)                ] [ call-next-method ] vv->v-op ; inline
 214 M: A v>                \ A-rep [ (simd-v>)                ] [ call-next-method ] vv->v-op ; inline
 215 M: A v>=               \ A-rep [ (simd-v>=)               ] [ call-next-method ] vv->v-op ; inline
 216 M: A vunordered?       \ A-rep [ (simd-vunordered?)       ] [ call-next-method ] vv->v-op ; inline
 217 M: A vany?             \ A-rep [ (simd-vany?)             ] [ call-next-method ] v->n-op  ; inline
 218 M: A vall?             \ A-rep [ (simd-vall?)             ] [ call-next-method ] v->n-op  ; inline
 219 M: A vnone?            \ A-rep [ (simd-vnone?)            ] [ call-next-method ] v->n-op  ; inline
 220
 221 ! SIMD high-level specializations
 222
 223 M: A vbroadcast [ swap nth ] keep simd-with ; inline
 224 M: A n+v [ simd-with ] keep v+ ; inline
 225 M: A n-v [ simd-with ] keep v- ; inline
 226 M: A n*v [ simd-with ] keep v* ; inline
 227 M: A n/v [ simd-with ] keep v/ ; inline
 228 M: A v+n over simd-with v+ ; inline
 229 M: A v-n over simd-with v- ; inline
 230 M: A v*n over simd-with v* ; inline
 231 M: A v/n over simd-with v/ ; inline
 232 M: A norm-sq dup v. assert-positive ; inline
 233 M: A norm      norm-sq sqrt ; inline
 234 M: A distance  v- norm ; inline
 235
 236 ! M: simd-128 >pprint-sequence ;
 237 ! M: simd-128 pprint* pprint-object ;
 238
 239 \ A-boa \ A new N {
 240     { 2 [ '[ _ [ (simd-gather-2) ] simd-construct-op ] ] }
 241     { 4 [ '[ _ [ (simd-gather-4) ] simd-construct-op ] ] }
 242     [ swap '[ _ _ nsequence ] ]
 243 } case BOA-EFFECT define-inline
 244
 245 M: A pprint-delims drop \ A{ \ } ;
 246 SYNTAX: A{ \ } [ >A ] parse-literal ;
 247
 248 c:<c-type>
 249     byte-array >>class
 250     A >>boxed-class
 251     [ A-rep alien-vector \ A boa ] >>getter
 252     [ [ underlying>> ] 2dip A-rep set-alien-vector ] >>setter
 253     16 >>size
 254     16 >>align
 255     A-rep >>rep
 256 \ A c:typedef
 257
 258 ;FUNCTOR
 259
 260 SYNTAX: SIMD-128:
 261     scan define-simd-128 ;
 262
 263 PRIVATE>
 264
 265 >>
 266
 267 INSTANCE: simd-128 sequence
 268
 269 ! SIMD constructors
 270
 271 : simd-with ( n seq -- v )
 272     [ (simd-with) ] simd-construct-op ; inline
 273
 274 MACRO: simd-boa ( class -- )
 275     new dup length {
 276         { 2 [ '[ _ [ (simd-gather-2) ] simd-construct-op ] ] }
 277         { 4 [ '[ _ [ (simd-gather-4) ] simd-construct-op ] ] }
 278         [ swap '[ _ _ nsequence ] ]
 279     } case ;
 280
 281 : simd-cast ( v seq -- v' )
 282     [ underlying>> ] dip new-underlying ; inline
 283
 284 ! SIMD instances
 285
 286 SIMD-128: char-16
 287 SIMD-128: uchar-16
 288 SIMD-128: short-8
 289 SIMD-128: ushort-8
 290 SIMD-128: int-4
 291 SIMD-128: uint-4
 292 SIMD-128: longlong-2
 293 SIMD-128: ulonglong-2
 294 SIMD-128: float-4
 295 SIMD-128: double-2
 296
 297 ! misc
 298
 299 M: simd-128 vshuffle ( u perm -- v )
 300     vshuffle-bytes ; inline
 301
 302 "mirrors" vocab [
 303     "math.vectors.simd.mirrors" require
 304 ] when