def visit_Subscript(self, node):
if isinstance(node.slice.value, ast.Tuple):
idxs = node.slice.value.elts
new_node = C.ArrayRef(node.value, idxs[0])
for idx in idxs[1:]:
new_node = C.ArrayRef(new_node, idx)
return new_node
return node
def visit_BinaryOp(self, node):
if isinstance(node.op, C.Op.ArrayRef):
if util.contains_symbol(node, self.loop_var):
idx = 0
curr_node = node
while not isinstance(curr_node.right, C.SymbolRef) or \
curr_node.right.name != self.loop_var:
idx += 1
curr_node = curr_node.left
while not isinstance(curr_node, C.SymbolRef):
curr_node = curr_node.left
self.vectorized_buffers[curr_node.name] = idx
if self.vectorize:
return simd_macros.mm256_load_ps(node)
else:
return C.ArrayRef(node,
C.SymbolRef("_neuron_index_1_inner"))
else:
if self.vectorize:
return simd_macros.mm256_set1_ps(node)
else:
return node
node.left = self.visit(node.left)
node.right = self.visit(node.right)
return node
def visit_AugAssign(self, node):
node.value = self.visit(node.value)
if util.contains_symbol(node.target, self.loop_var):
if not util.contains_symbol(node.target.right, self.loop_var):
target = self.visit(deepcopy(node.target))
curr_node = node.target
idx = 1
while curr_node.left.right.name != self.loop_var:
curr_node = curr_node.left
idx += 1
curr_node.left = curr_node.left.left
node.target = C.ArrayRef(node.target,
C.SymbolRef(self.loop_var))
while not isinstance(curr_node, C.SymbolRef):
curr_node = curr_node.left
if curr_node.name in self.transposed_buffers and self.transposed_buffers[
curr_node.name] != idx:
raise NotImplementedError()
self.transposed_buffers[curr_node.name] = idx
curr_node.name += "_transposed"
if isinstance(node.target.right,
C.Constant) and node.target.value == 0.0:
return store_ps(node.target.left,
C.BinaryOp(target, node.op, node.value))
else:
return store_ps(C.Ref(node.target),
C.BinaryOp(target, node.op, node.value))
else:
if isinstance(node.target.right,
C.Constant) and node.target.value == 0.0:
return store_ps(
node.target.left,
C.BinaryOp(self.visit(node.target), node.op,
node.value))
else:
return store_ps(
C.Ref(node.target),
C.BinaryOp(self.visit(node.target), node.op,
node.value))
elif isinstance(node.op, C.Op.Add) and isinstance(
node.value, C.FunctionCall):
# TODO: Verfiy it's a vector intrinsic
return C.Assign(
node.target,
C.FunctionCall(C.SymbolRef("_mm256_add_ps"),
[node.value, node.target]))
elif isinstance(node.target, C.BinaryOp) and isinstance(
node.target.op, C.Op.ArrayRef):
raise NotImplementedError(node)
node.target = self.visit(node.target)
return node
def visit_FunctionDecl(self, node):
new_body = []
count = 0
for statement in node.defn:
if isinstance(statement, ast.For) or isinstance(statement, C.For):
pre = C.SubAssign(C.ArrayRef(C.SymbolRef('times'), C.Constant(count)),C.FunctionCall('omp_get_wtime', []))
post = C.AddAssign(C.ArrayRef(C.SymbolRef('times'), C.Constant(count)),C.FunctionCall('omp_get_wtime', []))
new_body.append(pre)
new_body.append(statement)
new_body.append(post)
count = count + 1
else:
new_body.append(statement)
memset = C.Assign(C.SymbolRef('times'), C.FunctionCall(C.SymbolRef('calloc_doubles'),[C.Constant(count)]))
new_body.insert(0, memset)
new_body.insert(0, C.Assign(C.SymbolRef("*times", ctypes.c_double()), C.Constant(0)))
for i in range(0,count):
print_stmt = C.FunctionCall(C.SymbolRef('printf'),[C.String("\ttimes[%d] = %g\\n"), C.Constant(i), C.ArrayRef(C.SymbolRef('times'), C.Constant(i))])
new_body.append(print_stmt)
node.defn = new_body
return node
def visit_BinaryOp(self, node):
if isinstance(node.op, C.Op.ArrayRef) and \
contains_name(node, self.idx):
idx = 0
curr_node = node
while not contains_name(curr_node.right, self.idx):
idx += 1
curr_node = curr_node.left
idx_expr = deepcopy(curr_node.right)
while not isinstance(curr_node, ast.Name):
curr_node = curr_node.left
self.tiled_buffers[curr_node.id] = idx
idx_expr = AppendAllNames("_inner").visit(idx_expr)
return C.ArrayRef(node, idx_expr)
node.left = self.visit(node.left)
node.right = self.visit(node.right)
return node
def gen_index_expr(target, idxs):
node = C.ArrayRef(target, idxs[0])
for idx in idxs[1:]:
node = C.ArrayRef(node, idx)
return node
def visit_Return(self, node):
node.value = self.visit(node.value)
return C.Assign(C.ArrayRef(C.SymbolRef(self.retval_name),
C.SymbolRef(self.loopvar)),
node.value)
def visit_SymbolRef(self, node):
if node.name in self.param_dict:
return C.ArrayRef(node, C.SymbolRef(self.loopvar))
return node
def visit_BinaryOp(self, node):
if isinstance(node.op, C.Op.ArrayRef):
if util.contains_symbol(node, self.loop_var):
if not util.contains_symbol(node.right, self.loop_var):
curr_node = node
idx = 1
while curr_node.left.right.name != self.loop_var:
curr_node = curr_node.left
idx += 1
curr_node.left = curr_node.left.left
node = C.ArrayRef(node, C.SymbolRef(self.loop_var))
while not isinstance(curr_node, C.SymbolRef):
curr_node = curr_node.left
if curr_node.name in self.transposed_buffers and self.transposed_buffers[
curr_node.name] != idx:
raise NotImplementedError()
self.transposed_buffers[curr_node.name] = idx
curr_node.name += "_transposed"
if isinstance(node.right,
C.Constant) and node.target.value == 0.0:
return load_ps(node.left)
else:
return load_ps(C.Ref(node))
else:
return broadcast_ss(C.Ref(node))
elif isinstance(node.op, C.Op.Assign):
node.right = self.visit(node.right)
if isinstance(node.right, C.FunctionCall) and \
("load_ps" in node.right.func.name or
"broadcast_ss" in node.right.func.name) and \
isinstance(node.left, C.SymbolRef) and node.left.type is not None:
node.left.type = get_simd_type()()
self.symbol_table[node.left.name] = node.left.type
return node
elif isinstance(node.left, C.BinaryOp) and util.contains_symbol(
node.left, self.loop_var):
if node.left.right.name != self.loop_var:
curr_node = node
idx = 1
while curr_node.left.right.name != self.loop_var:
curr_node = curr_node.left
idx += 1
curr_node.left = curr_node.left.left
node = C.ArrayRef(node, C.SymbolRef(self.loop_var))
while not isinstance(curr_node, C.SymbolRef):
curr_node = curr_node.left
if curr_node.name in self.transposed_buffers and self.transposed_buffers[
curr_node.name] != idx:
raise NotImplementedError()
self.transposed_buffers[curr_node.name] = idx
curr_node.name += "_transposed"
is_float = self.get_type(node.left)
if isinstance(is_float, ctypes.c_float):
if isinstance(node.left.right,
C.Constant) and node.target.value == 0.0:
return store_ps(node.left.left, node.right)
else:
return store_ps(C.Ref(node.left), node.right)
elif isinstance(is_float, ctypes.c_int):
if isinstance(node.left.right,
C.Constant) and node.target.value == 0.0:
return store_epi32(node.left.left, node.right)
else:
return store_epi32(C.Ref(node.left), node.right)
else:
if isinstance(node.left.right,
C.Constant) and node.target.value == 0.0:
return store_ps(node.left.left, node.right)
else:
return store_ps(C.Ref(node.left), node.right)
node.left = self.visit(node.left)
return node
node.left = self.visit(node.left)
node.right = self.visit(node.right)
return node