Skip to content

Commit 3e2c6e4

Browse files
author
hartsantler
committed
webCLGL backend: fixed subroutines.
1 parent 721584d commit 3e2c6e4

3 files changed

Lines changed: 93 additions & 52 deletions

File tree

pythonjs/python_to_pythonjs.py

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,10 @@ def __init__(self, source=None, module=None, module_path=None, dart=False, coffe
153153
self._with_rpc = None
154154
self._with_rpc_name = None
155155
self._with_direct_keys = False
156+
156157
self._with_glsl = False
158+
self._in_gpu_main = False
159+
self._gpu_return_type = 'array' ## 'array' or float32 or array of 'vec4' float32's.
157160

158161
self._source = source.splitlines()
159162
self._classes = dict() ## class name : [method names]
@@ -1054,7 +1057,7 @@ def visit_If(self, node):
10541057
elif isinstance(node.test, ast.List):
10551058
writer.write('if %s.length:' % self.visit(node.test))
10561059

1057-
elif self._with_ll:
1060+
elif self._with_ll or self._with_glsl:
10581061
writer.write('if %s:' % self.visit(node.test))
10591062
elif isinstance(node.test, ast.Compare):
10601063
writer.write('if %s:' % self.visit(node.test))
@@ -1131,17 +1134,22 @@ def visit_Return(self, node):
11311134
elif isinstance(node.value, Name) and node.value.id == 'self' and 'self' in self._instances:
11321135
self._return_type = self._instances['self']
11331136

1134-
## cached property is DEPRECATED
1135-
#if self._cached_property:
1136-
# writer.write('self["__dict__"]["%s"] = %s' %(self._cached_property, self.visit(node.value)) )
1137-
# writer.write('return self["__dict__"]["%s"]' %self._cached_property)
1138-
#else:
1139-
if self._with_glsl:
1140-
writer.write('out_float = %s' %self.visit(node.value))
1137+
1138+
if self._with_glsl and self._in_gpu_main:
1139+
## _id_ is inserted into all function headers by pythonjs.py for glsl functions.
1140+
if self._gpu_return_type == 'array':
1141+
writer.write('out_float = %s' %self.visit(node.value))
1142+
elif self._gpu_return_type == 'vec4':
1143+
writer.write('out_float4 = %s' %self.visit(node.value))
1144+
else:
1145+
raise NotImplementedError(node)
1146+
1147+
11411148
elif self._inline:
11421149
writer.write('__returns__%s = %s' %(self._inline[-1], self.visit(node.value)) )
11431150
if self._inline_breakout:
11441151
writer.write('break')
1152+
11451153
elif isinstance(node.value, ast.Lambda):
11461154
self.visit( node.value )
11471155
writer.write( 'return __lambda__' )
@@ -2393,6 +2401,9 @@ def visit_FunctionDef(self, node):
23932401
inline = False
23942402
threaded = self._with_webworker
23952403
jsfile = None
2404+
2405+
gpu = False
2406+
gpu_main = False
23962407
gpu_vectorize = False
23972408

23982409
## deprecated?
@@ -2404,7 +2415,10 @@ def visit_FunctionDef(self, node):
24042415

24052416
for decorator in reversed(node.decorator_list):
24062417
log('@decorator: %s' %decorator)
2407-
if isinstance(decorator, Name) and decorator.id == 'inline':
2418+
if isinstance(decorator, Name) and decorator.id == 'gpu':
2419+
gpu = True
2420+
2421+
elif isinstance(decorator, Name) and decorator.id == 'inline':
24082422
inline = True
24092423
self._with_inline = True
24102424

@@ -2420,10 +2434,15 @@ def visit_FunctionDef(self, node):
24202434
return_type = decorator.args[0].id
24212435

24222436
elif isinstance(decorator, Attribute) and isinstance(decorator.value, Name) and decorator.value.id == 'gpu':
2423-
assert decorator.attr == 'vectorize'
2424-
gpu_vectorize = True
2425-
restore_with_glsl = self._with_glsl
2426-
self._with_glsl = True
2437+
gpu = True
2438+
if decorator.attr == 'vectorize':
2439+
gpu_vectorize = True
2440+
elif decorator.attr == 'main':
2441+
gpu_main = True
2442+
elif decorator.attr == 'typedef':
2443+
pass
2444+
else:
2445+
raise NotImplementedError(decorator)
24272446

24282447
elif self._with_dart:
24292448
with_dart_decorators.append( self.visit(decorator) )
@@ -2484,6 +2503,15 @@ def visit_FunctionDef(self, node):
24842503
decorators.append( decorator )
24852504

24862505

2506+
if gpu:
2507+
restore_with_glsl = self._with_glsl
2508+
self._with_glsl = True
2509+
if gpu_main: ## sets float
2510+
self._in_gpu_main = True
2511+
writer.write('@gpu.main')
2512+
2513+
2514+
24872515
if threaded:
24882516
if not jsfile: jsfile = 'worker.js'
24892517
writer_main.write('%s = "%s"' %(node.name, jsfile))
@@ -2613,8 +2641,6 @@ def visit_FunctionDef(self, node):
26132641
writer.write('def %s(args, kwargs):' % node.name)
26142642
writer.push()
26152643

2616-
2617-
26182644
## write local typedefs and var scope ##
26192645
a = ','.join( vars )
26202646
if local_typedefs:
@@ -2897,8 +2923,10 @@ def visit_FunctionDef(self, node):
28972923

28982924
writer.pull() ## end function body
28992925

2900-
if gpu_vectorize:
2926+
if gpu:
29012927
self._with_glsl = restore_with_glsl
2928+
if gpu_main:
2929+
self._in_gpu_main = False
29022930

29032931
self._typedef_vars = dict() ## clear typed variables
29042932

pythonjs/pythonjs.py

Lines changed: 37 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ def visit_Module(self, node):
105105
lines.append('}) //end requirejs define')
106106

107107
if self._has_glsl:
108+
header.append( 'var __shader_header__ = []' )
108109
header.append( 'var __shader__ = []' )
109110

110111
lines = header + lines
@@ -200,6 +201,7 @@ def _visit_call_helper_var_glsl(self, node):
200201

201202

202203
def _visit_function(self, node):
204+
is_main = node.name == 'main'
203205
return_type = None
204206
glsl = False
205207
glsl_wrapper_name = False
@@ -217,67 +219,85 @@ def _visit_function(self, node):
217219
return_type = decor.args[0].id
218220

219221
elif isinstance(decor, Attribute) and isinstance(decor.value, Name) and decor.value.id == 'gpu':
220-
assert decor.attr == 'vectorize'
221-
gpu_vectorize = True
222+
if decor.attr == 'vectorize':
223+
gpu_vectorize = True
224+
elif decor.attr == 'main':
225+
is_main = True
222226

223227

224228
args = self.visit(node.args)
225229

226230
if glsl:
227-
is_main = node.name == 'main'
228231
self._has_glsl = True ## writes `__shader__ = []` in header
229232
lines = []
230233
x = []
231234
for i,arg in enumerate(args):
232235
if gpu_vectorize and arg not in args_typedefs:
233236
x.append( 'float* %s' %arg )
234237
else:
235-
assert arg in args_typedefs
236-
x.append( '%s %s' %(args_typedefs[arg].replace('POINTER', '*'), arg) )
238+
if arg in args_typedefs:
239+
x.append( '%s %s' %(args_typedefs[arg].replace('POINTER', '*'), arg) )
240+
else:
241+
x.append( 'float* %s' %arg )
237242

238-
if is_main or gpu_vectorize:
243+
if is_main:
239244
lines.append( '__shader__.push("void main( %s ) {");' %', '.join(x) )
240245
elif return_type:
241-
lines.append( '__shader__.push("%s %s( %s ) {");' %(return_type, node.name, ', '.join(x)) )
246+
lines.append( '__shader_header__.push("%s %s( %s ) {");' %(return_type, node.name, ', '.join(x)) )
242247
else:
243-
lines.append( '__shader__.push("void %s( %s ) {");' %(node.name, ', '.join(x)) )
248+
lines.append( '__shader_header__.push("void %s( %s ) {");' %(node.name, ', '.join(x)) )
244249

245250
self.push()
246-
if gpu_vectorize:
251+
# `_id_` always write out an array of floats or array of vec4floats
252+
if is_main:
247253
lines.append( '__shader__.push("vec2 _id_ = get_global_id();");')
254+
else:
255+
lines.append( '__shader_header__.push("vec2 _id_ = get_global_id();");')
248256

249257
self._glsl = True
250258
for child in node.body:
251259
if isinstance(child, Str):
252260
continue
253261
else:
254262
for sub in self.visit(child).splitlines():
255-
lines.append( '__shader__.push("%s");' %(self.indent()+sub) )
263+
if is_main:
264+
lines.append( '__shader__.push("%s");' %(self.indent()+sub) )
265+
else:
266+
lines.append( '__shader_header__.push("%s");' %(self.indent()+sub) )
256267
self._glsl = False
257268
#buffer += '\n'.join(body)
258269
self.pull()
259-
lines.append('__shader__.push("%s}");' %self.indent())
270+
if is_main:
271+
lines.append('__shader__.push("%s}");' %self.indent())
272+
else:
273+
lines.append('__shader_header__.push("%s}");' %self.indent())
260274

261-
if is_main or gpu_vectorize:
275+
lines.append(';')
276+
277+
if is_main:
262278
if not glsl_wrapper_name:
263279
glsl_wrapper_name = node.name
264280
lines.append('function %s( %s, __offset ) {' %(glsl_wrapper_name, ','.join(args)) )
265-
lines.append(' __offset = __offset || 0') ## note by default: 0 allows 0-1.0
281+
lines.append(' __offset = __offset || 0') ## note by default: 0 allows 0-1.0 ## TODO this needs to be set per-buffer
282+
266283
lines.append(' var __webclgl = new WebCLGL()')
267-
lines.append(' var __kernel = __webclgl.createKernel( "\\n".join(__shader__) );')
268-
lines.append(' var __return_length = 1')
284+
lines.append(' var header = "\\n".join(__shader_header__)')
285+
lines.append(' var shader = "\\n".join(__shader__)')
286+
lines.append(' var __kernel = __webclgl.createKernel( shader, header );')
287+
288+
lines.append(' var __return_length = 64') ## minimum size is 64
269289

270290
for i,arg in enumerate(args):
271291
lines.append(' if (%s instanceof Array) {' %arg)
272-
lines.append(' __return_length = %s.length' %arg)
292+
lines.append(' __return_length = %s.length==2 ? %s : %s.length' %(arg,arg, arg) )
273293
lines.append(' var %s_buffer = __webclgl.createBuffer(%s.length, "FLOAT", __offset)' %(arg,arg))
274294
lines.append(' __webclgl.enqueueWriteBuffer(%s_buffer, %s)' %(arg, arg))
275295
lines.append(' __kernel.setKernelArg(%s, %s_buffer)' %(i, arg))
276296
lines.append(' } else { __kernel.setKernelArg(%s, %s) }' %(i, arg))
277297

278298
lines.append(' var return_buffer = __webclgl.createBuffer(__return_length, "FLOAT", __offset)')
279-
280299
lines.append(' __kernel.compile()')
300+
281301
lines.append(' __webclgl.enqueueNDRangeKernel(__kernel, return_buffer)')
282302
lines.append(' return __webclgl.enqueueReadBuffer_Float( return_buffer )')
283303
lines.append('} // end of wrapper')

regtests/webclgl/subroutine.py

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,19 @@
11
"""subroutine"""
22

33
def main():
4-
with glsl as myfunc:
5-
## glsl compiler requires subroutines defined first
6-
float def mysub(x,y):
7-
float x
8-
float y
9-
return x+y
4+
@gpu
5+
float def mysub(x,y):
6+
float x
7+
float y
8+
return x-y
109

11-
def main(buffA, buffB, num):
12-
float* buffA
13-
float* buffB
14-
float num
15-
vec2 n = get_global_id() ## WebCL API
16-
float result = 0.0
17-
for i in range(1000):
18-
result = sqrt(result + A[n] + B[n] + float(i))
19-
return mysub( result, num )
10+
#@gpu.returns.vec4
11+
@gpu.main
12+
def myfunc(a):
13+
float* a
14+
return mysub( 1.1 * a[_id_], 2.2 )
2015

2116

22-
23-
A = [1,2,3]
24-
B = [4,5,6]
25-
res = myfunc( A, B, 2.0 )
17+
A = [1.3 for i in range(64)]
18+
res = myfunc( A )
2619
print(res)

0 commit comments

Comments
 (0)