StarPilot
This commit is contained in:
firestar5683
+82
-17
@@ -20,39 +20,104 @@ if arch == "Darwin":
|
||||
else:
|
||||
libs += ['OpenCL']
|
||||
|
||||
# Set path definitions
|
||||
for pathdef, fn in {'TRANSFORM': 'transforms/transform.cl', 'LOADYUV': 'transforms/loadyuv.cl'}.items():
|
||||
# Set path definitions to repository-relative paths so prebuilt artifacts stay relocatable.
|
||||
for pathdef, fn in {'TRANSFORM': 'selfdrive/modeld/transforms/transform.cl', 'LOADYUV': 'selfdrive/modeld/transforms/loadyuv.cl'}.items():
|
||||
for xenv in (lenv, lenvCython):
|
||||
xenv['CXXFLAGS'].append(f'-D{pathdef}_PATH=\\"{File(fn).abspath}\\"')
|
||||
xenv['CXXFLAGS'].append(f'-D{pathdef}_PATH=\\"{fn}\\"')
|
||||
|
||||
# Compile cython
|
||||
cython_libs = envCython["LIBS"] + libs
|
||||
commonmodel_lib = lenv.Library('commonmodel', common_src)
|
||||
lenvCython.Program('models/commonmodel_pyx.so', 'models/commonmodel_pyx.pyx', LIBS=[commonmodel_lib, *cython_libs], FRAMEWORKS=frameworks)
|
||||
tinygrad_files = ["#"+x for x in glob.glob(env.Dir("#tinygrad_repo").relpath + "/**", recursive=True, root_dir=env.Dir("#").abspath) if 'pycache' not in x]
|
||||
skip_dm_tinygrad_pkl = os.getenv("SP_SKIP_DM_TINYGRAD_PKL", "").lower() in {"1", "true", "yes", "on"}
|
||||
allow_host_tinygrad_pkl = os.getenv("SP_ALLOW_HOST_TINYGRAD_PKL", "").lower() in {"1", "true", "yes", "on"}
|
||||
build_model_tinygrad_pkl = os.getenv("SP_BUILD_MODEL_TINYGRAD_PKL", "").lower() in {"1", "true", "yes", "on"}
|
||||
has_qcom_gpu = os.path.exists('/dev/kgsl-3d0')
|
||||
|
||||
# Compile camera warp artifacts used by modeld/dmonitoringmodeld on C4.
|
||||
if arch == 'larch64':
|
||||
compile_warp_flags = 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0' \
|
||||
if os.path.exists('/dev/kgsl-3d0') else 'DEV=CPU CPU=1 THREADS=0 NOLOCALS=1 DEBUG=0 HOME=/tmp'
|
||||
else:
|
||||
compile_warp_flags = {
|
||||
'Darwin': f'DEV=CPU THREADS=0 NOLOCALS=1 DEBUG=0 HOME={os.path.expanduser("~")}',
|
||||
}.get(arch, 'DEV=CPU CPU_LLVM=1 THREADS=0 NOLOCALS=1 DEBUG=0')
|
||||
|
||||
build_warp_artifacts = os.getenv("SP_BUILD_WARP_ARTIFACTS", "").lower() in {"1", "true", "yes", "on"}
|
||||
if build_warp_artifacts:
|
||||
default_warp_resolutions = [(1928, 1208), (1344, 760)]
|
||||
raw_warp_resolutions = os.getenv("SP_WARP_RESOLUTIONS", "").strip()
|
||||
if raw_warp_resolutions:
|
||||
selected_warp_resolutions = []
|
||||
seen_warp_resolutions = set()
|
||||
for token in raw_warp_resolutions.replace(";", ",").split(","):
|
||||
token = token.strip().lower()
|
||||
if not token:
|
||||
continue
|
||||
w_str, h_str = token.split("x", 1)
|
||||
wh = (int(w_str), int(h_str))
|
||||
if wh not in seen_warp_resolutions:
|
||||
seen_warp_resolutions.add(wh)
|
||||
selected_warp_resolutions.append(wh)
|
||||
if not selected_warp_resolutions:
|
||||
selected_warp_resolutions = default_warp_resolutions
|
||||
else:
|
||||
selected_warp_resolutions = default_warp_resolutions
|
||||
|
||||
compile_warp_script = "#selfdrive/modeld/compile_warp.py"
|
||||
warp_targets = []
|
||||
# Camera resolutions required by modeld compile_warp.py
|
||||
for w, h in selected_warp_resolutions:
|
||||
warp_targets += [File(f"models/warp_{w}x{h}_tinygrad.pkl").abspath, File(f"models/dm_warp_{w}x{h}_tinygrad.pkl").abspath]
|
||||
lenv.Command(
|
||||
warp_targets,
|
||||
tinygrad_files + [compile_warp_script],
|
||||
f'{compile_warp_flags} python3 selfdrive/modeld/compile_warp.py',
|
||||
)
|
||||
else:
|
||||
print("Skipping model warp precompile (set SP_BUILD_WARP_ARTIFACTS=1 to enable)")
|
||||
|
||||
# Get model metadata
|
||||
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
|
||||
fn = File(f"models/{model_name}").abspath
|
||||
script_files = [File(Dir("#selfdrive/modeld").File("get_model_metadata.py").abspath)]
|
||||
cmd = f'python3 {Dir("#selfdrive/modeld").abspath}/get_model_metadata.py {fn}.onnx'
|
||||
lenv.Command(fn + "_metadata.pkl", [fn + ".onnx"] + tinygrad_files + script_files, cmd)
|
||||
model_rel_path = f"selfdrive/modeld/models/{model_name}"
|
||||
model_node_path = f"#{model_rel_path}"
|
||||
metadata_script = "#selfdrive/modeld/get_model_metadata.py"
|
||||
cmd = f'python3 selfdrive/modeld/get_model_metadata.py {model_rel_path}.onnx'
|
||||
lenv.Command(model_node_path + "_metadata.pkl", [model_node_path + ".onnx"] + tinygrad_files + [metadata_script], cmd)
|
||||
|
||||
def tg_compile(flags, model_name):
|
||||
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + '"'
|
||||
fn = File(f"models/{model_name}").abspath
|
||||
tinygrad_root = "tinygrad_repo"
|
||||
compile_script = "#tinygrad_repo/examples/openpilot/compile3.py"
|
||||
pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + tinygrad_root + '"'
|
||||
model_rel_path = f"selfdrive/modeld/models/{model_name}"
|
||||
model_node_path = f"#{model_rel_path}"
|
||||
model_input = f"./{model_rel_path}.onnx"
|
||||
return lenv.Command(
|
||||
fn + "_tinygrad.pkl",
|
||||
[fn + ".onnx"] + tinygrad_files,
|
||||
f'{pythonpath_string} {flags} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {fn}_tinygrad.pkl'
|
||||
model_node_path + "_tinygrad.pkl",
|
||||
[model_node_path + ".onnx"] + tinygrad_files + [compile_script],
|
||||
f'{pythonpath_string} {flags} python3 tinygrad_repo/examples/openpilot/compile3.py {model_input} {model_rel_path}_tinygrad.pkl'
|
||||
)
|
||||
|
||||
compile_model_tinygrad_pkl = build_model_tinygrad_pkl or allow_host_tinygrad_pkl
|
||||
if not compile_model_tinygrad_pkl:
|
||||
print("Skipping tinygrad model PKL compile (set SP_BUILD_MODEL_TINYGRAD_PKL=1 or SP_ALLOW_HOST_TINYGRAD_PKL=1 to enable)")
|
||||
|
||||
# Compile small models
|
||||
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
|
||||
flags = {
|
||||
'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 IMAGE=2 JIT_BATCH_SIZE=0',
|
||||
'Darwin': f'DEV=CPU HOME={os.path.expanduser("~")}', # tinygrad calls brew which needs a $HOME in the env
|
||||
}.get(arch, 'DEV=CPU CPU_LLVM=1')
|
||||
if model_name == 'dmonitoring_model' and skip_dm_tinygrad_pkl:
|
||||
print("Skipping dmonitoring_model_tinygrad.pkl compile (SP_SKIP_DM_TINYGRAD_PKL enabled)")
|
||||
continue
|
||||
if not compile_model_tinygrad_pkl:
|
||||
continue
|
||||
if arch == 'larch64':
|
||||
# On real devices keep QCOM codegen.
|
||||
flags = 'DEV=QCOM QCOM=1 FLOAT16=1 NOLOCALS=1 IMAGE=2 JIT_BATCH_SIZE=0 DEBUG=0'
|
||||
else:
|
||||
# Opt-in host compile (for debugging only).
|
||||
flags = {
|
||||
'Darwin': 'DEV=CPU CPU=1 IMAGE=0 NOLOCALS=1 DEBUG=0 HOME=/tmp',
|
||||
}.get(arch, 'DEV=CPU CPU=1 IMAGE=0 NOLOCALS=1 DEBUG=0')
|
||||
tg_compile(flags, model_name)
|
||||
|
||||
# Compile BIG model if USB GPU is available
|
||||
@@ -62,7 +127,7 @@ if "USBGPU" in os.environ:
|
||||
devs = subprocess.check_output('python3 -c "from tinygrad import Device; print(list(Device.get_available_devices()))"', shell=True, cwd=env.Dir('#').abspath)
|
||||
if b"AMD" in devs:
|
||||
print("USB GPU detected... building")
|
||||
flags = "DEV=AMD AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0"
|
||||
flags = "DEV=AMD AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0 DEBUG=0"
|
||||
bp = tg_compile(flags, "big_driving_policy")
|
||||
bv = tg_compile(flags, "big_driving_vision")
|
||||
lenv.SideEffect('lock', [bp, bv]) # tg doesn't support multi-process so build serially
|
||||
|
||||
Reference in New Issue
Block a user