def optimize_for_gpu(sdfg: dace.SDFG, m: int, n: int, k: int):
""" Optimize the matrix multiplication example for GPUs. """
# Ensure integers are 32-bit by default
dace.Config.set('compiler', 'default_data_types', value='C')
# Fuse the map and reduce nodes
sdfg.apply_transformations(MapReduceFusion)
# Apply GPU transformation
sdfg.apply_gpu_transformations()
# Find multiplication map
entry = find_map_by_param(sdfg, 'k')
# Create a tiling strategy
divides_evenly = (m % 64 == 0) and (n % 64 == 0) and (k % 8 == 0)
xfutil.tile(sdfg, entry, divides_evenly, True, i=64, j=64, k=8)
xfutil.tile(sdfg, entry, divides_evenly, True, i=8, j=4)
# Create kernel schedule by collapsing and reordering maps
gtile_i = find_map_by_param(sdfg, 'tile_i')
gtile_j = find_map_by_param(sdfg, 'tile_j')
btile_i = find_map_by_param(sdfg, 'tile1_i')
btile_j = find_map_by_param(sdfg, 'tile1_j')
MapCollapse.apply_to(sdfg, outer_map_entry=gtile_i, inner_map_entry=gtile_j, permissive=True)
MapCollapse.apply_to(sdfg, outer_map_entry=btile_i, inner_map_entry=btile_j, permissive=True)
btile = find_map_by_param(sdfg, 'tile1_i')
btile.map.schedule = dace.ScheduleType.GPU_ThreadBlock
# Add local storage (shared memory) for A and B on GPU
ktile = find_map_by_param(sdfg, 'tile_k')
smem_a = InLocalStorage.apply_to(sdfg, dict(array='A'), node_a=ktile, node_b=btile)
smem_b = InLocalStorage.apply_to(sdfg, dict(array='B'), node_a=ktile, node_b=btile)
sdfg.arrays[smem_a.data].storage = dace.StorageType.GPU_Shared
sdfg.arrays[smem_b.data].storage = dace.StorageType.GPU_Shared
# Add local storage (registers) for A and B
ttile = find_map_by_param(sdfg, 'k')
warptile, ttile = xfutil.extract_map_dims(sdfg, ttile, [2])
InLocalStorage.apply_to(sdfg, dict(array='trans_gpu_A'), node_a=warptile, node_b=ttile)
InLocalStorage.apply_to(sdfg, dict(array='trans_gpu_B'), node_a=warptile, node_b=ttile)
# Add local storage (registers) for C
state = next(s for s in sdfg.nodes() if warptile in s.nodes())
warptile_exit = state.exit_node(warptile)
btile_exit = state.exit_node(btile)
AccumulateTransient.apply_to(sdfg, map_exit=warptile_exit, outer_map_exit=btile_exit)
# Set C tile to zero on allocation
c_access = next(n for n in state.data_nodes() if n.data == 'trans_gpu_C')
c_access.setzero = True
# Unroll microkernel maps
ttile.map.unroll = True
# Apply double-buffering on shared memory
DoubleBuffering.apply_to(sdfg, map_entry=ktile, transient=smem_a)
def program_for_node(program, sdfg: SDFG, state: SDFGState,
node: onnx_op.ONNXOp) -> SDFG:
""" Expand a function to a dace program.
The dtypes for the arguments will be extracted by matching the parameter names to edges.
"""
input_names = node.schema.non_variadic_inputs()
variadic_input_names = node.schema.variadic_inputs()
output_names = node.schema.non_variadic_outputs()
variadic_output_names = node.schema.variadic_outputs()
if set(input_names).intersection(output_names):
# this is currently the case for only one onnx op
raise ValueError(
"program_for_node cannot be applied on nodes of this type;"
" '{}' is both an input and an output".format(
next(input_names.intersection(output_names))))
params = inspect.signature(program).parameters
annotations = {}
for name, param in params.items():
if name in input_names or ("__" in name
and parse_variadic_param(name)[0]
in variadic_input_names):
annotations[name] = in_desc_with_name(node, state, sdfg, name)
elif name in output_names or ("__" in name
and parse_variadic_param(name)[0]
in variadic_output_names):
annotations[name] = out_desc_with_name(node, state, sdfg, name)
else:
raise ValueError(
"'{}' was not found as an input or output for {}".format(
name, node.schema.name))
program.__annotations__ = annotations
result = DaceProgram(program, (), {}, False, dace.DeviceType.CPU)
result.name = node.label + "_expansion"
sdfg = result.to_sdfg()
if node.schedule in [dtypes.ScheduleType.GPU_Default
] + dtypes.GPU_SCHEDULES:
sdfg.apply_gpu_transformations()
return sdfg
