def generate_subview(self, node: ast.Assign, view_name: str) -> cppast.DeclStmt:
subview_args: List[cppast.Expr] = [cppast.DeclRefExpr(view_name)]
slice_node = node.value
for dim in slice_node.slice.dims:
if isinstance(dim, ast.Index):
subview_args.append(self.visit(dim))
else:
if dim.lower is None and dim.upper is None:
subview_args.append(cppast.DeclRefExpr("Kokkos::ALL"))
elif dim.lower is not None and dim.upper is not None:
make_pair = cppast.CallExpr("std::make_pair",
[self.visit(dim.lower), self.visit(dim.upper)])
subview_args.append(make_pair)
else:
self.error(
slice_node, "Partial slice not supported, use [n:m] or [:]")
if len(node.targets) > 1:
self.error(node, "Multiple declarations of subview not supported")
auto = cppast.ClassType("auto")
target = node.targets[0]
target_ref = cppast.DeclRefExpr(target.id)
if target_ref in self.views:
self.error(
node, "Redeclaration of existing subview")
else:
self.views[target_ref] = None
self.subviews[target_ref.declname] = view_name
call = cppast.CallExpr("Kokkos::subview", subview_args)
decl = cppast.DeclStmt(cppast.VarDecl(auto, self.visit(target), call))
return decl
def visit_AnnAssign(self, node: ast.AnnAssign) -> cppast.Stmt:
if isinstance(node.value, ast.Call):
decltype: cppast.Type = visitors_util.get_type(
node.annotation, self.pk_import)
if decltype is None:
self.error(node, "Type not supported")
declname: cppast.DeclRefExpr = self.visit(node.target)
function_name: str = visitors_util.get_node_name(node.value.func)
# Call to a TeamMember method
if function_name in dir(TeamMember):
vardecl = cppast.VarDecl(decltype, declname,
self.visit(node.value))
return cppast.DeclStmt(vardecl)
# Nested parallelism
if function_name in ("parallel_reduce", "parallel_scan"):
args: List[cppast.Expr] = [
self.visit(a) for a in node.value.args
]
initial_value: cppast.Expr
if len(args) == 3:
initial_value = args[2]
else:
initial_value = cppast.IntegerLiteral(0)
vardecl = cppast.VarDecl(decltype, declname, initial_value)
declstmt = cppast.DeclStmt(vardecl)
work_unit: str = args[1].declname
function = cppast.DeclRefExpr(f"Kokkos::{function_name}")
call: cppast.CallExpr
if work_unit in self.nested_work_units:
call = cppast.CallExpr(
function,
[args[0], self.nested_work_units[work_unit], declname])
else:
call = cppast.CallExpr(
function, [args[0], f"pk_id_{work_unit}", declname])
callstmt = cppast.CallStmt(call)
return cppast.CompoundStmt([declstmt, callstmt])
return super().visit_AnnAssign(node)
def visit_For(self, node: ast.For) -> cppast.ForStmt:
if not isinstance(node.target, ast.Name):
self.error(node.target, "Must use single loop variable")
if node.orelse:
self.error(node.orelse, "Else clause not supported for translation")
if (
not isinstance(node.iter, ast.Call)
or node.iter.func.id != "range"
):
# TODO: support other iterators?
self.error(
node.iter, "Only range() iterator is supported for translation")
index: cppast.DeclRefExpr = self.visit(node.target)
start: cppast.Expr
end: cppast.Expr
step: cppast.Expr = cppast.IntegerLiteral(1)
op = cppast.BinaryOperatorKind.LT
args = node.iter.args
if len(args) == 1:
start = cppast.IntegerLiteral(0)
end = self.visit(args[0])
else:
start = self.visit(args[0])
end = self.visit(args[1])
if len(args) == 3:
step = self.visit(args[2])
# Negative step sizes are only handled correctly if they're
# written with a preceeding minus sign
if (
isinstance(args[2], ast.UnaryOp)
and isinstance(args[2].op, ast.USub)
):
op = cppast.BinaryOperatorKind.GT
body = cppast.CompoundStmt([self.visit(b) for b in node.body])
init = cppast.DeclStmt(cppast.VarDecl(
cppast.PrimitiveType(cppast.BuiltinType.INT), index, start))
condition = cppast.BinaryOperator(index, end, op)
increment = cppast.BinaryOperator(
index, step, cppast.BinaryOperatorKind.AddAssign)
forstmt = cppast.ForStmt(init, condition, increment, body)
return forstmt
def visit_AnnAssign(self, node: ast.AnnAssign) -> cppast.DeclStmt:
if not isinstance(node.target, ast.Name):
self.generic_error(node)
decltype: cppast.Type = visitors_util.get_type(node.annotation, self.pk_import)
if decltype is None:
self.error(node, "Type not supported")
target: cppast.DeclRefExpr = self.visit(node.target)
value: cppast.Expr = self.visit(node.value)
# Add the name of the list to self.lists
# and add the length of the list to its name
if isinstance(node.annotation, ast.Subscript):
self.lists.append(target.declname)
target.add_length(len(node.value.elts))
annassign = cppast.DeclStmt(cppast.VarDecl(decltype, target, value))
return annassign
def visit_Assign(self, node: ast.Assign) -> cppast.Stmt:
target = node.targets[0]
if isinstance(node.value, ast.Call):
name: str = visitors_util.get_node_name(node.value.func)
# Create Timer object
if name == "Timer":
decltype = cppast.ClassType("Kokkos::Timer")
declname = cppast.DeclRefExpr("timer")
return cppast.DeclStmt(cppast.VarDecl(decltype, declname,
None))
# Call Timer.seconds()
if name == "seconds":
target_name: str = visitors_util.get_node_name(target)
if target_name not in self.timer_result_queue:
self.timer_result_queue.append(target_name)
call = cppast.CallStmt(self.visit(node.value))
target_ref = cppast.DeclRefExpr(target_name)
target_view_ref = cppast.DeclRefExpr(
f"timer_result_{target_name}")
subscript = cppast.ArraySubscriptExpr(
target_view_ref, [cppast.IntegerLiteral(0)])
assign_op = cppast.BinaryOperatorKind.Assign
# Holds the result of the reduction temporarily
temp_ref = cppast.DeclRefExpr("pk_acc")
target_assign = cppast.AssignOperator([target_ref], temp_ref,
assign_op)
view_assign = cppast.AssignOperator([subscript], target_ref,
assign_op)
return cppast.CompoundStmt([call, target_assign, view_assign])
if name in ("BinSort", "BinOp1D", "BinOp3D"):
args: List = node.value.args
# if not isinstance(args[0], ast.Attribute):
# self.error(node.value, "First argument has to be a view")
view = cppast.DeclRefExpr(visitors_util.get_node_name(args[0]))
if view not in self.views:
self.error(args[0], "Undefined view")
view_type: cppast.ClassType = self.views[view]
is_subview: bool = view_type is None
if is_subview:
parent_view = cppast.DeclRefExpr(
self.subviews[view.declname])
view_type = self.views[parent_view]
view_type_str: str = visitors_util.cpp_view_type(view_type)
if name != "BinSort":
dimension: int = 1 if name == "BinOp1D" else 3
cpp_type = cppast.DeclRefExpr(
BinOp.get_type(dimension, view_type_str))
# Do not translate the first argument (view)
constructor = cppast.CallExpr(
cpp_type, [self.visit(a) for a in args[1:]])
else:
bin_op_type: str = f"decltype({visitors_util.get_node_name(args[1])})"
cpp_type = cppast.DeclRefExpr(
BinSort.get_type(view_type_str, bin_op_type))
binsort_args: List[cppast.DeclRefExpr] = [
self.visit(a) for a in args
]
constructor = cppast.CallExpr(cpp_type, binsort_args)
cpp_target: cppast.DeclRefExpr = self.visit(target)
auto_type = cppast.ClassType("auto")
return cppast.DeclStmt(
cppast.VarDecl(auto_type, cpp_target, constructor))
if name in ("get_bin_count", "get_bin_offsets",
"get_permute_vector"):
if not isinstance(target,
ast.Attribute) or target.value.id != "self":
self.error(
node,
"Views defined in pk.main must be an instance variable"
)
cpp_target: str = visitors_util.get_node_name(target)
cpp_device_target = f"pk_d_{cpp_target}"
cpp_target_ref = cppast.DeclRefExpr(cpp_device_target)
sorter: cppast.DeclRefExpr = self.visit(node.value.func.value)
initial_target_ref = cppast.DeclRefExpr(
f"_pk_{cpp_target_ref.declname}")
function = cppast.MemberCallExpr(sorter,
cppast.DeclRefExpr(name), [])
# Add to the dict of declarations made in pk.main
if name == "get_permute_vector":
# This occurs when a workload is executed multiple times
# Initially the view has not been defined in the workload,
# so it needs to be classified as a pkmain_view.
if cpp_target in self.views:
self.views[cpp_target_ref].add_template_param(
cppast.PrimitiveType(cppast.BuiltinType.INT))
return cppast.AssignOperator(
[cpp_target_ref], function,
cppast.BinaryOperatorKind.Assign)
# return f"{cpp_target} = {sorter}.{name}();"
self.pkmain_views[cpp_target_ref] = cppast.ClassType(
"View1D")
else:
self.pkmain_views[cpp_target_ref] = None
auto_type = cppast.ClassType("auto")
decl = cppast.DeclStmt(
cppast.VarDecl(auto_type, initial_target_ref, function))
# resize the workload's vector to match the generated vector
resize_call = cppast.CallStmt(
cppast.CallExpr(cppast.DeclRefExpr("Kokkos::resize"), [
cpp_target_ref,
cppast.MemberCallExpr(initial_target_ref,
cppast.DeclRefExpr("extent"),
[cppast.IntegerLiteral(0)])
]))
copy_call = cppast.CallStmt(
cppast.CallExpr(cppast.DeclRefExpr("Kokkos::deep_copy"),
[cpp_target_ref, initial_target_ref]))
# Assign to the functor after resizing
functor = cppast.DeclRefExpr("pk_f")
functor_access = cppast.MemberExpr(functor, cpp_target)
functor_assign = cppast.AssignOperator(
[functor_access], cpp_target_ref,
cppast.BinaryOperatorKind.Assign)
return cppast.CompoundStmt(
[decl, resize_call, copy_call, functor_assign])
# Assign result of parallel_reduce
if type(target) not in {ast.Name, ast.Subscript
} and target.value.id == "self":
target_name: str = visitors_util.get_node_name(target)
if target_name not in self.reduction_result_queue:
self.reduction_result_queue.append(target_name)
call = cppast.CallStmt(self.visit(node.value))
target_ref = cppast.DeclRefExpr(target_name)
target_view_ref = cppast.DeclRefExpr(
f"reduction_result_{target_name}")
subscript = cppast.ArraySubscriptExpr(target_view_ref,
[cppast.IntegerLiteral(0)])
assign_op = cppast.BinaryOperatorKind.Assign
# Holds the result of the reduction temporarily
temp_ref = cppast.DeclRefExpr("pk_acc")
target_assign = cppast.AssignOperator([target_ref], temp_ref,
assign_op)
view_assign = cppast.AssignOperator([subscript], target_ref,
assign_op)
return cppast.CompoundStmt([call, target_assign, view_assign])
return super().visit_Assign(node)