import re
import glob, os
from pathlib import Path
from os.path import exists
from sys import stderr
import math
import shutil
import argparse
parser = argparse.ArgumentParser("lc")
parser.add_argument("--output_dir", default=".")
parser.add_argument("--verbose", action="store_true")
parser.add_argument("--base_file", default="framework.h")
parser.add_argument("--main_file", default="framework.cu")
args = parser.parse_args()
shutil.copyfile(args.main_file, args.output_dir + "/lc.cu")
shutil.copyfile(args.base_file, args.output_dir + "/lc.h")
for i in ["/components/include", "/preprocessors/include", "/verifiers/include"]:
os.makedirs(args.output_dir + i, exist_ok=True)
def update_enum(filename, comps, item):
with open(filename, 'w') as f:
f.write("#ifndef LC_" + item + "_H\n")
f.write("#define LC_" + item + "_H\n\n")
f.write("enum {NUL_" + item)
for c in comps:
if item == 'VERIFIER':
c = c[0:]
else:
c = c[2:]
f.write(", " + str(c))
f.write("};\n\n")
def update_cpu_components(filename, comps, component_type):
with open(filename, 'a') as f:
for c in comps:
if c.startswith("v_"):
c = c[2:]
f.write("#include \"" + component_type + "/" + str(c) + ".h\"\n")
def update_gpu_components(filename, comps, component_type):
with open(filename, 'a') as f:
for c in comps:
if c.startswith("v_"):
c = c[2:]
f.write("#include \"" + component_type + "/" + str(c) + ".h\"\n")
f.write("\n#endif\n")
cfiles = next(os.walk('./components'))[2]
cpucomps = []
for f in cfiles:
if f.startswith("h_"):
cpucomps.append(f[:-2])
if args.verbose:
print("cpucomps \n", ', '.join(cpucomps), file=stderr)
gfiles = next(os.walk('./components'))[2]
gpucomps = []
for f in gfiles:
if f.startswith("d_"):
gpucomps.append(f[:-2])
if args.verbose:
print("\ngpucomps \n", ', '.join(gpucomps), file=stderr)
cpucomps.sort()
gpucomps.sort()
cname = []
gname = []
for c in cpucomps:
cname.append(c[2:])
for g in gpucomps:
gname.append(g[2:])
if set(cname) == set(gname):
print("\nComponents match.\n")
else:
diff = [x for x in cpucomps if x[2:] not in [y[2:] for y in gpucomps]] + [y for y in gpucomps if y[2:] not in [x[2:] for x in cpucomps]]
result = ', '.join(diff)
print("\nWarning: skipping unmatched components:", result, "\n")
cpucomps = [c for c in cpucomps if c[2:] in gname]
gpucomps = [c for c in gpucomps if c[2:] in cname]
with open(args.output_dir + '/include/consts.h', 'w') as f:
f.write("static const int CS = 1024 * 16; // chunk size (in bytes) [do not change]\n")
f.write("static const int TPB = 512; // threads per block [must be power of 2 and at least 128]\n")
f.write("#if defined(__AMDGCN_WAVEFRONT_SIZE) && (__AMDGCN_WAVEFRONT_SIZE == 64)\n")
f.write("#define WS 64\n")
f.write("#else\n")
f.write("#define WS 32\n")
f.write("#endif\n")
update_enum(args.output_dir + '/components/include/components.h', cpucomps, 'components')
update_cpu_components(args.output_dir + '/components/include/components.h', cpucomps, "components")
update_gpu_components(args.output_dir + '/components/include/components.h', gpucomps, "components")
gfiles = next(os.walk('./preprocessors'))[2]
gpupreprocess = []
for f in gfiles:
if f.startswith("d_"):
gpupreprocess.append(f[:-2])
if args.verbose:
print("\ngpupreprocess \n",', '.join(gpupreprocess), file=stderr)
cfiles = next(os.walk('./preprocessors'))[2]
cpupreprocess = []
for f in cfiles:
if f.startswith("h_"):
cpupreprocess.append(f[:-2])
if args.verbose:
print("\ncpupreprocess \n",', '.join(cpupreprocess), file=stderr)
cpupreprocess.sort()
gpupreprocess.sort()
cprepro = []
gprepro = []
for c in cpupreprocess:
cprepro.append(c[2:])
for g in gpupreprocess:
gprepro.append(g[2:])
if set(cprepro) == set(gprepro):
print("\npreprocessors match.\n")
else:
diff = [x for x in cpupreprocess if x[2:] not in [y[2:] for y in gpupreprocess]] + [y for y in gpupreprocess if y[2:] not in [x[2:] for x in cpupreprocess]]
result = ', '.join(diff)
print("\nWarning: skipping unmatched preprocessors:", result, "\n")
cpupreprocess = [c for c in cpupreprocess if c[2:] in gprepro]
gpupreprocess = [c for c in gpupreprocess if c[2:] in cprepro]
update_enum(args.output_dir + '/preprocessors/include/preprocessors.h', cpupreprocess, 'preprocessor')
update_cpu_components(args.output_dir + '/preprocessors/include/preprocessors.h', cpupreprocess, "preprocessors")
update_gpu_components(args.output_dir + '/preprocessors/include/preprocessors.h', gpupreprocess, "preprocessors")
cfiles = next(os.walk('./verifiers'))[2]
gpuverifier = []
for f in cfiles:
if f.endswith(".h"):
gpuverifier.append("v_" + f[:-2])
if args.verbose:
print("\nverifier \n", ', '.join(gpuverifier), file=stderr)
update_enum(args.output_dir + '/verifiers/include/verifiers.h', gpuverifier, 'VERIFIER')
update_gpu_components(args.output_dir + '/verifiers/include/verifiers.h', gpuverifier, "verifiers")
file = args.output_dir + "/lc.h"
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-host-encode-beg##[\s\S]*##switch-host-encode-end##", contents)
str_to_add = ''
for c in cpucomps:
c = c[2:]
str_to_add += " case " + str(c) + ": good = h_" + str(c) + "(csize, in, out); break;\n"
contents = contents[:m.span()[0]] + "##switch-host-encode-beg##*/\n" + str_to_add + " /*##switch-host-encode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-host-decode-beg##[\s\S]*##switch-host-decode-end##", contents)
str_to_add = ''
for c in cpucomps:
c = c[2:]
str_to_add += " case " + str(c) + ": h_i" + str(c) + "(csize, in, out); break;\n"
contents = contents[:m.span()[0]] + "##switch-host-decode-beg##*/\n" + str_to_add + " /*##switch-host-decode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-device-encode-beg##[\s\S]*##switch-device-encode-end##", contents)
str_to_add = ''
for c in gpucomps:
c = c[2:]
str_to_add += " case " + str(c) + ": good = d_" + str(c) + "(csize, in, out, temp); break;\n"
contents = contents[:m.span()[0]] + "##switch-device-encode-beg##*/\n" + str_to_add + " /*##switch-device-encode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-device-decode-beg##[\s\S]*##switch-device-decode-end##", contents)
str_to_add = ''
for c in gpucomps:
c = c[2:]
str_to_add += " case " + str(c) + ": d_i" + str(c) + "(csize, in, out, temp); break;\n"
contents = contents[:m.span()[0]] + "##switch-device-decode-beg##*/\n" + str_to_add + " /*##switch-device-decode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-pipeline-beg##[\s\S]*##switch-pipeline-end##", contents)
str_to_add = ''
for c in cpucomps:
c = c[2:]
str_to_add += " case " + str(c) + ": s += \" " + str(c) + "\"; break;\n"
contents = contents[:m.span()[0]] + "##switch-pipeline-beg##*/\n" + str_to_add + " /*##switch-pipeline-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-verify-beg##[\s\S]*##switch-verify-end##", contents)
str_to_add = ''
for c in gpuverifier:
c = c[2:]
str_to_add += " case v_" + str(c) + ": " + str(c) + "(size, recon, orig, params.size(), params.data()); break;\n"
contents = contents[:m.span()[0]] + "##switch-verify-beg##*/\n" + str_to_add + " /*##switch-verify-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-host-preprocess-encode-beg##[\s\S]*##switch-host-preprocess-encode-end##", contents)
str_to_add = ''
for c in cpupreprocess:
c = c[2:]
str_to_add += " case " + str(c) + ": h_" + str(c) + "(hpreencsize, hpreencdata, params.size(), params.data()); break;\n"
contents = contents[:m.span()[0]] + "##switch-host-preprocess-encode-beg##*/\n" + str_to_add + " /*##switch-host-preprocess-encode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-host-preprocess-decode-beg##[\s\S]*##switch-host-preprocess-decode-end##", contents)
str_to_add = ''
for c in cpupreprocess:
c = c[2:]
str_to_add += " case " + str(c) + ": h_i" + str(c) + "(hpredecsize, hpredecdata, params.size(), params.data()); break;\n"
contents = contents[:m.span()[0]] + "##switch-host-preprocess-decode-beg##*/\n" + str_to_add + " /*##switch-host-preprocess-decode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-device-preprocess-encode-beg##[\s\S]*##switch-device-preprocess-encode-end##", contents)
str_to_add = ''
for c in gpupreprocess:
c = c[2:]
str_to_add += " case " + str(c) + ": d_" + str(c) + "(dpreencsize, dpreencdata, params.size(), params.data()); break;\n"
contents = contents[:m.span()[0]] + "##switch-device-preprocess-encode-beg##*/\n" + str_to_add + " /*##switch-device-preprocess-encode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##switch-device-preprocess-decode-beg##[\s\S]*##switch-device-preprocess-decode-end##", contents)
str_to_add = ''
for c in gpupreprocess:
c = c[2:]
str_to_add += " case " + str(c) + ": d_i" + str(c) + "(dpredecsize, dpredecdata, params.size(), params.data()); break;\n"
contents = contents[:m.span()[0]] + "##switch-device-preprocess-decode-beg##*/\n" + str_to_add + " /*##switch-device-preprocess-decode-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##component-map-beg##[\s\S]*##component-map-end##", contents)
str_to_add = ''
i = 0
for c in cpucomps:
c = c[2:]
i += 1
str_to_add += " components[\"" + str(c) + "\"] = " + str(i) + ";\n"
contents = contents[:m.span()[0]] + "##component-map-beg##*/\n" + str_to_add + " /*##component-map-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##preprocessor-map-beg##[\s\S]*##preprocessor-map-end##", contents)
str_to_add = ''
for c in cpupreprocess:
c = c[2:]
str_to_add += " preprocessors[\"" + str(c) + "\"] = " + str(c) + ";\n"
contents = contents[:m.span()[0]] + "##preprocessor-map-beg##*/\n" + str_to_add + " /*##preprocessor-map-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
with open(file, "r+") as f:
contents = f.read()
m = re.search(r"##verifier-map-beg##[\s\S]*##verifier-map-end##", contents)
str_to_add = ''
for c in gpuverifier:
c = c[2:]
str_to_add += " verifs[\"" + str(c) + "\"] = v_" + str(c) + ";\n"
contents = contents[:m.span()[0]] + "##verifier-map-beg##*/\n" + str_to_add + " /*##verifier-map-end##" + contents[m.span()[1]:]
f.seek(0)
f.truncate()
f.write(contents)
print("\nCompile with\nnvcc -O3 -arch=sm_70 -fmad=false -DUSE_CPU -DUSE_GPU -Xcompiler \"-O3 -march=native -fopenmp -mno-fma -ffp-contract=off\" -I. -std=c++17 -o lc lc.cu\n")
print("Run the following command to see the usage message\n./lc")