def get_kernel_params(
members: PyKokkosMembers,
is_hierarchical: bool,
is_workload: bool,
real: Optional[str]
) -> Dict[str, str]:
"""
Get the parameters of the kernel. The parameters include the fields, the views,
the views holding the reduction results, and the view holding the timer results.
Also add parameters for the parameters of the execution policy.
:param members: an object containing the fields and views
:param is_hierarchical: does the workunit use hierarchical parallelism
:param real: the precision for which to generate a binding
:returns: a dict mapping from argument name to type
"""
s = cppast.Serializer()
params: Dict[str, str] = {}
for n, t in members.fields.items():
params[n.declname] = s.serialize(t)
for n, t in members.views.items():
# skip subviews
if t is None:
continue
layout: str = f"{Keywords.DefaultExecSpace.value}::array_layout"
params[n.declname] = cpp_view_type(t, space=Keywords.ArgMemSpace.value, layout=layout, real=real)
if not is_workload:
params[Keywords.KernelName.value] = "const std::string&"
if is_hierarchical:
params[Keywords.LeagueSize.value] = "int"
params[Keywords.TeamSize.value] = "int"
params[Keywords.VectorLength.value] = "int"
else:
params[Keywords.ThreadsBegin.value] = "int"
params[Keywords.ThreadsEnd.value] = "int"
for result in members.reduction_result_queue:
view_name = f"reduction_result_{result}"
view_type = cppast.ClassType("View1D")
view_type.add_template_param(cppast.DeclRefExpr("double"))
view_type.add_template_param(cppast.DeclRefExpr("HostSpace"))
params[view_name] = cpp_view_type(view_type, space="Kokkos::HostSpace", layout="Kokkos::LayoutRight")
for result in members.timer_result_queue:
view_name = f"timer_result_{result}"
view_type = cppast.ClassType("View1D")
view_type.add_template_param(cppast.DeclRefExpr("double"))
view_type.add_template_param(cppast.DeclRefExpr("HostSpace"))
params[view_name] = cpp_view_type(view_type, space="Kokkos::HostSpace", layout="Kokkos::LayoutRight")
return params
def generate_functor(
name: str,
members: PyKokkosMembers,
workunits: Dict[cppast.DeclRefExpr, Tuple[str, cppast.MethodDecl]],
functions: List[cppast.MethodDecl],
) -> cppast.RecordDecl:
"""
Generate the functor source
:param name: the functor class name
:param members: an object containing the fields and views
:param workunits: a dict mapping from workunit name to a tuple of operation type and source
:param functions: a list of KOKKOS_FUNCTIONS defined in the functor
:param has_real: whether the function contains a pk.real datatype
:returns: the cppast representation of the functor
"""
fields: Dict[cppast.DeclRefExpr, cppast.PrimitiveType] = members.fields
views: Dict[cppast.DeclRefExpr, cppast.ClassType] = members.views
decls: List[cppast.DeclStmt] = []
# Create the tags needed to call individual workunits. Tags in Kokkos are empty structs.
for n in workunits:
tag = cppast.RecordDecl(cppast.ClassType(n.declname), [])
tag.is_definition = True
decls.append(tag)
for n, t in fields.items():
decls.append(cppast.DeclStmt(cppast.FieldDecl(t, n)))
for n, t in views.items():
# skip subviews
if t is None:
continue
view_type: str = get_view_type(t)
decls.append(cppast.DeclStmt(cppast.FieldDecl(view_type, n)))
decls.append(generate_constructor(name, fields, views))
for _, s in workunits.values():
decls.append(s)
for f in functions:
decls.append(cppast.DeclStmt(f))
struct = cppast.RecordDecl(cppast.ClassType(name), decls)
struct.is_definition = True
if members.has_real:
struct.add_template_param(Keywords.RealPrecision.value)
s = cppast.Serializer()
return struct
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_FunctionDef(
self, node: ast.FunctionDef
) -> Union[cppast.ConstructorDecl, cppast.MethodDecl]:
name: str = node.name
return_type: Optional[cppast.ClassType]
if name == "__init__":
name = node.parent.name
return_type = None
elif self.is_void_function(node):
return_type = cppast.ClassType("void")
else:
return_type = visitors_util.get_type(node.returns, self.pk_import)
if len(node.args.args) == 0 or node.args.args[0].arg != "self":
self.error(node, "Static functions are not supported")
params: List[cppast.ParmVarDecl] = self.visit(node.args)
body = cppast.CompoundStmt([self.visit(b) for b in node.body])
attributes: str = "KOKKOS_FUNCTION"
if return_type is None:
return cppast.ConstructorDecl(attributes, name, params, body)
else:
return cppast.MethodDecl(attributes, return_type, name, params,
body)
def visit_ClassDef(self, node: ast.ClassDef) -> cppast.RecordDecl:
name: str = node.name
# Add class as allowed type
visitors_util.allowed_types[name] = name
member_variables: Dict[cppast.DeclRefExpr,
cppast.Type] = self.get_member_variables(node)
decls: List[cppast.DeclStmt] = []
if len(member_variables) == 0:
self.error(node,
"Missing constructor or no member variables detected")
for n, t in member_variables.items():
decls.append(cppast.DeclStmt(cppast.FieldDecl(t, n)))
for b in node.body:
decls.append(self.visit(b))
if not self.has_default_constructor(node):
decls.append(
cppast.ConstructorDecl("KOKKOS_FUNCTION", name, [], None))
classdef = cppast.RecordDecl(cppast.ClassType(name), decls)
classdef.is_definition = True
return classdef
def visit_Subscript(self, node: ast.Subscript) -> Union[cppast.ArraySubscriptExpr, cppast.CallExpr]:
current_node: ast.Subscript = node
slices: List = []
dim: int = 0
while isinstance(current_node, ast.Subscript):
index = current_node.slice
if sys.version_info.minor <= 8:
# In Python >= 3.9, ast.Index is deprecated
# (see # https://docs.python.org/3/whatsnew/3.9.html)
# Instead of ast.Index, value will be used directly
if not isinstance(index, ast.Index):
self.error(
current_node, "Slices not supported, use simple indices")
slices.insert(0, index)
current_node = current_node.value
dim += 1
name: str = visitors_util.get_node_name(current_node)
ref = cppast.DeclRefExpr(name)
if ref not in self.views and name not in self.lists:
self.error(current_node, "Unknown view or list")
dim_map: List = [cppast.ClassType("View1D"),
cppast.ClassType("View2D"),
cppast.ClassType("View3D"),
cppast.ClassType("View4D"),
cppast.ClassType("View5D"),
cppast.ClassType("View6D"),
cppast.ClassType("View7D"),
cppast.ClassType("View8D")]
if name in self.lists:
indices: List[cppast.Expr] = [self.visit(s) for s in slices]
subscript = cppast.ArraySubscriptExpr(ref, indices)
return subscript
if (
ref in self.views
and (
self.views[ref] is None # For views added in @pk.main
or self.views[ref].typename == dim_map[dim - 1].typename
)
):
args: List[cppast.Expr] = [self.visit(s) for s in slices]
subscript = cppast.CallExpr(ref, args)
return subscript
self.error(node, f"'{name}' is not a View{dim}D")
def visit_FunctionDef(
self, node: ast.FunctionDef
) -> Union[str, Tuple[str, cppast.MethodDecl]]:
if self.is_nested_call(node):
params: List[cppast.ParmVarDecl] = [
a for a in self.visit(node.args)
]
body = cppast.CompoundStmt([self.visit(b) for b in node.body])
workunit = cppast.LambdaExpr("[&]", params, body)
self.nested_work_units[node.name] = workunit
return ""
else:
operation: Optional[str] = self.get_operation_type(node)
if operation is None:
self.error(node.args, "Incorrect types in workunit definition")
tag_type = cppast.ClassType(f"const {node.name}")
tag_type.is_reference = True
tag = cppast.ParmVarDecl(tag_type, cppast.DeclRefExpr(""))
params: List[cppast.ParmVarDecl] = [tag]
params.extend(self.visit(node.args))
body = cppast.CompoundStmt([self.visit(b) for b in node.body])
attributes: str = "KOKKOS_FUNCTION"
decltype = cppast.ClassType("void")
declname: str = "operator()"
method = cppast.MethodDecl(attributes, decltype, declname, params,
body)
method.is_const = True
return (operation, method)
def visit_FunctionDef(self, node: ast.FunctionDef) -> cppast.MethodDecl:
if not self.is_valid_kokkos_function(node):
self.error(node, "Invalid Kokkos function")
return_type: cppast.ClassType
if self.is_void_function(node):
return_type = cppast.ClassType("void")
else:
return_type = visitors_util.get_type(node.returns, self.pk_import)
if return_type is None:
self.error(node, "Return type is not supported for translation")
params: List[cppast.ParmVarDecl] = self.visit(node.args)
name: str = node.name
body = cppast.CompoundStmt([self.visit(b) for b in node.body])
attributes: str = "KOKKOS_FUNCTION"
method = cppast.MethodDecl(attributes, return_type, name, params, body)
method.is_const = True
return method
def get_type(annotation: Union[ast.Attribute, ast.Name, ast.Subscript],
pk_import: str) -> Optional[cppast.Type]:
if isinstance(annotation, ast.Attribute):
if annotation.value.id == pk_import:
type_name: str = get_node_name(annotation)
if type_name in view_dtypes:
return cppast.PrimitiveType(view_dtypes[type_name])
if type_name in allowed_types:
type_name = allowed_types[type_name]
return cppast.ClassType(type_name)
if isinstance(annotation, ast.Name):
type_name: str = annotation.id
if type_name == "int":
return cppast.PrimitiveType(cppast.BuiltinType.INT)
if type_name == "float":
return cppast.PrimitiveType(cppast.BuiltinType.DOUBLE)
if type_name == "bool":
return cppast.PrimitiveType(cppast.BuiltinType.BOOL)
if type_name in allowed_types:
return cppast.ClassType(type_name)
if type_name == "List":
return None
if isinstance(annotation, ast.Subscript):
value: Union[ast.Name, ast.Attribute] = annotation.value
subscript: ast.Index = annotation.slice
id: str = ""
if isinstance(value, ast.Name):
id = value.id
elif isinstance(value, ast.Attribute):
id = value.value.id
if id == "List":
if sys.version_info.minor <= 8:
# In Python >= 3.9, ast.Index is deprecated
# (see # https://docs.python.org/3/whatsnew/3.9.html)
value = subscript.value
else:
value = subscript
member_type: cppast.Type = get_type(value, pk_import)
return member_type
if id == pk_import:
type_name: str = get_node_name(value)
if sys.version_info.minor <= 8:
# In Python >= 3.9, ast.Index is deprecated
# (see # https://docs.python.org/3/whatsnew/3.9.html)
dtype_node: ast.Attribute = subscript.value
else:
dtype_node: ast.Attribute = subscript
if type_name == "Acc":
return get_type(dtype_node, pk_import)
dtype: cppast.PrimitiveType = get_type(dtype_node, pk_import)
if dtype is None:
return None
view_type = cppast.ClassType(type_name)
view_type.add_template_param(dtype)
return view_type
return None
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)
