def clSetKernelArg(kernel, arg_index, arg_size, arg_value):
if isinstance(kernel, str): kernel = SymbolRef(kernel)
if isinstance(arg_index, int): arg_index = Constant(arg_index)
if isinstance(arg_size, int): arg_size = Constant(arg_size)
if isinstance(arg_value, str): arg_value = Ref(SymbolRef(arg_value))
return FunctionCall(SymbolRef("clSetKernelArg"),
[kernel, arg_index, arg_size, arg_value])
def printf(fmt, *args):
"""
Makes a printf call. Args must be CtreeNodes.
"""
for arg in args:
assert isinstance(arg, CtreeNode)
return FunctionCall(SymbolRef("printf"), [String(fmt)] + list(args))
def visit_Call(self, node):
args = [self.visit(a) for a in node.args]
fn = self.visit(node.func)
if node.starargs is not None:
node.func = fn
node.args = args
node.starargs = self.visit(node.starargs)
return node
return FunctionCall(fn, args)
def clEnqueueReadBuffer(queue, buf, blocking, offset, cb, ptr, num_events=0, evt_list_ptr=None, evt=None):
if isinstance(buf, str): buf = SymbolRef(buf)
if isinstance(blocking, bool): blocking = Constant(int(blocking))
if isinstance(ptr, str): ptr = SymbolRef(ptr)
if not isinstance(offset, ast.AST): offset = Constant(offset)
if not isinstance(cb, ast.AST): cb = Constant(cb)
if not isinstance(num_events, ast.AST): num_events = Constant(num_events)
if not isinstance(evt_list_ptr, ast.AST): event_list_ptr = NULL()
if not isinstance(evt, ast.AST): evt = NULL()
return FunctionCall(SymbolRef('clEnqueueReadBuffer'), [
queue, buf, blocking, offset, cb, ptr, num_events, event_list_ptr, evt])
def visit_Call(self, node):
args = [self.visit(a) for a in node.args]
fn = self.visit(node.func)
if getattr(node, 'starargs', None) is not None:
node.func = fn
node.args = args
node.starargs = self.visit(node.starargs)
return node
if hasattr(fn, "name") and fn.name == "float":
return Cast(c_float(), args[0])
return FunctionCall(fn, args)
def clEnqueueCopyBuffer(queue, src_buf, dst_buf, src_offset=0, dst_offset=0, cb=0):
if isinstance(src_buf, str): src_buf = SymbolRef(src_buf)
if isinstance(dst_buf, str): dst_buf = SymbolRef(dst_buf)
if isinstance(src_offset, int): src_offset = Constant(src_offset)
if isinstance(dst_offset, int): dst_offset = Constant(dst_offset)
if isinstance(cb, int): cb = Constant(cb)
num_events = Constant(0)
event_list_ptr = NULL()
evt = NULL()
return FunctionCall(SymbolRef('clEnqueueCopyBuffer'), [
queue, src_buf, dst_buf, src_offset, dst_offset, cb, num_events, event_list_ptr, evt])
def clEnqueueNDRangeKernel(queue, kernel, work_dim=1, work_offset=0, global_size=0, local_size=0):
assert isinstance(queue, SymbolRef)
assert isinstance(kernel, SymbolRef)
global_size_sym = SymbolRef('global_size', c_size_t())
local_size_sym = SymbolRef('local_size', c_size_t())
call = FunctionCall(SymbolRef("clEnqueueNDRangeKernel"), [
queue, kernel,
work_dim, work_offset,
Ref(global_size_sym.copy()), Ref(local_size_sym.copy()),
0, NULL(), NULL()
])
return Block([
Assign(global_size_sym, Constant(global_size)),
Assign(local_size_sym, Constant(local_size)),
call
])
def get_local_size(id):
return FunctionCall(SymbolRef('get_local_size'), [Constant(id)])
def get_global_id(id):
return FunctionCall(SymbolRef('get_global_id'), [Constant(id)])
def barrier(arg):
return FunctionCall(SymbolRef('barrier'), [arg])
def _the_call(*args):
assert len(args) == nArgs, \
"Macro expected %d args, got %d." % (nArgs, len(args))
return FunctionCall(SymbolRef(name), args)
def visit_Call(self, node):
args = [self.visit(a) for a in node.args]
fn = self.visit(node.func)
return FunctionCall(fn, args)
def omp_get_num_threads():
return FunctionCall(SymbolRef("omp_get_num_threads"), [])
def omp_get_thread_num():
return FunctionCall(SymbolRef("omp_get_thread_num"), [])
def get_num_groups(id):
return FunctionCall(SymbolRef('get_num_groups'), [Constant(id)])
def clReleaseMemObject(arg):
return FunctionCall(SymbolRef('clReleaseMemObject'), [arg])
def omp_get_wtime():
return FunctionCall(SymbolRef("omp_get_wtime"), [])
