def rocm_link(in_file, out_file, lld=None):
"""Link relocatable ELF object to shared ELF object using lld
Parameters
----------
in_file : str
Input file name (relocatable ELF object file)
out_file : str
Output file name (shared ELF object file)
lld : str, optional
The lld linker, if not specified,
we will try to guess the matched clang version.
"""
args = [
lld if lld is not None else find_lld()[0], "-shared", in_file, "-o",
out_file
]
proc = subprocess.Popen(args,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT)
(out, _) = proc.communicate()
if proc.returncode != 0:
msg = "Linking error using ld.lld:\n"
msg += py_str(out)
raise RuntimeError(msg)
def create_tvmjs_wasm(output, objects, options=None, cc="emcc"):
"""Create wasm that is supposed to run with the tvmjs.
Parameters
----------
output : str
The target shared library.
objects : list
List of object files.
options : str
The additional options.
cc : str, optional
The compile string.
"""
cmd = [cc]
cmd += ["-O3"]
cmd += ["-std=c++17"]
cmd += ["--no-entry"]
cmd += ["-s", "ERROR_ON_UNDEFINED_SYMBOLS=0"]
cmd += ["-s", "STANDALONE_WASM=1"]
cmd += ["-s", "ALLOW_MEMORY_GROWTH=1"]
objects = [objects] if isinstance(objects, str) else objects
with_runtime = False
for obj in objects:
if obj.find("wasm_runtime.bc") != -1:
with_runtime = True
if not with_runtime:
objects += [find_lib_path("wasm_runtime.bc")[0]]
objects += [find_lib_path("tvmjs_support.bc")[0]]
objects += [find_lib_path("webgpu_runtime.bc")[0]]
cmd += ["-o", output]
cmd += objects
if options:
cmd += options
proc = subprocess.Popen(cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT)
(out, _) = proc.communicate()
if proc.returncode != 0:
msg = "Compilation error:\n"
msg += py_str(out)
raise RuntimeError(msg)
def create_llvm(inputs, output=None, options=None, cc=None):
"""Create llvm text ir.
Parameters
----------
inputs : list of str
List of input files name or code source.
output : str, optional
Output file, if it is none
a temporary file is created
options : list
The list of additional options string.
cc : str, optional
The clang compiler, if not specified,
we will try to guess the matched clang version.
Returns
-------
code : str
The generated llvm text IR.
"""
cc = cc if cc else find_clang()[0]
cmd = [cc]
cmd += ["-S", "-emit-llvm"]
temp = util.tempdir()
output = output if output else temp.relpath("output.ll")
inputs = [inputs] if isinstance(inputs, str) else inputs
input_files = []
for i, code in enumerate(inputs):
if util.is_source_path(code):
input_files.append(code)
else:
temp_path = temp.relpath("input%d.cc" % i)
with open(temp_path, "w") as output_file:
output_file.write(code)
input_files.append(temp_path)
if options:
cmd += options
cmd += ["-o", output]
cmd += input_files
proc = subprocess.Popen(cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT)
(out, _) = proc.communicate()
if proc.returncode != 0:
msg = "Compilation error:\n"
msg += py_str(out)
raise RuntimeError(msg)
return open(output).read()
def test_ndarray_reflection():
x = np.random.uniform(size=(10, 2)).astype("float32")
xx = tvm.nd.array(x)
xnode = tvm.make.node("NDArrayWrapper", name="xx", array=xx)
xnode2 = tvm.make.node("NDArrayWrapper", name="x2", array=xx)
assert xnode.array.same_as(xx)
json_str = tvm.save_json([xnode, xnode2])
json_dict = json.loads(json_str)
b64_str = json_dict["b64ndarrays"][0]
decoded = py_str(base64.b64encode(base64.b64decode(b64_str)))
assert b64_str == decoded
xlist = tvm.load_json(json_str)
np.testing.assert_equal(xlist[0].array.asnumpy(), xx.asnumpy())
assert xlist[1].array == xlist[0].array
def test_ndarray_reflection():
x = np.random.uniform(size=(10, 2)).astype("float32")
xx = tvm.nd.array(x)
xnode = tvm.make.node("NDArrayWrapper", name="xx", array=xx)
xnode2 = tvm.make.node("NDArrayWrapper", name="x2", array=xx)
assert xnode.array.same_as(xx)
json_str = tvm.save_json([xnode, xnode2])
json_dict = json.loads(json_str)
b64_str = json_dict["b64ndarrays"][0]
decoded = py_str(base64.b64encode(base64.b64decode(b64_str)))
assert b64_str == decoded
xlist = tvm.load_json(json_str)
np.testing.assert_equal(xlist[0].array.asnumpy(), xx.asnumpy())
assert xlist[1].array == xlist[0].array
def _connect_proxy_loop(addr, key, load_library):
key = "server:" + key
retry_count = 0
max_retry = 5
retry_period = 5
while True:
try:
sock = socket.socket(base.get_addr_family(addr),
socket.SOCK_STREAM)
sock.connect(addr)
sock.sendall(struct.pack("<i", base.RPC_MAGIC))
sock.sendall(struct.pack("<i", len(key)))
sock.sendall(key.encode("utf-8"))
magic = struct.unpack("<i", base.recvall(sock, 4))[0]
if magic == base.RPC_CODE_DUPLICATE:
raise RuntimeError("key: %s has already been used in proxy" %
key)
if magic == base.RPC_CODE_MISMATCH:
logger.warning("RPCProxy do not have matching client key %s",
key)
elif magic != base.RPC_CODE_SUCCESS:
raise RuntimeError("%s is not RPC Proxy" % str(addr))
keylen = struct.unpack("<i", base.recvall(sock, 4))[0]
remote_key = py_str(base.recvall(sock, keylen))
opts = _parse_server_opt(remote_key.split()[1:])
logger.info("connected to %s", str(addr))
process = multiprocessing.Process(target=_serve_loop,
args=(sock, addr, load_library))
process.deamon = True
process.start()
sock.close()
process.join(opts.get("timeout", None))
if process.is_alive():
logger.info("Timeout in RPC session, kill..")
process.terminate()
retry_count = 0
except (socket.error, IOError) as err:
retry_count += 1
logger.warning("Error encountered %s, retry in %g sec", str(err),
retry_period)
if retry_count > max_retry:
raise RuntimeError("Maximum retry error: last error: %s" %
str(err))
time.sleep(retry_period)
def rocm_link(in_file, out_file, lld=None):
"""Link relocatable ELF object to shared ELF object using lld
Parameters
----------
in_file : str
Input file name (relocatable ELF object file)
out_file : str
Output file name (shared ELF object file)
lld : str, optional
The lld linker, if not specified,
we will try to guess the matched clang version.
"""
# if our result has undefined symbols, it will fail to load
# (hipModuleLoad/hipModuleLoadData), but with a somewhat opaque message
# so we have ld.lld check this here.
# If you get a complaint about missing symbols you might want to check the
# list of bitcode files below.
args = [
lld if lld is not None else find_lld()[0],
"--no-undefined",
"-shared",
in_file,
"-o",
out_file,
]
proc = subprocess.Popen(args,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT)
(out, _) = proc.communicate()
if proc.returncode != 0:
msg = "Linking error using ld.lld:\n"
msg += py_str(out)
raise RuntimeError(msg)
def _accept_conn(listen_sock, tracker_conn, ping_period=2):
"""Accept connection from the other places.
Parameters
----------
listen_sock: Socket
The socket used by listening process.
tracker_conn : connnection to tracker
Tracker connection
ping_period : float, optional
ping tracker every k seconds if no connection is accepted.
"""
old_keyset = set()
# Report resource to tracker
if tracker_conn:
matchkey = base.random_key(rpc_key + ":")
base.sendjson(
tracker_conn,
[TrackerCode.PUT, rpc_key, (port, matchkey), custom_addr])
assert base.recvjson(tracker_conn) == TrackerCode.SUCCESS
else:
matchkey = rpc_key
unmatch_period_count = 0
unmatch_timeout = 4
# Wait until we get a valid connection
while True:
if tracker_conn:
trigger = select.select([listen_sock], [], [], ping_period)
if not listen_sock in trigger[0]:
base.sendjson(tracker_conn,
[TrackerCode.GET_PENDING_MATCHKEYS])
pending_keys = base.recvjson(tracker_conn)
old_keyset.add(matchkey)
# if match key not in pending key set
# it means the key is acquired by a client but not used.
if matchkey not in pending_keys:
unmatch_period_count += 1
else:
unmatch_period_count = 0
# regenerate match key if key is acquired but not used for a while
if unmatch_period_count * ping_period > unmatch_timeout + ping_period:
logger.info(
"no incoming connections, regenerate key ...")
matchkey = base.random_key(rpc_key + ":", old_keyset)
base.sendjson(tracker_conn, [
TrackerCode.PUT, rpc_key,
(port, matchkey), custom_addr
])
assert base.recvjson(
tracker_conn) == TrackerCode.SUCCESS
unmatch_period_count = 0
continue
conn, addr = listen_sock.accept()
magic = struct.unpack("<i", base.recvall(conn, 4))[0]
if magic != base.RPC_MAGIC:
conn.close()
continue
keylen = struct.unpack("<i", base.recvall(conn, 4))[0]
key = py_str(base.recvall(conn, keylen))
arr = key.split()
expect_header = "client:" + matchkey
server_key = "server:" + rpc_key
if arr[0] != expect_header:
conn.sendall(struct.pack("<i", base.RPC_CODE_MISMATCH))
conn.close()
logger.warning("mismatch key from %s", addr)
continue
conn.sendall(struct.pack("<i", base.RPC_CODE_SUCCESS))
conn.sendall(struct.pack("<i", len(server_key)))
conn.sendall(server_key.encode("utf-8"))
return conn, addr, _parse_server_opt(arr[1:])
