modeld: group npy -> qcom copies to avoid graph breaks (#37866)

* modeld: group npy -> qcom copies to avoid graph breaks

* batch realize

* dm as well
This commit is contained in:
Armand du Parc Locmaria
2026-04-21 13:50:20 -07:00
committed by GitHub
parent 4988a62b31
commit 859bd215bf
+15 -8
View File
@@ -148,7 +148,7 @@ def make_warp_dm(cam_w, cam_h, dm_w, dm_h):
stride_pad = stride - cam_w
def warp_dm(input_frame, M_inv):
M_inv = M_inv.to(Device.DEFAULT)
M_inv = M_inv.to(Device.DEFAULT).realize()
result = warp_perspective_tinygrad(input_frame[:cam_h*stride], M_inv, (dm_w, dm_h), (cam_h, cam_w), stride_pad).reshape(-1, dm_h * dm_w)
return result
return warp_dm
@@ -162,8 +162,14 @@ def make_run_policy(vision_runner, policy_runner, cam_w, cam_h,
sample_desire_fn = partial(sample_desire, frame_skip=frame_skip)
def run_policy(img_q, big_img_q, feat_q, desire_q, desire, traffic_convention, tfm, big_tfm, frame, big_frame):
img = shift_and_sample(img_q, frame_prepare(frame, tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn)
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm.to(Device.DEFAULT)).unsqueeze(0), sample_skip_fn)
tfm = tfm.to(Device.DEFAULT)
big_tfm = big_tfm.to(Device.DEFAULT)
desire = desire.to(Device.DEFAULT)
traffic_convention = traffic_convention.to(Device.DEFAULT)
Tensor.realize(tfm, big_tfm, desire, traffic_convention)
img = shift_and_sample(img_q, frame_prepare(frame, tfm).unsqueeze(0), sample_skip_fn)
big_img = shift_and_sample(big_img_q, frame_prepare(big_frame, big_tfm).unsqueeze(0), sample_skip_fn)
if prepare_only:
return img, big_img
@@ -172,9 +178,9 @@ def make_run_policy(vision_runner, policy_runner, cam_w, cam_h,
new_feat = vision_out[:, vision_features_slice].reshape(1, -1).unsqueeze(0)
feat_buf = shift_and_sample(feat_q, new_feat, sample_skip_fn)
desire_buf = shift_and_sample(desire_q, desire.to(Device.DEFAULT).reshape(1, 1, -1), sample_desire_fn)
desire_buf = shift_and_sample(desire_q, desire.reshape(1, 1, -1), sample_desire_fn)
inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention.to(Device.DEFAULT)}
inputs = {'features_buffer': feat_buf, 'desire_pulse': desire_buf, 'traffic_convention': traffic_convention}
policy_out = next(iter(policy_runner(inputs).values())).cast('float32')
return vision_out, policy_out
@@ -210,10 +216,11 @@ def compile_modeld(cam_w, cam_h, prepare_only, pkl_path):
def random_inputs_run_fn(fn, seed, test_val=None, test_buffers=None, expect_match=True):
input_queues, npy = make_input_queues(vision_input_shapes, policy_input_shapes, frame_skip)
np.random.seed(seed)
Tensor.manual_seed(seed)
for i in range(N_RUNS):
frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize()
big_frame = Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize()
frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize()
big_frame = Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize()
for v in npy.values():
v[:] = np.random.randn(*v.shape).astype(v.dtype)
Device.default.synchronize()
@@ -263,7 +270,7 @@ def compile_dm_warp(cam_w, cam_h, pkl_path):
warp_dm_jit = TinyJit(warp_dm, prune=True)
for i in range(10):
inputs = [Tensor(np.random.randint(0, 256, yuv_size, dtype=np.uint8)).realize(),
inputs = [Tensor.randint(yuv_size, low=0, high=256, dtype='uint8').realize(),
Tensor(Tensor.randn(3, 3).mul(8).realize().numpy(), device='NPY')]
Device.default.synchronize()
st = time.perf_counter()