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 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
15 SYMBOL: cuda-functions
17 SYMBOL: cuda-libraries
18 cuda-libraries [ H{ } clone ] initialize
20 SYMBOL: current-cuda-library
22 : cuda-param-size ( function n -- )
23 cuParamSetSize cuda-error ; inline
25 : cuda-vector ( function offset ptr n -- )
26 cuParamSetv cuda-error ; inline
28 : launch-function-grid ( function width height -- )
29 cuLaunchGrid cuda-error ; inline
31 : function-block-shape ( function x y z -- )
32 cuFuncSetBlockShape cuda-error ; inline
34 : function-shared-size ( function n -- )
35 cuFuncSetSharedSize cuda-error ; inline
38 { dim-grid read-only }
39 { dim-block read-only }
40 { shared-size read-only initial: 0 }
41 { stream read-only } ;
43 : <grid> ( dim-grid dim-block -- grid )
46 : <grid-shared> ( dim-grid dim-block shared-size -- grid )
49 : <grid-shared-stream> ( dim-grid dim-block shared-size stream -- grid )
53 GENERIC: block-dim ( block-size -- x y z ) foldable
54 M: integer block-dim 1 1 ; inline
63 GENERIC: grid-dim ( grid-size -- x y ) foldable
64 M: integer grid-dim 1 ; inline
73 : load-module ( path -- module )
75 '[ _ cuModuleLoad cuda-error ] with-out-parameters ;
77 : unload-module ( module -- )
78 cuModuleUnload cuda-error ;
80 : load-cuda-library ( library -- handle )
83 ERROR: no-cuda-library name ;
85 : lookup-cuda-library ( name -- cuda-library )
86 cuda-libraries get ?at [ no-cuda-library ] unless ;
88 : remove-cuda-library ( name -- library )
89 cuda-libraries get ?delete-at [ no-cuda-library ] unless ;
91 : unload-cuda-library ( name -- )
92 remove-cuda-library handle>> unload-module ;
94 : launch-function ( function -- ) cuLaunch cuda-error ; inline
96 : run-grid ( grid function -- )
99 [ dim-block>> block-dim function-block-shape ]
100 [ shared-size>> function-shared-size ]
103 [ grid-dim launch-function-grid ]
104 [ launch-function ] if*
109 : make-param-buffer ( function size -- buffer size )
110 [ cuda-param-size ] [ (byte-array) ] [ ] tri ; inline
112 : fill-param-buffer ( values... buffer quots... n -- )
113 [ cleave-curry ] [ spread* ] bi ; inline
115 : pointer-argument-type? ( c-type -- ? )
116 { [ c:void* = ] [ CUdeviceptr = ] [ c:pointer? ] } 1|| ;
118 : abi-pointer-type ( abi -- type )
120 { cuda32 [ c:uint ] }
121 { cuda64 [ CUulonglong ] }
124 : >argument-type ( c-type abi -- c-type' )
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 ] }
133 : >argument-struct-slot ( c-type abi -- slot )
134 >argument-type "cuda-arg" swap { } <struct-slot-spec> ;
136 : [cuda-arguments] ( c-types abi -- quot )
137 '[ _ >argument-struct-slot ] map
138 [ compute-struct-offsets ]
139 [ [ '[ _ write-struct-slot ] ] [ ] map-as ]
142 [ _ make-param-buffer [ drop @ _ fill-param-buffer ] 2keep ]
143 [ '[ _ 0 ] 2dip cuda-vector ] bi
147 MACRO: cuda-arguments ( c-types abi -- quot: ( args... function -- ) )
148 [ [ 0 cuda-param-size ] ] swap '[ _ [cuda-arguments] ] if-empty ;
150 : get-function-ptr ( module string -- function )
151 [ { CUfunction } ] 2dip
152 '[ _ _ cuModuleGetFunction cuda-error ] with-out-parameters ;
154 : cached-module ( module-name -- alien )
156 cuda-modules get-global [ load-cuda-library ] cache ;
158 : cached-function ( module-name function-name -- alien )
159 [ cached-module ] dip
160 2array cuda-functions get [ first2 get-function-ptr ] cache ;
162 MACRO: cuda-invoke ( module-name function-name arguments -- quot )
163 pick lookup-cuda-library abi>> '[
165 [ nip _ _ cuda-arguments ]
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
174 : cuda-global ( module-name symbol-name -- device-ptr )
175 cuda-global* drop ; inline
177 :: define-cuda-function ( word module-name function-name types names -- )
179 [ module-name function-name types cuda-invoke ]
180 names "grid" suffix c:void function-effect
183 : define-cuda-global ( word module-name symbol-name -- )
184 '[ _ _ cuda-global ] ( -- device-ptr ) define-inline ;
186 TUPLE: cuda-library name abi path handle ;
187 ERROR: bad-cuda-abi abi ;
189 : check-cuda-abi ( abi -- abi )
190 dup cuda-abi? [ bad-cuda-abi ] unless ; inline
192 : <cuda-library> ( name abi path -- obj )
195 swap check-cuda-abi >>abi
198 : add-cuda-library ( name abi path -- )
199 normalize-path <cuda-library>
200 dup name>> cuda-libraries get-global set-at ;
202 USE-WHEN-LOADED: cuda.prettyprint { "cuda.libraries" "prettyprint" }