extra/cuda/libraries/libraries.factor

   1 ! Copyright (C) 2010 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors alien.data alien.parser arrays assocs
   4 byte-arrays classes.struct combinators combinators.short-circuit
   5 cuda cuda.ffi fry generalizations io.backend kernel macros math
   6 namespaces sequences variants words ;
   7 FROM: classes.struct.private => compute-struct-offsets write-struct-slot ;
   8 QUALIFIED-WITH: alien.c-types c
   9 IN: cuda.libraries
  10
  11 VARIANT: cuda-abi
  12     cuda32 cuda64 ;
  13
  14 SYMBOL: cuda-modules
  15 SYMBOL: cuda-functions
  16
  17 SYMBOL: cuda-libraries
  18 cuda-libraries [ H{ } clone ] initialize
  19
  20 SYMBOL: current-cuda-library
  21
  22 : ?delete-at ( key assoc -- old/key ? )
  23     2dup delete-at* [ 2nip t ] [ 2drop f ] if ; inline
  24
  25 : cuda-param-size ( function n -- )
  26     cuParamSetSize cuda-error ; inline
  27
  28 : cuda-vector ( function offset ptr n -- )
  29     cuParamSetv cuda-error ; inline
  30
  31 : launch-function-grid ( function width height -- )
  32     cuLaunchGrid cuda-error ; inline
  33
  34 : function-block-shape ( function x y z -- )
  35     cuFuncSetBlockShape cuda-error ; inline
  36
  37 : function-shared-size ( function n -- )
  38     cuFuncSetSharedSize cuda-error ; inline
  39
  40 TUPLE: grid
  41     { dim-grid read-only }
  42     { dim-block read-only }
  43     { shared-size read-only initial: 0 }
  44     { stream read-only } ;
  45
  46 : <grid> ( dim-grid dim-block -- grid )
  47     0 f grid boa ; inline
  48
  49 : <grid-shared> ( dim-grid dim-block shared-size -- grid )
  50     f grid boa ; inline
  51
  52 : <grid-shared-stream> ( dim-grid dim-block shared-size stream -- grid )
  53     grid boa ; inline
  54
  55 <PRIVATE
  56 GENERIC: block-dim ( block-size -- x y z ) foldable
  57 M: integer block-dim 1 1 ; inline
  58 M: sequence block-dim
  59     dup length {
  60         { 0 [ drop 1 1 1 ] }
  61         { 1 [ first 1 1 ] }
  62         { 2 [ first2 1 ] }
  63         [ drop first3 ]
  64     } case ; inline
  65
  66 GENERIC: grid-dim ( grid-size -- x y ) foldable
  67 M: integer grid-dim 1 ; inline
  68 M: sequence grid-dim
  69     dup length {
  70         { 0 [ drop 1 1 ] }
  71         { 1 [ first 1 ] }
  72         [ drop first2 ]
  73     } case ; inline
  74 PRIVATE>
  75
  76 : load-module ( path -- module )
  77     [ CUmodule <c-object> ] dip
  78     [ cuModuleLoad cuda-error ] 2keep drop c:*void* ;
  79
  80 : unload-module ( module -- )
  81     cuModuleUnload cuda-error ;
  82
  83 : load-cuda-library ( library -- handle )
  84     path>> load-module ;
  85
  86 ERROR: no-cuda-library name ;
  87
  88 : lookup-cuda-library ( name -- cuda-library )
  89     cuda-libraries get ?at [ no-cuda-library ] unless ;
  90
  91 : remove-cuda-library ( name -- library )
  92     cuda-libraries get ?delete-at [ no-cuda-library ] unless ;
  93
  94 : unload-cuda-library ( name -- )
  95     remove-cuda-library handle>> unload-module ;
  96
  97 : launch-function ( function -- ) cuLaunch cuda-error ; inline
  98
  99 : run-grid ( grid function -- )
 100     swap
 101     {
 102         [ dim-block>> block-dim function-block-shape ]
 103         [ shared-size>> function-shared-size ]
 104         [
 105             dim-grid>>
 106             [ grid-dim launch-function-grid ]
 107             [ launch-function ] if*
 108         ]
 109     } 2cleave ; inline
 110
 111 <PRIVATE
 112 : make-param-buffer ( function size -- buffer size )
 113     [ cuda-param-size ] [ (byte-array) ] [ ] tri ; inline
 114
 115 : fill-param-buffer ( values... buffer quots... n -- )
 116     [ cleave-curry ] [ spread* ] bi ; inline
 117
 118 : pointer-argument-type? ( c-type -- ? )
 119     { [ c:void* = ] [ CUdeviceptr = ] [ c:pointer? ] } 1|| ;
 120
 121 : abi-pointer-type ( abi -- type )
 122     {
 123         { cuda32 [ c:uint ] }
 124         { cuda64 [ CUulonglong ] }
 125     } case ;
 126
 127 : >argument-type ( c-type abi -- c-type' )
 128     swap {
 129         { [ dup pointer-argument-type? ] [ drop abi-pointer-type ] }
 130         { [ dup c:double    = ] [ 2drop CUdouble ] }
 131         { [ dup c:longlong  = ] [ 2drop CUlonglong ] }
 132         { [ dup c:ulonglong = ] [ 2drop CUulonglong ] }
 133         [ nip ]
 134     } cond ;
 135
 136 : >argument-struct-slot ( c-type abi -- slot )
 137     >argument-type "cuda-arg" swap { } <struct-slot-spec> ;
 138
 139 : [cuda-arguments] ( c-types abi -- quot )
 140     '[ _ >argument-struct-slot ] map
 141     [ compute-struct-offsets ]
 142     [ [ '[ _ write-struct-slot ] ] [ ] map-as ]
 143     [ length ] tri
 144     '[
 145         [ _ make-param-buffer [ drop @ _ fill-param-buffer ] 2keep ]
 146         [ '[ _ 0 ] 2dip cuda-vector ] bi
 147     ] ;
 148 PRIVATE>
 149
 150 MACRO: cuda-arguments ( c-types abi -- quot: ( args... function -- ) )
 151     [ [ 0 cuda-param-size ] ] swap '[ _ [cuda-arguments] ] if-empty ;
 152
 153 : get-function-ptr ( module string -- function )
 154     [ CUfunction <c-object> ] 2dip
 155     [ cuModuleGetFunction cuda-error ] 3keep 2drop c:*void* ;
 156
 157 : cached-module ( module-name -- alien )
 158     lookup-cuda-library
 159     cuda-modules get-global [ load-cuda-library ] cache ;
 160
 161 : cached-function ( module-name function-name -- alien )
 162     [ cached-module ] dip
 163     2array cuda-functions get [ first2 get-function-ptr ] cache ;
 164
 165 MACRO: cuda-invoke ( module-name function-name arguments -- )
 166     pick lookup-cuda-library abi>> '[
 167         _ _ cached-function
 168         [ nip _ _ cuda-arguments ]
 169         [ run-grid ] 2bi
 170     ] ;
 171
 172 : cuda-global* ( module-name symbol-name -- device-ptr size )
 173     [ CUdeviceptr <c-object> c:uint <c-object> ] 2dip
 174     [ cached-module ] dip
 175     '[ _ _ cuModuleGetGlobal cuda-error ] 2keep [ c:*uint ] bi@ ; inline
 176
 177 : cuda-global ( module-name symbol-name -- device-ptr )
 178     cuda-global* drop ; inline
 179
 180 : define-cuda-function ( word module-name function-name arguments -- )
 181     [ '[ _ _ _ cuda-invoke ] ]
 182     [ 2nip \ grid suffix c:void function-effect ]
 183     3bi define-inline ;
 184
 185 : define-cuda-global ( word module-name symbol-name -- )
 186     '[ _ _ cuda-global ] (( -- device-ptr )) define-inline ;
 187
 188 TUPLE: cuda-library name abi path handle ;
 189 ERROR: bad-cuda-abi abi ;
 190
 191 : check-cuda-abi ( abi -- abi )
 192     dup cuda-abi? [ bad-cuda-abi ] unless ; inline
 193
 194 : <cuda-library> ( name abi path -- obj )
 195     \ cuda-library new
 196         swap >>path
 197         swap check-cuda-abi >>abi
 198         swap >>name ; inline
 199
 200 : add-cuda-library ( name abi path -- )
 201     normalize-path <cuda-library>
 202     dup name>> cuda-libraries get-global set-at ;
 203