use upstreams

sunnypilot/modeld_v2/SConscript

@@ -1,5 +1,6 @@
 import os
 import glob
+from tinygrad import Device
 
 Import('env', 'arch')
 lenv = env.Clone()
@@ -21,10 +22,19 @@ if PC:
   if outputs:
     lenv.Command(outputs, inputs, cmd)
 
-tg_flags = {
-  'larch64': 'DEV=QCOM FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1',
-  'Darwin': f'DEV=CPU THREADS=0 HOME={os.path.expanduser("~")}',
-}.get(arch, 'DEV=CPU:LLVM THREADS=0')
+available = set(Device.get_available_devices())
+if 'CUDA' in available:
+  tg_backend = 'CUDA'
+  tg_flags = f'DEV={tg_backend}'
+elif 'QCOM' in available:
+  tg_backend = 'QCOM'
+  tg_flags = f'DEV={tg_backend} FLOAT16=1 NOLOCALS=1 JIT_BATCH_SIZE=0 OPENPILOT_HACKS=1'
+else:
+  tg_backend = 'CPU' if arch == 'Darwin' else 'CPU:LLVM'
+  # THREADS=0 is need to prevent bug: https://github.com/tinygrad/tinygrad/issues/14689
+  tg_flags = f'DEV={tg_backend} THREADS=0'
+
+mac_brew_string = f'HOME={os.path.expanduser("~")}' if arch == 'Darwin' else ''
 
 image_flag = {
   'larch64': 'IMAGE=2',
@@ -38,7 +48,7 @@ def tg_compile(flags, model_name):
   return lenv.Command(
     out,
     [fn + ".onnx"] + tinygrad_files,
-    f'{pythonpath_string} {flags} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {out}'
+    f'{pythonpath_string} {tg_flags} {mac_brew_string} {image_flag} python3 {Dir("#tinygrad_repo").abspath}/examples/openpilot/compile3.py {fn}.onnx {out}'
   )
 
 # Compile models
@@ -48,7 +58,7 @@ for model_name in ['supercombo', 'driving_vision', 'driving_off_policy', 'drivin
 
 script_files = [File("warp.py"), File(Dir("#selfdrive/modeld").File("compile_modeld.py").abspath)]
 pythonpath_string = 'PYTHONPATH="${PYTHONPATH}:' + env.Dir("#tinygrad_repo").abspath + ':' + env.Dir("#").abspath + '"'
-compile_warp_cmd = f'{pythonpath_string} {tg_flags} {image_flag} python3 -m sunnypilot.modeld_v2.warp'
+compile_warp_cmd = f'{pythonpath_string} {tg_flags} {mac_brew_string} {image_flag} python3 -m sunnypilot.modeld_v2.warp'
 
 from openpilot.common.transformations.camera import _ar_ox_fisheye, _os_fisheye
 warp_targets = []