optimize openpilot vision on qcom

examples/openpilot/compile3.py

@@ -65,6 +65,30 @@ def compile(onnx_file):
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
     assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_RPT", 1):
+    from extra.gemm.qcom_openpilot_vision_fp16 import patch_fp32_rpt
+    if (patched:=patch_fp32_rpt(run_onnx_jit)): print(f"repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_schedule_projection import patch_projection
+    if (patched:=patch_projection(run_onnx_jit)): print(f"rescheduled {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_FULL_RPT", 1):
+    from extra.gemm.qcom_openpilot_inverse_full_rpt import patch_model as patch_full_rpt
+    if (patched:=patch_full_rpt(run_onnx_jit)): print(f"fully repeat-packed {patched} QCOM vision kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_DEDUPE", 1):
+    from extra.gemm.qcom_openpilot_dedupe_head import dedupe_identical_calls
+    if (removed:=dedupe_identical_calls(run_onnx_jit)): print(f"deduplicated {len(removed)} QCOM kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_PACK_CONV", 1):
+    from extra.gemm.qcom_openpilot_pack_conv_weights import patch_conv
+    if (patched:=patch_conv(run_onnx_jit)): print(f"packed weights for {patched} QCOM convolution kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_LEVEL_SCHEDULE", 1):
+    from extra.gemm.qcom_openpilot_level_schedule import schedule_levels
+    if (moved:=schedule_levels(run_onnx_jit)): print(f"rescheduled {moved} QCOM kernels by dependency level")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_BATCH_HEAD", 1):
+    from extra.gemm.qcom_openpilot_batch_head import batch_head
+    if (combined:=batch_head(run_onnx_jit)): print(f"batched {combined} groups of QCOM head kernels")
+  if Device.DEFAULT.startswith("QCOM") and getenv("OPENPILOT_QCOM_INPUT_PACK", 1):
+    from extra.gemm.qcom_openpilot_input_pack import patch_input_pack
+    if (patched:=patch_input_pack(run_onnx_jit)): print(f"vectorized {patched} QCOM input kernel")
   with open(OUTPUT, "wb") as f:
     pickle.dump(run_onnx_jit, f)
   mdl_sz = os.path.getsize(onnx_file)
@@ -72,7 +96,7 @@ def compile(onnx_file):
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
   print(f"pkl size is {pkl_sz/1e6:.2f}M")
   print("**** compile done ****")
-  return inputs, test_val
+  return run_onnx_jit, inputs, test_val
 
 def test_vs_compile(run, inputs, test_val=None):
 
@@ -142,9 +166,10 @@ if __name__ == "__main__":
     test_vs_compile(pickle_loaded, inputs)
   else:
     onnx_file = fetch(OPENPILOT_MODEL)
-    inputs, outputs = compile(onnx_file)
+    pickle_loaded, inputs, outputs = compile(onnx_file)
 
-    with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+    if OUTPUT != os.devnull:
+      with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
 
     test_vs_compile(pickle_loaded, inputs, outputs)
     if getenv("SELFTEST"):

extra/gemm/qcom_openpilot_batch_head.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""Batch adjacent independent openpilot head kernels into one QCOM launch."""
+import argparse, pickle, re
+from dataclasses import replace
+
+from tinygrad import Device
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+MAX_BATCH={"r_256_4_128_4":4,"r_128_16_4_16_4":4,"r_128_16_4_32_4":4,
+           "r_8_16_4_8_4":4,"r_8_4_8_4":4,"r_8_4_8_4n1":4}
+MAX_BATCH.update({"r_16_16_4_8_4":4,"r_4_16_4_8_4":4,"r_16_16_4_4":4,
+                  "r_4_4_4_4":4,"r_16_16_4_4n1":4,"r_4_4_4_4n1":4})
+
+
+def batched_source(source:str, name:str, batch_count:int) -> str:
+  match=re.search(r"__kernel void \w+\((.*?)\) \{",source,re.S)
+  if match is None: raise RuntimeError("kernel signature not found")
+  declarations=[x.strip() for x in match.group(1).split(",")]
+  arg_names=[x.rsplit(" ",1)[1] for x in declarations]
+  renamed=[]
+  bodies=[]
+  body=source[match.end():source.rfind("}")]
+  local_decls=re.findall(r"__attribute__\s*\(\(aligned \(\d+\)\)\)\s*__local\s+[^;]+;",body)
+  hoisted=[]
+  for batch in range(batch_count):
+    mapping={arg:f"{arg}_{batch}" for arg in arg_names}
+    renamed.extend(decl.rsplit(" ",1)[0]+" "+mapping[arg] for decl,arg in zip(declarations,arg_names))
+    branch=body
+    for declaration in local_decls:
+      local_match=re.search(r"(\w+)(\[[^;]+;)$",declaration)
+      if local_match is None: raise RuntimeError(f"local declaration not understood: {declaration}")
+      old=local_match.group(1)
+      new=f"{old}_{batch}"
+      hoisted.append(declaration[:local_match.start(1)]+new+local_match.group(2))
+      branch=branch.replace(declaration,"")
+      branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
+    for old,new in mapping.items(): branch=re.sub(rf"\b{re.escape(old)}\b",new,branch)
+    bodies.append(branch)
+  prefix=source[:match.start()]
+  count=len(declarations)
+  order=tuple(batch*count for batch in range(batch_count))+tuple(
+    batch*count+arg for batch in range(batch_count) for arg in range(1,count))
+  branches=" else ".join((f"if (get_group_id(1)=={batch}) " if batch < batch_count-1 else "")+f"{{{body}}}"
+                         for batch,body in enumerate(bodies))
+  return f"{prefix}__kernel void {name}_batch{batch_count}({','.join(renamed[i] for i in order)}) {{\n" \
+         f"{''.join(hoisted)}\n{branches}\n}}"
+
+
+def independent(calls:list) -> bool:
+  outputs={call.src[out+1] for call in calls for out in call.src[0].arg.outs}
+  return not any(arg in outputs for call in calls for i,arg in enumerate(call.src[1:]) if i not in call.src[0].arg.outs)
+
+
+def batch_head(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  new_batch=[]
+  combined=0
+  index=0
+  cache={}
+  while index < len(batch):
+    first=batch[index]
+    name=plain_name(first.src[0].arg.name) if first.op is Ops.CALL and first.src[0].op is Ops.PROGRAM else ""
+    if index+1 < len(batch) and name in MAX_BATCH:
+      calls=[first]
+      while index+len(calls) < len(batch) and len(calls) < MAX_BATCH[name]:
+        candidate=batch[index+len(calls)]
+        candidate_name=plain_name(candidate.src[0].arg.name) if candidate.op is Ops.CALL and candidate.src[0].op is Ops.PROGRAM else ""
+        if candidate_name != name or first.src[0].src[3].arg != candidate.src[0].src[3].arg: break
+        calls.append(candidate)
+      if len(calls) > 1 and independent(calls):
+        batch_count=len(calls)
+        program=first.src[0]
+        source=batched_source(program.src[2].arg,name,batch_count)
+        if source not in cache: cache[source]=Device["QCOM"].compiler.compile_cached(source)
+        aux0=program.arg.aux[0]
+        count=len(aux0)
+        ordered_aux=tuple(aux0[0] for _ in calls)+tuple(entry for _ in calls for entry in aux0[1:])
+        combined_aux=tuple(tuple((new_index,dtype,shape) for _old_index,dtype,shape in entry)
+                           for new_index,entry in enumerate(ordered_aux))
+        info=replace(program.arg,name=f"{name}_batch{batch_count}",global_size=(program.arg.global_size[0],batch_count,1),
+                     globals=tuple(range(count*batch_count)),outs=tuple(range(batch_count)),
+                     ins=tuple(range(batch_count,count*batch_count)),aux=(combined_aux,))
+        program=program.replace(arg=info,src=program.src[:2]+
+          (program.src[2].replace(arg=source),program.src[3].replace(arg=cache[source])))
+        new_batch.append(first.replace(src=(program,*[call.src[1] for call in calls],
+                                                *[arg for call in calls for arg in call.src[2:]])))
+        combined+=1
+        index+=batch_count
+        continue
+    new_batch.append(first)
+    index+=1
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return combined
+
+
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("combined",batch_head(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_dedupe_head.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Remove byte-identical duplicate linear chains in the driving-vision head."""
+import argparse, hashlib, pickle
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops, UOp
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGETS = {"r_128_16_4_32_4", "r_256_4_128_4", "r_128_16_4_16_4"}
+
+
+def dedupe_identical_calls(model, all_calls:bool=True) -> list[tuple[int, str]]:
+  """Alias calls with identical programs, inputs, and byte-identical constants."""
+  outer = model.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  produced:dict[UOp, UOp] = {}
+  static_hash:dict[UOp, str] = {}
+  seen:dict[tuple, tuple[UOp, ...]] = {}
+  new_batch, removed = [], []
+
+  def representative(buf:UOp) -> UOp:
+    while buf in produced and produced[buf] is not buf: buf = produced[buf]
+    return buf
+
+  def content_hash(buf:UOp) -> str:
+    if buf not in static_hash:
+      static_hash[buf] = hashlib.sha256(memoryview(buf.buffer.numpy()).cast("B")).hexdigest()
+    return static_hash[buf]
+
+  for index, original in enumerate(batch):
+    call = original.replace(src=tuple(representative(x) if x in produced else x for x in original.src))
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM or (not all_calls and plain_name(call.src[0].arg.name) not in TARGETS):
+      new_batch.append(call)
+      if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM:
+        for out_index in call.src[0].arg.outs: produced[original.src[out_index+1]] = call.src[out_index+1]
+      continue
+    program = call.src[0]
+    output_indices = set(program.arg.outs)
+    signature_args = []
+    for arg_index, (before, after) in enumerate(zip(original.src[1:], call.src[1:])):
+      if arg_index in output_indices: continue
+      if before.op is Ops.PARAM:
+        signature_args.append(("param", before.arg))
+      elif before in produced:
+        signature_args.append(("dynamic", representative(before)))
+      else:
+        signature_args.append((str(after.dtype), after.buffer.size, content_hash(after)))
+    signature = (plain_name(program.arg.name), program.src[3].arg, tuple(signature_args))
+    outputs = tuple(original.src[i+1] for i in program.arg.outs)
+    if signature in seen:
+      canonical_outputs = seen[signature]
+      for output, canonical in zip(outputs, canonical_outputs): produced[output] = representative(canonical)
+      removed.append((index, plain_name(program.arg.name)))
+    else:
+      new_batch.append(call)
+      canonical_outputs = tuple(call.src[i+1] for i in program.arg.outs)
+      seen[signature] = canonical_outputs
+      for output, canonical in zip(outputs, canonical_outputs): produced[output] = canonical
+
+  # Apply aliases to consumers which occur after the duplicate chains.
+  new_batch = [call.replace(src=tuple(representative(x) if x in produced else x for x in call.src)) for call in new_batch]
+  model.captured._linear = model.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+  model.captured.__dict__.pop("linear", None)
+  return removed
+
+
+def main() -> None:
+  parser = argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  parser.add_argument("--all", action="store_true", help="deduplicate every program family, not only the head linears")
+  args = parser.parse_args()
+  with open(args.input, "rb") as f: model = pickle.load(f)
+  removed = dedupe_identical_calls(model, args.all)
+  with open(args.output, "wb") as f: pickle.dump(model, f)
+  print(f"removed {len(removed)} duplicate head calls: {removed}")
+
+
+if __name__ == "__main__": main()

extra/gemm/qcom_openpilot_input_pack.py

@@ -0,0 +1,49 @@
+"""Vectorize the driving-vision uint8 input normalization kernel on QCOM."""
+from dataclasses import replace
+
+from tinygrad import Device
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGET = "E_8192_3_4_2_4"
+SOURCE = r"""#pragma OPENCL EXTENSION cl_khr_fp16 : enable
+__kernel void E_8192_3_4_2_4(write_only image2d_t O,__global uchar *A,__global uchar *B,
+                             __global half *MEAN,__global half *STD) {
+  int c=get_global_id(0),i=get_global_id(1),off=(c<<17)+(i<<2),mc=c<<2;
+  uchar4 a0=vload4(0,A+off),a1=vload4(0,A+off+32768);
+  uchar4 a2=vload4(0,A+off+65536),a3=vload4(0,A+off+98304);
+  uchar4 b0=vload4(0,B+off),b1=vload4(0,B+off+32768);
+  uchar4 b2=vload4(0,B+off+65536),b3=vload4(0,B+off+98304);
+  half4 ma=vload4(0,MEAN+mc),mb=vload4(0,MEAN+mc+12);
+  half4 ia=(half4)(1)/vload4(0,STD+mc),ib=(half4)(1)/vload4(0,STD+mc+12);
+  int x=c+(i&7)*24,y=i>>3;
+  write_imagef(O,(int2)(x,y),convert_float4(((half4)(a0.x,a1.x,a2.x,a3.x)-ma)*ia));
+  write_imagef(O,(int2)(x+3,y),convert_float4(((half4)(b0.x,b1.x,b2.x,b3.x)-mb)*ib));
+  write_imagef(O,(int2)(x+6,y),convert_float4(((half4)(a0.y,a1.y,a2.y,a3.y)-ma)*ia));
+  write_imagef(O,(int2)(x+9,y),convert_float4(((half4)(b0.y,b1.y,b2.y,b3.y)-mb)*ib));
+  write_imagef(O,(int2)(x+12,y),convert_float4(((half4)(a0.z,a1.z,a2.z,a3.z)-ma)*ia));
+  write_imagef(O,(int2)(x+15,y),convert_float4(((half4)(b0.z,b1.z,b2.z,b3.z)-mb)*ib));
+  write_imagef(O,(int2)(x+18,y),convert_float4(((half4)(a0.w,a1.w,a2.w,a3.w)-ma)*ia));
+  write_imagef(O,(int2)(x+21,y),convert_float4(((half4)(b0.w,b1.w,b2.w,b3.w)-mb)*ib));
+}"""
+
+
+def patch_input_pack(jit) -> int:
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  lib, new_batch, replaced = None, [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name == TARGET:
+      if lib is None: lib = Device["QCOM"].compiler.compile_cached(SOURCE)
+      program = call.src[0]
+      program = program.replace(arg=replace(program.arg, global_size=(1, 64, 1), local_size=(3, 128, 1)),
+                                src=program.src[:2] +
+                                (program.src[2].replace(arg=SOURCE), program.src[3].replace(arg=lib)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced

extra/gemm/qcom_openpilot_inverse_full_rpt.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""Pack all four FP32 accumulator vectors in the OpenPilot inverse projection."""
+import argparse, hashlib, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import branch as BR, mad_f32 as MAD_F32, mov_f32 as MOV_F32, nop as NOP, plain_name
+
+TARGET="r_32_64_4_4_192_4"
+INVERSE_W_TARGETS={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+SAFE_DONORS={"d4c281a1","1fe26758","e34e7e58"}
+
+
+def replace_src2(ins:bytes, src2:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xff00ffff)|(src2<<16),hi)
+
+
+def replace_low_src(ins:bytes, src:int) -> bytes:
+  lo,hi=struct.unpack("<II",ins)
+  return struct.pack("<II",(lo&0xffffff00)|src,hi)
+
+
+def pack_inverse_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(instrs)!=175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # Move loop control from r13.x into the existing r12.w zero register. This
+  # makes r13-r16 four contiguous accumulator vectors without growing the
+  # shader's declared register file.
+  out=instrs[:11]
+  for acc in ("r13.x","r14.x","r15.x","r16.x"): out.append(MOV_F32(acc,"r12.w",rpt=3))
+  loop_start=len(out)
+  body=list(instrs[16:32])
+  body[0]=replace_low_src(body[0],51)   # add r8.z, r12.w, 192
+  body[1]=replace_low_src(body[1],51)   # add r9.x, r12.w, 384
+  body[2]=replace_src2(body[2],51)      # add r9.z, c28.y, r12.w
+  body[3]=replace_low_src(body[3],51)   # mov r10.x, r12.w
+  out+=body
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r13.x",f"r7.{component}",weight,"r13.x",rpt=3,r=True,sy=component=="x"))
+  out+=instrs[48:60]
+  control=list(instrs[60:65])
+  control[0]=replace_low_src(control[0],51)  # increment r12.w
+  control[2]=replace_low_src(control[2],51)  # compare r12.w
+  control[3]=MOV_F32("r12.w","r0.x")
+  out+=control
+  out.append(BR(loop_start-len(out)))
+  tail=list(instrs[66:131])
+  # The first residual moved from r12.w to r13.x; y/z/w were already in r13.
+  tail[90-66]=replace_src2(tail[90-66],52)
+  out+=tail
+  out += [NOP()]*(len(instrs)-len(out))
+  if len(out)!=len(instrs): raise RuntimeError(f"packed image has {len(out)} instructions")
+  return lib[:off]+b"".join(out)+lib[off+size:]
+
+
+def with_fregs(lib:bytes, count:int) -> bytes:
+  out=bytearray(lib)
+  regoff=struct.unpack_from("<I",out,0x34)[0]+0x14
+  regs=struct.unpack_from("<I",out,regoff)[0]
+  struct.pack_into("<I",out,regoff,(regs&0x80000000)|max(regs&0x7fffffff,count))
+  return bytes(out)
+
+
+def pack_inverse_w_full(lib:bytes) -> bytes:
+  off,size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  ins=[lib[i:i+8] for i in range(off,off+size,8)]
+  if len(ins) not in (174,178): raise RuntimeError(f"expected 174/178 inverse-W instructions, got {len(ins)}")
+  out=ins[:19]+[MOV_F32("r17.x","r12.z",rpt=3)]
+  loop_start=len(out)
+  out+=ins[19:35]
+  for component,weight in zip("xyzw",("r5.x","r2.x","r3.x","r4.x")):
+    out.append(MAD_F32("r17.x",f"r7.{component}",weight,"r17.x",rpt=3,r=True,sy=component=="x"))
+  out+=ins[51:68]
+  out.append(BR(loop_start-len(out)))
+  tail=list(ins[69:])
+  mapping={50:68,52:69,53:70,54:71}
+  first_store=next(i for i,x in enumerate(tail) if struct.unpack_from("<I",x,4)[0]>>24==0xc0)
+  for i in range(first_store):
+    lo,_=struct.unpack("<II",tail[i])
+    src2=(lo>>16)&0xff
+    if src2 in mapping: tail[i]=replace_src2(tail[i],mapping[src2])
+  out+=tail
+  out += [NOP()]*(len(ins)-len(out))
+  return with_fregs(lib[:off]+b"".join(out)+lib[off+size:],18)
+
+
+def patch_model(model,names:set[str]|None=None) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  cache,patched={},0
+  for index,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM: continue
+    name=plain_name(call.src[0].arg.name)
+    if name not in (names if names is not None else {TARGET}|INVERSE_W_TARGETS): continue
+    program=call.src[0]
+    old=program.src[3].arg
+    # These transforms relocate fixed compiler registers. A different QCOM compiler allocation can
+    # have the same instruction count but different live values and must not be patched by index.
+    if hashlib.sha1(old).hexdigest()[:8] not in SAFE_DONORS: continue
+    if old not in cache:
+      cache[old]=pack_inverse_full(old) if name==TARGET else pack_inverse_w_full(old)
+    program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+    batch[index]=call.replace(src=(program,*call.src[1:]))
+    patched+=1
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return patched
+
+
+def main() -> None:
+  ap=argparse.ArgumentParser()
+  ap.add_argument("input")
+  ap.add_argument("output")
+  ap.add_argument("--names",help="comma-separated program families")
+  args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_model(model,set(args.names.split(",")) if args.names else None))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()

extra/gemm/qcom_openpilot_ir3.py

@@ -0,0 +1,34 @@
+"""Small IR3 encoding helpers used by the OpenPilot QCOM graph patches."""
+import re, struct
+
+ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
+
+def plain_name(name:str) -> str: return ANSI_RE.sub("", name)
+
+def _freg(name:str|int) -> int:
+  if isinstance(name, int): return name
+  register, component = name.replace("r", "").split(".")
+  return int(register) * 4 + "xyzw".index(component)
+
+def _pack(lo:int, hi:int) -> bytes: return struct.pack("<II", lo & 0xffffffff, hi & 0xffffffff)
+
+def nop() -> bytes: return _pack(0, 0)
+
+def branch(offset:int) -> bytes: return struct.pack("<iI", offset, 0x00900000)
+
+def mov_f32(dst:str, src:str, rpt:int=0, sy:bool=False, ss:bool=False, r:bool=False) -> bytes:
+  return _pack(_freg(src), (0x30044000 if sy else 0x20044000) | (0x1000 if ss else 0) |
+               (0x800 if r else 0) | ((rpt & 0x7f) << 8) | _freg(dst))
+
+def add_s(dst:str, src:str, imm:int, ss:bool=False) -> bytes:
+  lo = ((0x27 if imm < 0 else 0x20) << 24) | ((imm & 0xff) << 16) | _freg(src)
+  return _pack(lo, 0x42300000 | (0x1000 if ss else 0) | _freg(dst))
+
+def mad_f32(dst:str, src1:str, src2:str, src3:str, rpt:int=0, sy:bool=False, r:bool=False) -> bytes:
+  d, s1, s2, s3 = _freg(dst), _freg(src1), _freg(src2), _freg(src3)
+  hi = ((0x73 if sy else 0x63) << 24) | (0x80 << 16) | ((s2 >> 1) << 16) | (((s2 & 1) << 7 | (rpt & 0x7f)) << 8) | d
+  lo = (0x20000000 if r else 0) | (s3 << 16) | (0x8000 if r else 0) | s1
+  return _pack(lo, hi)
+
+def isam_f32(dst:str, coord:str, tex:int=0, samp:int=0) -> bytes:
+  return _pack((tex * 2) << 24 | ((samp & 7) << 21) | (_freg(coord) * 2 + 1), 0xa0001f00 | _freg(dst))

extra/gemm/qcom_openpilot_level_schedule.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python3
+"""Group ready OpenPilot graph calls by program family without crossing dependency levels."""
+import argparse, pickle
+from collections import defaultdict
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+
+def schedule_levels(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=list(outer.src[0].src[0].src)
+  writer,levels={},{}
+  grouped=defaultdict(list)
+  for sequence,call in enumerate(batch):
+    if call.op is not Ops.CALL or call.src[0].op is not Ops.PROGRAM:
+      grouped[sequence].append((sequence,call))
+      continue
+    deps={writer[arg] for arg in call.src[1:] if arg in writer}
+    level=1+max((levels[dep] for dep in deps),default=-1)
+    levels[sequence]=level
+    grouped[level].append((sequence,call))
+    for output in call.src[0].arg.outs: writer[call.src[output+1]]=sequence
+  scheduled=[]
+  moved=0
+  for entries in grouped.values():
+    ordered=sorted(entries,key=lambda item:(plain_name(item[1].src[0].arg.name),item[0]))
+    scheduled.extend(call for _index,call in ordered)
+    moved+=sum(old_index!=entries[new_index][0] for new_index,(old_index,_call) in enumerate(ordered))
+  if scheduled != batch:
+    model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(scheduled)},walk=True)
+    model.captured.__dict__.pop("linear",None)
+  return moved
+
+
+def main() -> None:
+  ap=argparse.ArgumentParser()
+  ap.add_argument("input")
+  ap.add_argument("output")
+  args=ap.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("moved",schedule_levels(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__=="__main__":main()

extra/gemm/qcom_openpilot_pack_conv_weights.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Prepack static weights for the slow stride-2 openpilot 7x7 convolutions."""
+import argparse, pickle, re
+
+import numpy as np
+from tinygrad import Device
+from tinygrad.device import Buffer
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops, UOp
+from extra.gemm.qcom_openpilot_ir3 import plain_name
+
+TARGETS={"r_16_8_16_2_4_4_7_7", "r_8_4_32_2_4_4_7_7", "r_4_2_64_2_4_4_7_7"}
+
+
+def packed_source(source:str) -> str:
+  start=source.index("    half val0")
+  end=source.index("    int alu20",start)
+  weight_name=re.search(r"__global half\* (data2_\d+)",source).group(1)  # type: ignore[union-attr]
+  block="    int wp=(((alu0*2+alu2)*7+Ridx0)*28);\n"+"\n".join(
+    f"    float4 w{i}=convert_float4(vload4(0,{weight_name}+wp+{i*4}));" for i in range(7))+"\n"
+  source=source[:start]+block+source[end:]
+  accum_start=source.index("    *(buf0+0)",start)
+  loop_prefix=source[start:accum_start]
+  casts=re.findall(r"    float (cast\d+) = \(\(float\)\(val(\d+)\)\);\n",loop_prefix)
+  assert len(casts) == 28
+  source=source[:start]+re.sub(r"    float cast\d+ = \(\(float\)\(val\d+\)\);\n", "", loop_prefix)+source[accum_start:]
+  mapping={cast:("w0.x" if int(val) == 27 else f"w{int(val)%7+1}.x" if int(val) < 6 else
+                       f"w{(int(val)-6)%7}.{'yzw'[(int(val)-6)//7]}") for cast,val in casts}
+  for old,new in sorted(mapping.items(),key=lambda item:-len(item[0])):
+    source=re.sub(rf"\b{old}\b",new,source)
+  return source
+
+
+def pack_weight_buffer(weight:UOp) -> UOp:
+  original=np.asarray(weight.buffer.numpy()).reshape(-1)
+  outputs=original.size//896
+  assert outputs*896 == original.size
+  packed=np.empty((outputs,2,7,7,4),dtype=np.float16)
+  for output in range(outputs):
+    for parity in range(2):
+      for row in range(7):
+        base=output*896+parity+row*28
+        for tap in range(7):
+          for component in range(4): packed[output,parity,row,tap,component]=original[base+component*224+tap*4]
+  buf=Buffer("QCOM",packed.size,weight.dtype,initial_value=bytearray(packed.tobytes()))
+  return UOp.from_buffer(buf)
+
+
+def patch_conv(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      source=packed_source(program.src[2].arg)
+      lib=Device["QCOM"].compiler.compile_cached(source)
+      program=program.replace(src=program.src[:2]+(program.src[2].replace(arg=source),program.src[3].replace(arg=lib)))
+      call=call.replace(src=(program,call.src[1],call.src[2],pack_weight_buffer(call.src[3]),*call.src[4:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+
+
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_conv(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_schedule_projection.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""Reschedule independent texture addresses in the dominant vision projection."""
+import argparse, pickle, struct
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import add_s as ADD_S, branch as BR, isam_f32 as ISAM_F32, mov_f32 as MOV_F32, nop as NOP, plain_name
+
+TARGET="r_32_192_4_4_64_4"
+FIRST_CONV_TARGET="r_64_32_16_4_4_6_3_3_4"
+FORWARD_STYLE={TARGET,"r_32_64_4_4_64_4","r_8_384_4_4_128_4"}
+GAP_STYLE={"r_512_16_4_4_16_4","r_512_48_4_4_16_4","r_128_32_4_4_32_4","r_128_96_4_4_32_4"}
+INVERSE_W_STYLE={"r_512_16_4_4_48_4","r_128_32_4_4_96_4"}
+INVERSE_STYLE={"r_32_64_4_4_192_4"}
+TARGETS=FORWARD_STYLE|GAP_STYLE|INVERSE_W_STYLE|INVERSE_STYLE|{FIRST_CONV_TARGET}
+
+
+def schedule_first_conv(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  # Use registers that the subsequent texture loads overwrite, allowing all
+  # eight independent input/weight coordinates to precede the texture reads.
+  addresses=[]
+  for index,(register,offset) in enumerate(zip(("r0","r2","r3","r4"),(-36,-24,-12,0))):
+    addresses.append(MOV_F32(f"{register}.x","r16.w",ss=index > 0) if offset == 0 else
+                     ADD_S(f"{register}.x","r16.w",offset,ss=index > 0))
+    addresses.append(MOV_F32(f"{register}.y","r16.z"))
+  addresses.extend(instrs[i] for i in (48,51,54,57))
+  loads=[ISAM_F32(dst,f"{coord}.x",tex=0) for dst,coord in zip(("r7.x","r6.x","r1.x","r0.x"),("r0","r2","r3","r4"))]
+  loads.extend(ISAM_F32(dst,coord,tex=1) for dst,coord in zip(("r2.x","r3.x","r4.x","r5.x"),("r8.x","r8.z","r9.x","r9.z")))
+  out=instrs[:32]+addresses+loads+instrs[60:86]
+  out.append(BR(31-len(out)))
+  out.extend(instrs[87:92])
+  out.append(BR(26-len(out)))
+  out.extend(instrs[93:99])
+  out.append(BR(24-len(out)))
+  out.extend(instrs[100:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled first conv has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_native_f16(lib:bytes) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if len(instrs) != 222: raise RuntimeError(f"expected 222 native-FP16 instructions, got {len(instrs)}")
+  addresses=(21,24,27,30,33,40,47,54)
+  loads=(23,26,29,32,35,42,49,56)
+  mads=tuple(range(36,40))+tuple(range(43,47))+tuple(range(50,54))+tuple(range(57,61))
+  out=instrs[:21]+[instrs[i] for i in addresses]+[instrs[i] for i in loads]+[instrs[i] for i in mads]+instrs[61:67]
+  out.append(BR(21-len(out)))
+  out.extend(instrs[68:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled native FP16 has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def schedule_loads(lib:bytes, name:str) -> bytes:
+  image_offset, image_size=struct.unpack_from("<I",lib,0xc0)[0],struct.unpack_from("<I",lib,0x100)[0]
+  instrs=[lib[i:i+8] for i in range(image_offset,image_offset+image_size,8)]
+  if name == TARGET and len(instrs) == 222: return schedule_native_f16(lib)
+  if len(instrs) < 160: raise RuntimeError(f"expected projection shader, got {len(instrs)} instructions")
+  # The compiler emits address, rpt5 nop, texture-read eight times. Calculate
+  # every independent address first, then issue the reads as one contiguous run.
+  if name in FORWARD_STYLE: start,body_end=20,66
+  elif name in GAP_STYLE: start,body_end=26,84
+  elif name in INVERSE_W_STYLE: start,body_end=19,76
+  elif name in INVERSE_STYLE: start,body_end=16,73
+  else: raise RuntimeError(f"unsupported projection {name}")
+  address_indices=tuple(start+3*i for i in range(8))
+  load_indices=tuple(start+3*i+2 for i in range(8))
+  out=instrs[:start]+[instrs[i] for i in address_indices]+[instrs[i] for i in load_indices]+instrs[start+24:body_end]
+  branch_index=len(out)
+  out.append(BR(start-branch_index))
+  out.extend(instrs[body_end+1:])
+  out.extend([NOP()]*(len(instrs)-len(out)))
+  if len(out) != len(instrs): raise RuntimeError(f"scheduled image has {len(out)} instructions")
+  return lib[:image_offset]+b"".join(out)+lib[image_offset+image_size:]
+
+
+def patch_projection(model) -> int:
+  outer=model.captured.linear.src[0]
+  batch=outer.src[0].src[0].src
+  new_batch=[]
+  cache={}
+  replaced=0
+  for call in batch:
+    name=plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in TARGETS:
+      program=call.src[0]
+      old=program.src[3].arg
+      if old not in cache: cache[old]=schedule_first_conv(old) if name == FIRST_CONV_TARGET else schedule_loads(old,name)
+      program=program.replace(src=program.src[:3]+(program.src[3].replace(arg=cache[old]),))
+      call=call.replace(src=(program,*call.src[1:]))
+      replaced+=1
+    new_batch.append(call)
+  model.captured._linear=model.captured.linear.substitute({outer:create_graph_call(new_batch)},walk=True)
+  model.captured.__dict__.pop("linear",None)
+  return replaced
+def main() -> None:
+  parser=argparse.ArgumentParser()
+  parser.add_argument("input")
+  parser.add_argument("output")
+  args=parser.parse_args()
+  with open(args.input,"rb") as f:model=pickle.load(f)
+  print("patched",patch_projection(model))
+  with open(args.output,"wb") as f:pickle.dump(model,f)
+
+
+if __name__ == "__main__":main()

extra/gemm/qcom_openpilot_vision_fp16.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+"""Replace selected driving_vision 1x1 convolutions with vector FP16-acc kernels."""
+import struct
+from dataclasses import replace
+
+from tinygrad.engine.jit import create_graph_call
+from tinygrad.uop.ops import Ops
+from extra.gemm.qcom_openpilot_ir3 import branch as BR, mad_f32 as MAD_F32, nop as NOP, plain_name
+
+TARGET = "r_32_192_4_4_64_4"
+INVERSE_TARGET = "r_32_64_4_4_192_4"
+FIRST_CONV_TARGET = "r_64_32_16_4_4_6_3_3_4"
+FULL_Y_TARGETS = {TARGET, "r_32_64_4_4_64_4"}
+FULL_Z_TARGETS = {"r_8_384_4_4_128_4"}
+GAP_Y_TARGETS = {"r_512_16_4_4_16_4", "r_512_48_4_4_16_4", "r_128_32_4_4_32_4", "r_128_96_4_4_32_4"}
+INVERSE_W_TARGETS = {"r_512_16_4_4_48_4", "r_128_32_4_4_96_4"}
+OTHER_INVERSE_TARGETS: set[str] = set()
+FP32_TARGETS = FULL_Y_TARGETS | FULL_Z_TARGETS | GAP_Y_TARGETS | INVERSE_W_TARGETS | OTHER_INVERSE_TARGETS | {INVERSE_TARGET, FIRST_CONV_TARGET}
+
+def pack_fp32_mads(lib:bytes, component:str="y") -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 116: raise RuntimeError(f"expected FP32 matmul loop through instruction 115, got {len(instrs)}")
+  out = instrs[:44]
+  for k_component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(tuple(f"r{reg}.{component}" for reg in range(13, 17)), ("r7", "r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{k_component}", weight, acc, rpt=3,
+                         sy=len(out) == 44, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed FP32 image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_gap_y_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 237: raise RuntimeError(f"expected at least 237 gap-Y instructions, got {len(instrs)}")
+  out = instrs[:66]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[114:120]
+  branch_index = len(out)
+  out.append(BR(26-branch_index))
+  out += instrs[121:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed gap-Y image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_w_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) < 174: raise RuntimeError(f"expected at least 174 inverse-W instructions, got {len(instrs)}")
+  out = instrs[:59]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.w", "r14.w", "r15.w"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[107:112]
+  branch_index = len(out)
+  out.append(BR(19-branch_index))
+  out += instrs[113:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse-W image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_other_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 179: raise RuntimeError(f"expected 179 other-inverse instructions, got {len(instrs)}")
+  # The first output vector is split by loop-control registers. Keep it scalar;
+  # the remaining three vectors are contiguous from r14.y through r17.x.
+  out = instrs[:60]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.y", "r15.y", "r16.y"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[108:114]
+  branch_index = len(out)
+  out.append(BR(20-branch_index))
+  out += instrs[115:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed other-inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_first_conv_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 262: raise RuntimeError(f"expected 262 first-conv instructions, got {len(instrs)}")
+  out = instrs[:60]
+  first = True
+  for component, weight in zip("xyzw", ("r5", "r2", "r3", "r4")):
+    out.append(MAD_F32("r11.w", f"r7.{component}", f"{weight}.x", "r11.w", sy=first, r=True))
+    first = False
+    out.append(MAD_F32("r12.y", f"r7.{component}", f"{weight}.y", "r12.y", rpt=2, r=True))
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r13.x", "r14.x", "r15.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[124:130]
+  branch_index = len(out)
+  out.append(BR(31-branch_index))
+  out += instrs[131:]
+  # Compacting the innermost loop also relocates the two enclosing-loop branches.
+  # Their targets remain in the untouched prologue, so rebuild their relative offsets.
+  out[92] = BR(26-92)
+  out[99] = BR(24-99)
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed first-conv image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def pack_inverse_fp32_mads(lib:bytes) -> bytes:
+  image_offset, image_size = struct.unpack_from("<I", lib, 0xc0)[0], struct.unpack_from("<I", lib, 0x100)[0]
+  image = lib[image_offset:image_offset+image_size]
+  instrs = [image[i:i+8] for i in range(0, len(image), 8)]
+  if len(instrs) != 175: raise RuntimeError(f"expected 175 inverse instructions, got {len(instrs)}")
+  # The first output vector straddles loop-control r13.x, so retain its scalar
+  # instructions. The remaining r14/r15/r16 accumulator vectors are contiguous.
+  out = instrs[:56]
+  for component, weight in zip("xyzw", ("r5.x", "r2.x", "r3.x", "r4.x")):
+    for acc, activation in zip(("r14.x", "r15.x", "r16.x"), ("r6", "r1", "r0")):
+      out.append(MAD_F32(acc, f"{activation}.{component}", weight, acc, rpt=3, r=True))
+  out += instrs[104:109]
+  branch_index = len(out)
+  out.append(BR(16-branch_index))
+  out += instrs[110:]
+  out += [NOP()] * (len(instrs)-len(out))
+  if len(out) != len(instrs): raise RuntimeError(f"packed inverse image has {len(out)} instructions")
+  return lib[:image_offset] + b"".join(out) + lib[image_offset+image_size:]
+
+
+def patch_fp32_rpt(jit, names:set[str]|None=None) -> int:
+  """Apply the verified FP32-accumulate repeat packing to a captured vision JIT."""
+  outer = jit.captured.linear.src[0]
+  batch = outer.src[0].src[0].src
+  new_batch, replaced = [], 0
+  for call in batch:
+    name = plain_name(call.src[0].arg.name) if call.op is Ops.CALL and call.src[0].op is Ops.PROGRAM else ""
+    if name in FP32_TARGETS and (names is None or name in names):
+      program = call.src[0]
+      patch = (pack_fp32_mads if name in FULL_Y_TARGETS else
+               (lambda lib:pack_fp32_mads(lib, "z")) if name in FULL_Z_TARGETS else
+               pack_gap_y_fp32_mads if name in GAP_Y_TARGETS else
+               pack_inverse_w_fp32_mads if name in INVERSE_W_TARGETS else
+               pack_other_inverse_fp32_mads if name in OTHER_INVERSE_TARGETS else
+               pack_first_conv_fp32_mads if name == FIRST_CONV_TARGET else pack_inverse_fp32_mads)
+      program = program.replace(src=program.src[:3] + (program.src[3].replace(arg=patch(program.src[3].arg)),))
+      if name == INVERSE_TARGET:
+        program = program.replace(arg=replace(program.arg, global_size=(8, 2, 1), local_size=(8, 16, 1)))
+      call, replaced = call.replace(src=(program, *call.src[1:])), replaced+1
+    new_batch.append(call)
+  if replaced:
+    jit.captured._linear = jit.captured.linear.substitute({outer:create_graph_call(new_batch)}, walk=True)
+    jit.captured.__dict__.pop("linear", None)
+  return replaced

tinygrad/runtime/ops_qcom.py

@@ -84,7 +84,6 @@ class QCOMComputeQueue(HWQueue):
     self.cmd(mesa.CP_WAIT_FOR_IDLE)
     if self.dev.gpu_id[:2] < (7, 3):
       self.cmd(mesa.CP_EVENT_WRITE, qreg.cp_event_write_0(event=mesa.CACHE_FLUSH_TS), *data64_le(signal.value_addr), lo32(value))
-      self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
     else:
       # TODO: support devices starting with 8 Gen 1. Also, 700th series have convenient CP_GLOBAL_TIMESTAMP and CP_LOCAL_TIMESTAMP
       raise RuntimeError('CP_EVENT_WRITE7 is not supported')
@@ -135,7 +134,7 @@ class QCOMComputeQueue(HWQueue):
     self.reg(mesa.REG_A6XX_SP_PERFCTR_SHADER_MASK, qreg.a6xx_sp_perfctr_shader_mask(cs=True))
     self.reg(mesa.REG_A6XX_TPL1_MODE_CNTL, qreg.a6xx_tpl1_mode_cntl(isammode=mesa.ISAMMODE_GL if prg.NIR else mesa.ISAMMODE_CL))
     self.reg(mesa.REG_A6XX_TPL1_DBG_ECO_CNTL, 0)
-    self.cmd(mesa.CP_WAIT_FOR_IDLE)
+    # CP_RUN_OPENCL snapshots the programmed compute state, so graph command streams can program the next dispatch without draining prior shader work.
 
     self.reg(mesa.REG_A6XX_SP_CS_NDRANGE_0,
              qreg.a6xx_sp_cs_ndrange_0(kerneldim=3, localsizex=local_size[0] - 1, localsizey=local_size[1] - 1, localsizez=local_size[2] - 1),
@@ -192,7 +191,7 @@ class QCOMComputeQueue(HWQueue):
                qreg.cp_exec_cs_1(ngroups_x=global_size[0]), qreg.cp_exec_cs_2(ngroups_y=global_size[1]), qreg.cp_exec_cs_3(_ngroups_z=global_size[2]))
     else: self.cmd(mesa.CP_RUN_OPENCL, 0)
 
-    self._cache_flush(write_back=True, invalidate=False, sync=False, memsync=False)
+    # The final queue signal writes back data needed by host and cross-queue consumers.
     return self
 
 class QCOMArgsState(HCQArgsState):