mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-07 21:13:41 +08:00
firestar5683
32 lines
1.4 KiB
Python
32 lines
1.4 KiB
Python
"""Protect one pinned stock QCOM capture; no artifact selection or metadata bypass."""
|
|
from collections import Counter
|
|
import hashlib
|
|
from tinygrad.uop.ops import Ops
|
|
from tinygrad.runtime.support.hcq2 import HCQInfo
|
|
|
|
|
|
def protect_capture(jit, protection):
|
|
cap = getattr(jit, 'captured', None)
|
|
if cap is None or 'linear' in cap.__dict__:
|
|
raise ValueError('Expected a never-executed stored capture')
|
|
for entry in cap._linear.src:
|
|
call = entry.without_after
|
|
if call.op is not Ops.CALL or not (call.body.op is Ops.COPY or isinstance(call.arg.aux, HCQInfo)):
|
|
raise ValueError('Unknown stock capture schedule')
|
|
def inventory(root):
|
|
cpu = {u.src[-1] for u in root.toposort() if u.op is Ops.PROGRAM and u.arg.target.device == 'CPU'}
|
|
return Counter(hashlib.sha256(u.arg).hexdigest() for u in root.toposort() if u.op is Ops.BINARY and u not in cpu)
|
|
replacements = {}
|
|
for u in cap._linear.toposort():
|
|
if u.op is Ops.CALL and isinstance(u.arg.aux, HCQInfo):
|
|
if u.arg.aux.device != ('QCOM',):
|
|
raise ValueError('Stock capture is not the declared QCOM backend')
|
|
replacements[u] = protection.protect_call(u)
|
|
if not replacements:
|
|
raise ValueError('No native QCOM submission')
|
|
proposed = cap._linear.substitute(replacements, enter_calls=True)
|
|
if inventory(cap._linear) != inventory(proposed):
|
|
raise ValueError('Stock capture opaque device payload changed')
|
|
cap._linear = proposed
|
|
return jit
|