extra/cuda/cuda.factor

   1 ! Copyright (C) 2010 Doug Coleman.
   2 ! See http://factorcode.org/license.txt for BSD license.
   3 USING: accessors alien alien.c-types alien.data alien.parser
   4 alien.strings arrays assocs byte-arrays classes.struct
   5 combinators continuations cuda.ffi destructors fry io
   6 io.backend io.encodings.string io.encodings.utf8 kernel lexer
   7 locals math math.parser namespaces opengl.gl.extensions
   8 prettyprint quotations sequences ;
   9 IN: cuda
  10
  11 SYMBOL: cuda-device
  12 SYMBOL: cuda-context
  13 SYMBOL: cuda-module
  14 SYMBOL: cuda-function
  15 SYMBOL: cuda-launcher
  16 SYMBOL: cuda-memory-hashtable
  17
  18 ERROR: throw-cuda-error n ;
  19
  20 : cuda-error ( n -- )
  21     dup CUDA_SUCCESS = [ drop ] [ throw-cuda-error ] if ;
  22
  23 : cuda-version ( -- n )
  24     int <c-object> [ cuDriverGetVersion cuda-error ] keep *int ;
  25
  26 : init-cuda ( -- )
  27     0 cuInit cuda-error ;
  28
  29 TUPLE: launcher
  30 { device integer initial: 0 }
  31 { device-flags initial: 0 }
  32 path block-shape shared-size grid ;
  33
  34 : with-cuda-context ( flags device quot -- )
  35     [
  36         [ CUcontext <c-object> ] 2dip
  37         [ cuCtxCreate cuda-error ] 3keep 2drop *void*
  38     ] dip
  39     [ '[ _ @ ] ]
  40     [ drop '[ _ cuCtxDestroy cuda-error ] ] 2bi
  41     [ ] cleanup ; inline
  42
  43 : with-cuda-module ( path quot -- )
  44     [
  45         normalize-path
  46         [ CUmodule <c-object> ] dip
  47         [ cuModuleLoad cuda-error ] 2keep drop *void*
  48     ] dip
  49     [ '[ _ @ ] ]
  50     [ drop '[ _ cuModuleUnload cuda-error ] ] 2bi
  51     [ ] cleanup ; inline
  52
  53 : with-cuda-program ( flags device path quot -- )
  54     [ dup cuda-device set ] 2dip
  55     '[
  56         cuda-context set
  57         _ [
  58             cuda-module set
  59             _ call
  60         ] with-cuda-module
  61     ] with-cuda-context ; inline
  62
  63 : with-cuda ( launcher quot -- )
  64     [
  65         init-cuda
  66         H{ } clone cuda-memory-hashtable
  67     ] 2dip '[
  68         _
  69         [ cuda-launcher set ]
  70         [ [ device>> ] [ device-flags>> ] [ path>> ] tri ] bi
  71         _ with-cuda-program
  72     ] with-variable ; inline
  73
  74 <PRIVATE
  75
  76 : #cuda-devices ( -- n )
  77     int <c-object> [ cuDeviceGetCount cuda-error ] keep *int ;
  78
  79 : n>cuda-device ( n -- device )
  80     [ CUdevice <c-object> ] dip [ cuDeviceGet cuda-error ] 2keep drop *int ;
  81
  82 : enumerate-cuda-devices ( -- devices )
  83     #cuda-devices iota [ n>cuda-device ] map ;
  84
  85 : cuda-device-properties ( device -- properties )
  86     [ CUdevprop <c-object> ] dip
  87     [ cuDeviceGetProperties cuda-error ] 2keep drop
  88     CUdevprop memory>struct ;
  89
  90 PRIVATE>
  91
  92 : cuda-devices ( -- assoc )
  93     enumerate-cuda-devices [ dup cuda-device-properties ] { } map>assoc ;
  94
  95 : cuda-device-name ( n -- string )
  96     [ 256 [ <byte-array> ] keep ] dip
  97     [ cuDeviceGetName cuda-error ]
  98     [ 2drop utf8 alien>string ] 3bi ;
  99
 100 : cuda-device-capability ( n -- pair )
 101     [ int <c-object> int <c-object> ] dip
 102     [ cuDeviceComputeCapability cuda-error ]
 103     [ drop [ *int ] bi@ ] 3bi 2array ;
 104
 105 : cuda-device-memory ( n -- bytes )
 106     [ uint <c-object> ] dip
 107     [ cuDeviceTotalMem cuda-error ]
 108     [ drop *uint ] 2bi ;
 109
 110 : get-cuda-function* ( module string -- function )
 111     [ CUfunction <c-object> ] 2dip
 112     [ cuModuleGetFunction cuda-error ] 3keep 2drop *void* ;
 113
 114 : get-cuda-function ( string -- function )
 115     [ cuda-module get ] dip get-cuda-function* ;
 116
 117 : with-cuda-function ( string quot -- )
 118     [
 119         get-cuda-function cuda-function set
 120     ] dip call ; inline
 121
 122 : launch-function* ( function -- ) cuLaunch cuda-error ;
 123
 124 : launch-function ( -- ) cuda-function get cuLaunch cuda-error ;
 125
 126 : launch-function-grid* ( function width height -- )
 127     cuLaunchGrid cuda-error ;
 128
 129 : launch-function-grid ( width height -- )
 130     [ cuda-function get ] 2dip
 131     cuLaunchGrid cuda-error ;
 132
 133 TUPLE: cuda-memory < disposable ptr length ;
 134
 135 : <cuda-memory> ( ptr length -- obj )
 136     cuda-memory new-disposable
 137         swap >>length
 138         swap >>ptr ;
 139
 140 : add-cuda-memory ( obj -- obj )
 141     dup dup ptr>> cuda-memory-hashtable get set-at ;
 142
 143 : delete-cuda-memory ( obj -- )
 144     cuda-memory-hashtable delete-at ;
 145
 146 ERROR: invalid-cuda-memory ptr ;
 147
 148 : cuda-memory-length ( cuda-memory -- n )
 149     ptr>> cuda-memory-hashtable get ?at [
 150         length>>
 151     ] [
 152         invalid-cuda-memory
 153     ] if ;
 154
 155 M: cuda-memory byte-length length>> ;
 156
 157 : cuda-malloc ( n -- ptr )
 158     [ CUdeviceptr <c-object> ] dip
 159     [ cuMemAlloc cuda-error ] 2keep
 160     [ *int ] dip <cuda-memory> add-cuda-memory ;
 161
 162 : cuda-free* ( ptr -- )
 163     cuMemFree cuda-error ;
 164
 165 M: cuda-memory dispose ( ptr -- )
 166     ptr>> cuda-free* ;
 167
 168 : host>device ( dest-ptr src-ptr -- )
 169     [ ptr>> ] dip dup length cuMemcpyHtoD cuda-error ;
 170
 171 :: device>host ( ptr -- seq )
 172     ptr byte-length <byte-array>
 173     [ ptr [ ptr>> ] [ byte-length ] bi cuMemcpyDtoH cuda-error ] keep ;
 174
 175 : memcpy-device>device ( dest-ptr src-ptr count -- )
 176     cuMemcpyDtoD cuda-error ;
 177
 178 : memcpy-device>array ( dest-array dest-index src-ptr count -- )
 179     cuMemcpyDtoA cuda-error ;
 180
 181 : memcpy-array>device ( dest-ptr src-array src-index count -- )
 182     cuMemcpyAtoD cuda-error ;
 183
 184 : memcpy-array>host ( dest-ptr src-array src-index count -- )
 185     cuMemcpyAtoH cuda-error ;
 186
 187 : memcpy-host>array ( dest-array dest-index src-ptr count -- )
 188     cuMemcpyHtoA cuda-error ;
 189
 190 : memcpy-array>array ( dest-array dest-index src-array src-ptr count -- )
 191     cuMemcpyAtoA cuda-error ;
 192
 193 : cuda-int* ( function offset value -- )
 194     cuParamSeti cuda-error ;
 195
 196 : cuda-int ( offset value -- )
 197     [ cuda-function get ] 2dip cuda-int* ;
 198
 199 : cuda-float* ( function offset value -- )
 200     cuParamSetf cuda-error ;
 201
 202 : cuda-float ( offset value -- )
 203     [ cuda-function get ] 2dip cuda-float* ;
 204
 205 : cuda-vector* ( function offset ptr n -- )
 206     cuParamSetv cuda-error ;
 207
 208 : cuda-vector ( offset ptr n -- )
 209     [ cuda-function get ] 3dip cuda-vector* ;
 210
 211 : param-size* ( function n -- )
 212     cuParamSetSize cuda-error ;
 213
 214 : param-size ( n -- )
 215     [ cuda-function get ] dip param-size* ;
 216
 217 : malloc-device-string ( string -- n )
 218     utf8 encode
 219     [ length cuda-malloc ] keep
 220     [ host>device ] [ drop ] 2bi ;
 221
 222 ERROR: bad-cuda-parameter parameter ;
 223
 224 :: set-parameters ( seq -- )
 225     cuda-function get :> function
 226     0 :> offset!
 227     seq [
 228         [ offset ] dip
 229         {
 230             { [ dup cuda-memory? ] [ ptr>> cuda-int ] }
 231             { [ dup float? ] [ cuda-float ] }
 232             { [ dup integer? ] [ cuda-int ] }
 233             [ bad-cuda-parameter ]
 234         } cond
 235         offset 4 + offset!
 236     ] each
 237     offset param-size ;
 238
 239 : cuda-device-attribute ( attribute dev -- n )
 240     [ int <c-object> ] 2dip
 241     [ cuDeviceGetAttribute cuda-error ]
 242     [ 2drop *int ] 3bi ;
 243
 244 : function-block-shape* ( function x y z -- )
 245     cuFuncSetBlockShape cuda-error ;
 246
 247 : function-block-shape ( x y z -- )
 248     [ cuda-function get ] 3dip
 249     cuFuncSetBlockShape cuda-error ;
 250
 251 : function-shared-size* ( function n -- )
 252     cuFuncSetSharedSize cuda-error ;
 253
 254 : function-shared-size ( n -- )
 255     [ cuda-function get ] dip
 256     cuFuncSetSharedSize cuda-error ;
 257
 258 : launch ( -- )
 259     cuda-launcher get {
 260         [ block-shape>> first3 function-block-shape ]
 261         [ shared-size>> function-shared-size ]
 262         [
 263             grid>> [
 264                 launch-function
 265             ] [
 266                 first2 launch-function-grid
 267             ] if-empty
 268         ]
 269     } cleave ;
 270
 271 : cuda-device. ( n -- )
 272     {
 273         [ "Device: " write number>string print ]
 274         [ "Name: " write cuda-device-name print ]
 275         [ "Memory: " write cuda-device-memory number>string print ]
 276         [
 277             "Capability: " write
 278             cuda-device-capability [ number>string ] map " " join print
 279         ]
 280         [ "Properties: " write cuda-device-properties . ]
 281         [
 282             "CU_DEVICE_ATTRIBUTE_GPU_OVERLAP: " write
 283             CU_DEVICE_ATTRIBUTE_GPU_OVERLAP swap
 284             cuda-device-attribute number>string print
 285         ]
 286     } cleave ;
 287
 288 : cuda. ( -- )
 289     "CUDA Version: " write cuda-version number>string print nl
 290     #cuda-devices iota [ nl ] [ cuda-device. ] interleave ;
 291
 292
 293 : test-cuda0 ( -- )
 294     T{ launcher
 295         { path "vocab:cuda/hello.ptx" }
 296         { block-shape { 6 6 6 } }
 297         { shared-size 2 }
 298         { grid { 2 6 } }
 299     } [
 300         "helloWorld" [
 301             "Hello World!" [ - ] map-index
 302             malloc-device-string &dispose
 303
 304             [ 1array set-parameters ]
 305             [ drop launch ]
 306             [ device>host utf8 alien>string . ] tri
 307         ] with-cuda-function
 308     ] with-cuda ;