def build_hard_fusion_kernel(base_loop, fuse_loop, fusion_map, loop_chain_index):
"""
Build AST and :class:`Kernel` for two loops suitable to hard fusion.
The AST consists of three functions: fusion, base, fuse. base and fuse
are respectively the ``base_loop`` and the ``fuse_loop`` kernels, whereas
fusion is the orchestrator that invokes, for each ``base_loop`` iteration,
base and, if still to be executed, fuse.
The orchestrator has the following structure: ::
fusion (buffer, ..., executed):
base (buffer, ...)
for i = 0 to arity:
if not executed[i]:
additional pointer staging required by kernel2
fuse (sub_buffer, ...)
insertion into buffer
The executed array tracks whether the i-th iteration (out of /arity/)
adjacent to the main kernel1 iteration has been executed.
"""
finder = Find((ast.FunDecl, ast.PreprocessNode))
base = base_loop.kernel
base_ast = dcopy(base._ast)
base_info = finder.visit(base_ast)
base_headers = base_info[ast.PreprocessNode]
base_fundecl = base_info[ast.FunDecl]
assert len(base_fundecl) == 1
base_fundecl = base_fundecl[0]
fuse = fuse_loop.kernel
fuse_ast = dcopy(fuse._ast)
fuse_info = finder.visit(fuse_ast)
fuse_headers = fuse_info[ast.PreprocessNode]
fuse_fundecl = fuse_info[ast.FunDecl]
assert len(fuse_fundecl) == 1
fuse_fundecl = fuse_fundecl[0]
# Create /fusion/ arguments and signature
body = ast.Block([])
fusion_name = '%s_%s' % (base_fundecl.name, fuse_fundecl.name)
fusion_args = dcopy(base_fundecl.args + fuse_fundecl.args)
fusion_fundecl = ast.FunDecl(base_fundecl.ret, fusion_name, fusion_args, body)
# Make sure kernel and variable names are unique
base_fundecl.name = "%s_base" % base_fundecl.name
fuse_fundecl.name = "%s_fuse" % fuse_fundecl.name
for i, decl in enumerate(fusion_args):
decl.sym.symbol += '_%d' % i
# Filter out duplicate arguments, and append extra arguments to the fundecl
binding = WeakFilter().kernel_args([base_loop, fuse_loop], fusion_fundecl)
fusion_args += [ast.Decl('int*', 'executed'),
ast.Decl('int*', 'fused_iters'),
ast.Decl('int', 'i')]
# Which args are actually used in /fuse/, but not in /base/ ? The gather for
# such arguments is moved to /fusion/, to avoid usless memory LOADs
base_dats = set(a.data for a in base_loop.args)
fuse_dats = set(a.data for a in fuse_loop.args)
unshared = OrderedDict()
for arg, decl in binding.items():
if arg.data in fuse_dats - base_dats:
unshared.setdefault(decl, arg)
# Track position of Args that need a postponed gather
# Can't track Args themselves as they change across different parloops
fargs = {fusion_args.index(i): ('postponed', False) for i in unshared.keys()}
fargs.update({len(set(binding.values())): ('onlymap', True)})
# Add maps for arguments that need a postponed gather
for decl, arg in unshared.items():
decl_pos = fusion_args.index(decl)
fusion_args[decl_pos].sym.symbol = arg.c_arg_name()
if arg._is_indirect:
fusion_args[decl_pos].sym.rank = ()
fusion_args.insert(decl_pos + 1, ast.Decl('int*', arg.c_map_name(0, 0)))
# Append the invocation of /base/; then, proceed with the invocation
# of the /fuse/ kernels
base_funcall_syms = [binding[a].sym.symbol for a in base_loop.args]
body.children.append(ast.FunCall(base_fundecl.name, *base_funcall_syms))
for idx in range(fusion_map.arity):
fused_iter = ast.Assign('i', ast.Symbol('fused_iters', (idx,)))
fuse_funcall = ast.FunCall(fuse_fundecl.name)
if_cond = ast.Not(ast.Symbol('executed', ('i',)))
if_update = ast.Assign(ast.Symbol('executed', ('i',)), 1)
if_body = ast.Block([fuse_funcall, if_update], open_scope=True)
if_exec = ast.If(if_cond, [if_body])
body.children.extend([ast.FlatBlock('\n'), fused_iter, if_exec])
# Modify the /fuse/ kernel
# This is to take into account that many arguments are shared with
# /base/, so they will only staged once for /base/. This requires
# tweaking the way the arguments are declared and accessed in /fuse/.
# For example, the shared incremented array (called /buffer/ in
# the pseudocode in the comment above) now needs to take offsets
# to be sure the locations that /base/ is supposed to increment are
# actually accessed. The same concept apply to indirect arguments.
init = lambda v: '{%s}' % ', '.join([str(j) for j in v])
for i, fuse_loop_arg in enumerate(fuse_loop.args):
fuse_kernel_arg = binding[fuse_loop_arg]
buffer_name = '%s_vec' % fuse_kernel_arg.sym.symbol
fuse_funcall_sym = ast.Symbol(buffer_name)
# What kind of temporaries do we need ?
if fuse_loop_arg.access == INC:
op, lvalue, rvalue = ast.Incr, fuse_kernel_arg.sym.symbol, buffer_name
stager = lambda b, l: b.children.extend(l)
indexer = lambda indices: [(k, j) for j, k in enumerate(indices)]
pointers = []
elif fuse_loop_arg.access == READ:
op, lvalue, rvalue = ast.Assign, buffer_name, fuse_kernel_arg.sym.symbol
stager = lambda b, l: [b.children.insert(0, j) for j in reversed(l)]
indexer = lambda indices: [(j, k) for j, k in enumerate(indices)]
pointers = list(fuse_kernel_arg.pointers)
# Now gonna handle arguments depending on their type and rank ...
if fuse_loop_arg._is_global:
# ... Handle global arguments. These can be dropped in the
# kernel without any particular fiddling
fuse_funcall_sym = ast.Symbol(fuse_kernel_arg.sym.symbol)
elif fuse_kernel_arg in unshared:
# ... Handle arguments that appear only in /fuse/
staging = unshared[fuse_kernel_arg].c_vec_init(False).split('\n')
rvalues = [ast.FlatBlock(j.split('=')[1]) for j in staging]
lvalues = [ast.Symbol(buffer_name, (j,)) for j in range(len(staging))]
staging = [ast.Assign(j, k) for j, k in zip(lvalues, rvalues)]
# Set up the temporary
buffer_symbol = ast.Symbol(buffer_name, (len(staging),))
buffer_decl = ast.Decl(fuse_kernel_arg.typ, buffer_symbol,
qualifiers=fuse_kernel_arg.qual,
pointers=list(pointers))
# Update the if-then AST body
stager(if_exec.children[0], staging)
if_exec.children[0].children.insert(0, buffer_decl)
elif fuse_loop_arg._is_mat:
# ... Handle Mats
staging = []
for b in fused_inc_arg._block_shape:
for rc in b:
lvalue = ast.Symbol(lvalue, (idx, idx),
((rc[0], 'j'), (rc[1], 'k')))
rvalue = ast.Symbol(rvalue, ('j', 'k'))
staging = ItSpace(mode=0).to_for([(0, rc[0]), (0, rc[1])],
('j', 'k'),
[op(lvalue, rvalue)])[:1]
# Set up the temporary
buffer_symbol = ast.Symbol(buffer_name, (fuse_kernel_arg.sym.rank,))
buffer_init = ast.ArrayInit(init([init([0.0])]))
buffer_decl = ast.Decl(fuse_kernel_arg.typ, buffer_symbol, buffer_init,
qualifiers=fuse_kernel_arg.qual, pointers=pointers)
# Update the if-then AST body
stager(if_exec.children[0], staging)
if_exec.children[0].children.insert(0, buffer_decl)
elif fuse_loop_arg._is_indirect:
cdim = fuse_loop_arg.data.cdim
if cdim == 1 and fuse_kernel_arg.sym.rank:
# [Special case]
# ... Handle rank 1 indirect arguments that appear in both
# /base/ and /fuse/: just point into the right location
rank = (idx,) if fusion_map.arity > 1 else ()
fuse_funcall_sym = ast.Symbol(fuse_kernel_arg.sym.symbol, rank)
else:
# ... Handle indirect arguments. At the C level, these arguments
# are of pointer type, so simple pointer arithmetic is used
# to ensure the kernel accesses are to the correct locations
fuse_arity = fuse_loop_arg.map.arity
base_arity = fuse_arity*fusion_map.arity
size = fuse_arity*cdim
# Set the proper storage layout before invoking /fuse/
ofs_vals = [[base_arity*j + k for k in range(fuse_arity)]
for j in range(cdim)]
ofs_vals = [[fuse_arity*j + k for k in flatten(ofs_vals)]
for j in range(fusion_map.arity)]
ofs_vals = list(flatten(ofs_vals))
indices = [ofs_vals[idx*size + j] for j in range(size)]
staging = [op(ast.Symbol(lvalue, (j,)), ast.Symbol(rvalue, (k,)))
for j, k in indexer(indices)]
# Set up the temporary
buffer_symbol = ast.Symbol(buffer_name, (size,))
if fuse_loop_arg.access == INC:
buffer_init = ast.ArrayInit(init([0.0]))
else:
buffer_init = ast.EmptyStatement()
pointers.pop()
buffer_decl = ast.Decl(fuse_kernel_arg.typ, buffer_symbol, buffer_init,
qualifiers=fuse_kernel_arg.qual,
pointers=pointers)
# Update the if-then AST body
stager(if_exec.children[0], staging)
if_exec.children[0].children.insert(0, buffer_decl)
else:
# Nothing special to do for direct arguments
pass
# Finally update the /fuse/ funcall
fuse_funcall.children.append(fuse_funcall_sym)
fused_headers = set([str(h) for h in base_headers + fuse_headers])
fused_ast = ast.Root([ast.PreprocessNode(h) for h in fused_headers] +
[base_fundecl, fuse_fundecl, fusion_fundecl])
return Kernel([base, fuse], fused_ast, loop_chain_index), fargs
