def expansion(node: 'Allreduce', state: SDFGState, sdfg: SDFG, **kwargs):
node.validate(sdfg, state)
for edge in state.in_edges(node):
if edge.dst_conn == '_inbuffer':
input_edge = edge
for edge in state.out_edges(node):
if edge.src_conn == '_outbuffer':
output_edge = edge
input_dims = input_edge.data.subset.size_exact()
output_dims = output_edge.data.subset.size_exact()
input_data = sdfg.arrays[input_edge.data.data]
output_data = sdfg.arrays[output_edge.data.data]
# Verify that data is on the GPU
if input_data.storage is not dtypes.StorageType.GPU_Global:
raise ValueError('Input of NCCL Send must reside '
' in global GPU memory.')
if output_data.storage is not dtypes.StorageType.GPU_Global:
raise ValueError('Output of NCCL Recv must reside '
' in global GPU memory.')
redtype = node.reduction_type
redtype = nutil.NCCL_SUPPORTED_OPERATIONS[redtype]
wcr_str = str(redtype)
wcr_str = wcr_str[wcr_str.find('.') + 1:] # Skip "NcclReductionType."
nccl_dtype_str = nutil.Nccl_dtypes(input_data.dtype.base_type)
count_str = "*".join(str(e) for e in input_dims)
if input_data.dtype.veclen > 1:
raise (NotImplementedError)
code = f"""ncclAllReduce(_inbuffer, _outbuffer, {count_str}, {nccl_dtype_str}, {wcr_str}, __state->ncclCommunicators->at(__dace_cuda_device), __dace_current_stream)"""
if Config.get('compiler', 'build_type') == 'Debug':
code = '''DACE_NCCL_CHECK(''' + code + ''');\n'''
else:
code = code + ''';\n'''
if Config.get_bool('debugprint'):
code = (
f'''printf("{str(node)}: begin; dev: %d\\n", __dace_cuda_device);\n'''
+ code +
f'''printf("{str(node)}: end; dev: %d\\n\\n", __dace_cuda_device);\n'''
)
code += """\ncudaStreamSynchronize(__dace_current_stream);"""
tasklet = nodes.Tasklet(node.name + "_" + wcr_str,
node.in_connectors,
node.out_connectors,
code,
location=node.location,
language=dtypes.Language.CPP)
return tasklet
def expansion(node: 'Send', state: SDFGState, sdfg: SDFG, **kwargs):
node.validate(sdfg, state)
for edge in state.in_edges(node):
if edge.dst_conn == '_inbuffer':
input_edge = edge
input_dims = input_edge.data.subset.size_exact()
input_data = sdfg.arrays[input_edge.data.data]
# Verify that data is on the GPU
if input_data.storage is not dtypes.StorageType.GPU_Global:
raise ValueError('Input of NCCL Send must reside '
' in global GPU memory.')
peer = node.peer
peerstr = str(peer)
for fs in peer.free_symbols:
if fs.name in sdfg.arrays:
sdfg.arrays[fs.name].lifetime = dtypes.AllocationLifetime.SDFG
if fs.name in sdfg.parent_sdfg.arrays:
sdfg.parent_sdfg.arrays[
fs.name].lifetime = dtypes.AllocationLifetime.SDFG
nccl_dtype_str = nutil.Nccl_dtypes(input_data.dtype.base_type)
count_str = "*".join(str(e) for e in input_dims)
if input_data.dtype.veclen > 1:
raise (NotImplementedError)
code = f"""ncclSend(_inbuffer, {count_str}, {nccl_dtype_str}, {peerstr}, __state->ncclCommunicators->at(__dace_cuda_device), __dace_current_stream)"""
if Config.get('compiler', 'build_type') == 'Debug':
code = '''DACE_NCCL_CHECK(''' + code + ''');\n'''
else:
code = code + ''';\n'''
if Config.get_bool('debugprint'):
code = (
f'''printf("{str(node)}: begin; dev,peer: %d, %d\\n", __dace_cuda_device, {peerstr});\n'''
+ code +
f'''printf("{str(node)}: end; dev,peer: %d, %d\\n\\n", __dace_cuda_device, {peerstr});\n'''
)
code = nutil.aggregate_calls(sdfg, state, node, code)
tasklet = nodes.Tasklet(node.name,
node.in_connectors,
node.out_connectors,
code,
location=node.location,
language=dtypes.Language.CPP,
library_expansion_symbols=set(
map(str, peer.free_symbols)))
return tasklet
def tmp_location(name: str) -> str:
""" Returns the absolute path to the temporary folder
:param name: name of the SDFG
:return: path to the tmp folder
"""
build_folder = Config.get('default_build_folder')
return os.path.abspath(os.path.join(build_folder, name, "map", "tmp.json"))
def create_cpp_map(code: str, name: str, target_name: str, build_folder: str,
sourceFiles: [str], made_with_api: bool):
""" Creates the mapping from the SDFG nodes to the C++ code lines.
The mapping gets saved at: <SDFG build folder>/map/map_cpp.json
:param code: C++ code containing the identifiers '////__DACE:0:0:0'
:param name: The name of the SDFG
:param target_name: The target type, example: 'cpu'
:param build_folder: The build_folder of the SDFG
:param sourceFiles: A list of source files of to the SDFG
:param made_with_api: true if the SDFG was created just with the API
"""
codegen_debug = Config.get_bool('compiler', 'codegen_lineinfo')
cpp_mapper = MapCpp(code, name, target_name)
cpp_mapper.mapper(codegen_debug)
folder = save("cpp", name, cpp_mapper.map, build_folder)
if codegen_debug:
save("codegen", name, cpp_mapper.codegen_map, build_folder)
# Send information about the SDFG to VSCode
send({
"type": "registerFunction",
"name": name,
"path_cache": folder,
"path_file": sourceFiles,
"target_name": target_name,
"made_with_api": made_with_api,
"codegen_map": codegen_debug
})
def create_py_map(sdfg):
""" Creates the mapping from the python source lines to the SDFG nodes.
The mapping gets saved at: <SDFG build folder>/map/map_py.json
:param sdfg: The SDFG for which the mapping will be created
"""
# If the cache setting is set to 'hash' then we don't create a
# mapping as we don't know where to save it to due to
# the hash value changing every time the state of the SDFG changes.
if Config.get('cache') != 'hash':
tmp = get_tmp(sdfg.name)
py_mapper = MapPython(sdfg.name)
py_mapper.mapper(sdfg, tmp)
folder = sdfg.build_folder
save("py", sdfg.name, py_mapper.map, folder)
# If the SDFG was made with the API we need to create tmp info
# as it doesn't have any
sourceFiles = [src for src in get_src_files(sdfg, set())]
# If tmp is None, then the SDFG was created with the API
if tmp is None:
temporaryInfo(
sdfg.name, {
'build_folder': folder,
'src_files': sourceFiles,
'made_with_api': True,
})
else:
tmp['src_files'] = sourceFiles
temporaryInfo(sdfg.name, tmp)
def create_cpp_map(code: str, name: str, target_name: str):
""" Creates the mapping from the SDFG nodes to the C++ code lines.
The mapping gets saved at: <SDFG build folder>/map/map_cpp.json
:param code: C++ code containing the identifiers '////__DACE:0:0:0'
:param name: The name of the SDFG
:param target_name: The target type, example: 'cpu'
"""
tmp = get_tmp(name)
if tmp is None:
return
remove_tmp(name, True)
# If the cache setting is set to 'hash' then we don't create a
# mapping as we don't know where to save it to due to
# the hash value changing every time the state of the SDFG changes.
# We send a message to the IDE in that case to notify the user of
# restricted features.
if Config.get('cache') != 'hash':
codegen_debug = Config.get_bool('compiler', 'codegen_lineinfo')
cpp_mapper = MapCpp(code, name, target_name)
cpp_mapper.mapper(codegen_debug)
folder = save("cpp", name, cpp_mapper.map, tmp.get("build_folder"))
api = tmp.get('made_with_api')
if codegen_debug:
save("codegen", name, cpp_mapper.codegen_map,
tmp.get("build_folder"))
# Send information about the SDFG to VSCode
send({
"type": "registerFunction",
"name": name,
"path_cache": folder,
"path_file": tmp.get("src_files"),
"target_name": target_name,
"made_with_api": api if api else False,
"codegen_map": codegen_debug
})
else:
send({'type': 'restrictedFeatures', 'reason': 'config.cache.hash'})
def create_cache(name: str, folder: str) -> str:
""" Creates the map folder in the build path if it
does not yet exist
:param name: name of the SDFG
:param folder: the build folder
:return: relative path to the created folder
"""
if (folder is not None):
create_folder(os.path.join(folder, "map"))
return folder
else:
build_folder = Config.get('default_build_folder')
cache_folder = os.path.join(build_folder, name)
create_folder(os.path.join(cache_folder, "map"))
return cache_folder
def remove_tmp(name: str, remove_cache: bool = False):
""" Remove the tmp created by "temporaryInfo"
:param name: name of the sdfg for which the tmp will be removed
:param remove_cache: If true, checks if the directory only contains
the tmp file, if this is the case, remove the SDFG cache directory.
"""
build_folder = Config.get('default_build_folder')
path = os.path.join(build_folder, name)
if not os.path.exists(path):
return
os.remove(os.path.join(path, 'map', 'tmp.json'))
if (remove_cache and len(os.listdir(path)) == 1
and len(os.listdir(os.path.join(path, 'map'))) == 0):
if os.path.exists(path):
os.rmdir(os.path.join(path, 'map'))
os.rmdir(path)