def apply(self, sdfg: SDFG):
graph = sdfg.nodes()[self.state_id]
map_entry = graph.nodes()[self.subgraph[Vectorization._map_entry]]
tasklet: nodes.Tasklet = graph.successors(map_entry)[0]
param = symbolic.pystr_to_symbolic(map_entry.map.params[-1])
# Create new vector size.
vector_size = self.vector_len
dim_from, dim_to, dim_skip = map_entry.map.range[-1]
# Determine whether to create preamble or postamble maps
if self.preamble is not None:
create_preamble = self.preamble
else:
create_preamble = not ((dim_from % vector_size == 0) == True
or dim_from == 0)
if self.postamble is not None:
create_postamble = self.postamble
else:
if isinstance(dim_to, symbolic.SymExpr):
create_postamble = (((dim_to.approx + 1) %
vector_size == 0) == False)
else:
create_postamble = (((dim_to + 1) % vector_size == 0) == False)
# Determine new range for vectorized map
if self.strided_map:
new_range = [dim_from, dim_to - vector_size + 1, vector_size]
else:
new_range = [
dim_from // vector_size, ((dim_to + 1) // vector_size) - 1,
dim_skip
]
# Create preamble non-vectorized map (replacing the original map)
if create_preamble:
old_scope = graph.scope_subgraph(map_entry, True, True)
new_scope: ScopeSubgraphView = replicate_scope(
sdfg, graph, old_scope)
new_begin = dim_from + (vector_size - (dim_from % vector_size))
map_entry.map.range[-1] = (dim_from, new_begin - 1, dim_skip)
# Replace map_entry with the replicated scope (so that the preamble
# will usually come first in topological sort)
map_entry = new_scope.entry
tasklet = new_scope.nodes()[old_scope.nodes().index(tasklet)]
new_range[0] = new_begin
# Create postamble non-vectorized map
if create_postamble:
new_scope: ScopeSubgraphView = replicate_scope(
sdfg, graph, graph.scope_subgraph(map_entry, True, True))
dim_to_ex = dim_to + 1
new_scope.entry.map.range[-1] = (dim_to_ex -
(dim_to_ex % vector_size), dim_to,
dim_skip)
# Change the step of the inner-most dimension.
map_entry.map.range[-1] = tuple(new_range)
# Vectorize connectors adjacent to the tasklet.
for edge in graph.all_edges(tasklet):
connectors = (tasklet.in_connectors
if edge.dst == tasklet else tasklet.out_connectors)
conn = edge.dst_conn if edge.dst == tasklet else edge.src_conn
if edge.data.data is None: # Empty memlets
continue
desc = sdfg.arrays[edge.data.data]
contigidx = desc.strides.index(1)
newlist = []
lastindex = edge.data.subset[contigidx]
if isinstance(lastindex, tuple):
newlist = [(rb, re, rs) for rb, re, rs in edge.data.subset]
symbols = set()
for indd in lastindex:
symbols.update(
symbolic.pystr_to_symbolic(indd).free_symbols)
else:
newlist = [(rb, rb, 1) for rb in edge.data.subset]
symbols = symbolic.pystr_to_symbolic(lastindex).free_symbols
oldtype = connectors[conn]
if oldtype is None or oldtype.type is None:
oldtype = desc.dtype
# Vector to scalar WCR edge: change connector and continue
lastedge = graph.memlet_path(edge)[-1]
if (lastedge.data.subset.num_elements() == 1
and edge.data.wcr is not None):
connectors[conn] = dtypes.vector(oldtype, vector_size)
continue
if str(param) not in map(str, symbols):
continue
# Vectorize connector, if not already vectorized
if isinstance(oldtype, dtypes.vector):
continue
connectors[conn] = dtypes.vector(oldtype, vector_size)
# Modify memlet subset to match vector length
if self.strided_map:
rb = newlist[contigidx][0]
if self.propagate_parent:
newlist[contigidx] = (rb / self.vector_len,
rb / self.vector_len, 1)
else:
newlist[contigidx] = (rb, rb + self.vector_len - 1, 1)
else:
rb = newlist[contigidx][0]
if self.propagate_parent:
newlist[contigidx] = (rb, rb, 1)
else:
newlist[contigidx] = (self.vector_len * rb,
self.vector_len * rb +
self.vector_len - 1, 1)
edge.data.subset = subsets.Range(newlist)
edge.data.volume = vector_size
# Vector length propagation using data descriptors, recursive traversal
# outwards
if self.propagate_parent:
for edge in graph.all_edges(tasklet):
cursdfg = sdfg
curedge = edge
while cursdfg is not None:
arrname = curedge.data.data
dtype = cursdfg.arrays[arrname].dtype
# Change type and shape to vector
if not isinstance(dtype, dtypes.vector):
cursdfg.arrays[arrname].dtype = dtypes.vector(
dtype, vector_size)
new_shape = list(cursdfg.arrays[arrname].shape)
contigidx = cursdfg.arrays[arrname].strides.index(1)
new_shape[contigidx] /= vector_size
try:
new_shape[contigidx] = int(new_shape[contigidx])
except TypeError:
pass
cursdfg.arrays[arrname].shape = new_shape
propagation.propagate_memlets_sdfg(cursdfg)
# Find matching edge in parent
nsdfg = cursdfg.parent_nsdfg_node
if nsdfg is None:
break
tstate = cursdfg.parent
curedge = ([
e
for e in tstate.in_edges(nsdfg) if e.dst_conn == arrname
] + [
e for e in tstate.out_edges(nsdfg)
if e.src_conn == arrname
])[0]
cursdfg = cursdfg.parent_sdfg
def apply(self, graph: SDFGState, sdfg: SDFG) -> nodes.MapEntry:
me = self.mapentry
# Add new map within map
mx = graph.exit_node(me)
new_me, new_mx = graph.add_map('warp_tile',
dict(__tid=f'0:{self.warp_size}'),
dtypes.ScheduleType.GPU_ThreadBlock)
__tid = symbolic.pystr_to_symbolic('__tid')
for e in graph.out_edges(me):
xfh.reconnect_edge_through_map(graph, e, new_me, True)
for e in graph.in_edges(mx):
xfh.reconnect_edge_through_map(graph, e, new_mx, False)
# Stride and offset all internal maps
maps_to_stride = xfh.get_internal_scopes(graph, new_me, immediate=True)
for nstate, nmap in maps_to_stride:
nsdfg = nstate.parent
nsdfg_node = nsdfg.parent_nsdfg_node
# Map cannot be partitioned across a warp
if (nmap.range.size()[-1] < self.warp_size) == True:
continue
if nsdfg is not sdfg and nsdfg_node is not None:
nsdfg_node.symbol_mapping['__tid'] = __tid
if '__tid' not in nsdfg.symbols:
nsdfg.add_symbol('__tid', dtypes.int32)
nmap.range[-1] = (nmap.range[-1][0], nmap.range[-1][1] - __tid,
nmap.range[-1][2] * self.warp_size)
subgraph = nstate.scope_subgraph(nmap)
subgraph.replace(nmap.params[-1], f'{nmap.params[-1]} + __tid')
inner_map_exit = nstate.exit_node(nmap)
# If requested, replicate maps with multiple dependent maps
if self.replicate_maps:
destinations = [
nstate.memlet_path(edge)[-1].dst
for edge in nstate.out_edges(inner_map_exit)
]
for dst in destinations:
# Transformation will not replicate map with more than one
# output
if len(destinations) != 1:
break
if not isinstance(dst, nodes.AccessNode):
continue # Not leading to access node
if not xfh.contained_in(nstate, dst, new_me):
continue # Memlet path goes out of map
if not nsdfg.arrays[dst.data].transient:
continue # Cannot modify non-transients
for edge in nstate.out_edges(dst)[1:]:
rep_subgraph = xfh.replicate_scope(
nsdfg, nstate, subgraph)
rep_edge = nstate.out_edges(
rep_subgraph.sink_nodes()[0])[0]
# Add copy of data
newdesc = copy.deepcopy(sdfg.arrays[dst.data])
newname = nsdfg.add_datadesc(dst.data,
newdesc,
find_new_name=True)
newaccess = nstate.add_access(newname)
# Redirect edges
xfh.redirect_edge(nstate,
rep_edge,
new_dst=newaccess,
new_data=newname)
xfh.redirect_edge(nstate,
edge,
new_src=newaccess,
new_data=newname)
# If has WCR, add warp-collaborative reduction on outputs
for out_edge in nstate.out_edges(inner_map_exit):
dst = nstate.memlet_path(out_edge)[-1].dst
if not xfh.contained_in(nstate, dst, new_me):
# Skip edges going out of map
continue
if dst.desc(nsdfg).storage == dtypes.StorageType.GPU_Global:
# Skip shared memory
continue
if out_edge.data.wcr is not None:
ctype = nsdfg.arrays[out_edge.data.data].dtype.ctype
redtype = detect_reduction_type(out_edge.data.wcr)
if redtype == dtypes.ReductionType.Custom:
raise NotImplementedError
credtype = ('dace::ReductionType::' +
str(redtype)[str(redtype).find('.') + 1:])
# One element: tasklet
if out_edge.data.subset.num_elements() == 1:
# Add local access between thread-local and warp reduction
name = nsdfg._find_new_name(out_edge.data.data)
nsdfg.add_scalar(
name,
nsdfg.arrays[out_edge.data.data].dtype,
transient=True)
# Initialize thread-local to global value
read = nstate.add_read(out_edge.data.data)
write = nstate.add_write(name)
edge = nstate.add_nedge(read, write,
copy.deepcopy(out_edge.data))
edge.data.wcr = None
xfh.state_fission(nsdfg,
SubgraphView(nstate, [read, write]))
newnode = nstate.add_access(name)
nstate.remove_edge(out_edge)
edge = nstate.add_edge(out_edge.src, out_edge.src_conn,
newnode, None,
copy.deepcopy(out_edge.data))
for e in nstate.memlet_path(edge):
e.data.data = name
e.data.subset = subsets.Range([(0, 0, 1)])
wrt = nstate.add_tasklet(
'warpreduce', {'__a'}, {'__out'},
f'__out = dace::warpReduce<{credtype}, {ctype}>::reduce(__a);',
dtypes.Language.CPP)
nstate.add_edge(newnode, None, wrt, '__a',
Memlet(name))
out_edge.data.wcr = None
nstate.add_edge(wrt, '__out', out_edge.dst, None,
out_edge.data)
else: # More than one element: mapped tasklet
# Could be a parallel summation
# TODO(later): Check if reduction
continue
# End of WCR to warp reduction
# Make nested SDFG out of new scope
xfh.nest_state_subgraph(sdfg, graph,
graph.scope_subgraph(new_me, False, False))
return new_me
