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
15 SYMBOL: cuda-functions
17 SYMBOL: cuda-libraries
18 cuda-libraries [ H{ } clone ] initialize
20 SYMBOL: current-cuda-library
22 : ?delete-at ( key assoc -- old/key ? )
23 2dup delete-at* [ 2nip t ] [ 2drop f ] if ; inline
25 : cuda-param-size ( function n -- )
26 cuParamSetSize cuda-error ; inline
28 : cuda-vector ( function offset ptr n -- )
29 cuParamSetv cuda-error ; inline
31 : launch-function-grid ( function width height -- )
32 cuLaunchGrid cuda-error ; inline
34 : function-block-shape ( function x y z -- )
35 cuFuncSetBlockShape cuda-error ; inline
37 : function-shared-size ( function n -- )
38 cuFuncSetSharedSize cuda-error ; inline
41 { dim-grid read-only }
42 { dim-block read-only }
43 { shared-size read-only initial: 0 }
44 { stream read-only } ;
46 : <grid> ( dim-grid dim-block -- grid )
49 : <grid-shared> ( dim-grid dim-block shared-size -- grid )
52 : <grid-shared-stream> ( dim-grid dim-block shared-size stream -- grid )
56 GENERIC: block-dim ( block-size -- x y z ) foldable
57 M: integer block-dim 1 1 ; inline
66 GENERIC: grid-dim ( grid-size -- x y ) foldable
67 M: integer grid-dim 1 ; inline
76 : load-module ( path -- module )
78 '[ _ cuModuleLoad cuda-error ] with-out-parameters ;
80 : unload-module ( module -- )
81 cuModuleUnload cuda-error ;
83 : load-cuda-library ( library -- handle )
86 ERROR: no-cuda-library name ;
88 : lookup-cuda-library ( name -- cuda-library )
89 cuda-libraries get ?at [ throw-no-cuda-library ] unless ;
91 : remove-cuda-library ( name -- library )
92 cuda-libraries get ?delete-at [ throw-no-cuda-library ] unless ;
94 : unload-cuda-library ( name -- )
95 remove-cuda-library handle>> unload-module ;
97 : launch-function ( function -- ) cuLaunch cuda-error ; inline
99 : run-grid ( grid function -- )
102 [ dim-block>> block-dim function-block-shape ]
103 [ shared-size>> function-shared-size ]
106 [ grid-dim launch-function-grid ]
107 [ launch-function ] if*
112 : make-param-buffer ( function size -- buffer size )
113 [ cuda-param-size ] [ (byte-array) ] [ ] tri ; inline
115 : fill-param-buffer ( values... buffer quots... n -- )
116 [ cleave-curry ] [ spread* ] bi ; inline
118 : pointer-argument-type? ( c-type -- ? )
119 { [ c:void* = ] [ CUdeviceptr = ] [ c:pointer? ] } 1|| ;
121 : abi-pointer-type ( abi -- type )
123 { cuda32 [ c:uint ] }
124 { cuda64 [ CUulonglong ] }
127 : >argument-type ( c-type abi -- c-type' )
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 ] }
136 : >argument-struct-slot ( c-type abi -- slot )
137 >argument-type "cuda-arg" swap { } <struct-slot-spec> ;
139 : [cuda-arguments] ( c-types abi -- quot )
140 '[ _ >argument-struct-slot ] map
141 [ compute-struct-offsets ]
142 [ [ '[ _ write-struct-slot ] ] [ ] map-as ]
145 [ _ make-param-buffer [ drop @ _ fill-param-buffer ] 2keep ]
146 [ '[ _ 0 ] 2dip cuda-vector ] bi
150 MACRO: cuda-arguments ( c-types abi -- quot: ( args... function -- ) )
151 [ [ 0 cuda-param-size ] ] swap '[ _ [cuda-arguments] ] if-empty ;
153 : get-function-ptr ( module string -- function )
154 [ { CUfunction } ] 2dip
155 '[ _ _ cuModuleGetFunction cuda-error ] with-out-parameters ;
157 : cached-module ( module-name -- alien )
159 cuda-modules get-global [ load-cuda-library ] cache ;
161 : cached-function ( module-name function-name -- alien )
162 [ cached-module ] dip
163 2array cuda-functions get [ first2 get-function-ptr ] cache ;
165 MACRO: cuda-invoke ( module-name function-name arguments -- quot )
166 pick lookup-cuda-library abi>> '[
168 [ nip _ _ cuda-arguments ]
172 : cuda-global* ( module-name symbol-name -- device-ptr size )
173 [ { CUdeviceptr { c:uint initial: 0 } } ] 2dip
174 [ cached-module ] dip
175 '[ _ _ cuModuleGetGlobal cuda-error ] with-out-parameters ; inline
177 : cuda-global ( module-name symbol-name -- device-ptr )
178 cuda-global* drop ; inline
180 : define-cuda-function ( word module-name function-name arguments -- )
181 [ '[ _ _ _ cuda-invoke ] ]
182 [ 2nip \ grid suffix c:void function-effect ]
185 : define-cuda-global ( word module-name symbol-name -- )
186 '[ _ _ cuda-global ] ( -- device-ptr ) define-inline ;
188 TUPLE: cuda-library name abi path handle ;
189 ERROR: bad-cuda-abi abi ;
191 : check-cuda-abi ( abi -- abi )
192 dup cuda-abi? [ throw-bad-cuda-abi ] unless ; inline
194 : <cuda-library> ( name abi path -- obj )
197 swap check-cuda-abi >>abi
200 : add-cuda-library ( name abi path -- )
201 normalize-path <cuda-library>
202 dup name>> cuda-libraries get-global set-at ;