extra/cuda/libraries/libraries.factor

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