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.data alien.parser alien.strings
   4 alien.syntax arrays assocs byte-arrays classes.struct
   5 combinators continuations cuda.ffi cuda.memory cuda.utils
   6 destructors fry init io io.backend io.encodings.string
   7 io.encodings.utf8 kernel lexer locals macros math math.parser
   8 namespaces opengl.gl.extensions parser prettyprint quotations
   9 sequences words cuda.libraries ;
  10 QUALIFIED-WITH: alien.c-types c
  11 IN: cuda
  12
  13 TUPLE: launcher
  14 { device integer initial: 0 }
  15 { device-flags initial: 0 } ;
  16
  17 : <launcher> ( device-id -- launcher )
  18     launcher new
  19         swap >>device ; inline
  20
  21 TUPLE: function-launcher
  22 dim-grid dim-block shared-size stream ;
  23
  24 : (set-up-cuda-context) ( flags device create-quot -- )
  25     H{ } clone cuda-modules set-global
  26     H{ } clone cuda-functions set
  27     call ; inline
  28
  29 : (with-cuda-context) ( context quot -- )
  30     [ '[ _ @ ] ]
  31     [ drop '[ [ sync-context ] ignore-errors _ destroy-context ] ] 2bi
  32     [ ] cleanup ; inline
  33
  34 : with-cuda-context ( flags device quot -- )
  35     [ [ create-context ] (set-up-cuda-context) ] dip (with-cuda-context) ; inline
  36
  37 : with-cuda-program ( flags device quot -- )
  38     [ dup cuda-device set ] 2dip
  39     '[ cuda-context set _ call ] with-cuda-context ; inline
  40
  41 : with-cuda ( launcher quot -- )
  42     init-cuda [
  43         [ cuda-launcher set ]
  44         [ [ device>> ] [ device-flags>> ] bi ] bi
  45     ] [ with-cuda-program ] bi* ; inline
  46
  47 : c-type>cuda-setter ( c-type -- n cuda-type )
  48     {
  49         { [ dup c:int = ] [ drop 4 [ cuda-int* ] ] }
  50         { [ dup c:uint = ] [ drop 4 [ cuda-int* ] ] }
  51         { [ dup c:float = ] [ drop 4 [ cuda-float* ] ] }
  52         { [ dup c:pointer? ] [ drop 4 [ cuda-int* ] ] }
  53         { [ dup c:void* = ] [ drop 4 [ cuda-int* ] ] }
  54     } cond ;
  55
  56 <PRIVATE
  57 : block-dim ( block -- x y z )
  58     dup sequence? [ 3 1 pad-tail first3 ] [ 1 1 ] if ; inline
  59 : grid-dim ( block -- x y )
  60     dup sequence? [ 2 1 pad-tail first2 ] [ 1 ] if ; inline
  61 PRIVATE>
  62
  63 : run-function-launcher ( function-launcher function -- )
  64     swap
  65     {
  66         [ dim-block>> block-dim function-block-shape* ]
  67         [ shared-size>> function-shared-size* ]
  68         [
  69             dim-grid>>
  70             [ grid-dim launch-function-grid* ]
  71             [ launch-function* ] if*
  72         ]
  73     } 2cleave ;
  74
  75 : cuda-argument-setter ( offset c-type -- offset' quot )
  76     c-type>cuda-setter
  77     [ over [ + ] dip ] dip
  78     '[ swap _ swap _ call ] ;
  79
  80 MACRO: cuda-arguments ( c-types -- quot: ( args... function -- ) )
  81     [ 0 ] dip [ cuda-argument-setter ] map reverse
  82     swap '[ _ param-size* ] suffix
  83     '[ _ cleave ] ;
  84
  85 : define-cuda-word ( word module-name function-name arguments -- )
  86     [
  87         '[
  88             _ _ cached-function
  89             [ nip _ cuda-arguments ]
  90             [ run-function-launcher ] 2bi
  91         ]
  92     ]
  93     [ 2nip \ function-launcher suffix c:void function-effect ]
  94     3bi define-declared ;