mawei/onepilot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

TR16 Model, fix radar routine (#211)

Browse Source 
* UV+DTR model

* DTR model.. again.

* fix naviGPS

* fix radar...

* fix..

* test

* fix..

* carrot serv

* fix..

* fix.. fleet

* fix.. radar

* fix atc

* Steam Powered model..

* fix.. radarLatFactor range.. 200->500

* fix.. dbc..

* side

* SP v2

* brake light

* fix brakelight

* fix..

* add datetime...

* fix..

* fix..

* fix..

* fix..

* blind spot

* fix tz

* fix..

* ff

* radarLatFactor

* fix.. bsd

* Revert "fix.. bsd"

This reverts commit 1d0d1434470e1b92c65eaffaeb8dd7cd779f85ee.

* fix.. bsd side..

* test

* fix.. e2e conditions

* Revert "test"

This reverts commit 0ce791dbd66c17260366ed1a4df2626c602dbb7d.

* TR16

* fix cut-in detect threshold  3.4 -> 2.6

* fix.. jerk_l limit 5->10

* fix..

* fix.. gm

* fix.. OPTIMA_H mass

* fix.. radar..

* fix radar..

* fix..

* Radar...

* fix..

* fix..

* fix..

* fix.. radartrack 3

* fix..

* fix..

* fix..

* merge..

* fix.. canfd

* fix..

* fix..

* fix..

* fix.. radard

* new cut_in

* Revert "new cut_in"

This reverts commit b9b6e9b33318fe1ce7d626468139b17848efcdcd.

* fix..

* new cut_in detect...

* fix.. disp..

* fix..

* fix..

* fix.. center radar..

* fix.. radar y_sane..

* fix..

* fix..

* hkg jerk 10 -> 5

* fix..

* fix..

* fix.. radar dbc..

* fix..

* fix.. jLead filter..

* test new radar interface..

* fix..

* fix..

* test time...

* Revert "test time..."

This reverts commit 63e9187736985c4dc4b4f3736674ba7cda6adc3f.

* fix radar..

* fix..

* FireHose model..

* tinygrad

* Update interface.py

* fix..

* fix.. nff toyota corolla_tss2

* fix..

* fix..

* fix.. radar

* fix..

* fix.. radar, y_gate

* fix.. radar..

* fix.. for clone..

* scc radar enable at low speed..

* fix.. settings..

* fix.

* fix..

* fix.. radarTimeStep.

* TR16 model again..

* RELEASE.md

* fix cut-in detection...

* fix.. registeration timeout 15sec..

* fix..

* fix.. radar processing.

* fix..

* fix..

* fix..

* fix..

* fix..

* fix..
This commit is contained in:
carrot
2025-09-05 15:43:10 +09:00
committed by GitHub
parent 9d35822092
commit 77a8919349
254 changed files with 8412 additions and 225959 deletions
Download Patch File Download Diff File Expand all files Collapse all files
RELEASES.md
+8
View File
@@ -1,3 +1,11 @@
Carrot2-v9 (2025-09-xx)
========================
* TR16 Model
* RadarTrack Option:3 (Cutin Detect, vision fail detection)
* RdarTrack Option: 2 (always use SCC radar)
* Brake light (CANFD)
Carrot2-v9 (2025-08-12)
========================
* TombRaider16 v2 model.
cereal/log.capnp
+2
View File
@@ -768,6 +768,8 @@ struct RadarState @0x9a185389d6fdd05f {
leadsCenter @15 : List(LeadData);
leadsLeft @16 : List(LeadData);
leadsRight @17 : List(LeadData);
leadsLeft2 @19 : List(LeadData);
leadsRight2 @20 : List(LeadData);
struct LeadData {
dRel @0 :Float32;
opendbc_repo/opendbc/car/car.capnp
+5
View File
@@ -254,6 +254,7 @@ struct CarState {
carrotCruise @73 : Int16;
leftLaneLine @74 : Int16; # -1: no lane, 0: dashed, 1: solid, +10: white, +20: yellow, ex) 21: solid yellow
rightLaneLine @75 : Int16; # -1: no lane, 0: dashed, 1: solid, +10: white, +20: yellow, ex) 21: solid yellow
datetime @76 :UInt64; # timestamp in milliseconds since epoch
struct Tpms {
fl @0 :Float32;
@@ -455,6 +456,10 @@ struct CarControl {
leadRightDist @19: Float32;
leadLeftLat @20: Float32;
leadRightLat @21: Float32;
leadLeftDist2 @22: Float32;
leadRightDist2 @23: Float32;
leadLeftLat2 @24: Float32;
leadRightLat2 @25: Float32;
# not used with the dash, TODO: separate structs for dash UI and device UI
audibleAlert @5: AudibleAlert;
opendbc_repo/opendbc/car/gm/carstate.py
+1 -1
View File
@@ -214,7 +214,7 @@ class CarState(CarStateBase):
if CP.transmissionType == TransmissionType.direct:
pt_messages += [
("EBCMRegenPaddle", 50),
("EVDriveMode", 0),
("EVDriveMode", float('nan')),
]
loopback_messages = [
("ASCMLKASteeringCmd", float('nan')),
opendbc_repo/opendbc/car/hyundai/carstate.py
+36 -3
View File
@@ -13,6 +13,9 @@ from opendbc.car.interfaces import CarStateBase
from openpilot.common.params import Params
from datetime import datetime
from zoneinfo import ZoneInfo
ButtonType = structs.CarState.ButtonEvent.Type
@@ -25,6 +28,20 @@ BUTTONS_DICT = {Buttons.RES_ACCEL: ButtonType.accelCruise, Buttons.SET_DECEL: Bu
GearShifter = structs.CarState.GearShifter
NUMERIC_TO_TZ = {
840: "America/New_York", # 미국 (US) → 동부 시간대
124: "America/Toronto", # 캐나다 (CA) → 동부 시간대
250: "Europe/Paris", # 프랑스 (FR)
276: "Europe/Berlin", # 독일 (DE)
826: "Europe/London", # 영국 (GB)
392: "Asia/Tokyo", # 일본 (JP)
156: "Asia/Shanghai", # 중국 (CN)
410: "Asia/Seoul", # 한국 (KR)
36: "Australia/Sydney", # 호주 (AU)
356: "Asia/Kolkata", # 인도 (IN)
}
class CarState(CarStateBase):
def __init__(self, CP):
super().__init__(CP)
@@ -140,9 +157,11 @@ class CarState(CarStateBase):
self.CAM_0x2a4 = True if 0x2a4 in fingerprints[alt_bus] else False
self.STEER_TOUCH_2AF = True if 0x2af in fingerprints[pt_bus] else False
self.TPMS = True if 0x3a0 in fingerprints[pt_bus] else False
self.LOCAL_TIME = True if 1264 in fingerprints[pt_bus] else False
self.cp_bsm = None
self.time_zone = "UTC"
self.controls_ready_count = 0
def update(self, can_parsers) -> structs.CarState:
@@ -421,7 +440,7 @@ class CarState(CarStateBase):
ret.vEgo, ret.aEgo = self.update_speed_kf(ret.vEgoRaw)
ret.standstill = ret.wheelSpeeds.fl <= STANDSTILL_THRESHOLD and ret.wheelSpeeds.rr <= STANDSTILL_THRESHOLD
ret.brakeLights = ret.brakePressed or (ret.aEgo < -0.5) ### TODO: �ӽ÷� brakeLight�� ������.
ret.brakeLights = ret.brakePressed or cp.vl["TCS"]["BrakeLight"] == 1
ret.steeringRateDeg = cp.vl["STEERING_SENSORS"]["STEERING_RATE"]
ret.steeringAngleDeg = cp.vl["STEERING_SENSORS"]["STEERING_ANGLE"] * -1
@@ -517,9 +536,13 @@ class CarState(CarStateBase):
if not self.is_metric:
speedLimit *= CV.MPH_TO_KPH
ret.speedLimit = speedLimit if speedLimit < 255 else 0
if int(self.hda_info_4a3["NEW_SIGNAL_4"]) == 17:
if int(self.hda_info_4a3["MapSource"]) == 2:
speed_limit_cam = True
if self.time_zone == "UTC":
country_code = int(self.hda_info_4a3["CountryCode"])
self.time_zone = ZoneInfo(NUMERIC_TO_TZ.get(country_code, "UTC"))
self.new_msg_4b4 = cp.vl["NEW_MSG_4B4"] if self.NEW_MSG_4B4 else None
self.tcs_info_373 = cp.vl["TCS"]
@@ -552,6 +575,16 @@ class CarState(CarStateBase):
if self.CP.flags & HyundaiFlags.EV:
ret.cruiseState.nonAdaptive = cp.vl["MANUAL_SPEED_LIMIT_ASSIST"]["MSLA_ENABLED"] == 1
if self.LOCAL_TIME and self.time_zone != "UTC":
lt = cp.vl["LOCAL_TIME"]
y, m, d, H, M, S = int(lt["YEAR"]) + 2000, int(lt["MONTH"]), int(lt["DATE"]), int(lt["HOURS"]), int(lt["MINUTES"]), int(lt["SECONDS"])
try:
dt_local = datetime(y, m, d, H, M, S, tzinfo=self.time_zone)
ret.datetime = int(dt_local.timestamp() * 1000)
except:
#print(f"Error parsing local time: {y}-{m}-{d} {H}:{M}:{S} in {self.time_zone}")
pass
prev_cruise_buttons = self.cruise_buttons[-1]
#self.cruise_buttons.extend(cp.vl_all[self.cruise_btns_msg_canfd]["CRUISE_BUTTONS"])
#carrot {{
opendbc_repo/opendbc/car/hyundai/hyundaicanfd.py
+22 -2
View File
@@ -539,14 +539,16 @@ def create_ccnc_messages(CP, packer, CAN, frame, CC, CS, hud_control, disp_angle
values["LANELINE_CURVATURE"] = (min(abs(curvature), 15) + (-1 if curvature < 0 else 0)) if lat_active else 0
values["LANELINE_CURVATURE_DIRECTION"] = 1 if curvature < 0 and lat_active else 0
# lane_color = 6 if lat_active else 2
lane_color = 2 # 6: green, 2: white, 4: yellow
if hud_control.leftLaneDepart:
values["LANELINE_LEFT"] = 4 if (frame // 50) % 2 == 0 else 1
else:
values["LANELINE_LEFT"] = 2 if hud_control.leftLaneVisible else 0
values["LANELINE_LEFT"] = lane_color if hud_control.leftLaneVisible else 0
if hud_control.rightLaneDepart:
values["LANELINE_RIGHT"] = 4 if (frame // 50) % 2 == 0 else 1
else:
values["LANELINE_RIGHT"] = 2 if hud_control.rightLaneVisible else 0
values["LANELINE_RIGHT"] = lane_color if hud_control.rightLaneVisible else 0
#values["LANELINE_LEFT_POSITION"] = 15
#values["LANELINE_RIGHT_POSITION"] = 15
@@ -575,6 +577,16 @@ def create_ccnc_messages(CP, packer, CAN, frame, CC, CS, hud_control, disp_angle
values['RF_DETECT'] = 3 if hud_control.leadRightDist > 30 else 4
values['RF_DETECT_DISTANCE'] = hud_control.leadRightDist
values['RF_DETECT_LATERAL'] = hud_control.leadRightLat
"""
if values['LR_DETECT'] == 0 and hud_control.leadLeftDist2 > 0:
values['LR_DETECT'] = 4
values['LR_DETECT_DISTANCE'] = 2
values['LR_DETECT_LATERAL'] = hud_control.leadLeftLat2
if values['RR_DETECT'] == 0 and hud_control.leadRightDist2 > 0:
values['RR_DETECT'] = 4
values['RR_DETECT_DISTANCE'] = 2
values['RR_DETECT_LATERAL'] = hud_control.leadRightLat2
"""
ret.append(packer.make_can_msg("ADRV_0x1ea", CAN.ECAN, values))
if CS.adrv_info_162 is not None:
@@ -596,6 +608,14 @@ def create_ccnc_messages(CP, packer, CAN, frame, CC, CS, hud_control, disp_angle
values['RF_DETECT'] = 3 if hud_control.leadRightDist > 30 else 4
values['RF_DETECT_DISTANCE'] = hud_control.leadRightDist
values['RF_DETECT_LATERAL'] = hud_control.leadRightLat
if values['LR_DETECT'] == 0 and hud_control.leadLeftDist2 > 0:
values['LR_DETECT'] = 4
values['LR_DETECT_DISTANCE'] = 2
values['LR_DETECT_LATERAL'] = hud_control.leadLeftLat2
if values['RR_DETECT'] == 0 and hud_control.leadRightDist2 > 0:
values['RR_DETECT'] = 4
values['RR_DETECT_DISTANCE'] = 2
values['RR_DETECT_LATERAL'] = hud_control.leadRightLat2
else:
sensors = [
('lf', 'LF_DETECT'),
opendbc_repo/opendbc/car/hyundai/interface.py
+1 -1
View File
@@ -180,7 +180,7 @@ class CarInterface(CarInterfaceBase):
#ret.radarUnavailable = False # TODO: canfd... carrot, hyundai cars have radar
ret.radarTimeStep = 0.05 if params.get_int("EnableRadarTracks") > 0 else 0.02
ret.radarTimeStep = 0.05 #if params.get_int("EnableRadarTracks") > 0 else 0.02
ret.pcmCruise = not ret.openpilotLongitudinalControl
ret.startingState = False # True # carrot
opendbc_repo/opendbc/car/hyundai/radar_interface.py
+137 -128
View File
@@ -8,6 +8,7 @@ from openpilot.common.params import Params
from opendbc.car.hyundai.hyundaicanfd import CanBus
from openpilot.common.filter_simple import MyMovingAverage
SCC_TID = 0
RADAR_START_ADDR = 0x500
RADAR_MSG_COUNT = 32
RADAR_START_ADDR_CANFD1 = 0x210
@@ -49,7 +50,7 @@ def get_radar_can_parser_scc(CP):
class RadarInterface(RadarInterfaceBase):
def __init__(self, CP):
super().__init__(CP)
self.updated_messages = set()
self.canfd = True if CP.flags & HyundaiFlags.CANFD else False
self.radar_group1 = False
if self.canfd:
@@ -63,176 +64,184 @@ class RadarInterface(RadarInterfaceBase):
else:
self.radar_start_addr = RADAR_START_ADDR
self.radar_msg_count = RADAR_MSG_COUNT
self.trigger_msg = self.radar_start_addr + self.radar_msg_count - 1
self.params = Params()
self.radar_tracks = self.params.get_int("EnableRadarTracks") >= 1
self.updated_tracks = set()
self.updated_scc = set()
self.rcp_tracks = get_radar_can_parser(CP, self.radar_tracks, self.radar_start_addr, self.radar_msg_count)
self.rcp_scc = get_radar_can_parser_scc(CP)
self.trigger_msg_scc = 416 if self.canfd else 0x420
self.trigger_msg_tracks = self.radar_start_addr + self.radar_msg_count - 1
self.track_id = 0
self.radar_off_can = CP.radarUnavailable
self.params = Params()
self.radar_tracks = self.params.get_int("EnableRadarTracks") >= 1
self.rcp = get_radar_can_parser(CP, self.radar_tracks, self.radar_start_addr, self.radar_msg_count)
if not self.radar_tracks:
self.rcp = get_radar_can_parser_scc(CP)
self.trigger_msg = 416 if self.canfd else 0x420
# 50Hz (SCC), 20Hz (RadarTracks)
self.vLead_filter = MyMovingAverage(13) # for SCC radar 0.1 unit
self.vRel_last = 0
self.dRel_last = 0
# Initialize pts
total_tracks = self.radar_msg_count * ( 2 if self.radar_group1 else 1)
for track_id in range(total_tracks):
t_id = track_id + 32
self.pts[t_id] = structs.RadarData.RadarPoint()
self.pts[t_id].measured = False
self.pts[t_id].trackId = t_id
self.pts[SCC_TID] = structs.RadarData.RadarPoint()
self.pts[SCC_TID].trackId = SCC_TID
self.frame = 0
def update(self, can_strings):
if self.radar_off_can or (self.rcp is None):
self.frame += 1
if self.radar_off_can or (self.rcp_tracks is None and self.rcp_scc is None):
return super().update(None)
vls = self.rcp.update(can_strings)
self.updated_messages.update(vls)
if self.rcp_scc is not None:
vls_s = self.rcp_scc.update(can_strings)
self.updated_scc.update(vls_s)
if not self.radar_tracks and self.frame % 5 == 0:
self._update_scc(self.updated_scc)
self.updated_scc.clear()
ret = structs.RadarData()
if not self.rcp_scc.can_valid:
ret.errors.canError = True
ret.points = list(self.pts.values())
return ret
if self.radar_tracks and self.rcp_tracks is not None:
vls_t = self.rcp_tracks.update(can_strings)
self.updated_tracks.update(vls_t)
if self.trigger_msg_tracks in self.updated_tracks:
self._update(self.updated_tracks)
self._update_scc(self.updated_scc)
self.updated_scc.clear()
self.updated_tracks.clear()
ret = structs.RadarData()
if not self.rcp_tracks.can_valid:
ret.errors.canError = True
ret.points = list(self.pts.values())
return ret
if self.trigger_msg not in self.updated_messages:
return None
rr = self._update(self.updated_messages) if self.radar_tracks else self._update_scc(self.updated_messages)
self.updated_messages.clear()
return rr
return None
def _update(self, updated_messages):
ret = structs.RadarData()
if self.rcp is None:
return ret
if not self.rcp.can_valid:
ret.errors.canError = True
t_id = 32
for addr in range(self.radar_start_addr, self.radar_start_addr + self.radar_msg_count):
msg = self.rcp.vl[f"RADAR_TRACK_{addr:x}"]
if addr not in self.pts:
self.pts[addr] = structs.RadarData.RadarPoint()
self.pts[addr].trackId = self.track_id
self.track_id += 1
msg = self.rcp_tracks.vl[f"RADAR_TRACK_{addr:x}"]
if self.radar_group1:
valid = msg['VALID_CNT1'] > 0
valid = msg['VALID_CNT1'] > 10
elif self.canfd:
valid = msg['VALID'] > 0
valid = msg['VALID_CNT'] > 10
else:
valid = msg['STATE'] in (3, 4)
if valid:
if self.radar_group1:
self.pts[addr].measured = True
self.pts[addr].dRel = msg['LONG_DIST1']
self.pts[addr].yRel = msg['LAT_DIST1']
self.pts[addr].vRel = msg['REL_SPEED1']
self.pts[addr].vLead = self.pts[addr].vRel + self.v_ego
self.pts[addr].aRel = msg['REL_ACCEL1']
self.pts[addr].yvRel = msg['LAT_SPEED1']
elif self.canfd:
self.pts[addr].measured = True
self.pts[addr].dRel = msg['LONG_DIST']
self.pts[addr].yRel = msg['LAT_DIST']
self.pts[addr].vRel = msg['REL_SPEED']
self.pts[addr].vLead = self.pts[addr].vRel + self.v_ego
self.pts[addr].aRel = msg['REL_ACCEL']
self.pts[addr].yvRel = msg['LAT_SPEED']
else:
azimuth = math.radians(msg['AZIMUTH'])
self.pts[addr].measured = True
self.pts[addr].dRel = math.cos(azimuth) * msg['LONG_DIST']
self.pts[addr].yRel = 0.5 * -math.sin(azimuth) * msg['LONG_DIST']
self.pts[addr].vRel = msg['REL_SPEED']
self.pts[addr].vLead = self.pts[addr].vRel + self.v_ego
self.pts[addr].aRel = msg['REL_ACCEL']
self.pts[addr].yvRel = 0.0
self.pts[t_id].measured = bool(valid)
if not valid:
self.pts[t_id].dRel = 0
self.pts[t_id].yRel = 0
self.pts[t_id].vRel = 0
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0
elif self.radar_group1:
self.pts[t_id].dRel = msg['LONG_DIST1']
self.pts[t_id].yRel = msg['LAT_DIST1']
self.pts[t_id].vRel = msg['REL_SPEED1']
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = msg['REL_ACCEL1']
self.pts[t_id].yvRel = msg['LAT_SPEED1']
elif self.canfd:
self.pts[t_id].dRel = msg['LONG_DIST']
self.pts[t_id].yRel = msg['LAT_DIST']
self.pts[t_id].vRel = msg['REL_SPEED']
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = msg['REL_ACCEL']
self.pts[t_id].yvRel = msg['LAT_SPEED']
else:
del self.pts[addr]
azimuth = math.radians(msg['AZIMUTH'])
self.pts[t_id].dRel = math.cos(azimuth) * msg['LONG_DIST']
self.pts[t_id].yRel = 0.5 * -math.sin(azimuth) * msg['LONG_DIST']
self.pts[t_id].vRel = msg['REL_SPEED']
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = msg['REL_ACCEL']
self.pts[t_id].yvRel = 0.0
t_id += 1
# radar group1은 하나의 msg에 2개의 레이더가 들어있음.
if self.radar_group1:
for addr in range(self.radar_start_addr, self.radar_start_addr + self.radar_msg_count):
msg = self.rcp.vl[f"RADAR_TRACK_{addr:x}"]
msg = self.rcp_tracks.vl[f"RADAR_TRACK_{addr:x}"]
addr += 16
if addr not in self.pts:
self.pts[addr] = structs.RadarData.RadarPoint()
self.pts[addr].trackId = self.track_id
self.track_id += 1
valid = msg['VALID_CNT2'] > 0
if valid:
self.pts[addr].measured = True
self.pts[addr].dRel = msg['LONG_DIST2']
self.pts[addr].yRel = msg['LAT_DIST2']
self.pts[addr].vRel = msg['REL_SPEED2']
self.pts[addr].vLead = self.pts[addr].vRel + self.v_ego
self.pts[addr].aRel = msg['REL_ACCEL2']
self.pts[addr].yvRel = msg['LAT_SPEED2']
valid = msg['VALID_CNT2'] > 10
self.pts[t_id].measured = bool(valid)
if not valid:
self.pts[t_id].dRel = 0
self.pts[t_id].yRel = 0
self.pts[t_id].vRel = 0
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0
else:
del self.pts[addr]
ret.points = list(self.pts.values())
return ret
self.pts[t_id].dRel = msg['LONG_DIST2']
self.pts[t_id].yRel = msg['LAT_DIST2']
self.pts[t_id].vRel = msg['REL_SPEED2']
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = msg['REL_ACCEL2']
self.pts[t_id].yvRel = msg['LAT_SPEED2']
t_id += 1
def _update_scc(self, updated_messages):
ret = structs.RadarData()
if self.rcp is None:
return ret
if not self.rcp.can_valid:
ret.errors.canError = True
cpt = self.rcp.vl
cpt = self.rcp_scc.vl
t_id = SCC_TID
if self.canfd:
dRel = cpt["SCC_CONTROL"]['ACC_ObjDist']
vRel = cpt["SCC_CONTROL"]['ACC_ObjRelSpd']
new_pts = abs(dRel - self.dRel_last) > 3 or abs(vRel - self.vRel_last) > 1
vLead = vRel + self.v_ego
valid = 0 < dRel < 150 #cpt["SCC_CONTROL"]['OBJ_STATUS'] and dRel < 150
for ii in range(1):
if valid:
if ii not in self.pts or new_pts:
self.pts[ii] = structs.RadarData.RadarPoint()
self.pts[ii].trackId = self.track_id
self.track_id = min(1 - self.track_id, 1)
self.vLead_filter.set_all(vLead)
self.pts[ii].dRel = dRel
self.pts[ii].yRel = 0
self.pts[ii].vRel = vRel
self.pts[ii].vLead = self.vLead_filter.process(vLead)
self.pts[ii].aRel = 0 #float('nan')
self.pts[ii].yvRel = 0 #float('nan')
self.pts[ii].measured = True
else:
if ii in self.pts:
del self.pts[ii]
valid = 0 < dRel < 150 and not new_pts #cpt["SCC_CONTROL"]['OBJ_STATUS'] and dRel < 150
self.pts[t_id].measured = bool(valid)
if not valid:
self.pts[t_id].dRel = 0
self.pts[t_id].yRel = 0
self.pts[t_id].vRel = 0
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0
else:
self.pts[t_id].dRel = dRel
self.pts[t_id].yRel = 0
self.pts[t_id].vRel = vRel
self.pts[t_id].vLead = vLead
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0 #float('nan')
else:
dRel = cpt["SCC11"]['ACC_ObjDist']
vRel = cpt["SCC11"]['ACC_ObjRelSpd']
new_pts = abs(dRel - self.dRel_last) > 3 or abs(vRel - self.vRel_last) > 1
vLead = vRel + self.v_ego
valid = cpt["SCC11"]['ACC_ObjStatus'] and dRel < 150
for ii in range(1):
if valid:
if ii not in self.pts or new_pts:
self.pts[ii] = structs.RadarData.RadarPoint()
self.pts[ii].trackId = self.track_id
self.track_id = min(1 - self.track_id, 1)
self.vLead_filter.set_all(vLead)
self.pts[ii].dRel = dRel
self.pts[ii].yRel = -cpt["SCC11"]['ACC_ObjLatPos'] # in car frame's y axis, left is negative
self.pts[ii].vRel = vRel
self.pts[ii].vLead = self.vLead_filter.process(vLead)
self.pts[ii].aRel = 0 #float('nan')
self.pts[ii].yvRel = 0 #float('nan')
self.pts[ii].measured = True
else:
if ii in self.pts:
del self.pts[ii]
valid = cpt["SCC11"]['ACC_ObjStatus'] and dRel < 150 and not new_pts
self.pts[t_id].measured = bool(valid)
if not valid:
self.pts[t_id].dRel = 0
self.pts[t_id].yRel = 0
self.pts[t_id].vRel = 0
self.pts[t_id].vLead = self.pts[t_id].vRel + self.v_ego
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0
else:
self.pts[t_id].dRel = dRel
self.pts[t_id].yRel = -cpt["SCC11"]['ACC_ObjLatPos'] # in car frame's y axis, left is negative
self.pts[t_id].vRel = vRel
self.pts[t_id].vLead = vLead
self.pts[t_id].aRel = float('nan')
self.pts[t_id].yvRel = 0 #float('nan')
self.dRel_last = dRel
self.vRel_last = vRel
ret.points = list(self.pts.values())
return ret
opendbc_repo/opendbc/car/hyundai/values.py
+1 -1
View File
@@ -559,7 +559,7 @@ class CAR(Platforms):
# TODO: may support adjacent years. may have a non-zero minimum steering speed
KIA_OPTIMA_H = HyundaiPlatformConfig(
[HyundaiCarDocs("Kia Optima Hybrid 2017", "Advanced Smart Cruise Control", car_parts=CarParts.common([CarHarness.hyundai_c]))],
CarSpecs(mass=3558 * CV.LB_TO_KG, wheelbase=2.8, steerRatio=13.75, tireStiffnessFactor=0.5),
CarSpecs(mass=3758 * CV.LB_TO_KG, wheelbase=2.8, steerRatio=13.75, tireStiffnessFactor=0.5),
flags=HyundaiFlags.HYBRID | HyundaiFlags.LEGACY,
)
KIA_OPTIMA_H_G4_FL = HyundaiPlatformConfig(
opendbc_repo/opendbc/car/interfaces.py
+52 -28
View File
@@ -220,38 +220,62 @@ class MyTrack:
self.jLead_avg = FirstOrderFilter(self.jLead, 0.4, self.dt)
self.yRel_avg = FirstOrderFilter(self.yRel, 0.02, self.dt)
self.yvRel_avg = FirstOrderFilter(self.yvRel, 0.02, self.dt)
def update(self, radar_point):
self.vLead = radar_point.vLead
"""
if abs(radar_point.dRel - self.dRel) > 3.0 or abs(self.vRel - radar_point.vRel) > 20.0 * self.dt:
self.cnt = 0
self.jLead = 0.0
self.aLead = 0.0
self.vLead_avg.x = self.vLead
self.aLead_avg.x = self.aLead
self.jLead_avg.x = self.jLead
self.v_lead_filtered_last = self.vLead
"""
self.cnt = 0
self.yRel = self.yRel_avg.update(radar_point.yRel)
self.yvRel = self.yvRel_avg.update(radar_point.yvRel)
v_lead_filtered = self.vLead_avg.update(self.vLead)
pseudo_stop = abs(v_lead_filtered) < 0.3 and abs(self.vLead - v_lead_filtered) < 0.05
a_raw = (v_lead_filtered - self.v_lead_filtered_last) / self.dt
self.v_lead_filtered_last = v_lead_filtered
a_lead = self.aLead_avg.update(a_raw if not pseudo_stop else 0.0)
j_lead = (a_lead - self.aLead) / self.dt
self.aLead = a_lead
self.jLead = self.jLead_avg.update(j_lead)
# Store latest values
def init_point(self, radar_point):
self.dRel = radar_point.dRel
self.vRel = radar_point.vRel
self.yRel = radar_point.yRel
self.yvRel = radar_point.yvRel
self.vLead = radar_point.vLead
self.v_lead_filtered_last = self.vLead
self.aLead = 0.0
self.jLead = 0.0
self.vLead_avg.x = self.vLead
self.aLead_avg.x = self.aLead
self.jLead_avg.x = self.jLead
self.yRel_avg.x = self.yRel
self.yvRel_avg.x = self.yvRel
def update(self, radar_point):
if not radar_point.measured:
if self.cnt > 0:
self.init_point(radar_point)
self.cnt = 0
elif self.cnt < 1:
self.init_point(radar_point)
self.cnt += 1
else:
self.vLead = radar_point.vLead
"""
if abs(radar_point.dRel - self.dRel) > 3.0 or abs(self.vRel - radar_point.vRel) > 20.0 * self.dt:
self.cnt = 0
self.jLead = 0.0
self.aLead = 0.0
self.vLead_avg.x = self.vLead
self.aLead_avg.x = self.aLead
self.jLead_avg.x = self.jLead
self.v_lead_filtered_last = self.vLead
"""
self.cnt += 1
self.yRel = self.yRel_avg.update(radar_point.yRel)
self.yvRel = self.yvRel_avg.update(radar_point.yvRel)
v_lead_filtered = self.vLead_avg.update(self.vLead)
pseudo_stop = abs(v_lead_filtered) < 0.3 and abs(self.vLead - v_lead_filtered) < 0.05
a_raw = (v_lead_filtered - self.v_lead_filtered_last) / self.dt
self.v_lead_filtered_last = v_lead_filtered
a_lead = self.aLead_avg.update(a_raw if not pseudo_stop else 0.0)
j_lead = (a_lead - self.aLead) / self.dt
self.aLead = a_lead
self.jLead = self.jLead_avg.update(j_lead if self.cnt > 2 else 0.0)
# Store latest values
self.dRel = radar_point.dRel
self.vRel = radar_point.vRel
self.cnt += 1
# generic car and radar interfaces
class RadarInterfaceBase(ABC):
opendbc_repo/opendbc/car/torque_data/lat_models/TOYOTA_COROLLA_TSS2.json
+120 -1
View File
@@ -1 +1,120 @@
{"input_std":[[9.754316],[1.5248336],[0.56031793],[0.049654935],[1.5070868],[1.5136857],[1.5187283],[1.4884274],[1.4468032],[1.3927733],[1.3395035],[0.049471103],[0.049509283],[0.049545065],[0.049528807],[0.049440794],[0.049183726],[0.048698023]],"model_test_loss":0.01838030107319355,"input_size":18,"current_date_and_time":"2023-08-11_23-49-46","input_mean":[[22.216816],[-0.068382844],[-0.0003495027],[-0.009731857],[-0.06888316],[-0.06969439],[-0.07004197],[-0.06833163],[-0.062542245],[-0.054059293],[-0.046969976],[-0.0098139495],[-0.009790095],[-0.009772296],[-0.009818108],[-0.009930541],[-0.010102344],[-0.010210523]],"input_vars":["v_ego","lateral_accel","lateral_jerk","roll","lateral_accel_m03","lateral_accel_m02","lateral_accel_m01","lateral_accel_p03","lateral_accel_p06","lateral_accel_p10","lateral_accel_p15","roll_m03","roll_m02","roll_m01","roll_p03","roll_p06","roll_p10","roll_p15"],"output_size":1,"layers":[{"dense_1_b":[[0.071733765],[-2.3017254],[-0.002747584],[-0.16260307],[-0.039528668],[1.2633095],[0.40724456]],"dense_1_W":[[0.14254493,0.7560598,0.0022538262,0.026828635,-1.0626036,1.6978638,-0.58277315,-0.532054,0.13709316,0.08024807,-0.01103483,0.7790654,-0.36199996,-0.27354154,-0.40135407,0.112546146,-0.12273542,0.17872572],[-0.75257075,1.8171842,-1.4895765,0.6856587,-1.047781,-2.1456847,1.0205821,-1.1760736,-0.51131254,-0.1572458,0.3351541,-0.8285507,-0.10711446,0.89037967,-0.5866782,0.046956923,-0.018645424,0.15313248],[0.04169559,-0.58989626,-6.2346396,0.21148093,1.0282952,-0.15679906,0.015086645,-0.4938231,-0.98810667,-0.8081533,1.2733475,-1.0567358,-0.6226471,-0.07101477,1.028134,0.7304045,0.7475682,-0.87618846],[0.00729714,-1.1409336,0.00029273634,0.021291457,0.47364047,-1.0528986,0.6670064,-0.31857133,-0.2512035,-0.14272158,0.29329458,-0.47188643,-0.18927035,0.7875639,0.17967238,0.47387323,-0.0324174,-0.19277702],[0.13785772,-0.30069488,0.002800617,0.31292388,0.408026,-1.1410295,0.61747766,0.056347255,-0.16441363,0.14121056,-0.08929794,-0.39708617,-0.2790936,0.44575593,0.28871247,-0.30545992,-0.063934155,0.045969535],[0.9263208,-0.0041096397,-1.2649248,-0.22491741,-0.7898359,-0.6432323,0.46669382,-0.69405776,-0.039997756,-0.059218984,0.16558072,-0.21184966,-0.06702091,0.2703531,0.22768503,0.39336708,-0.25291198,0.07505179],[-0.007783014,-2.9238482,0.010547396,-0.17914231,-0.9764414,-1.8973196,-2.0817938,-0.5377732,-0.8602919,-0.92122966,-1.1030004,-0.061209604,0.05769448,-0.2639501,-0.011335088,0.0692365,0.37218606,-0.34912306]],"activation":"σ"},{"dense_2_W":[[-1.2037164,0.82475835,0.58777946,-0.41806436,-0.07912497,-0.9752202,0.43038344],[-0.87383485,1.7625434,-0.9010436,0.04913966,-0.80236125,-0.8888718,-1.1724374],[-1.149093,0.588674,-0.79056436,0.5840478,-0.6800404,-2.3066614,-0.30252007],[-0.7044807,-0.3438525,0.18204984,0.9572669,0.17036141,-0.4475017,-0.10285365],[-0.35474542,-0.91775477,0.16575153,-0.6729545,0.004384356,-0.5627307,0.20499422],[-0.62687695,-0.89932716,-0.40928647,0.004326587,0.10205387,-0.38104382,0.05874869],[0.20604928,0.56815124,-0.031190628,-0.6709605,-0.70632714,0.3854373,-0.38627437],[-0.32655486,-0.48946378,-0.27688035,-0.28829703,-0.28283417,-1.2973477,-0.30639303],[0.95871663,0.015167555,-0.4939611,-0.82378834,-0.6367465,0.11321682,0.36064836],[0.17759234,-0.05241082,-0.9149439,-0.5095627,-0.028105808,-0.1622127,-0.56177574],[0.9313323,-0.7414045,-0.10210027,-0.6010315,-0.9564538,-0.379816,0.40230134],[-0.09108528,0.32599035,-0.026278114,0.36614916,0.5561161,0.84679973,0.47440732],[-1.5251282,1.6984508,0.8223704,-0.14148536,-0.21781565,-1.3052686,0.55815554]],"activation":"σ","dense_2_b":[[-0.5853772],[-0.1194652],[0.09660614],[-0.20516787],[-0.2301157],[-0.33061823],[-0.039316446],[-0.11710167],[-0.029613337],[-0.25951824],[0.089471936],[-0.0013098614],[-0.92236656]]},{"dense_3_W":[[-0.20270818,0.6552373,0.9410651,-0.7787618,-0.18410496,0.016181383,0.62408584,-0.3269968,0.7875554,-0.40091902,0.7076024,-0.21312849,-1.0302222],[0.27061725,-0.13254738,-0.19332989,-0.070887744,-0.23539513,-0.14211313,-0.56738764,-0.3196864,-0.56094,0.17847942,-0.31821814,0.57166994,0.60353655],[0.5228776,-0.36678898,-0.47331262,0.4382709,0.092253216,0.38762012,-0.44059324,-0.44392672,-0.19738796,-0.43091238,-0.8047946,0.5009608,-0.23722588]],"activation":"identity","dense_3_b":[[-0.08095956],[0.10173851],[0.10956224]]},{"dense_4_W":[[0.48155427,-0.78180456,-0.52694654]],"dense_4_b":[[-0.0919536]],"activation":"identity"}]}
{ "input_std": [
[ 10.799555 ],
[ 1.1120019 ],
[ 0.6616854 ],
[ 0.028482547 ],
[ 1.1156307 ],
[ 1.1159346 ],
[ 1.1148219 ],
[ 1.0780851 ],
[ 1.0371089 ],
[ 0.97595906 ],
[ 0.9065878 ],
[ 0.02847537 ],
[ 0.028477926 ],
[ 0.028481513 ],
[ 0.0284707 ],
[ 0.028458958 ],
[ 0.028488474 ],
[ 0.028483905 ]
],
"model_test_loss": 0.02256479673087597,
"input_size": 18,
"current_date_and_time": "2023-12-25_22-57-21",
"input_mean": [
[ 19.2235 ],
[ -0.06585429 ],
[ -0.0064189 ],
[ -0.005323282 ],
[ -0.067707 ],
[ -0.0671286 ],
[ -0.06649458 ],
[ -0.060063194 ],
[ -0.05377905 ],
[ -0.04412506 ],
[ -0.035479717 ],
[ -0.0053720735 ],
[ -0.0053506913 ],
[ -0.0053401864 ],
[ -0.0052933325 ],
[ -0.005350221 ],
[ -0.0054709557 ],
[ -0.005746576 ]
],
"input_vars": [ "v_ego", "lateral_accel", "lateral_jerk", "roll", "lateral_accel_m03", "lateral_accel_m02", "lateral_accel_m01", "lateral_accel_p03", "lateral_accel_p06", "lateral_accel_p10", "lateral_accel_p15", "roll_m03", "roll_m02", "roll_m01", "roll_p03", "roll_p06", "roll_p10", "roll_p15" ],
"output_size": 1,
"layers": [
{
"dense_1_b": [
[ -0.075336866 ],
[ -1.4379132 ],
[ 0.12846093 ],
[ -0.83898723 ],
[ -0.36279294 ],
[ -0.056873046 ],
[ 0.047107063 ]
],
"dense_1_W": [
[ -0.050702084, 0.23833075, 9.068456, 0.11729692, -2.24747, -1.9006892, -1.2767023, 2.3568945, 2.131596, 1.4344207, -1.032736, 0.49961105, -0.04982021, -0.3041553, -0.38732168, -0.08194013, 0.1400462, 0.06643751 ],
[ -0.7581488, -0.020518934, 0.5159834, -0.08114609, 0.7436459, -0.33153492, -0.49035493, 0.12532651, -0.12278706, 0.09580404, 0.12504527, 0.20612119, -0.45066848, 0.1813876, -0.24614686, 0.4758823, 0.2178751, -0.30032337 ],
[ -1.3949319, 0.67544717, -0.7753792, -0.15714955, -0.17681755, -0.25429842, -0.54722327, 0.5418026, 0.10016261, 0.0053565986, -0.39035755, -0.04587303, 0.36687377, -0.22805889, 0.093026236, 0.05298613, -0.23806311, 0.14899784 ],
[ -0.42220324, 0.04483509, -0.40907508, 0.11821204, 0.073940635, -0.017443, 0.19060583, -0.6515505, 0.34508246, -0.08026583, -0.018325916, -0.16471778, -0.05598569, 0.17481138, -0.14715311, -0.080125585, 0.3393668, -0.19069892 ],
[ -1.4330497, -0.64631116, 0.6935077, 0.2762588, -0.14516634, 0.15581381, 0.6870634, -0.014884118, -0.44512376, 0.12536626, 0.2644134, 0.14975125, -0.33620888, 0.1905219, -0.19605163, -0.1808306, -0.0016299848, 0.106743194 ],
[ -0.010526905, -0.7717455, 0.0070439703, -0.22383846, -0.415478, 0.0844275, 0.2170057, -0.069986425, 0.22479944, 0.2915719, -0.4006947, 0.07989299, -0.0019174323, 0.220215, -0.020739302, 0.010475229, -0.27987975, 0.21558948 ],
[ 0.00034062448, 0.6978409, -0.0499567, -0.16071245, 0.07709055, -0.11454262, -0.42469633, 0.14393206, 0.38380122, 0.15885039, -0.24375756, 0.1721643, 0.024053562, 0.219485, -0.079801545, -0.6738785, -0.18035688, 0.3041013 ]
],
"activation": "σ"
},
{
"dense_2_W": [
[ -0.47582296, -0.3389653, -0.25085196, 0.4886556, 0.6374338, 1.1124122, -1.06381 ],
[ 0.27025986, -1.2483587, 0.17614846, 0.6991092, -0.2602735, -0.120778754, -0.09193157 ],
[ 0.4948387, 0.7217855, 0.111300245, -0.9388221, -0.4064722, -0.17349482, 0.43001428 ],
[ 0.22625388, 0.74308527, 0.6891021, -0.26306373, -0.597923, -1.2967286, 0.22016492 ],
[ 0.021139862, -1.1632888, -1.0416571, 0.3494392, 0.6425689, 1.8975137, -0.8081407 ],
[ -0.52389294, 0.6372325, 0.6633958, -0.64765304, 0.20338675, -0.027452167, 0.506291 ],
[ 0.39088923, 0.7343456, 0.09120028, -1.1630384, 0.023677234, -0.19186972, 0.07007424 ],
[ -0.449708, 0.3169958, 0.4992463, -0.67045796, -0.47493598, -0.00015411235, 0.43920174 ],
[ -1.1875321, -0.68294007, -1.1207345, -0.65358764, -0.14696452, 0.3581226, 0.048447974 ],
[ -1.1105211, 0.01114233, -0.15045986, 0.6279203, 0.81262004, 0.6925559, -0.9141913 ],
[ 0.29771498, -0.085332215, -0.10415446, 0.30199984, -0.15799537, 0.36737302, -0.6503492 ],
[ -0.64377147, -0.65133274, -0.5873282, 1.0216556, 0.32842088, 1.4965141, -0.34098 ],
[ 0.14713718, -0.773232, -0.35718775, 0.5903775, -0.02301098, 0.48789024, 0.0144560225 ]
],
"activation": "σ",
"dense_2_b": [
[ -0.06163904 ],
[ -0.23192078 ],
[ -0.1009967 ],
[ -0.16059391 ],
[ 0.23239176 ],
[ 0.14465374 ],
[ -0.1871548 ],
[ 0.073026784 ],
[ -0.046453997 ],
[ -0.41197765 ],
[ -0.2929203 ],
[ 0.12238579 ],
[ -0.23292084 ]
]
},
{
"dense_3_W": [
[ 0.4508514, 0.08078731, 0.37861827, 0.28543904, 0.042243864, 0.20070161, 0.45177004, -0.23327906, 0.2789905, -0.30319306, -0.13606939, -0.2258465, 0.30893984 ],
[ 0.19797017, -0.54021436, 0.50532675, 0.6035604, -0.7982811, 0.510931, 0.10049933, 0.506153, 0.59682614, -0.4687595, -0.1153152, -0.7137997, -0.22665377 ],
[ 0.5610325, -0.23229897, -0.61403847, -0.3500231, 0.35477564, -0.08623257, -0.5712205, 0.10214144, -0.10332114, 0.71979856, 0.07346336, 0.16925354, -0.10480455 ]
],
"activation": "identity",
"dense_3_b": [
[ -0.05611255 ],
[ 0.059826545 ],
[ -0.05803037 ]
]
},
{
"dense_4_W": [ [ -0.15714681, 1.1904545, -1.1465737 ] ],
"dense_4_b": [ [ 0.058684155 ] ],
"activation": "identity"
}
]
}
opendbc_repo/opendbc/dbc/generator/hyundai/hyundai_canfd.dbc
+12 -10
View File
@@ -235,7 +235,7 @@ BO_ 373 TCS: 24 XXX
SG_ ACCEnable : 67|2@0+ (1,0) [0|3] "" XXX
SG_ ACC_REQ : 68|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_5 : 72|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_2 : 74|1@0+ (1,0) [0|1] "" XXX
SG_ BrakeLight : 74|2@1+ (1,0) [0|3] "" XXX
SG_ NEW_SIGNAL_3 : 76|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_1 : 80|1@0+ (1,0) [0|1] "" XXX
SG_ DriverBraking : 81|1@0+ (1,0) [0|1] "" XXX
@@ -337,11 +337,11 @@ BO_ 354 CCNC_0x162: 32 CCNC
SG_ RF_DETECT_DISTANCE : 141|11@1+ (0.1,0) [0|204.7] "m" XXX
SG_ RF_DETECT_LATERAL : 152|7@1+ (0.1,0) [0|127] "m" XXX
SG_ LR_DETECT : 163|5@1+ (1,0) [0|31] "" XXX
SG_ LR_DETECT_DISTANCE : 168|9@1+ (0.1,0) [0|204.7] "m" XXX
SG_ LR_DETECT_LATERAL : 177|7@1+ (0.1,0) [0|127] "m" XXX
SG_ LR_DETECT_DISTANCE : 175|8@0+ (0.1,0) [0|255] "" XXX
SG_ LR_DETECT_LATERAL : 176|7@1+ (0.1,0) [0|127] "" XXX
SG_ RR_DETECT : 192|5@1+ (1,0) [0|31] "" XXX
SG_ RR_DETECT_DISTANCE : 197|9@1+ (0.1,0) [0|204.7] "m" XXX
SG_ RR_DETECT_LATERAL : 206|7@1+ (0.1,0) [0|127] "m" XXX
SG_ RR_DETECT_DISTANCE : 197|8@1+ (0.1,0) [0|255] "" XXX
SG_ RR_DETECT_LATERAL : 205|7@1+ (0.1,0) [0|127] "" XXX
SG_ FAULT_FSS : 213|3@1+ (1,0) [0|7] "" XXX
SG_ FAULT_FCA : 216|3@1+ (1,0) [0|7] "" XXX
SG_ FAULT_LSS : 219|3@1+ (1,0) [0|7] "" XXX
@@ -834,12 +834,14 @@ BO_ 1151 HVAC_TOUCH_BUTTONS: 8 XXX
SG_ HEAT_BUTTON : 52|1@0+ (1,0) [0|1] "" XXX
BO_ 1187 HDA_INFO_4A3: 8 XXX
SG_ SIGNAL_0 : 0|4@1+ (1,0) [0|15] "" XXX
SG_ NEW_SIGNAL_1 : 4|4@1+ (1,0) [0|15] "" XXX
SG_ LinkClass : 0|3@1+ (1,0) [0|7] "" XXX
SG_ Frwinfo : 3|3@1+ (1,0) [0|7] "" XXX
SG_ SpeedUnit : 6|2@1+ (1,0) [0|3] "" XXX
SG_ SPEED_LIMIT : 15|8@0+ (1,0) [0|255] "" XXX
SG_ NEW_SIGNAL_3 : 23|8@0+ (1,0) [0|255] "" XXX
SG_ NEW_SIGNAL_4 : 31|8@0+ (1,0) [0|255] "" XXX
SG_ NEW_SIGNAL_5 : 39|8@0+ (1,0) [0|255] "" XXX
SG_ CountryCode : 16|10@1+ (1,0) [0|1023] "" XXX
SG_ MapSource : 27|3@1+ (1,0) [0|7] "" XXX
SG_ TollExist : 30|2@1+ (1,0) [0|3] "" XXX
SG_ TunnelExist : 38|2@1+ (1,0) [0|3] "" XXX
SG_ NEW_SIGNAL_6 : 47|24@0+ (1,0) [0|16777215] "" XXX
BO_ 1204 NEW_MSG_4B4: 8 XXX
opendbc_repo/opendbc/dbc/generator/hyundai/hyundai_canfd_radar.py
+6 -7
View File
@@ -79,13 +79,12 @@ BO_ {a} RADAR_TRACK_{a:x}: 32 RADAR
for a in range(0x3a5, 0x3a5 + 32):
f.write(f"""
BO_ {a} RADAR_TRACK_{a:x}: 24 RADAR
SG_ NEW_SIGNAL_9 : 24|1@0+ (1,0) [0|1] "" XXX
SG_ VALID : 25|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_5 : 30|10@1+ (1,0) [0|1023] "" XXX
SG_ NEW_SIGNAL_3 : 40|10@1+ (1,0) [0|1023] "" XXX
SG_ NEW_SIGNAL_6 : 50|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_8 : 51|1@0+ (1,0) [0|1] "" XXX
SG_ NEW_SIGNAL_7 : 55|4@0+ (1,0) [0|15] "" XXX
SG_ VALID : 25|2@0+ (1,0) [0|3] "" XXX
SG_ VALID2 : 28|2@0+ (1,0) [0|3] "" XXX
SG_ PROB : 30|10@1+ (1,0) [0|1023] "" XXX
SG_ VALID_CNT : 47|8@0+ (1,0) [0|255] "" XXX
SG_ NEW_SIGNAL_7 : 51|4@0+ (1,0) [0|15] "" XXX
SG_ NEW_SIGNAL_6 : 55|4@0+ (1,0) [0|15] "" XXX
SG_ NEW_SIGNAL_2 : 62|7@0- (1,0) [0|127] "" XXX
SG_ LONG_DIST : 63|13@1+ (0.05,0) [0|8191] "" XXX
SG_ LAT_DIST : 76|12@1- (0.05,0) [0|4095] "" XXX
selfdrive/car/cruise.py
+4 -4
View File
@@ -406,12 +406,12 @@ class VCruiseCarrot:
elif not b.pressed and self.button_cnt > 0 and bt == self.button_prev:
if bt == ButtonType.cancel:
button_type = bt
elif not self.long_pressed:
elif not self.long_pressed:
if bt == ButtonType.accelCruise:
button_kph += SPEED_UP_UNIT if is_metric else SPEED_UP_UNIT * CV.MPH_TO_KPH
unit = SPEED_UP_UNIT if is_metric else SPEED_UP_UNIT * CV.MPH_TO_KPH
button_kph = math.ceil((button_kph + 0.01) / unit) * unit
elif bt == ButtonType.decelCruise:
#button_kph -= SPEED_DOWN_UNIT if is_metric else SPEED_DOWN_UNIT * CV.MPH_TO_KPH
unit = SPEED_DOWN_UNIT * 1 if is_metric else CV.MPH_TO_KPH
unit = SPEED_DOWN_UNIT if is_metric else SPEED_DOWN_UNIT * CV.MPH_TO_KPH
button_kph = math.floor((button_kph - 0.01) / unit) * unit
button_type = bt
self.long_pressed = False
selfdrive/carrot/carrot_functions.py
+4 -4
View File
@@ -368,9 +368,9 @@ class CarrotPlanner:
v_cruise_kph = self.cruise_eco_control(v_ego_cluster_kph, v_cruise_kph)
v_cruise_kph, atc_active = self._update_carrot_man(sm, v_ego_kph, v_cruise_kph)
if atc_active and not self.atc_active and self.xState not in [XState.e2eStop, XState.e2eStopped, XState.lead]:
if self.atcType in ["turn left", "turn right", "atc left", "atc right"]:
self.xState = XState.e2ePrepare
#if atc_active and not self.atc_active and self.xState not in [XState.e2eStop, XState.e2eStopped, XState.lead]:
# if self.atcType in ["turn left", "turn right", "atc left", "atc right"]:
# self.xState = XState.e2ePrepare
self.atc_active = atc_active
v_cruise = v_cruise_kph * CV.KPH_TO_MS
@@ -407,7 +407,7 @@ class CarrotPlanner:
self.events.add(EventName.trafficSignChanged)
elif self.xState == XState.e2eStopped:
if carstate.gasPressed:
self.xState = XState.e2ePrepare
self.xState = XState.e2eCruise #XState.e2ePrepare
elif lead_detected and (radarstate.leadOne.dRel - stop_model_x) < 2.0:
self.xState = XState.lead
elif self.stopping_count == 0:
selfdrive/carrot/carrot_man.py
+2 -1100
View File
File diff suppressed because it is too large Load Diff
selfdrive/carrot/carrot_serv.py
+1145
View File
File diff suppressed because it is too large Load Diff
selfdrive/carrot_settings.json
+4 -4
View File
@@ -1371,12 +1371,12 @@
"group": "시작",
"name": "EnableRadarTracks",
"title": "레이더트랙,SCC레이더설정(0)",
"descr": "-1: 레이더트랙 미사용(SCC 항시사용)\n 0: 레이더트랙 미사용\n 1: 레이더트랙 사용\n 2: 레이더트랙 사용(SCC 항시사용)",
"descr": "-1: SCC레이더사용(SCC 항시사용), 0: SCC레이더사용\n 1: 레이더트랙 사용, 2: 레이더트랙 사용(SCC 항시사용)\n 3:레이더트랙 사용(끼어들기,스텔스차량검출)",
"egroup": "START",
"etitle": "Enable radar tracks(0)",
"edescr": "-1: Radar Track Disabled, 0: Radar Track Disabled (SCC Always On), 1: Radar Track Enabled (SCC Always On), 2: Radar Track Enabled",
"edescr": "-1: SCC Radar(SCC Always On), 0: SCC Radar\n, 1: Radar Track Enabled, 2: Radar Track Enabled(SCC Always On)\n 3:Radar Track Eanbled(Cutin detect and etc)",
"min": -1,
"max": 2,
"max": 3,
"default": 0,
"unit": 1
},
@@ -1389,7 +1389,7 @@
"etitle": "RadarTrack Cutin Factor(0)",
"edescr": "Higher values increase sensitivity to cutting-in or exiting vehicles",
"min": 0,
"max": 200,
"max": 1000,
"default": 0,
"unit": 10
},
selfdrive/controls/controlsd.py
+41 -41
View File
@@ -56,10 +56,10 @@ class Controls:
self.desired_curvature = 0.0
self.yStd = 0.0
self.lead_left_dRel = None
self.lead_left_Lat = None
self.lead_right_dRel = None
self.lead_right_Lat = None
self.side_state = {
"left": {"main": {"dRel": None, "lat": None}, "sub": {"dRel": None, "lat": None}},
"right": {"main": {"dRel": None, "lat": None}, "sub": {"dRel": None, "lat": None}},
}
self.LoC = LongControl(self.CP)
self.VM = VehicleModel(self.CP)
@@ -187,6 +187,41 @@ class Controls:
return CC, lac_log
def _update_side(self, side: str, leads2, road_edge, bsd_state, hudControl):
def ema(prev, curr, a=0.02):
return curr if prev is None else prev * (1 - a) + curr * a
def set_hud(side_cap, name, val):
setattr(hudControl, f"lead{side_cap}{name}", float(val if val is not None else 0.0))
st = self.side_state[side]
if road_edge <= 2.0 or not leads2:
st["main"] = {"dRel": None, "lat": None}
st["sub"] = {"dRel": None, "lat": None}
if not bsd_state:
return
lead_main = leads2[0] if len(leads2) > 0 else None
side_cap = side.capitalize()
if bsd_state:
set_hud(side_cap, "Dist2", 1)
set_hud(side_cap, "Lat2", 3.2)
# 첫 번째가 10m 이내라면 sub 업데이트 + 두 번째를 main으로
elif len(leads2) > 1 and lead_main.dRel < 10:
st["sub"]["dRel"] = ema(st["sub"]["dRel"], lead_main.dRel)
st["sub"]["lat"] = ema(st["sub"]["lat"], abs(lead_main.dPath))
set_hud(side_cap, "Dist2", st["sub"]["dRel"])
set_hud(side_cap, "Lat2", st["sub"]["lat"])
lead_main = leads2[1]
if len(leads2) > 0:
st["main"]["dRel"] = ema(st["main"]["dRel"], lead_main.dRel)
st["main"]["lat"] = ema(st["main"]["lat"], abs(lead_main.dPath))
set_hud(side_cap, "Dist", st["main"]["dRel"])
set_hud(side_cap, "Lat", st["main"]["lat"])
def publish(self, CC, lac_log):
CS = self.sm['carState']
@@ -251,43 +286,8 @@ class Controls:
meta = self.sm['modelV2'].meta
hudControl.modelDesire = 1 if meta.desire == log.Desire.turnLeft else 2 if meta.desire == log.Desire.turnRight else 0
road_edge_left = meta.distanceToRoadEdgeLeft
road_edge_right = meta.distanceToRoadEdgeRight
def _find_closest_lead(leads, road_edge):
if road_edge < 2.0:
return None
valid_leads = [
lead for lead in leads
if lead.status and abs(lead.dPath) < 4.2 and lead.vLead > 2.0 and 2 < lead.dRel < 130
#if lead.status and abs(lead.dPath) < 4.2 and ((lead.vLead >= 2.0 and 5 < lead.dRel < 100) or (lead.vLead < 2.0 and 3 < lead.dRel < 30))
]
return min(valid_leads, key=lambda l: l.dRel) if valid_leads else None
lead_left = _find_closest_lead(radarState.leadsLeft, road_edge_left)
lead_right = _find_closest_lead(radarState.leadsRight, road_edge_right)
if lead_left is not None:
if self.lead_left_dRel is None:
self.lead_left_dRel = lead_left.dRel
self.lead_left_Lat = abs(lead_left.dPath)
else:
self.lead_left_dRel = self.lead_left_dRel * 0.98 + lead_left.dRel * 0.02
self.lead_left_Lat = self.lead_left_Lat * 0.98 + abs(lead_left.dPath) * 0.02
hudControl.leadLeftDist = self.lead_left_dRel
hudControl.leadLeftLat = self.lead_left_Lat
else:
self.lead_left_dRel = None
if lead_right is not None:
if self.lead_right_dRel is None:
self.lead_right_dRel = lead_right.dRel
self.lead_right_Lat = abs(lead_right.dPath)
else:
self.lead_right_dRel = self.lead_right_dRel * 0.98 + lead_right.dRel * 0.02
self.lead_right_Lat = self.lead_right_Lat * 0.98 + abs(lead_right.dPath) * 0.02
hudControl.leadRightDist = self.lead_right_dRel
hudControl.leadRightLat = self.lead_right_Lat
else:
self.lead_right_dRel = None
self._update_side("left", radarState.leadsLeft2, meta.distanceToRoadEdgeLeft, CS.leftBlindspot, hudControl)
self._update_side("right", radarState.leadsRight2, meta.distanceToRoadEdgeRight, CS.rightBlindspot, hudControl)
hudControl.rightLaneVisible = True
hudControl.leftLaneVisible = True
selfdrive/controls/lib/desire_helper.py
+15 -9
View File
@@ -130,7 +130,9 @@ class DesireHelper:
self.lane_available_last = False
self.edge_available_last = False
self.object_detected_count = 0
self.object_detected_count = 0
self.lane_available_trigger = False
self.lane_appeared = False
self.laneChangeNeedTorque = 0
self.laneChangeBsd = 0
@@ -147,6 +149,7 @@ class DesireHelper:
self.desire_disable_count = 0
self.blindspot_detected_counter = 0
self.auto_lane_change_enable = False
self.next_lane_change = False
def check_lane_state(self, modeldata):
lane_width_left, self.distance_to_road_edge_left, self.distance_to_road_edge_left_far, lane_prob_left = calculate_lane_width(modeldata.laneLines[0], modeldata.laneLineProbs[0],
@@ -264,7 +267,7 @@ class DesireHelper:
lane_exist_counter = self.lane_exist_left_count.counter if blinker_state == BLINKER_LEFT else self.lane_exist_right_count.counter
lane_available = self.available_left_lane if blinker_state == BLINKER_LEFT else self.available_right_lane
edge_available = self.available_left_edge if blinker_state == BLINKER_LEFT else self.available_right_edge
lane_appeared = lane_exist_counter == int(0.2 / DT_MDL)
self.lane_appeared = self.lane_appeared or lane_exist_counter == int(0.2 / DT_MDL)
radar = radarState.leadLeft if blinker_state == BLINKER_LEFT else radarState.leadRight
side_object_dist = radar.dRel + radar.vLead * 4.0 if radar.status else 255
@@ -275,20 +278,21 @@ class DesireHelper:
lane_exist_counter = 0
lane_available = True
edge_available = True
lane_appeared = False
self.lane_appeared = False
self.lane_available_trigger = False
self.object_detected_count = 0
#lane_available_trigger = not self.lane_available_last and lane_available
lane_change_available = (lane_available or edge_available) and lane_line_info < 20 # lane_line_info가 20보다 작으면 흰색라인임.
lane_available_trigger = False
self.lane_available_trigger = False
lane_width_diff = self.lane_width_left_diff if atc_blinker_state == BLINKER_LEFT else self.lane_width_right_diff
distance_to_road_edge = self.distance_to_road_edge_left if atc_blinker_state == BLINKER_LEFT else self.distance_to_road_edge_right
lane_width_side = self.lane_width_left if atc_blinker_state == BLINKER_LEFT else self.lane_width_right
if lane_width_diff > 0.8 and (lane_width_side < distance_to_road_edge):
lane_available_trigger = True
self.lane_available_trigger = True
edge_availabled = not self.edge_available_last and edge_available
side_object_detected = self.object_detected_count > -0.3 / DT_MDL
lane_appeared = lane_appeared and distance_to_road_edge < 4.0
self.lane_appeared = self.lane_appeared and distance_to_road_edge < 4.0
if self.carrot_lane_change_count > 0:
@@ -297,8 +301,8 @@ class DesireHelper:
else:
auto_lane_change_blocked = ((atc_blinker_state == BLINKER_LEFT) and (driver_blinker_state != BLINKER_LEFT))
#auto_lane_change_trigger = not auto_lane_change_blocked and edge_available and (lane_available_trigger or edge_availabled or lane_appeared) and not side_object_detected
auto_lane_change_trigger = self.auto_lane_change_enable and not auto_lane_change_blocked and edge_available and (lane_available_trigger or lane_appeared) and not side_object_detected
self.desireLog = f"L:{self.auto_lane_change_enable},{auto_lane_change_blocked},E:{lane_available},{edge_available},A:{lane_available_trigger},{lane_appeared},{lane_width_diff:.1f},{lane_width_side:.1f},{distance_to_road_edge:.1f}={auto_lane_change_trigger}"
auto_lane_change_trigger = self.auto_lane_change_enable and not auto_lane_change_blocked and edge_available and (self.lane_available_trigger or self.lane_appeared) and not side_object_detected
self.desireLog = f"L:{self.auto_lane_change_enable},{auto_lane_change_blocked},E:{lane_available},{edge_available},A:{self.lane_available_trigger},{self.lane_appeared},{lane_width_diff:.1f},{lane_width_side:.1f},{distance_to_road_edge:.1f}={auto_lane_change_trigger}"
if not lateral_active or self.lane_change_timer > LANE_CHANGE_TIME_MAX:
#print("Desire canceled")
@@ -328,6 +332,7 @@ class DesireHelper:
# 맨끝차선이 아니면(측면에 차선이 있으면), ATC 자동작동 안함.
#self.auto_lane_change_enable = False if lane_exist_counter > 0 else True
self.auto_lane_change_enable = False if lane_exist_counter > 0 or lane_change_available else True
self.next_lane_change = False
# LaneChangeState.preLaneChange
@@ -360,7 +365,7 @@ class DesireHelper:
if self.blindspot_detected_counter > 0 and not ignore_bsd: # BSD검출시
if torque_applied and not block_lanechange_bsd:
self.lane_change_state = LaneChangeState.laneChangeStarting
elif self.laneChangeNeedTorque > 0: # 조향토크필요
elif self.laneChangeNeedTorque > 0 or self.next_lane_change: # 조향토크필요
if torque_applied:
self.lane_change_state = LaneChangeState.laneChangeStarting
elif driver_desire_enabled:
@@ -388,6 +393,7 @@ class DesireHelper:
self.lane_change_direction = LaneChangeDirection.none
if desire_enabled:
self.lane_change_state = LaneChangeState.preLaneChange
self.next_lane_change = True
else:
self.lane_change_state = LaneChangeState.off
selfdrive/controls/radard.py
+228 -146
View File
@@ -3,6 +3,7 @@ import math
import numpy as np
from collections import deque
from typing import Any
import heapq
import capnp
from cereal import messaging, log, car
@@ -34,6 +35,8 @@ class Track:
self.is_stopped_car_count = 0
self.selected_count = 0
self.cut_in_count = 0
self.measured = False
def update(self, md, pt, ready, radar_reaction_factor):
@@ -48,6 +51,8 @@ class Track:
self.yvLead = pt.yvRel
self.measured = pt.measured # measured or estimate
if not self.measured:
self.cnt = 0
if ready:
self.dPath = self.yRel + np.interp(self.dRel, md.position.x, md.position.y)
@@ -56,8 +61,8 @@ class Track:
self.yRel_filtered = self.yRel
self.yvLead_filtered = self.yvLead
else:
self.yRel_filtered = self.yRel_filtered * 0.95 + self.yRel * 0.05
self.yvLead_filtered = self.yvLead_filtered * 0.95 + self.yvLead * 0.05
self.yRel_filtered = self.yRel_filtered * 0.9 + self.yRel * 0.1
self.yvLead_filtered = self.yvLead_filtered * 0.9 + self.yvLead * 0.1
a_lead_threshold = 0.5 * radar_reaction_factor
if abs(self.aLead) < a_lead_threshold and abs(self.jLead) < 0.5:
@@ -107,12 +112,80 @@ def laplacian_pdf(x: float, mu: float, b: float):
def match_vision_to_track(v_ego: float, lead: capnp._DynamicStructReader, tracks: dict[int, Track], radar_lat_factor = 0.0):
offset_vision_dist = lead.x[0] - RADAR_TO_CAMERA
#vel_tolerance = 25.0 if lead.prob > 0.99 else 10.0
max_offset_vision_dist = max(offset_vision_dist * 0.35, 5.0)
max_vision_dist = max(offset_vision_dist * 1.25, 5.0)
min_vision_dist = max(offset_vision_dist * 0.6, 1.0)
max_offset_vision_vel = max(lead.v[0] * np.interp(lead.prob, [0.8, 0.98], [0.3, 0.5]), 5.0) # 확률이 낮으면 속도오차를 줄임.
def prob(c):
prob_d = laplacian_pdf(c.dRel, offset_vision_dist, lead.xStd[0])
prob_y = laplacian_pdf(c.yRel + c.yvLead_filtered * radar_lat_factor, -lead.y[0], lead.yStd[0])
prob_y = laplacian_pdf(c.yRel, -lead.y[0], lead.yStd[0])
prob_v = laplacian_pdf(c.vLead, lead.v[0], lead.vStd[0])
weight_v = np.interp(c.vLead, [0, 10], [0.3, 1])
return prob_d * prob_y * prob_v * weight_v
best_track = max(tracks.values(), key=prob)
# 끼어드는차량을 간헐적 멀리있는 차량으로 검출하는 문제가 있음..
y_gate = 2.0 #min(1.7, lead.yStd[0] * 2.0)
v_gate = max(5.0, lead.vStd[0] * 2.0)
yv_candidates = [
c for c in tracks.values()
if (min_vision_dist < c.dRel < max_vision_dist)
and (abs(c.yRel + lead.y[0]) < y_gate)
and (abs(c.vLead - lead.v[0]) < v_gate)
]
if False and yv_candidates:
best_track = min(yv_candidates, key=lambda c: c.dRel)
dist_sane = min_vision_dist < best_track.dRel < max_vision_dist #abs(best_track.dRel - offset_vision_dist) < max([(offset_vision_dist)*.35, 5.0])
vel_sane = (abs(best_track.vLead - lead.v[0]) < 10) or (best_track.vLead > 3)
y_sane = abs(best_track.yRel + lead.y[0]) < y_gate
if dist_sane and y_sane:
if vel_sane and lead.prob < 0.45: # 근처에 달리고 있는차를 오감지 했을수 있음
best_track = None
elif not vel_sane or lead.prob < 0.5: # 속도가 안맞거나 희미하게 감지된 차인경우
if best_track.selected_count < 1: # 이전에 선택된 경우에는 그냥 통과함.
best_track.is_stopped_car_count += 1
if best_track.is_stopped_car_count < int(1.0/DT_MDL): # 2초 -> 1초
best_track = None
else:
best_track.is_stopped_car_count = max(0, best_track.is_stopped_car_count - 1)
best_track = None
for c in tracks.values():
if c is best_track:
best_track.selected_count += 1
c.is_stopped_car_count = 0
else:
c.selected_count = 0
return best_track
def match_vision_to_track_old(v_ego: float, lead: capnp._DynamicStructReader, tracks: dict[int, Track], radar_lat_factor = 0.0):
offset_vision_dist = lead.x[0] - RADAR_TO_CAMERA
#vel_tolerance = 25.0 if lead.prob > 0.99 else 10.0
max_offset_vision_dist = max(offset_vision_dist * 0.35, 5.0)
max_offset_vision_vel = max(lead.v[0] * np.interp(lead.prob, [0.8, 0.98], [0.3, 0.5]), 5.0) # 확률이 낮으면 속도오차를 줄임.
def prob(c):
#if abs(offset_vision_dist - c.dRel) > max_offset_vision_dist:
# return -1e6
#if abs(lead.v[0] - c.vLead) > max_offset_vision_vel:
# return -1e6
#if abs(c.yRel + c.yvLead * radar_lat_factor + lead.y[0]) > 3.0: # lead.y[0]는 반대..
# return -1e6
prob_d = laplacian_pdf(c.dRel, offset_vision_dist, lead.xStd[0])
prob_y = laplacian_pdf(c.yRel + c.yvLead * radar_lat_factor, -lead.y[0], lead.yStd[0])
prob_v = laplacian_pdf(c.vLead, lead.v[0], lead.vStd[0])
weight_v = np.interp(c.vLead, [0, 10], [0.3, 1])
@@ -135,7 +208,7 @@ def match_vision_to_track(v_ego: float, lead: capnp._DynamicStructReader, tracks
#if best_track is not None and lead.v[0] - best_track.vLead > max_offset_vision_vel:
# best_track = None
if best_track is not None and abs(best_track.yRel + best_track.yvLead_filtered * radar_lat_factor + lead.y[0]) > 3.0: # lead.y[0]는 반대..
if best_track is not None and abs(best_track.yRel + best_track.yvLead * radar_lat_factor + lead.y[0]) > 3.0: # lead.y[0]는 반대..
best_track = None
if best_track is not None:
@@ -179,110 +252,6 @@ def get_RadarState_from_vision(md, lead_msg: capnp._DynamicStructReader, v_ego:
"radarTrackId": -1,
}
def get_lead_side(v_ego, tracks, md, lane_width, model_v_ego, radar_lat_factor = 0.0):
lead_msg = md.leadsV3[0]
leadCenter = {'status': False}
leadLeft = {'status': False}
leadRight = {'status': False}
leadCutIn = {'status': False}
## SCC레이더는 일단 보관하고 리스트에서 삭제...
track_scc = tracks.get(0)
#if track_scc is not None:
# del tracks[0]
#if len(tracks) == 0:
# return [[],[],[],leadLeft,leadRight]
if md is not None and len(md.position.x) == 33: #ModelConstants.IDX_N:
md_y = md.position.y
md_x = md.position.x
else:
return [[],[],[],leadCenter,leadLeft,leadRight,leadCutIn]
leads_center = {}
leads_left = {}
leads_right = {}
next_lane_y = 1e6 #lane_width / 2 + lane_width * 0.8
for c in tracks.values():
# d_y : path_y - traks_y 의 diff값
# yRel값은 왼쪽이 +값, lead.y[0]값은 왼쪽이 -값
d_y = c.yRel_filtered + np.interp(c.dRel, md_x, md_y) + c.yvLead_filtered * radar_lat_factor
if abs(d_y) < lane_width / 2 * 0.8:
if c.cnt > 6:
ld = c.get_RadarState(lead_msg.prob, float(-lead_msg.y[0]))
leads_center[c.dRel] = ld
elif -next_lane_y < d_y < 0:
ld = c.get_RadarState(0, 0)
leads_right[c.dRel] = ld
elif 0 < d_y < next_lane_y:
ld = c.get_RadarState(0, 0)
leads_left[c.dRel] = ld
# 레이더가 3.4m 차폭보다 적으면,
if abs(d_y) < 3.4/2 and 4 < c.dRel < 20.0 and c.vLead > 4.0 and c.cnt > int(2.0/DT_MDL) and d_y * c.yvLead_filtered < 0:
if leadCutIn['status'] is False or c.dRel < leadCutIn['dRel']:
leadCutIn = c.get_RadarState(lead_msg.prob)
if False: #lead_msg.prob > 0.5: # center에 비젼데이터 안넣음..
ld = get_RadarState_from_vision(md, lead_msg, v_ego, model_v_ego)
leads_center[ld['dRel']] = ld
#ll,lr = [[l[k] for k in sorted(list(l.keys()))] for l in [leads_left,leads_right]]
#lc = sorted(leads_center.values(), key=lambda c:c["dRel"])
ll = list(leads_left.values())
lr = list(leads_right.values())
lc = list(leads_center.values())
#if leads_center:
# dRel_min = min(leads_center.keys())
# lc = [leads_center[dRel_min]]
#else:
# lc = {}
leadLeft = min((lead for dRel, lead in leads_left.items() if lead['dRel'] > 5.0 and abs(lead['dPath']) < 3.5), key=lambda x: x['dRel'], default=leadLeft)
leadRight = min((lead for dRel, lead in leads_right.items() if lead['dRel'] > 5.0 and abs(lead['dPath']) < 3.5), key=lambda x: x['dRel'], default=leadRight)
leadCenter = min((lead for dRel, lead in leads_center.items() if lead['vLead'] > 5 and lead['radar']), key=lambda x: x['dRel'], default=leadCenter)
#filtered_leads_left = {dRel: lead for dRel, lead in leads_left.items() if lead['dRel'] > 5.0}
#if filtered_leads_left:
# dRel_min = min(filtered_leads_left.keys())
# leadLeft = filtered_leads_left[dRel_min]
#filtered_leads_right = {dRel: lead for dRel, lead in leads_right.items() if lead['dRel'] > 5.0}
#if filtered_leads_right:
# dRel_min = min(filtered_leads_right.keys())
# leadRight = filtered_leads_right[dRel_min]
return [ll, lc, lr, leadCenter, leadLeft, leadRight, leadCutIn]
def get_lead(v_ego: float, ready: bool, tracks: dict[int, Track], lead_msg: capnp._DynamicStructReader,
model_v_ego: float, low_speed_override: bool = True) -> dict[str, Any]:
# Determine leads, this is where the essential logic happens
if len(tracks) > 0 and ready and lead_msg.prob > .5:
track = match_vision_to_track(v_ego, lead_msg, tracks)
else:
track = None
lead_dict = {'status': False}
if track is not None:
lead_dict = track.get_RadarState(lead_msg.prob)
elif (track is None) and ready and (lead_msg.prob > .5):
lead_dict = get_RadarState_from_vision(lead_msg, v_ego, model_v_ego)
if low_speed_override:
low_speed_tracks = [c for c in tracks.values() if c.potential_low_speed_lead(v_ego)]
if len(low_speed_tracks) > 0:
closest_track = min(low_speed_tracks, key=lambda c: c.dRel)
# Only choose new track if it is actually closer than the previous one
if (not lead_dict['status']) or (closest_track.dRel < lead_dict['dRel']):
lead_dict = closest_track.get_RadarState()
return lead_dict
class VisionTrack:
def __init__(self, radar_ts):
self.radar_ts = radar_ts
@@ -474,42 +443,22 @@ class RadarD:
if len(leads_v3) > 1:
md = sm['modelV2']
if model_updated:
if self.radar_detected:
self.vision_tracks[0].cnt = 0
self.vision_tracks[1].cnt = 0
self.vision_tracks[0].update(leads_v3[0], model_v_ego, self.v_ego, sm['modelV2'])
self.vision_tracks[1].update(leads_v3[1], model_v_ego, self.v_ego, sm['modelV2'])
self.vision_tracks[0].update(leads_v3[0], model_v_ego, self.v_ego, md)
self.vision_tracks[1].update(leads_v3[1], model_v_ego, self.v_ego, md)
self.radar_state.leadOne, self.radar_detected = self.get_lead(sm['carState'], sm['modelV2'], self.tracks, 0, leads_v3[0], model_v_ego, low_speed_override=False)
self.radar_state.leadTwo, _ = self.get_lead(sm['carState'], sm['modelV2'], self.tracks, 1, leads_v3[1], model_v_ego, low_speed_override=False)
ll, lc, lr, leadCenter, self.radar_state.leadLeft, self.radar_state.leadRight, leadCutIn = get_lead_side(self.v_ego, self.tracks, sm['modelV2'], 3.2, model_v_ego, self.radar_lat_factor)
if leadCutIn is not None and leadCutIn["status"] and self.detect_cut_in:
if self.radar_state.leadOne.status:
if leadCutIn["dRel"] < self.radar_state.leadOne.dRel:
leadCutIn["modelProb"] = 0.03
self.radar_state.leadOne = leadCutIn
self.radar_detected = True
else:
self.radar_detected = True
leadCutIn["modelProb"] = 0.03
self.radar_state.leadOne = leadCutIn
elif leadCenter is not None and leadCenter["status"]:
if self.radar_detected:
if leadCenter["dRel"] < self.radar_state.leadOne.dRel:
leadCenter["modelProb"] = 0.01
self.radar_state.leadOne = leadCenter
elif True: #가끔 다리교랑이 검출됨.. 커브길..
self.radar_detected = True
leadCenter["modelProb"] = 0.02
self.radar_state.leadOne = leadCenter
self.radar_state.leadsLeft = list(ll)
self.radar_state.leadsCenter = list(lc)
self.radar_state.leadsRight = list(lr)
alive_tracks = {tid: trk for tid, trk in self.tracks.items() if trk.measured }
self.radar_state.leadOne, self.radar_detected = self.get_lead(sm['carState'], md, alive_tracks, 0, leads_v3[0], model_v_ego, low_speed_override=False)
self.radar_state.leadTwo, _ = self.get_lead(sm['carState'], md, alive_tracks, 1, leads_v3[1], model_v_ego, low_speed_override=False)
self.lane_line_available = md.laneLineProbs[1] > 0.5 and md.laneLineProbs[2] > 0.5
self.compute_leads(self.v_ego, alive_tracks, md, lane_width=3.2, model_v_ego=model_v_ego)
if self.enable_radar_tracks == 3:
self._pick_lead_one_from_state()
def publish(self, pm: messaging.PubMaster):
assert self.radar_state is not None
@@ -526,18 +475,19 @@ class RadarD:
ready = self.ready
## backup SCC radar(0, 1 trackid)
track_scc = tracks.get(0)
if track_scc is None:
track_scc = tracks.get(1)
if self.enable_radar_tracks <= 0:
track_scc = tracks.get(0)
else:
track_scc = tracks.pop(0, None)
# Determine leads, this is where the essential logic happens
if len(tracks) > 0 and ready and lead_msg.prob > .5:
if len(tracks) > 0 and ready and lead_msg.prob > .3:
track = match_vision_to_track(v_ego, lead_msg, tracks, self.radar_lat_factor)
else:
track = None
if self.enable_radar_tracks in [-1, 2]:
if track_scc is not None and track is None:
if track is None and track_scc is not None and track_scc.measured:
if self.enable_radar_tracks in [-1, 2] or model_v_ego < 5 or track_scc.vLead < 5.0:
track = track_scc
lead_dict = {'status': False}
@@ -545,7 +495,7 @@ class RadarD:
if track is not None:
lead_dict = track.get_RadarState(lead_msg.prob, self.vision_tracks[0].yRel)
radar = True
elif (track is None) and ready and (lead_msg.prob > .8):
elif (track is None) and ready and (lead_msg.prob > .5):
lead_dict = self.vision_tracks[index].get_lead(md)
if self.enable_corner_radar > 0:
@@ -562,6 +512,138 @@ class RadarD:
lead_dict = closest_track.get_RadarState(lead_msg.prob, self.vision_tracks[0].yRel)
return lead_dict, radar
def compute_leads(self, v_ego, tracks, md, lane_width=3.2, model_v_ego=0.0):
lead_msg = md.leadsV3[0] if (md is not None and len(md.position.x) == 33) else None
self.leadCutIn = {'status': False}
if lead_msg is None:
# reset
self.radar_state.leadsLeft = []
self.radar_state.leadsCenter = []
self.radar_state.leadsRight = []
self.radar_state.leadLeft = {'status': False}
self.radar_state.leadRight = {'status': False}
return
md_x, md_y = md.position.x, md.position.y
lane_xs = md.laneLines[1].x
left_ys = md.laneLines[1].y
right_ys = md.laneLines[2].y
left_list, right_list, center_list = [], [], []
for c in tracks.values():
dy = c.yRel_filtered + np.interp(c.dRel, md_x, md_y) # + c.yvLead_filtered * self.radar_lat_factor
dy_with_vel = dy + c.yvLead_filtered * self.radar_lat_factor
y_with_vel_neg = -(c.yRel_filtered + c.yvLead_filtered * self.radar_lat_factor)
left_y = np.interp(c.dRel, lane_xs, left_ys)
right_y = np.interp(c.dRel, lane_xs, right_ys)
y_rel_neg = - c.yRel
# center
if left_y < y_rel_neg < right_y:
if c.cnt > 6:
ld = c.get_RadarState(lead_msg.prob, float(-lead_msg.y[0]))
center_list.append(ld)
# left/right
elif y_rel_neg < left_y:
ld = c.get_RadarState(0, 0)
left_list.append(ld)
else:
ld = c.get_RadarState(0, 0)
right_list.append(ld)
# cut-in
#cut_in_width = 3.0 #3.4 # 끼어들기 차폭
#if self.lane_line_available and left_y < y_with_vel_neg < right_y and (3 < c.dRel < 20 and c.vLead > 4 and c.cnt > int(2.0/DT_MDL) and c.yRel_filtered * c.yvLead_filtered < 0):
if self.lane_line_available and 3 < c.dRel < 50 and c.vLead > 4 and c.cnt > int(2.0/DT_MDL):
if (y_rel_neg < left_y and y_with_vel_neg > left_y) or (y_rel_neg > right_y and y_with_vel_neg < right_y):
if not self.leadCutIn['status'] or c.dRel < self.leadCutIn['dRel']:
c.cut_in_count += 1
else:
c.cut_in_count = 0
if c.cut_in_count > int(0.5/DT_MDL):
self.leadCutIn = c.get_RadarState(lead_msg.prob)
else:
c.cut_in_count = 0
else:
c.cut_in_count = 0
self.radar_state.leadsLeft = left_list
self.radar_state.leadsRight = right_list
self.radar_state.leadsCenter = center_list
self.radar_state.leadLeft = min(
(ld for ld in left_list if ld['dRel'] > 5 and abs(ld['dPath']) < 3.5),
key=lambda d: d['dRel'],
default={'status': False}
)
self.radar_state.leadRight = min(
(ld for ld in right_list if ld['dRel'] > 5 and abs(ld['dPath']) < 3.5),
key=lambda d: d['dRel'],
default={'status': False}
)
if self.lane_line_available:
self.leadCenter = min(
(ld for ld in center_list if ld['vLead'] > 5 and ld['radar'] and abs(ld['yRel']) < 5.0 and ld['dRel'] > 3.5),
key=lambda d: d['dRel'],
default={'status': False}
)
else:
self.leadCenter = None
def _ok(ld):
return (ld.get('vLead', 0) > 2 and
abs(ld.get('dPath', 0)) < 4.2 and
ld.get('dRel', 0) > 2)
def _pick_two_with_gap(cands, min_gap=5.0):
xs = sorted((ld for ld in cands if _ok(ld)), key=lambda d: d['dRel'])
if not xs:
return []
first = xs[0]
second = None
for ld in xs[1:]:
# 5m 이상 떨어진 후보만 허용 (>= 5.0)
if (ld['dRel'] - first['dRel']) >= min_gap:
second = ld
break
return [first] if second is None else [first, second]
self.radar_state.leadsLeft2 = _pick_two_with_gap(left_list, min_gap=5.0)
self.radar_state.leadsRight2 = _pick_two_with_gap(right_list, min_gap=5.0)
def _pick_lead_one_from_state(self):
chosen = None
detected = self.radar_detected
if self.leadCutIn and self.leadCutIn.get("status") and self.detect_cut_in:
if self.radar_state.leadOne.status:
if self.leadCutIn["dRel"] < self.radar_state.leadOne.dRel:
chosen = self.leadCutIn
chosen["modelProb"] = 0.03
detected = True
else:
chosen = self.leadCutIn
chosen["modelProb"] = 0.03
detected = True
elif self.leadCenter and self.leadCenter["status"]:
if self.radar_detected:
if self.radar_state.leadOne.status and self.leadCenter["dRel"] < self.radar_state.leadOne.dRel:
chosen = self.leadCenter
chosen["modelProb"] = 0.01
else:
chosen = self.leadCenter
chosen["modelProb"] = 0.02
detected = True
if chosen is not None:
self.radar_state.leadOne = chosen
self.radar_detected = detected
def corner_radar(self, CS, lead_dict):
lat_dist = 1e6
selfdrive/frogpilot/fleetmanager/helpers.py
+6 -2
View File
@@ -266,9 +266,13 @@ def get_locations():
return data
def preload_favs():
raw_json = params.get("ApiCache_NavDestinations", encoding='utf8')
if raw_json is None:
return (None, None, None, None, None)
try:
nav_destinations = json.loads(params.get("ApiCache_NavDestinations", encoding='utf8'))
except TypeError:
nav_destinations = json.loads(raw_json)
except (TypeError, json.JSONDecodeError):
return (None, None, None, None, None)
locations = {"home": None, "work": None, "fav1": None, "fav2": None, "fav3": None}
selfdrive/modeld/SConscript
+4 -4
View File
@@ -50,9 +50,9 @@ def tg_compile(flags, model_name):
# Compile small models
for model_name in ['driving_vision', 'driving_policy', 'dmonitoring_model']:
flags = {
'larch64': 'QCOM=1',
'Darwin': 'CPU=1 IMAGE=0 JIT=2',
}.get(arch, 'LLVM=1 LLVMOPT=1 BEAM=0 IMAGE=0 JIT=2')
'larch64': 'DEV=QCOM',
'Darwin': 'DEV=CPU IMAGE=0',
}.get(arch, 'DEV=LLVM IMAGE=0')
tg_compile(flags, model_name)
# Compile BIG model if USB GPU is available
@@ -62,7 +62,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 = "AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0"
flags = "DEV=AMD AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=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
selfdrive/modeld/modeld.py
+2 -7
View File
@@ -1,16 +1,11 @@
#!/usr/bin/env python3
import os
from openpilot.system.hardware import TICI
os.environ['DEV'] = 'QCOM' if TICI else 'LLVM'
USBGPU = "USBGPU" in os.environ
if USBGPU:
os.environ['AMD'] = '1'
os.environ['DEV'] = 'AMD'
os.environ['AMD_IFACE'] = 'USB'
elif TICI:
from openpilot.selfdrive.modeld.runners.tinygrad_helpers import qcom_tensor_from_opencl_address
os.environ['QCOM'] = '1'
else:
os.environ['LLVM'] = '1'
os.environ['JIT'] = '2'
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes
import time
selfdrive/modeld/models/driving_policy.onnx
BIN
View File
Binary file not shown.
selfdrive/modeld/models/driving_vision.onnx
BIN
View File
Binary file not shown.
selfdrive/modeld/parse_model_outputs.py
+8 -7
View File
@@ -22,9 +22,10 @@ class Parser:
self.ignore_missing = ignore_missing
def check_missing(self, outs, name):
if name not in outs and not self.ignore_missing:
missing = name not in outs
if missing and not self.ignore_missing:
raise ValueError(f"Missing output {name}")
return name not in outs
return missing
def parse_categorical_crossentropy(self, name, outs, out_shape=None):
if self.check_missing(outs, name):
@@ -90,22 +91,22 @@ class Parser:
self.parse_mdn('road_transform', outs, in_N=0, out_N=0, out_shape=(ModelConstants.POSE_WIDTH,))
self.parse_mdn('lane_lines', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_LANE_LINES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
self.parse_mdn('road_edges', outs, in_N=0, out_N=0, out_shape=(ModelConstants.NUM_ROAD_EDGES,ModelConstants.IDX_N,ModelConstants.LANE_LINES_WIDTH))
self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION,
out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH))
for k in ['lead_prob', 'lane_lines_prob']:
self.parse_binary_crossentropy(k, outs)
self.parse_binary_crossentropy('lane_lines_prob', outs)
self.parse_categorical_crossentropy('desire_pred', outs, out_shape=(ModelConstants.DESIRE_PRED_LEN,ModelConstants.DESIRE_PRED_WIDTH))
self.parse_binary_crossentropy('meta', outs)
return outs
def parse_policy_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
self.parse_mdn('plan', outs, in_N=ModelConstants.PLAN_MHP_N, out_N=ModelConstants.PLAN_MHP_SELECTION,
out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH))
out_shape=(ModelConstants.IDX_N,ModelConstants.PLAN_WIDTH))
if 'lat_planner_solution' in outs:
self.parse_mdn('lat_planner_solution', outs, in_N=0, out_N=0, out_shape=(ModelConstants.IDX_N,ModelConstants.LAT_PLANNER_SOLUTION_WIDTH))
if 'desired_curvature' in outs:
self.parse_mdn('desired_curvature', outs, in_N=0, out_N=0, out_shape=(ModelConstants.DESIRED_CURV_WIDTH,))
self.parse_categorical_crossentropy('desire_state', outs, out_shape=(ModelConstants.DESIRE_PRED_WIDTH,))
self.parse_binary_crossentropy('lead_prob', outs)
self.parse_mdn('lead', outs, in_N=ModelConstants.LEAD_MHP_N, out_N=ModelConstants.LEAD_MHP_SELECTION,
out_shape=(ModelConstants.LEAD_TRAJ_LEN,ModelConstants.LEAD_WIDTH))
return outs
def parse_outputs(self, outs: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
selfdrive/ui/carrot.cc
+9 -3
View File
@@ -676,7 +676,7 @@ protected:
return true;
};
bool isLeadSCC() {
return radarTrackId < 2;
return radarTrackId < 1;
}
bool isRadarDetected() {
return radarTrackId >= 0;
@@ -1120,8 +1120,14 @@ protected:
ui_draw_text(s, bx, by + 20, str, 35, COLOR_WHITE, BOLD);
break;
}
if (xDistToTurn < 1000) sprintf(str, "%d m", xDistToTurn);
else sprintf(str, "%.1f km", xDistToTurn / 1000.f);
if (s->scene.is_metric) {
if (xDistToTurn < 1000) sprintf(str, "%d m", xDistToTurn);
else sprintf(str, "%.1f km", xDistToTurn / 1000.f);
}
else {
if (xDistToTurn < 1609) sprintf(str, "%d ft", (int)(xDistToTurn * 3.28084));
else sprintf(str, "%.1f mi", xDistToTurn / 1609.344f);
}
ui_draw_text(s, bx, by + 120, str, 40, COLOR_WHITE, BOLD);
}
nvgTextAlign(s->vg, NVG_ALIGN_LEFT | NVG_ALIGN_BOTTOM);
selfdrive/ui/qt/offroad/settings.cc
+1 -1
View File
@@ -813,7 +813,7 @@ CarrotPanel::CarrotPanel(QWidget* parent) : QWidget(parent) {
startToggles->addItem(selectCarBtn);
startToggles->addItem(new CValueControl("HyundaiCameraSCC", "HYUNDAI: CAMERA SCC", "1:Connect the SCC's CAN line to CAM, 2:Sync Cruise state, 3:StockLong", 0, 3, 1));
startToggles->addItem(new CValueControl("CanfdHDA2", "CANFD: HDA2 mode", "1:HDA2,2:HDA2+BSM", 0, 2, 1));
startToggles->addItem(new CValueControl("EnableRadarTracks", "Enable Radar Track", "1:Enable RadarTrack, -1,2:Disable use HKG SCC radar at all times", -1, 2, 1));
startToggles->addItem(new CValueControl("EnableRadarTracks", "Enable Radar Track", "1:Enable RadarTrack, -1,2:Disable use HKG SCC radar at all times", -1, 3, 1));
startToggles->addItem(new CValueControl("AutoCruiseControl", "Auto Cruise control", "Softhold, Auto Cruise ON/OFF control", 0, 3, 1));
startToggles->addItem(new CValueControl("CruiseOnDist", "CRUISE: Auto ON distance(0cm)", "When GAS/Brake is OFF, Cruise ON when the lead car gets closer.", 0, 2500, 50));
startToggles->addItem(new CValueControl("AutoEngage", "Auto Engage control on start", "1:SteerEnable, 2:Steer/Cruise Engage", 0, 2, 1));
selfdrive/ui/translations/main_ko.ts
+2 -2
View File
@@ -774,7 +774,7 @@ Firehose Mode allows you to maximize your training data uploads to improve openp
<name>PrimeAdWidget</name>
<message>
<source>Upgrade Now</source>
<translation>지금 업그레이드하세요</translation>
<translation>절대 업그레이드마세요</translation>
</message>
<message>
<source>Become a comma prime member at connect.comma.ai</source>
@@ -782,7 +782,7 @@ Firehose Mode allows you to maximize your training data uploads to improve openp
</message>
<message>
<source>PRIME FEATURES:</source>
<translation>PRIME 기능:</translation>
<translation>PRIME 기능(주행성능과 전혀관련 없음):</translation>
</message>
<message>
<source>Remote access</source>
system/athena/registration.py
+5
View File
@@ -91,6 +91,11 @@ def register(show_spinner=False) -> str | None:
backoff = min(backoff + 1, 15)
time.sleep(backoff)
if time.monotonic() - start_time > 14:
cloudlog.error("pilotauth timed out; continuing as UNREGISTERED")
dongle_id = UNREGISTERED_DONGLE_ID
break
if time.monotonic() - start_time > 60 and show_spinner:
spinner.update(f"registering device - serial: {serial}, IMEI: ({imei1}, {imei2})")
tinygrad_repo/autogen_stubs.sh
-24
View File
@@ -198,11 +198,7 @@ generate_amd() {
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
extra/hip_gpu_driver/nvd.h \
extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
extra/hip_gpu_driver/soc21_enum.h \
extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
extra/hip_gpu_driver/gc_11_0_0_offset.h \
extra/hip_gpu_driver/gc_10_3_0_offset.h \
extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
--clang-args="-I/opt/rocm/include -x c++" \
-o $BASE/amd_gpu.py
@@ -376,26 +372,6 @@ generate_am() {
-o $BASE/am/pm4_nv.py
fixup $BASE/am/pm4_nv.py
clang2py -k cdefstum \
$AMKERN_INC/vega10_enum.h \
-o $BASE/am/vega10.py
fixup $BASE/am/vega10.py
clang2py -k cdefstum \
$AMKERN_INC/navi10_enum.h \
-o $BASE/am/navi10.py
fixup $BASE/am/navi10.py
clang2py -k cdefstum \
$AMKERN_INC/soc21_enum.h \
-o $BASE/am/soc21.py
fixup $BASE/am/soc21.py
clang2py -k cdefstum \
$AMKERN_INC/soc24_enum.h \
-o $BASE/am/soc24.py
fixup $BASE/am/soc24.py
clang2py -k cdefstum \
extra/hip_gpu_driver/sdma_registers.h \
$AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
tinygrad_repo/docs/developer/layout.md
+2 -2
View File
@@ -18,11 +18,11 @@ Group UOps into kernels.
---
## tinygrad/opt
## tinygrad/codegen/opt
Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
::: tinygrad.opt.get_optimized_ast
::: tinygrad.codegen.opt.get_optimized_ast
options:
members: false
show_labels: false
tinygrad_repo/docs/ramp.py
+1 -1
View File
@@ -126,7 +126,7 @@ print(t_log_grad.uop)
"""
void E_(float* restrict data0, float* restrict data1) {
float val0 = *(data1+0);
*(data0+0) = (0.6931471805599453f*(1/(val0*0.6931471805599453f)));
*(data0+0) = (1/val0);
}
"""
# the derivative is close to 1/3
tinygrad_repo/docs/tensor/elementwise.md
+1
View File
@@ -78,6 +78,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
::: tinygrad.Tensor.minimum
::: tinygrad.Tensor.where
::: tinygrad.Tensor.copysign
::: tinygrad.Tensor.logaddexp
## Casting Ops
tinygrad_repo/docs/tinybox.md
+4 -4
View File
@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra
## Welcome
Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.
Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.
@@ -47,8 +47,8 @@ Reboot after making these changes or restart the `displayservice.service` servic
The [default tinybox image](https://github.com/tinygrad/tinyos) ships with tinygrad and PyTorch. While we develop tinygrad, the box is universal hardware. Use whatever framework you desire, run notebooks, download demos, install more things, train, inference, live, laugh, love, you aren't paying per hour for this box so the only limit is your imagination.
## tinychat
## Building the OS image
Since LLMs are so popular, we ship with a built in tinygrad based chatbot using a LLaMA-3 finetune. Visit the IP (not the BMC IP) of your tinybox in a web browser on your computer or phone, and you'll find a friendly looking chat interface. This chatbot also provides an OpenAI compatible LLM API on that port, so you can script it.
The OS image is built using `ubuntu-image` from <https://github.com/tinygrad/tinyos>.
The conversations you have with this chatbot are between you and your tinybox. Also, the history in the web app is saved on the client, not the tinybox.
After cloning, run `make green` or `make red` to build a tinybox green or tinybox red image respectively.
tinygrad_repo/examples/beautiful_mnist.py
+4 -6
View File
@@ -1,12 +1,12 @@
# model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
from typing import List, Callable
from typing import Callable
from tinygrad import Tensor, TinyJit, nn, GlobalCounters
from tinygrad.helpers import getenv, colored, trange
from tinygrad.nn.datasets import mnist
class Model:
def __init__(self):
self.layers: List[Callable[[Tensor], Tensor]] = [
self.layers: list[Callable[[Tensor], Tensor]] = [
nn.Conv2d(1, 32, 5), Tensor.relu,
nn.Conv2d(32, 32, 5), Tensor.relu,
nn.BatchNorm(32), Tensor.max_pool2d,
@@ -21,17 +21,15 @@ if __name__ == "__main__":
X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
model = Model()
opt = nn.optim.Adam(nn.state.get_parameters(model))
opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
@TinyJit
@Tensor.train()
def train_step() -> Tensor:
opt.zero_grad()
samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
# TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
opt.step()
return loss
return loss.realize(*opt.schedule_step())
@TinyJit
def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
tinygrad_repo/examples/handcode_opt.py
-134
View File
@@ -1,134 +0,0 @@
from extra.models.resnet import ResNet50
from extra.mcts_search import mcts_search
from examples.mlperf.helpers import get_mlperf_bert_model
from tinygrad import Tensor, Device, dtypes, nn
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.uop.ops import Ops, sym_infer
from tinygrad.device import Compiled
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
from extra.optimization.helpers import time_linearizer
from tinygrad.engine.realize import get_program
def get_sched_resnet():
mdl = ResNet50()
optim = (nn.optim.LARS if getenv("LARS") else nn.optim.SGD)(nn.state.get_parameters(mdl))
BS = getenv("BS", 64)
# run model twice to get only what changes, these are the kernels of the model
for _ in range(2):
out = mdl(Tensor.empty(BS, 3, 224, 224))
targets = [out]
if getenv("BACKWARD"):
optim.zero_grad()
out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
targets += [x for x in optim.schedule_step()]
sched = Tensor.schedule(*targets)
print(f"schedule length {len(sched)}")
return sched
def get_sched_bert():
mdl = get_mlperf_bert_model()
optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
# fake data
BS = getenv("BS", 9)
input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)
masked_positions = Tensor.empty((BS, 76), dtype=dtypes.float32)
masked_lm_ids = Tensor.empty((BS, 76), dtype=dtypes.float32)
masked_lm_weights = Tensor.empty((BS, 76), dtype=dtypes.float32)
next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
# run model twice to get only what changes, these are the kernels of the model
for _ in range(2):
lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
targets = [lm_logits, seq_relationship_logits]
if getenv("BACKWARD"):
optim.zero_grad()
loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
# ignore grad norm and loss scaler for now
loss.backward()
targets += [x for x in optim.schedule_step()]
sched = Tensor.schedule(*targets)
print(f"schedule length {len(sched)}")
return sched
if __name__ == "__main__":
if getenv("HALF", 1):
dtypes.default_float = dtypes.half
# the device we are optimizing for
device: Compiled = Device[Device.DEFAULT]
if getenv("BACKWARD"): Tensor.training = True
print(f"optimizing for {Device.DEFAULT}")
sched = globals()[f"get_sched_{getenv('MODEL', 'resnet')}"]()
sched = [x for x in sched if x.ast.op is Ops.SINK]
# focus on one kernel
if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
# work with the schedule
total_tm = 0
running_gflops = 0
usage = {}
for i,si in enumerate(sched):
if DEBUG >= 3: print(si.ast)
rawbufs = bufs_from_lin(Kernel(si.ast))
# "linearize" the op into uops in different ways
lins: list[tuple[Kernel, str]] = []
# always try hand coded opt
lin = Kernel(si.ast, opts=device.renderer)
lin.apply_opts(hand_coded_optimizations(lin))
lins.append((lin, "HC"))
# maybe try tensor cores
lin = Kernel(si.ast, opts=device.renderer)
if lin.apply_tensor_cores():
lins.append((lin, "TC"))
# try a beam search
if beam:=getenv("BEAM"):
lin = Kernel(si.ast, opts=device.renderer)
lin = beam_search(lin, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
lins.append((lin, "BEAM"))
# try MCTS
if mcts:=getenv("MCTS"):
lin = Kernel(si.ast, opts=device.renderer)
lin = mcts_search(lin, rawbufs, mcts)
lins.append((lin, "MCTS"))
# benchmark the programs
choices = []
for lin, nm in lins:
tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
choices.append((tm, gflops, lin, prg, nm))
sorted_choices = sorted(choices, key=lambda x: x[0])
if DEBUG >= 1: # print all kernels
for tm, gflops, lin, prg, nm in choices:
print(f" kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
tm, gflops, lin, prg, nm = sorted_choices[0]
if getenv("SRC"):
print(si.ast)
print(lin.applied_opts)
print(get_program(lin.get_optimized_ast(), lin.opts).src)
total_tm += tm
running_gflops += gflops * tm
if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
usage[key] = (usage[key][0] + tm, usage[key][1] + 1)
print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS {[repr(m) if TRACEMETA >= 2 else str(m) for m in si.metadata]}")
print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")
print("usage:")
for k in sorted(usage, key=lambda x: -usage[x][0])[:10]:
print(f"{usage[k][0]*1000:.2f} ms: {k} ({usage[k][1]} times)")
tinygrad_repo/examples/hlb_cifar10.py
+1 -1
View File
@@ -118,7 +118,7 @@ class SpeedyResNet:
# hyper-parameters were exactly the same as the original repo
bias_scaler = 58
hyp = {
'seed' : 200,
'seed' : 201,
'opt': {
'bias_lr': 1.76 * bias_scaler/512,
'non_bias_lr': 1.76 / 512,
tinygrad_repo/examples/mlperf/dataloader.py
+261 -1
View File
@@ -1,4 +1,4 @@
import os, random, pickle, queue
import os, random, pickle, queue, struct, math, functools, hashlib, time
from typing import List
from pathlib import Path
from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -6,6 +6,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu
import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX
from tinygrad.nn.state import TensorIO
### ResNet
@@ -510,6 +511,253 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
# happens with BENCHMARK set
pass
# llama3
class BinIdxDataset:
def __init__(self, base_path:Path):
self.idx_t = Tensor(base_path.with_name(f"{base_path.name}.idx"))
self.idx = TensorIO(self.idx_t)
# parse idx file
magic = self.idx.read(9)
assert magic == b"MMIDIDX\x00\x00", "invalid index file format"
version, = struct.unpack("<Q", self.idx.read(8))
assert version == 1, "unsupported index version"
dtype_code, = struct.unpack("<B", self.idx.read(1))
self.dtype = {1:dtypes.uint8, 2:dtypes.int8, 3:dtypes.int16, 4:dtypes.int32, 5:dtypes.int64, 6:dtypes.float64, 7:dtypes.double, 8:dtypes.uint16}[dtype_code]
self.count, = struct.unpack("<Q", self.idx.read(8))
doc_count, = struct.unpack("<Q", self.idx.read(8))
start = self.idx.tell()
end = start + self.count * dtypes.int32.itemsize
self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
start = end
end = start + self.count * dtypes.int64.itemsize
self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
start = end
end = start + doc_count * dtypes.int64.itemsize
self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
# bin file
self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
def _index(self, idx) -> tuple[int, int]:
return int(self.pointers[idx]), int(self.sizes[idx])
def get(self, idx, offset:int=0, length:int|None=None):
ptr, size = self._index(idx)
if length is None: length = size - offset
ptr += offset * self.dtype.itemsize
return self.bin_t[ptr:ptr+length*self.dtype.itemsize].bitcast(self.dtype).to(None)
# https://docs.nvidia.com/megatron-core/developer-guide/latest/api-guide/datasets.html
class GPTDataset:
def __init__(self, base_path:Path, samples:int, seqlen:int, seed:int, shuffle:bool):
self.samples, self.seqlen = samples, seqlen
self.shuffle = shuffle
self.rng = np.random.RandomState(seed)
self.indexed_dataset = BinIdxDataset(base_path)
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
print(f"try loading GPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.doc_idx = self._build_doc_idx()
self.sample_idx = self._build_sample_idx()
self.shuffle_idx = self._build_shuffle_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.doc_idx, self.sample_idx, self.shuffle_idx), f)
def __getitem__(self, idx):
if idx is None:
text = self._get(0)
else:
text = self._get(idx)
return text
def _get(self, idx):
idx = self.shuffle_idx[idx]
doc_idx_beg, doc_idx_beg_offset = self.sample_idx[idx]
doc_idx_end, doc_idx_end_offset = self.sample_idx[idx + 1]
doc_ids, sample_parts = [], []
if doc_idx_beg == doc_idx_end:
doc_ids.append(self.doc_idx[doc_idx_beg])
sample_parts.append(
self.indexed_dataset.get(
int(self.doc_idx[doc_idx_beg]), offset=int(doc_idx_beg_offset), length=int(doc_idx_end_offset - doc_idx_beg_offset + 1)))
else:
for i in range(doc_idx_beg, doc_idx_end + 1):
doc_ids.append(self.doc_idx[i])
offset = 0 if i > doc_idx_beg else doc_idx_beg_offset
length = None if i < doc_idx_end else int(doc_idx_end_offset + 1)
sample_parts.append(self.indexed_dataset.get(int(self.doc_idx[i]), offset=int(offset), length=length))
# concat all parts
text = Tensor.cat(*sample_parts)
return text
@functools.cached_property
def tokens_per_epoch(self) -> int:
return sum(self.indexed_dataset.sizes.tolist())
@functools.cached_property
def num_epochs(self) -> int:
# we need enough epochs to cover the requested amount of tokens
num_epochs = 1
num_tokens = self.tokens_per_epoch
while num_tokens < self.samples * self.seqlen:
num_epochs += 1
num_tokens += self.tokens_per_epoch
return num_epochs
# https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
def _build_doc_idx(self):
print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
st = time.perf_counter()
# doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
doc_idx = doc_idx.astype(np.int32)
at = time.perf_counter()
if self.shuffle: self.rng.shuffle(doc_idx)
print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
return doc_idx
def _build_sample_idx(self):
print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
sample_idx_idx += 1
for _ in tqdm(range(1, self.samples + 1)):
remaining_seqlen = self.seqlen + 1
while remaining_seqlen > 0:
doc_idx = int(self.doc_idx[doc_idx_idx])
doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
remaining_seqlen -= doc_len
if remaining_seqlen <= 0:
doc_offset += remaining_seqlen + doc_len - 1
remaining_seqlen = 0
else:
if doc_idx_idx == len(self.doc_idx) - 1:
assert sample_idx_idx == self.samples
doc_idx = int(self.doc_idx[doc_idx_idx])
doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
break
doc_idx_idx += 1
doc_offset = 0
sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
sample_idx_idx += 1
return sample_idx
def _build_shuffle_idx(self):
print(f"building shuffle_idx for {self.samples=}")
st = time.perf_counter()
shuffle_idx = np.arange(self.samples, dtype=np.int32)
at = time.perf_counter()
if self.shuffle: self.rng.shuffle(shuffle_idx)
print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
return shuffle_idx
class BlendedGPTDataset:
def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
self.shuffle = shuffle
self.rng = np.random.RandomState(seed)
# normalize weights
total_weight = sum(weights)
self.weights = [w / total_weight for w in weights]
self.samples = samples
surplus = 0.005
samples_per_blend = [math.ceil(math.ceil(self.samples * w) * (1 + surplus)) for w in self.weights]
self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
# check for cache
cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
print(f"try loading BlendedGPTDataset from {cache_path}...")
if cache_path.exists():
print("cache found, loading...")
with open(cache_path, "rb") as f:
self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
else:
print("cache not found, building index...")
self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
# save cache
with open(cache_path, "wb") as f:
pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
def get(self, idx:int):
tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
return tokens
def _build_blend_idx(self):
dataset_idx = np.zeros(self.samples, dtype=np.int16)
dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
unspent_datasets = set(range(len(self.datasets)))
dataset_sample_counts = [0] * len(self.datasets)
for i in tqdm(range(self.samples)):
error_argmax, error_max = 0, 0.0
for di in unspent_datasets:
error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
if error > error_max:
error_max = error
error_argmax = di
dataset_idx[i] = error_argmax
dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
dataset_sample_counts[error_argmax] += 1
return dataset_idx, dataset_sample_idx
def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
if val:
dataset = BlendedGPTDataset([
base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
], [
1.0
], samples, seqlen, seed, False)
else:
dataset = BlendedGPTDataset([
base_dir / "c4-train.en_6_text_document",
base_dir / "c4-train.en_7_text_document",
], [
1.0, 1.0
], samples, seqlen, seed, True)
for b in range(math.ceil(samples / bs)):
batch = []
for i in range(bs):
tokens = dataset.get(b * bs + i)
batch.append(tokens)
yield Tensor.stack(batch, dim=0)
if __name__ == "__main__":
def load_unet3d(val):
assert not val, "validation set is not supported due to different sizes on inputs"
@@ -538,6 +786,18 @@ if __name__ == "__main__":
for x in batch_load_retinanet(dataset, val, base_dir):
pbar.update(x[0].shape[0])
def load_llama3(val):
bs = 24
samples = 5760 if val else 1_200_000 * 1152
seqlen = 8192
max_, min_ = 0, math.inf
for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
max_ = max(max_, tokens.shape[1])
min_ = min(min_, tokens.shape[1])
print(f"max seq length: {max_}")
print(f"min seq length: {min_}")
load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
if load_fn_name in globals():
globals()[load_fn_name](getenv("VAL", 1))
tinygrad_repo/examples/mlperf/model_eval.py
+29 -1
View File
@@ -1,4 +1,4 @@
import time
import time, math
start = time.perf_counter()
from pathlib import Path
import numpy as np
@@ -241,6 +241,34 @@ def eval_mrcnn():
evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
evaluate_predictions_on_coco(mask_output, iou_type='segm')
def eval_llama3():
from extra.models.llama import Transformer
from examples.llama3 import MODEL_PARAMS
from tinygrad.helpers import tqdm
bs = 4
sequence_length = 512
model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
@TinyJit
def eval_step(model, tokens):
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten()
from examples.mlperf.dataloader import batch_load_llama3
iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
losses = []
for tokens in tqdm(iter, total=5760//bs):
GlobalCounters.reset()
losses += eval_step(model, tokens).tolist()
tqdm.write(f"loss: {np.mean(losses)}")
log_perplexity = Tensor(losses).mean()
print(f"Log Perplexity: {log_perplexity.item()}")
if __name__ == "__main__":
# inference only
Tensor.training = False
tinygrad_repo/examples/mlperf/model_train.py
+129 -20
View File
@@ -1290,9 +1290,19 @@ def train_llama3():
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
config = {}
BS = config["BS"] = getenv("BS", 4)
BS = config["BS"] = getenv("BS", 16)
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
SEED = config["SEED"] = getenv("SEED", 5760)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
TRAIN_ON_VAL = config["TRAIN_ON_VAL"] = getenv("TRAIN_ON_VAL", 0)
SAMPLES = config["SAMPLES"] = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
EVAL_FREQ = config["EVAL_FREQ"] = getenv("EVAL_FREQ", 46080)
EVAL_BS = config["EVAL_BS"] = getenv("EVAL_BS", 16)
EVAL_TARGET = config["EVAL_TARGET"] = getenv("EVAL_TARGET", 5.6)
# LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
# trains to 7
opt_adamw_beta_1 = 0.9
opt_adamw_beta_2 = 0.95
@@ -1300,7 +1310,6 @@ def train_llama3():
opt_adamw_weight_decay = 0.1
opt_gradient_clip_norm = 1.0
sequence_length = 8192
opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152) # NOTE: cannot change for benchmark
@@ -1308,7 +1317,37 @@ def train_llama3():
# TODO: confirm weights are in bf16
# vocab_size from the mixtral tokenizer
model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
if getenv("FAKEDATA"):
for v in get_parameters(model):
v = v.assign(Tensor.empty(v.shape))
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
for v in get_parameters(model):
v.shard_(device, axis=None)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
for k,v in get_state_dict(model).items():
if 'scale' in k: v.shard_(device, axis=None) # from quantized
elif '.attention.wq' in k: v.shard_(device, axis=0)
elif '.attention.wk' in k: v.shard_(device, axis=0)
elif '.attention.wv' in k: v.shard_(device, axis=0)
elif '.attention.wo' in k: v.shard_(device, axis=1)
elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
elif 'output.weight' in k: v.shard_(device, axis=0)
else:
# attention_norm, ffn_norm, norm
v.shard_(device, axis=None)
# prevents memory spike on device 0
v.realize()
optim = AdamW(get_parameters(model), lr=0.0,
b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
@@ -1316,12 +1355,20 @@ def train_llama3():
@TinyJit
@Tensor.train()
def train_step(model, x, y):
def train_step(model, tokens:Tensor, grad_acc:int):
optim.zero_grad()
logits:Tensor = model(x, start_pos=0, temperature=math.nan)
loss = logits.cross_entropy(y)
loss.backward()
# grad acc
for batch in tokens.split(tokens.shape[0]//grad_acc):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
batch = batch.shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
batch = batch.shard(device)
logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
loss.backward()
Tensor.realize(*[p.grad for p in optim.params])
# L2 norm grad clip
# https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
# https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
@@ -1331,7 +1378,7 @@ def train_llama3():
total_norm += p.grad.float().square().sum()
total_norm = total_norm.sqrt().contiguous()
for p in optim.params:
p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
optim.step()
scheduler.step()
@@ -1340,19 +1387,81 @@ def train_llama3():
loss.realize(lr)
return loss, lr
# overfitting this example should give cross_entropy log(BS)
fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
fake_label = Tensor(list(range(BS)), dtype="int16")
@TinyJit
@Tensor.train(False)
def eval_step(model, tokens:Tensor):
if (DP := getenv("DP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
tokens = tokens.shard(device, 0)
if (MP := getenv("MP", 1)) > 1:
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
tokens = tokens.shard(device)
logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
return loss.flatten().float()
for _ in range(100):
# ** data iters **
def fake_data(bs, samples):
for _ in range(samples // bs):
yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
def get_train_iter():
if getenv("FAKEDATA", 0):
return fake_data(GBS, SAMPLES)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
def get_eval_iter():
if getenv("FAKEDATA", 0):
return fake_data(EVAL_BS, 5760)
else:
from examples.mlperf.dataloader import batch_load_llama3
return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
iter = get_train_iter()
i, sequences_seen = 0, 0
for tokens in tqdm(iter, total=SAMPLES//GBS):
t = time.perf_counter()
GlobalCounters.reset()
loss, lr = train_step(model, fake_input, fake_label)
# BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
# uses 43% ~= 83GB
# 8B bf16 = 16GB. model + grad + optim m and v = 64GB
# TODO: this OOM
# BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
loss, lr = train_step(model, tokens, grad_acc)
loss = loss.float().item()
# above as tqdm.write f-string
tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
if (fname:=getenv("LOSS_FILE", "")):
with open(fname, "a") as f:
f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
if getenv("CKPT") and (i % 200 == 0 or i == 10):
tqdm.write("saving checkpoint")
if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
fn = f"{ckpt_dir}/llama3_{i}.safe"
safe_save(get_state_dict(model), fn)
i += 1
sequences_seen += tokens.shape[0]
if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
tqdm.write(f"evaluating after {sequences_seen} sequences")
# run eval
eval_losses = []
eval_iter = get_eval_iter()
tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
eval_losses += eval_step(model, tokens).tolist()
log_perplexity = Tensor(eval_losses).mean().float().item()
tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
if log_perplexity < EVAL_TARGET:
tqdm.write(f"target achieved after {sequences_seen} sequences")
if getenv("CKPT"):
if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
fn = f"{ckpt_dir}/llama3.safe"
safe_save(get_state_dict(model), fn)
break
if __name__ == "__main__":
multiprocessing.set_start_method('spawn')
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_1xMI300X/dev_beam.sh
+2
View File
@@ -4,6 +4,8 @@ export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
# export BEAM_LOG_SURPASS_MAX=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh
+2
View File
@@ -5,6 +5,8 @@ export MODEL="bert"
export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_run.sh
+2
View File
@@ -8,6 +8,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+2
View File
@@ -11,6 +11,8 @@ export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
export TRAIN_STEPS=3900
export IGNORE_OOB=1
export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
export BASEDIR="/raid/datasets/wiki"
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_beam.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/dev_run.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." NV=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+2 -2
View File
@@ -5,9 +5,9 @@ set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." NV=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_green"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_beam.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/dev_run.sh
+2 -2
View File
@@ -2,9 +2,9 @@
export PYTHONPATH="." AMD=1
export MODEL="bert"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+2 -2
View File
@@ -5,9 +5,9 @@ set -o pipefail # Make pipeline fail if any command fails
export PYTHONPATH="." AMD=1
export MODEL="bert"
export SUBMISSION_PLATFORM="tinybox_red"
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
export IGNORE_OOB=1
export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
export IGNORE_JIT_FIRST_BEAM=1
tinygrad_repo/examples/olmoe.py
+3 -4
View File
@@ -1,8 +1,7 @@
# https://arxiv.org/pdf/2409.02060
import time
import time, functools
import numpy as np
np.set_printoptions(suppress=True, linewidth=1000)
import functools
from tinygrad import Tensor, nn, Device, GlobalCounters
from tinygrad.helpers import Timing, getenv
from extra.models.llama import Transformer, convert_from_huggingface
@@ -17,7 +16,7 @@ class MixtureFeedForward:
def __call__(self, x:Tensor) -> Tensor:
assert x.shape[0] == 1, "only BS=1"
assert x.shape[1] == 1, "only length=1"
g = self.gate(x).float().softmax(-1)
g = self.gate(x).softmax(-1)
g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
probs, sel = g.topk(self.activated_experts)
@@ -25,7 +24,7 @@ class MixtureFeedForward:
# run MoE
x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
return (x_down * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
# model is bf16, 1.3B active, 6.9B total
# M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
tinygrad_repo/examples/yolov3.py
+9 -9
View File
@@ -71,8 +71,8 @@ def bbox_iou(box1, box2):
# get the coordinates of the intersection rectangle
inter_rect_x1 = np.maximum(b1_x1, b2_x1)
inter_rect_y1 = np.maximum(b1_y1, b2_y1)
inter_rect_x2 = np.maximum(b1_x2, b2_x2)
inter_rect_y2 = np.maximum(b1_y2, b2_y2)
inter_rect_x2 = np.minimum(b1_x2, b2_x2)
inter_rect_y2 = np.minimum(b1_y2, b2_y2)
#Intersection area
inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, 99999) * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, 99999)
#Union Area
@@ -297,13 +297,13 @@ class Darknet:
# Get the number of weights of batchnorm
num_bn_biases = math.prod(bn.bias.shape)
# Load weights
bn_biases = Tensor(weights[ptr:ptr + num_bn_biases])
bn_biases = Tensor(weights[ptr:ptr + num_bn_biases].astype(np.float32))
ptr += num_bn_biases
bn_weights = Tensor(weights[ptr:ptr+num_bn_biases])
bn_weights = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
ptr += num_bn_biases
bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases])
bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
ptr += num_bn_biases
bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases])
bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
ptr += num_bn_biases
# Cast the loaded weights into dims of model weights
bn_biases = bn_biases.reshape(shape=tuple(bn.bias.shape))
@@ -319,7 +319,7 @@ class Darknet:
# load biases of the conv layer
num_biases = math.prod(conv.bias.shape)
# Load weights
conv_biases = Tensor(weights[ptr: ptr+num_biases])
conv_biases = Tensor(weights[ptr: ptr+num_biases].astype(np.float32))
ptr += num_biases
# Reshape
conv_biases = conv_biases.reshape(shape=tuple(conv.bias.shape))
@@ -327,7 +327,7 @@ class Darknet:
conv.bias = conv_biases
# Load weighys for conv layers
num_weights = math.prod(conv.weight.shape)
conv_weights = Tensor(weights[ptr:ptr+num_weights])
conv_weights = Tensor(weights[ptr:ptr+num_weights].astype(np.float32))
ptr += num_weights
conv_weights = conv_weights.reshape(shape=tuple(conv.weight.shape))
conv.weight = conv_weights
@@ -371,7 +371,7 @@ class Darknet:
if __name__ == "__main__":
model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
print("Loading weights file (237MB). This might take a while…")
model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
model.load_weights('https://github.com/shadiakiki1986/yolov3.weights/releases/download/3.0.1/yolov3.weights')
if len(sys.argv) > 1:
url = sys.argv[1]
else:
tinygrad_repo/extra/amdpci/proclogs.py
+16 -11
View File
@@ -1,6 +1,16 @@
import re, ctypes, sys
import re, ctypes, sys, importlib
from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
class AMDFake(AMDev):
def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
self._run_discovery()
self._build_regs()
amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
amdev.AMDev = AMDFake
from tinygrad.runtime.ops_amd import PCIIface
def parse_amdgpu_logs(log_content, register_names=None):
register_map = register_names
@@ -23,16 +33,11 @@ def parse_amdgpu_logs(log_content, register_names=None):
return processed_log
def main():
regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
reg_names = {}
def _prepare_registers(modules):
for base, m in modules:
for k, regval in m.__dict__.items():
if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
_prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
dev = PCIIface(None, 0)
for x, y in dev.dev_impl.__dict__.items():
if isinstance(y, AMRegister):
for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"
with open(sys.argv[1], 'r') as f:
log_content = log_content_them = f.read()
tinygrad_repo/extra/assembly/assembly.py
+1 -1
View File
@@ -1,5 +1,5 @@
from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
from tinygrad.opt.kernel import Ops, MemOp, UOp
from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
from tinygrad.uop.ops import BinaryOps, UnaryOps
from tinygrad.dtype import DType, dtypes
from tinygrad.helpers import DEBUG
tinygrad_repo/extra/assembly/assembly_arm64.py
+1 -1
View File
@@ -3,7 +3,7 @@ from platform import system
from typing import Tuple, Dict, List, Optional
from tinygrad import dtypes
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.opt.kernel import Ops, UOp
from tinygrad.codegen.opt.kernel import Ops, UOp
from tinygrad.helpers import CI
from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
tinygrad_repo/extra/assembly/assembly_ptx.py
+1 -1
View File
@@ -1,7 +1,7 @@
from typing import List
import struct
from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
from tinygrad.opt.kernel import Ops, UOp
from tinygrad.codegen.opt.kernel import Ops, UOp
from tinygrad import dtypes
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.runtime.ops_cuda import arch
tinygrad_repo/extra/assembly/assembly_rdna.py
+1 -1
View File
@@ -2,7 +2,7 @@ import yaml
from typing import Tuple, Set, Dict
from tinygrad import dtypes
from tinygrad.codegen.assembly import AssemblyCodegen, Register
from tinygrad.opt.kernel import Ops
from tinygrad.codegen.opt.kernel import Ops
from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
tinygrad_repo/extra/backends/triton.py
+1 -1
View File
@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
from collections import defaultdict
from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
from tinygrad.opt.kernel import UOp, Ops
from tinygrad.codegen.opt.kernel import UOp, Ops
from triton.compiler import compile as triton_compile
import linecache
import math
tinygrad_repo/extra/gemm/amd_uop_matmul.py
+25 -4
View File
@@ -2,11 +2,11 @@ from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
from tinygrad.dtype import AddrSpace
from tinygrad.schedule.kernelize import merge_views, view_left
from tinygrad.helpers import getenv, colored, prod, unwrap
from tinygrad.shape.shapetracker import ShapeTracker, View
from tinygrad.shape.view import strides_for_shape
from tinygrad.opt.kernel import axis_colors
from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
from tinygrad.codegen.opt.swizzler import merge_views, view_left
def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
@@ -44,6 +44,21 @@ pm = PatternMatcher([
(UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
])
def rangeify_kernel3():
a = Tensor.empty(N,N)
b = Tensor.empty(N,N)
c = a@b
#c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
with Context(RANGEIFY=1):
sink = c.schedule()[-1].ast
#print(sink)
opts = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
opts += [Opt(OptOps.UNROLL, 0, 8)]
return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
def top_spec_kernel3():
a = Tensor.empty(N,N)
b = Tensor.empty(N,N)
@@ -87,6 +102,7 @@ def hl_spec_kernel3():
# this makes all the global loads match
# this can also be more simply done by rebinding the RANGEs
# but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
permute_a[17:20] = [11,12,13]
permute_a[11:14] = [17,18,19]
permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
@@ -308,10 +324,15 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
if __name__ == "__main__":
HL = getenv("HL")
if HL == 2: hprg = top_spec_kernel3()
if HL == 3: hprg = rangeify_kernel3()
elif HL == 2: hprg = top_spec_kernel3()
elif HL == 1: hprg = hl_spec_kernel3()
else: hprg = hand_spec_kernel3()
prg = get_program(hprg, Device.default.renderer)
if HL == 3:
with Context(RANGEIFY=1, BLOCK_REORDER=0):
prg = get_program(hprg, Device.default.renderer)
else:
prg = get_program(hprg, Device.default.renderer)
print(prg.src)
if getenv("SRC"): exit(0)
hrunner = CompiledRunner(prg)
tinygrad_repo/extra/gemm/max_matmul.py
+2 -2
View File
@@ -5,9 +5,9 @@ from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
from tinygrad.engine.realize import get_program
# for copied uops
from tinygrad.opt.kernel import Kernel, KernelOptError
from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
from tinygrad.opt.search import Opt, OptOps
from tinygrad.codegen.opt.search import Opt, OptOps
from tinygrad import Device, dtypes, Tensor
from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
from tinygrad.shape.shapetracker import ShapeTracker
tinygrad_repo/extra/gemm/simple_matmul.py
+1 -1
View File
@@ -2,7 +2,7 @@ import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, get_single_element
from tinygrad.dtype import _to_np_dtype
from tinygrad.opt.kernel import OptOps
from tinygrad.codegen.opt.kernel import OptOps
from tinygrad.engine.realize import lower_schedule
dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py
+1 -1
View File
@@ -1,6 +1,6 @@
from tinygrad import Tensor, dtypes, Device
from tinygrad.helpers import getenv, DEBUG
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
from dataclasses import replace
tinygrad_repo/extra/gemm/tvm_gemm.py
+1 -1
View File
@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
sched = C.schedule()
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.device import CompilerOptions
lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
lin.to_program()
tinygrad_repo/extra/hcq/hcq_smi.py
+122
View File
@@ -0,0 +1,122 @@
#!/usr/bin/env python3
from tinygrad.runtime.support.system import System
import argparse, glob, os, re, time, subprocess, sys
def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
devs = []
for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
dev_id = dev[8:-5]
if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
return devs
def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
def cmd_remove_module(args):
modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
to_unload = [m for m in modules if _is_module_loaded(m)]
if not to_unload: print("Kernel modules are not loaded")
else:
print("Removing kernel modules:", ", ".join(to_unload))
try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
except subprocess.CalledProcessError as e:
print("Failed to unload all modules — they may be in use.", file=sys.stderr)
sys.exit(e.returncode)
def cmd_insert_module(args):
cmd_remove_module(args)
cmd_reset_devices(args)
module = "nvidia" if args.backend == "nv" else "amdgpu"
if _is_module_loaded(module):
print(f"{module} kernel module already loaded")
return
print(f"Inserting kernel module: {module}")
if args.backend == "nv":
subprocess.run(["nvidia-smi"], check=True)
elif args.backend == "amd":
subprocess.run(["sudo", "modprobe", "amdgpu"], check=True)
def cmd_reset_devices(args):
devs = scan_devs_based_on_lock({"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
print(f"Resetting device {dev}")
if args.backend != "amd": _do_reset_device(dev)
time.sleep(0.2)
def cmd_show_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
try:
pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
print(f"{dev}: {pid}")
except subprocess.CalledProcessError: print(f"{dev}: No processes found using this device")
def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
for dev in devs:
for i in range(128):
if i > 0: time.sleep(0.2)
try:
try: pid = subprocess.check_output(['sudo', 'lsof', f'/tmp/{prefix}_{dev}.lock']).decode('utf-8').strip().split('\n')[1].split()[1]
except subprocess.CalledProcessError: break
print(f"Killing process {pid} (which uses {dev})")
subprocess.run(['sudo', 'kill', '-9', pid], check=True)
except subprocess.CalledProcessError as e:
print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
def add_common_commands(parent_subparsers):
p_insmod = parent_subparsers.add_parser("insmod", help="Insert a kernel module")
p_insmod.set_defaults(func=cmd_insert_module)
p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
p_rmmod.set_defaults(func=cmd_remove_module)
p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device to reset")
p_reset.set_defaults(func=cmd_reset_devices)
p_reset = parent_subparsers.add_parser("pids", help="Show pids of processes using the device")
p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
p_reset.set_defaults(func=cmd_show_pids)
p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
p_reset.set_defaults(func=cmd_kill_pids)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
nv_commands = nv_parser.add_subparsers(dest="command", required=True)
add_common_commands(nv_commands)
amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
amd_commands = amd_parser.add_subparsers(dest="command", required=True)
add_common_commands(amd_commands)
args = parser.parse_args()
if args.command is None:
parser.print_help(sys.stderr)
sys.exit(1)
args.func(args)
tinygrad_repo/extra/hcqfuzz/tests/bert.py
-1
View File
@@ -8,7 +8,6 @@ bert_train_params = {
"BS": 96,
"EVAL_BS": 96,
"FUSE_ARANGE": 1,
"FUSE_ARANGE_UINT": 0,
"BASEDIR": "/raid/datasets/wiki",
}
tinygrad_repo/extra/huggingface_onnx/README.md
+61
View File
@@ -0,0 +1,61 @@
# HuggingFace ONNX
Tool for discovering, downloading, and validating ONNX models from HuggingFace.
## Extra Dependencies
```bash
pip install huggingface_hub pyyaml requests onnx onnxruntime numpy
```
## Huggingface Manager (discovering and downloading)
The `huggingface_manager.py` script discovers top ONNX models from HuggingFace, collects metadata, and optionally downloads them.
```bash
# Download top 50 models sorted by downloads
python huggingface_manager.py --limit 50 --download
# Just collect metadata (no download)
python huggingface_manager.py --limit 100
# Sort by likes instead of downloads
python huggingface_manager.py --limit 20 --sort likes --download
# Custom output file
python huggingface_manager.py --limit 10 --output my_models.yaml
```
### Output Format
The tool generates a YAML file with the following structure:
```yaml
repositories:
"model-name":
url: "https://huggingface.co/model-name"
download_path: "/path/to/models/..." # when --download used
files:
- file: "model.onnx"
size: "90.91MB"
total_size: "2.45GB"
created_at: "2024-01-15T10:30:00Z"
```
## Run Models (validation)
The `run_models.py` script validates ONNX models against ONNX Runtime for correctness.
```bash
# Validate models from a YAML configuration file
python run_models.py --validate huggingface_repos.yaml
# Debug specific repository (downloads and validates all ONNX models)
python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2
# Debug specific model file
python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx
# Debug with model truncation for debugging and validating intermediate results
DEBUGONNX=1 python run_models.py --debug sentence-transformers/all-MiniLM-L6-v2/onnx/model.onnx --truncate 10
```
tinygrad_repo/extra/huggingface_onnx/collect_metadata.py
-85
View File
@@ -1,85 +0,0 @@
import yaml, time, requests, argparse
from pathlib import Path
from huggingface_hub import list_models, HfApi
from tinygrad.helpers import tqdm
HUGGINGFACE_URL = "https://huggingface.co"
SKIPPED_FILES = [
"fp16", "int8", "uint8", "quantized", # numerical accuracy issues
"avx2", "arm64", "avx512", "avx512_vnni", # numerical accuracy issues
"q4", "q4f16", "bnb4", # unimplemented quantization
"model_O4", # requires non cpu ort runner and MemcpyFromHost op
"merged", # TODO implement attribute with graph type and Loop op
]
SKIPPED_REPO_PATHS = [
# Invalid model-index
"AdamCodd/vit-base-nsfw-detector",
# TODO: implement attribute with graph type and Loop op
"minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
"HuggingFaceTB/SmolLM2-360M-Instruct",
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
"HuggingFaceTB/SmolLM2-1.7B-Instruct",
# TODO: implmement RandomNormalLike
"stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
# TODO: implement NonZero
"mangoapps/fb_zeroshot_mnli_onnx",
# TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
"briaai/RMBG-2.0",
]
def get_top_repos(n: int, sort: str) -> list[str]: # list["FacebookAI/xlm-roberta-large", ...]
print(f"** Getting top {n} models sorted by {sort} **")
repos = []
i = 0
for model in list_models(filter="onnx", sort=sort):
if model.id in SKIPPED_REPO_PATHS: continue
print(f"{i+1}/{n}: {model.id} ({getattr(model, sort)})")
repos.append(model.id)
i += 1
if i == n: break
return repos
def get_metadata(repos:list[str]) -> dict:
api = HfApi()
repos_metadata = {"repositories": {}}
total_size = 0
# TODO: speed head requests up with async?
for repo in tqdm(repos, desc="Getting metadata"):
files_metadata = []
model_info = api.model_info(repo)
for file in model_info.siblings:
filename = file.rfilename
if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')): continue
if any(skip_str in filename for skip_str in SKIPPED_FILES): continue
head = requests.head(f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}", allow_redirects=True)
file_size = file.size or int(head.headers.get('Content-Length', 0))
files_metadata.append({"file": filename, "size": f"{file_size/1e6:.2f}MB"})
total_size += file_size
repos_metadata["repositories"][repo] = {
"url": f"{HUGGINGFACE_URL}/{repo}",
"download_path": None,
"files": files_metadata,
}
repos_metadata['total_size'] = f"{total_size/1e9:.2f}GB"
repos_metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
return repos_metadata
if __name__ == "__main__":
sort = "downloads" # recent 30 days downloads
huggingface_onnx_dir = Path(__file__).parent
parser = argparse.ArgumentParser(description="Produces a YAML file with metadata of top huggingface onnx models")
parser.add_argument("--limit", type=int, required=True, help="Number of top repositories to process (e.g., 100)")
parser.add_argument("--output", type=str, default="huggingface_repos.yaml", help="Output YAML file name to save the report")
args = parser.parse_args()
top_repos = get_top_repos(args.limit, sort)
metadata = get_metadata(top_repos)
yaml_path = huggingface_onnx_dir / args.output
with open(yaml_path, 'w') as f:
yaml.dump(metadata, f, sort_keys=False)
print(f"YAML saved to: {str(yaml_path)}")
tinygrad_repo/extra/huggingface_onnx/download_models.py
-29
View File
@@ -1,29 +0,0 @@
import yaml, argparse
from pathlib import Path
from huggingface_hub import snapshot_download
def download_models(yaml_file: str, download_dir: str) -> None:
with open(yaml_file, 'r') as f: metadata = yaml.safe_load(f)
n = len(metadata["repositories"])
for i, (model_id, model_data) in enumerate(metadata["repositories"].items()):
print(f"Downloading {i+1}/{n}: {model_id}...")
allow_patterns = [file_info["file"] for file_info in model_data["files"]]
root_path = Path(snapshot_download(repo_id=model_id, allow_patterns=allow_patterns, cache_dir=download_dir))
# download configs too (the sizes are small)
snapshot_download(repo_id=model_id, allow_patterns=["*config.json"], cache_dir=download_dir)
print(f"Downloaded model files to: {root_path}")
model_data["download_path"] = str(root_path)
# Save the updated metadata back to the YAML file
with open(yaml_file, 'w') as f: yaml.dump(metadata, f, sort_keys=False)
print("Download completed according to YAML file.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Download models from Huggingface Hub based on a YAML configuration file.")
parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
args = parser.parse_args()
models_folder = Path(__file__).parent / "models"
models_folder.mkdir(parents=True, exist_ok=True)
download_models(args.input, str(models_folder))
tinygrad_repo/extra/huggingface_onnx/huggingface_manager.py
+230
View File
@@ -0,0 +1,230 @@
import yaml
import time
import requests
import argparse
from pathlib import Path
from huggingface_hub import list_models, HfApi, snapshot_download
from tinygrad.helpers import _ensure_downloads_dir
DOWNLOADS_DIR = _ensure_downloads_dir() / "models"
from tinygrad.helpers import tqdm
def snapshot_download_with_retry(*, repo_id: str, allow_patterns: list[str]|tuple[str, ...]|None=None, cache_dir: str|Path|None=None,
tries: int=2, **kwargs) -> Path:
for attempt in range(tries):
try:
return Path(snapshot_download(
repo_id=repo_id,
allow_patterns=allow_patterns,
cache_dir=str(cache_dir) if cache_dir is not None else None,
**kwargs
))
except Exception as e:
if attempt == tries-1: raise
time.sleep(1)
# Constants for filtering models
HUGGINGFACE_URL = "https://huggingface.co"
SKIPPED_FILES = [
"fp16", "int8", "uint8", "quantized", # numerical accuracy issues
"avx2", "arm64", "avx512", "avx512_vnni", # numerical accuracy issues
"q4", "q4f16", "bnb4", # unimplemented quantization
"model_O4", # requires non cpu ort runner and MemcpyFromHost op
"merged", # TODO implement attribute with graph type and Loop op
]
SKIPPED_REPO_PATHS = [
# Invalid model-index
"AdamCodd/vit-base-nsfw-detector",
# TODO: implement attribute with graph type and Loop op
"minishlab/potion-base-8M", "minishlab/M2V_base_output", "minishlab/potion-retrieval-32M",
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, GroupQueryAttention
"HuggingFaceTB/SmolLM2-360M-Instruct",
# TODO: implement SimplifiedLayerNormalization, SkipSimplifiedLayerNormalization, RotaryEmbedding, MultiHeadAttention
"HuggingFaceTB/SmolLM2-1.7B-Instruct",
# TODO: implement RandomNormalLike
"stabilityai/stable-diffusion-xl-base-1.0", "stabilityai/sdxl-turbo", 'SimianLuo/LCM_Dreamshaper_v7',
# TODO: implement NonZero
"mangoapps/fb_zeroshot_mnli_onnx",
# TODO huge Concat in here with 1024 (1, 3, 32, 32) Tensors, and maybe a MOD bug with const folding
"briaai/RMBG-2.0",
]
class HuggingFaceONNXManager:
def __init__(self):
self.base_dir = Path(__file__).parent
self.models_dir = DOWNLOADS_DIR
self.api = HfApi()
def discover_models(self, limit: int, sort: str = "downloads") -> list[str]:
print(f"Discovering top {limit} ONNX models sorted by {sort}...")
repos = []
i = 0
for model in list_models(filter="onnx", sort=sort):
if model.id in SKIPPED_REPO_PATHS:
continue
print(f" {i+1}/{limit}: {model.id} ({getattr(model, sort)})")
repos.append(model.id)
i += 1
if i == limit:
break
print(f"Found {len(repos)} suitable ONNX models")
return repos
def collect_metadata(self, repos: list[str]) -> dict:
print(f"Collecting metadata for {len(repos)} repositories...")
metadata = {"repositories": {}}
total_size = 0
for repo in tqdm(repos, desc="Collecting metadata"):
try:
files_metadata = []
model_info = self.api.model_info(repo)
for file in model_info.siblings:
filename = file.rfilename
if not (filename.endswith('.onnx') or filename.endswith('.onnx_data')):
continue
if any(skip_str in filename for skip_str in SKIPPED_FILES):
continue
# Get file size from API or HEAD request
try:
head = requests.head(
f"{HUGGINGFACE_URL}/{repo}/resolve/main/{filename}",
allow_redirects=True,
timeout=10
)
file_size = file.size or int(head.headers.get('Content-Length', 0))
except requests.RequestException:
file_size = file.size or 0
files_metadata.append({
"file": filename,
"size": f"{file_size/1e6:.2f}MB"
})
total_size += file_size
if files_metadata: # Only add repos with valid ONNX files
metadata["repositories"][repo] = {
"url": f"{HUGGINGFACE_URL}/{repo}",
"download_path": None,
"files": files_metadata,
}
except Exception as e:
print(f"WARNING: Failed to collect metadata for {repo}: {e}")
continue
metadata['total_size'] = f"{total_size/1e9:.2f}GB"
metadata['created_at'] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
print(f"Collected metadata for {len(metadata['repositories'])} repositories")
print(f"Total estimated download size: {metadata['total_size']}")
return metadata
def download_models(self, metadata: dict) -> dict:
self.models_dir.mkdir(parents=True, exist_ok=True)
repos = metadata["repositories"]
n = len(repos)
print(f"Downloading {n} repositories to {self.models_dir}...")
for i, (model_id, model_data) in enumerate(repos.items()):
print(f" Downloading {i+1}/{n}: {model_id}...")
try:
# Download ONNX model files
allow_patterns = [file_info["file"] for file_info in model_data["files"]]
root_path = snapshot_download_with_retry(
repo_id=model_id,
allow_patterns=allow_patterns,
cache_dir=str(self.models_dir)
)
# Download config files (usually small)
snapshot_download_with_retry(
repo_id=model_id,
allow_patterns=["*config.json"],
cache_dir=str(self.models_dir)
)
model_data["download_path"] = str(root_path)
print(f" Downloaded to: {root_path}")
except Exception as e:
print(f" ERROR: Failed to download {model_id}: {e}")
model_data["download_path"] = None
continue
successful_downloads = sum(1 for repo in repos.values() if repo["download_path"] is not None)
print(f"Successfully downloaded {successful_downloads}/{n} repositories")
print(f"All models saved to: {self.models_dir}")
return metadata
def save_metadata(self, metadata: dict, output_file: str):
yaml_path = self.base_dir / output_file
with open(yaml_path, 'w') as f:
yaml.dump(metadata, f, sort_keys=False)
print(f"Metadata saved to: {yaml_path}")
def discover_and_download(self, limit: int, output_file: str = "huggingface_repos.yaml",
sort: str = "downloads", download: bool = True):
print(f"Starting HuggingFace ONNX workflow...")
print(f" Limit: {limit} models")
print(f" Sort by: {sort}")
print(f" Download: {'Yes' if download else 'No'}")
print(f" Output: {output_file}")
print("-" * 50)
repos = self.discover_models(limit, sort)
metadata = self.collect_metadata(repos)
if download:
metadata = self.download_models(metadata)
self.save_metadata(metadata, output_file)
print("-" * 50)
print("Workflow completed successfully!")
if download:
successful = sum(1 for repo in metadata["repositories"].values()
if repo["download_path"] is not None)
print(f"{successful}/{len(metadata['repositories'])} models downloaded")
return metadata
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="HuggingFace ONNX Model Manager - Discover, collect metadata, and download ONNX models",
)
parser.add_argument("--limit", type=int, help="Number of top repositories to process")
parser.add_argument("--output", type=str, default="huggingface_repos.yaml",
help="Output YAML file name (default: huggingface_repos.yaml)")
parser.add_argument("--sort", type=str, default="downloads",
choices=["downloads", "likes", "created", "modified"],
help="Sort criteria for model discovery (default: downloads)")
parser.add_argument("--download", action="store_true", default=False,
help="Download models after collecting metadata")
args = parser.parse_args()
if not args.limit: parser.error("--limit is required")
manager = HuggingFaceONNXManager()
manager.discover_and_download(
limit=args.limit,
output_file=args.output,
sort=args.sort,
download=args.download
)
tinygrad_repo/extra/huggingface_onnx/run_models.py
+23 -50
View File
@@ -1,10 +1,11 @@
import onnx, yaml, tempfile, time, collections, pprint, argparse, json
import onnx, yaml, tempfile, time, argparse, json
from pathlib import Path
from typing import Any
from tinygrad.frontend.onnx import OnnxRunner
from extra.onnx import get_onnx_ops
from extra.onnx_helpers import validate, get_example_inputs
from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
def get_config(root_path: Path):
def get_config(root_path: Path) -> dict[str, Any]:
ret = {}
for path in root_path.rglob("*config.json"):
config = json.load(path.open())
@@ -12,19 +13,19 @@ def get_config(root_path: Path):
ret.update(config)
return ret
def run_huggingface_validate(onnx_model_path, config, rtol, atol):
onnx_runner = OnnxRunner(onnx_model_path)
inputs = get_example_inputs(onnx_runner.graph_inputs, config)
validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
def get_tolerances(file_name): # -> rtol, atol
def get_tolerances(file_name: str) -> tuple[float, float]:
# TODO very high rtol atol
if "fp16" in file_name: return 9e-2, 9e-2
if any(q in file_name for q in ["int8", "uint8", "quantized"]): return 4, 4
return 4e-3, 3e-2
def run_huggingface_validate(onnx_model_path: str | Path, config: dict[str, Any], rtol: float, atol: float):
onnx_runner = OnnxRunner(onnx_model_path)
inputs = get_example_inputs(onnx_runner.graph_inputs, config)
validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
def validate_repos(models:dict[str, tuple[Path, Path]]):
print(f"** Validating {len(model_paths)} models **")
print(f"** Validating {len(models)} models **")
for model_id, (root_path, relative_path) in models.items():
print(f"validating model {model_id}")
model_path = root_path / relative_path
@@ -36,25 +37,6 @@ def validate_repos(models:dict[str, tuple[Path, Path]]):
et = time.time() - st
print(f"passed, took {et:.2f}s")
def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
ret = {}
op_counter = collections.Counter()
unsupported_ops = collections.defaultdict(set)
supported_ops = get_onnx_ops()
print(f"** Retrieving stats from {len(model_paths)} models **")
for model_id, (root_path, relative_path) in models.items():
print(f"examining {model_id}")
model_path = root_path / relative_path
onnx_runner = OnnxRunner(model_path)
for node in onnx_runner.graph_nodes:
op_counter[node.op] += 1
if node.op not in supported_ops:
unsupported_ops[node.op].add(model_id)
del onnx_runner
ret["unsupported_ops"] = {k:list(v) for k, v in unsupported_ops.items()}
ret["op_counter"] = op_counter.most_common()
return ret
def debug_run(model_path, truncate, config, rtol, atol):
if truncate != -1:
model = onnx.load(model_path)
@@ -71,12 +53,9 @@ def debug_run(model_path, truncate, config, rtol, atol):
run_huggingface_validate(model_path, config, rtol, atol)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator and Ops Checker")
parser.add_argument("input", type=str, help="Path to the input YAML configuration file containing model information.")
parser.add_argument("--check_ops", action="store_true", default=False,
help="Check support for ONNX operations in models from the YAML file")
parser.add_argument("--validate", action="store_true", default=False,
help="Validate correctness of models from the YAML file")
parser = argparse.ArgumentParser(description="Huggingface ONNX Model Validator")
parser.add_argument("--validate", type=str, default="",
help="Validate correctness of models from the specified YAML configuration file")
parser.add_argument("--debug", type=str, default="",
help="""Validates without explicitly needing a YAML or models pre-installed.
provide repo id (e.g. "minishlab/potion-base-8M") to validate all onnx models inside the repo
@@ -85,13 +64,13 @@ if __name__ == "__main__":
parser.add_argument("--truncate", type=int, default=-1, help="Truncate the ONNX model so intermediate results can be validated")
args = parser.parse_args()
if not (args.check_ops or args.validate or args.debug):
parser.error("Please provide either --validate, --check_ops, or --debug.")
if not (args.validate or args.debug):
parser.error("Please provide either --validate <yaml_file> or --debug <repo_id>.")
if args.truncate != -1 and not args.debug:
parser.error("--truncate and --debug should be used together for debugging")
if args.check_ops or args.validate:
with open(args.input, 'r') as f:
if args.validate:
with open(args.validate, 'r') as f:
data = yaml.safe_load(f)
assert all(repo["download_path"] is not None for repo in data["repositories"].values()), "please run `download_models.py` for this yaml"
model_paths = {
@@ -101,22 +80,16 @@ if __name__ == "__main__":
if model["file"].endswith(".onnx")
}
if args.check_ops:
pprint.pprint(retrieve_op_stats(model_paths))
if args.validate:
validate_repos(model_paths)
validate_repos(model_paths)
if args.debug:
from huggingface_hub import snapshot_download
download_dir = Path(__file__).parent / "models"
path:list[str] = args.debug.split("/")
if len(path) == 2:
# repo id
# validates all onnx models inside repo
repo_id = "/".join(path)
root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=download_dir))
snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*.onnx", "*.onnx_data"], cache_dir=DOWNLOADS_DIR)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
config = get_config(root_path)
for onnx_model in root_path.rglob("*.onnx"):
rtol, atol = get_tolerances(onnx_model.name)
@@ -128,8 +101,8 @@ if __name__ == "__main__":
onnx_model = path[-1]
assert path[-1].endswith(".onnx")
repo_id, relative_path = "/".join(path[:2]), "/".join(path[2:])
root_path = Path(snapshot_download(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=download_dir))
snapshot_download(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=download_dir)
root_path = snapshot_download_with_retry(repo_id=repo_id, allow_patterns=[relative_path], cache_dir=DOWNLOADS_DIR)
snapshot_download_with_retry(repo_id=repo_id, allow_patterns=["*config.json"], cache_dir=DOWNLOADS_DIR)
config = get_config(root_path)
rtol, atol = get_tolerances(onnx_model)
print(f"validating {relative_path} with truncate={args.truncate}, {rtol=}, {atol=}")
tinygrad_repo/extra/mcts_search.py
+2 -2
View File
@@ -4,9 +4,9 @@ import numpy as np
np.set_printoptions(suppress=True)
import math, functools, time, random, statistics
from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.device import Buffer, Device, CompileError
from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
from tinygrad.engine.realize import get_program
class MCTSNode:
tinygrad_repo/extra/models/llama.py
+7 -7
View File
@@ -99,7 +99,9 @@ class FeedForward:
self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
def __call__(self, x:Tensor) -> Tensor:
return self.w2(self.w1(x).silu() * self.w3(x)) # SwiGLU [arxiv/2002.05202, eq (5)]
w1 = self.w1(x).silu()
w3 = self.w3(x.contiguous_backward()) # this fixes a strange fusion that makes tensor cores miss
return self.w2(w1 * w3)
class TransformerBlock:
def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int, norm_eps:float, max_context:int, linear=nn.Linear,
@@ -111,7 +113,7 @@ class TransformerBlock:
def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
# standard openai sampling
def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
@@ -179,16 +181,14 @@ class Transformer:
def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
_bsz, seqlen = tokens.shape
h = self.tok_embeddings(tokens)
self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
freqs_cis = self.freqs_cis.cast(h.dtype)[:, start_pos:start_pos+seqlen, :, :, :]
mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
logits = self.output(self.norm(h)).float()[:, -1, :]
logits = self.output(self.norm(h))
if math.isnan(temperature): return logits
return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
return sample(logits[:, -1, :].flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
# TODO: better way to handle the first call v.s. the rest?
tinygrad_repo/extra/nvpci/gpu_reset.sh
-2
View File
@@ -1,2 +0,0 @@
GPU="$1"
echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null
tinygrad_repo/extra/nvpci/nv_smi.py
-65
View File
@@ -1,65 +0,0 @@
#!/usr/bin/env python3
from tinygrad.runtime.support.system import System
import argparse, glob, os, re, time, subprocess, sys
def scan_devs_based_on_lock(prefix:str) -> list[str]:
devs = []
for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
dev_id = dev[8:-5]
if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
return devs
def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
def cmd_remove_module(args):
to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
if not to_unload:
print("NVIDIA kernel modules are not loaded")
else:
print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
except subprocess.CalledProcessError as e:
print("Failed to unload all modules — they may be in use.", file=sys.stderr)
sys.exit(e.returncode)
def cmd_insert_module(args):
cmd_remove_module(args)
cmd_reset_devices(args)
if not os.path.exists("/sys/module/nvidia"):
print("Inserting nvidia kernel module")
subprocess.run(["nvidia-smi"], check=True)
else: print("Nvidia kernel module already loaded")
def cmd_reset_devices(args):
devs = scan_devs_based_on_lock("nv")
dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
for dev in devs:
if dev.startswith(dev_to_reset):
print(f"Resetting device {dev}")
_do_reset_device(dev)
time.sleep(0.2)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
subparsers = parser.add_subparsers(required=True, dest="cmd")
parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
parser_insmod.set_defaults(func=cmd_insert_module)
parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
parser_rmmod.set_defaults(func=cmd_remove_module)
parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
parser_reset.set_defaults(func=cmd_reset_devices)
args = parser.parse_args()
if args.cmd is None:
parser.print_help(sys.stderr)
sys.exit(1)
args.func(args)
tinygrad_repo/extra/onnx.py
-954
View File
@@ -1,954 +0,0 @@
# mypy: disable-error-code="misc, list-item, assignment, operator, index, arg-type"
from types import SimpleNamespace
from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element
from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
from tinygrad.device import is_dtype_supported, Device
from extra.onnx_parser import onnx_load
# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
data_types: dict[int, DType] = {
1:dtypes.float32, 2:dtypes.uint8, 3:dtypes.int8, 4:dtypes.uint16, 5:dtypes.int16, 6:dtypes.int32, 7:dtypes.int64,
9:dtypes.bool, 10:dtypes.float16, 11:dtypes.double, 12:dtypes.uint32, 13:dtypes.uint64, 16:dtypes.bfloat16,
}
# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L128-L145
attribute_types: dict[int, Callable] = {
1: lambda a: float(a.f),
2: lambda a: int(a.i),
3: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
4: lambda a: buffer_parse(a.t),
6: lambda a: tuple(float(x) for x in a.floats),
7: lambda a: tuple(int(x) for x in a.ints),
8: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
}
# ***** protobuf parsing ******
from onnx import AttributeProto, TensorProto, TypeProto
def has_field(onnx_type: TypeProto|SimpleNamespace, field):
if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
return hasattr(onnx_type, field)
def dtype_parse(onnx_dtype: int, fallback_context: str | None = None) -> DType:
if onnx_dtype not in data_types: raise NotImplementedError(f"onnx dtype id {onnx_dtype} is not supported")
if is_dtype_supported(dtype := data_types[onnx_dtype]): return dtype
# if fallback_context is provided, we can fall back to a default dtype
if fallback_context is not None:
default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
return default_dtype
raise RuntimeError(f"dtype {dtype} on device {Device.DEFAULT} is not supported")
def attribute_parse(onnx_attribute: AttributeProto):
if onnx_attribute.type not in attribute_types: raise NotImplementedError(f"attribute type {onnx_attribute.type} is not supported")
return attribute_types[onnx_attribute.type](onnx_attribute)
def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
to_dtype, true_dtype = dtype_parse(onnx_tensor.data_type, "buffer parse"), data_types[onnx_tensor.data_type]
shape = tuple(onnx_tensor.dims)
keys = ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', "raw_data"]
data = next((val for k in keys if (val := getattr(onnx_tensor, k)) is not None), None)
if data is None: raise RuntimeError("empty buffer")
if not isinstance(data, Tensor): return Tensor(data, dtype=to_dtype).reshape(shape)
assert data.dtype is dtypes.uint8, data.dtype
data = data.bitcast(true_dtype).reshape(shape)
data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
if shape == ():
if data.dtype is dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
return Tensor(data.item(), dtype=to_dtype).reshape(shape)
return data
def type_parse(onnx_type: TypeProto):
elem_type = onnx_type
if has_field(elem_type, "map_type") or has_field(elem_type, "sparse_tensor_type") or has_field(elem_type, "opaque_type"):
raise NotImplementedError("parsing for map_type, sparse_tensor_type and opaque_type are not implemented")
if is_optional := has_field(elem_type, "optional_type"): elem_type = elem_type.optional_type.elem_type
if is_sequence := has_field(elem_type, "sequence_type"): elem_type = elem_type.sequence_type.elem_type
if has_field(elem_type, "tensor_type"):
shape = tuple(getattr(d, "dim_param", None) or getattr(d, "dim_value") for d in elem_type.tensor_type.shape.dim) \
if has_field(elem_type.tensor_type, "shape") else None # test_identity_sequence_cpu
dtype = data_types[elem_type.tensor_type.elem_type]
return OnnxValue(shape, dtype, is_optional, is_sequence)
raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
# ***** onnx spec *****
@dataclasses.dataclass(frozen=True)
class OnnxValue:
shape: tuple[str|int, ...]
dtype: DType
is_optional: bool
is_sequence: bool
class Domain(enum.Enum):
ONNX = "ai.onnx"
ONNX_ML = "ai.onnx.ml"
AI_ONNX_TRAINING = "ai.onnx.training"
AI_ONNX_PREVIEW_TRAINING = "ai.onnx.preview.training"
MICROSOFT_CONTRIB_OPS = "com.microsoft"
@classmethod
def from_onnx(cls, domain: str | None) -> "Domain": return cls.ONNX if domain is None or domain == "" else cls(domain)
class OpSetId(NamedTuple):
domain: Domain
version: int
@dataclasses.dataclass(frozen=True)
class OnnxNode:
num: int
op: str
opset_id: OpSetId
inputs: tuple[str, ...]
outputs: tuple[str, ...]
opts: dict[str, Any]
# ***** python const *****
required_input_python_consts: dict[str, tuple[int, ...]] = {
"Tile": (1,), "Range": (0,1,2), "Expand": (1,), "Reshape": (1,), "Squeeze": (1,), "Unsqueeze": (1,), "Trilu": (1,), "ConstantOfShape": (0,),
"CumSum": (1,), "TopK": (1,), "Pad": (1,2,3), "MaxUnpool": (2,), "Dropout": (1,2), "CenterCropPad": (1,), "OneHot": (1,), "Compress": (1,),
"ImageDecoder": (0,), "AffineGrid": (1,), "Resize": (1,2,3), "Upsample": (1,), "Split": (1,), "Slice": (1,2,3,4),
**{"Reduce"+r: (1,) for r in ("Max", "Min", "Sum", "Mean", "SumSquare", "Prod", "L1", "L2", "LogSum", "LogSumExp")},
**{optim: (1,) for optim in ("Adam", "Adagrad", "Momentum")}
}
cache_misses = 0
@functools.cache
def _cached_to_python_const(t:Tensor):
if t.dtype is dtypes.uint8: return t.data().tobytes()
if 0 in t.shape: return []
return t.tolist()
# Tensor -> python value cache for parameters
def to_python_const(t:Any, op:str, idx:int) -> list[ConstType]|ConstType|bytes:
if idx not in required_input_python_consts.get(op, ()) or not isinstance(t, Tensor): return t
global cache_misses
ret = _cached_to_python_const(t)
if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3:
print(f"Cache miss for {t}")
cache_misses = info.misses
return ret
# ***** runner ******
debug = int(getenv("DEBUGONNX", "0"))
limit = int(getenv("ONNXLIMIT", "-1"))
class OnnxRunner:
"""
`OnnxRunner` executes an ONNX model using Tinygrad.
Args:
model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
"""
def __init__(self, model_path: Tensor | str | pathlib.Path):
model = onnx_load(model_path)
self.is_training = any(n.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in model.graph.node)
self.old_training = Tensor.training
Tensor.training = True if self.is_training else False
self.graph_values = {"": None, **{x.name:buffer_parse(x) for x in model.graph.initializer}}
self.graph_inputs = {x.name:type_parse(x.type) for x in model.graph.input if x.name not in self.graph_values}
self.graph_outputs = tuple(x.name for x in model.graph.output)
opset_imports = {Domain.from_onnx(getattr(x, "domain", "")):x.version for x in model.opset_import}
self.graph_nodes = []
for num, n in enumerate(model.graph.node):
domain = Domain.from_onnx(n.domain)
opset_id = OpSetId(domain, opset_imports.get(domain, 1))
self.graph_nodes.append(OnnxNode(num, n.op_type, opset_id, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute}))
self.graph_nodes = tuple(self.graph_nodes)
self.variable_dims: dict[str, int] = {}
self.onnx_ops = onnx_ops
def _parse_input(self, name: str, value: Any, spec: OnnxValue):
if spec.is_optional and value is None: return None
if spec.is_sequence:
if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
return sequence
dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) else spec.dtype
tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
if isinstance(onnx_dim, str):
onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
return tensor
def _select_op(self, op:str, required_opset:OpSetId) -> types.FunctionType:
if op not in self.onnx_ops: raise NotImplementedError(f"{op=} is not supported")
# return default implementation if no opset_id is specified
if isinstance(impl := self.onnx_ops[op], types.FunctionType): return impl
# match domain and select implementation with latest compatible version
eligible_ops = {impl_opset.version:impl_fxn for impl_opset,impl_fxn in impl.items()
if impl_opset.domain == required_opset.domain and impl_opset.version <= required_opset.version}
if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
return eligible_ops[max(eligible_ops.keys())]
def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
def to(self, device:str|None):
self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
self.graph_nodes = tuple(OnnxNode(n.num, n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
{k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
return self
def __call__(self, inputs:dict[str, Any], debug=debug):
for name, input_spec in self.graph_inputs.items():
if name not in inputs: raise RuntimeError(f"Please provide input data for {name}")
self.graph_values[name] = self._parse_input(name, inputs[name], input_spec)
for node in self.graph_nodes:
inps = [to_python_const(self.graph_values[name], node.op, i) for i,name in enumerate(node.inputs)]
opts = node.opts
# provide additional opts
if node.op == "Split" and 'num_outputs' not in opts: opts['num_outputs'] = len(node.outputs)
if node.op == "Gradient": opts['intermediate_tensors'] = self.graph_values
if debug >= 1: print(f"{node.num}: op '{node.op}' opt {opts}")
if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
ret = ret if isinstance(ret, tuple) else (ret,)
if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{x} - {o!r}" for x,o in zip(node.outputs, ret)))
self.graph_values.update(dict(zip(node.outputs, ret[:len(node.outputs)], strict=True)))
if node.num == limit:
Tensor.training = self.old_training
return {name:self.graph_values[name] for name in node.outputs}
Tensor.training = self.old_training
return {name:self.graph_values[name] for name in self.graph_outputs}
####################
##### ONNX OPS #####
####################
def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
# ***** helper functions *****
def _resolve_const(x: Sequence[ConstType]|ConstType): return x if isinstance(x, get_args(ConstType)) else get_single_element(x)
def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)
# (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
def _onnx_pads_to_tiny_pads(pads): return tuple(flatten(reversed(list(zip(pads, pads[len(pads)//2:])))))
AUTO_PAD_OPTIONS = Literal["NOTSET", "SAME_UPPER", "SAME_LOWER", "VALID"]
# (padding_height, padding_width) -> (padding_top, padding_left, padding_bottom, padding_right)
def _auto_pad(pads, auto_pad: AUTO_PAD_OPTIONS):
if auto_pad == "SAME_UPPER": return [pads[i]//2 for i in range(len(pads))] + [pads[i]-pads[i]//2 for i in range(len(pads))]
return [pads[i]-pads[i]//2 for i in range(len(pads))] + [pads[i]//2 for i in range(len(pads))]
def _resolve_pool_pads(x:Tensor, p_, k_, d_, s_, auto_pad:AUTO_PAD_OPTIONS):
if auto_pad == "VALID": return [0]*(len(k_)*2)
i_, (s_,d_,p_) = x.shape[-len(k_):], (make_tuple(x, len(k_)*2) for x in (s_, d_, p_))
if auto_pad == "NOTSET": return _onnx_pads_to_tiny_pads(p_ if len(p_)==len(k_)*2 else p_*2)
o_ = [((i - (1 if auto_pad in ("SAME_UPPER", "SAME_LOWER") else k)) // s + 1) for i,k,s in zip(i_, k_, s_)]
return _onnx_pads_to_tiny_pads(_auto_pad([(o-1)*s+k-i for o,i,k,s in zip(o_, i_, k_, s_)], auto_pad))
def _clamp_cast(x:Tensor, dtype:DType): return x.clamp(dtypes.min(dtype), dtypes.max(dtype)).cast(dtype)
def _prepare_quantize(x:Tensor, scale:Tensor, zero_point:Tensor|int, axis=1, block_size=0):
if axis < 0: axis += x.ndim
# https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_quantize_linear.py#L31
def reshape(val:Tensor):
if val.numel() == 1: return val
if block_size == 0: return val.reshape([val.shape[0] if dim == axis else 1 for dim in range(x.ndim)])
return val.repeat_interleave(block_size, axis)
return (reshape(scale), reshape(zero_point) if isinstance(zero_point, Tensor) else zero_point)
def _op_integer(op, inputs:list[Tensor], zero_points:list[Tensor], **opts):
adjusted_inputs = [inp.int() - zp for inp, zp in zip(inputs, zero_points)]
return op(*adjusted_inputs, **opts)
def _qlinearop_quantized(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
# op execution is done in quantized int
out = _op_integer(op, inputs, zero_points, **opts)
assert dtypes.is_int(out.dtype), "quantized op should've done math in int"
out_quantized = (out * prod(scales) / out_scale).round() + out_zero_point
return _clamp_cast(out_quantized, out_zero_point.dtype)
def _qlinearop_float(op, inputs:list[Tensor], zero_points:list[Tensor], scales:list[Tensor], out_scale:Tensor, out_zero_point:Tensor, **opts):
# op execution is done in float32
dequantized_inputs = [(inp.int() - zp) * scale for inp, zp, scale in zip(inputs, zero_points, scales)]
out = op(*dequantized_inputs, **opts)
assert dtypes.is_float(out.dtype), "op should've done math in float"
out_quantized = (out / out_scale).round() + out_zero_point
return _clamp_cast(out_quantized, out_zero_point.dtype)
def _onnx_training(input_group_size):
def __decorator(func):
def ___wrapper(R:Tensor, T:int, *inputs:Tensor, **kwargs):
R = R.detach()
groups = len(inputs) // input_group_size
ret = [func(R, T, *inps, **kwargs) for inps in (inputs[i::groups] for i in range(groups))]
return tuple(flatten(zip(*ret)))
return ___wrapper
return __decorator
# ***** Property/Graph Ops *****
def Identity(x:Tensor): return x
def Constant(sparse_value:Tensor|None=None, value:Tensor|None=None, value_float:float|None=None, value_floats:list[float]|None=None,
value_int:int|None=None, value_ints:list[int]|None=None, value_string:str|None=None, value_strings:list[str]|None=None):
if value is not None: return value
if value_float is not None: return Tensor(value_float, dtype=dtypes.float32, requires_grad=False)
if value_floats is not None: return Tensor(list(value_floats), dtype=dtypes.float32, requires_grad=False)
if value_int is not None: return Tensor(value_int, dtype=dtypes.int64, requires_grad=False)
if value_ints is not None: return Tensor(list(value_ints), dtype=dtypes.int64, requires_grad=False)
if value_string is not None or value_strings is not None and sparse_value is not None:
raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value')
def Range(start:float|int|list[float|int], limit:float|int|list[float|int], delta:float|int|list[float|int]):
return Tensor.arange(start=_resolve_const(start), stop=_resolve_const(limit), step=_resolve_const(delta))
def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
try: import PIL.Image
except ImportError as e: raise ImportError("Pillow must be installed for the ImageDecoder operator") from e
img = PIL.Image.open(io.BytesIO(encoded_stream))
if pixel_format == "BGR": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3).flip(-1)
if pixel_format == "RGB": return Tensor(img.tobytes(), dtype=dtypes.uint8).reshape(*img.size, 3)
if pixel_format == "Grayscale": return Tensor(img.convert("L").tobytes(), dtype=dtypes.uint8).reshape(*img.size, 1)
raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype, "EyeLike op") if dtype is not None else x.dtype)
return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
def OptionalGetElement(x:Tensor|None=None): return x if x is not None else Tensor([])
def ConstantOfShape(shape:list[int], value:Tensor|None=None):
if value is None: value = Tensor(0, dtype=dtypes.float32)
if shape == [0]: return Tensor([], dtype=value.dtype)
return value.expand(shape)
def Size(data:Tensor): return data.numel()
def Shape(data:Tensor, end:int|None=None, start:int=0): return Tensor(data.shape[start:end], dtype=dtypes.int64)
# ***** Unary Ops (math) *****
def Not(x:Tensor): return x.logical_not()
def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max) # noqa: A002
def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))
# ***** Unary Ops (activation) *****
def softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
def FastGelu(x:Tensor, bias:Tensor|None=None): return (x + bias).gelu() if bias is not None else x.gelu() # this is tanh approximated
def PRelu(X:Tensor, slope:Tensor): return (X > 0).where(X, X * slope)
def LeakyRelu(X:Tensor, alpha:float=0.01): return X.leaky_relu(alpha)
def ThresholdedRelu(X:Tensor, alpha:float=1.0): return (X > alpha).where(X, 0)
def LogSoftmax(x: Tensor, axis:int=-1): return x.log_softmax(axis)
def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
# ***** Unary Ops (broadcasted) *****
def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
def Less(x:Tensor,y:Tensor): return x < y
def LessOrEqual(x:Tensor,y:Tensor): return x <= y
def Greater(x:Tensor,y:Tensor): return x > y
def GreaterOrEqual(x:Tensor,y:Tensor): return x >= y
def Equal(x:Tensor,y:Tensor): return x == y
def And(x:Tensor,y:Tensor): return (x==y).where(x, False)
def Or(x:Tensor,y:Tensor): return (x==y).where(x, True)
def Xor(x:Tensor,y:Tensor): return x.bool().bitwise_xor(y.bool())
def BitwiseAnd(x:Tensor,y:Tensor): return x & y
def BitwiseOr(x:Tensor,y:Tensor): return x | y
def BitwiseXor(x:Tensor,y:Tensor): return x ^ y
def BitwiseNot(x:Tensor): return ~x
def Mod(x:Tensor,y:Tensor,fmod=0):
if fmod: return x - x.div(y, rounding_mode="trunc") * y
return x % y
# ***** Casting Ops *****
# TODO: saturate
def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to, "Cast op"))
def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
# ***** Reduce Ops *****
def Max(*data_0:Tensor): return functools.reduce(Tensor.maximum, data_0)
def Min(*data_0:Tensor): return functools.reduce(Tensor.minimum, data_0)
def Sum(*data_0:Tensor): return functools.reduce(Tensor.add, data_0)
def Mean(*data_0:Tensor): return Sum(*data_0) / len(data_0)
def ReduceMax(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return data.max(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
def ReduceMin(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return data.min(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
def ReduceSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
def ReduceMean(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
def ReduceSumSquare(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes)
def ReduceProd(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims)
def ReduceL1(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes)
def ReduceL2(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return ReduceSumSquare(data, axes, keepdims, noop_with_empty_axes).sqrt()
def ReduceLogSum(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log()
def ReduceLogSumExp(data:Tensor, axes:list[int]|None=None, keepdims:int=1, noop_with_empty_axes:int=0):
return ReduceSum(data.exp(), axes, keepdims, noop_with_empty_axes).log()
def ArgMax(x:Tensor, axis:int=0, keepdims:int=1, select_last_index:int=0):
if select_last_index: return ((x.shape[axis]-1) - x.flip(axis).argmax(axis, keepdim=keepdims)).cast(dtypes.int64)
return x.argmax(axis, keepdim=keepdims).cast(dtypes.int64)
def ArgMin(x, axis:int=0, keepdims:int=1, select_last_index:int=0):
return ArgMax(-x, axis=axis, keepdims=keepdims, select_last_index=select_last_index)
# ***** Movement Ops *****
def Reshape(data:Tensor, shape:list[int], allowzero:int=0):
return data.reshape([x if x != 0 else (0 if allowzero else data.shape[i]) for i,x in enumerate(shape)])
def Flatten(x:Tensor, axis:int=1): return x.reshape(prod(x.shape[0:axis]), -1)
def Expand(x:Tensor, shape:list[int]): return x.expand(_broadcast_shape(x.shape, tuple(shape)))
def Shrink(x:Tensor, bias:float=0.0, lambd:float=0.5): return (x < -lambd)*(x+bias) + (x > lambd)*(x-bias)
def Transpose(x:Tensor, perm:list[int]|None=None): return x.permute(order=perm or list(range(x.ndim)[::-1]))
def Squeeze(data:Tensor, axes:list[int]|None=None):
return data.squeeze() if axes is None else functools.reduce(lambda d, dim: d.squeeze(dim), sorted(axes, reverse=True), data)
def Unsqueeze(data:Tensor, axes:list[int]): return functools.reduce(lambda d, dim: d.unsqueeze(dim), sorted(axes), data)
def Tile(x:Tensor, repeats:list[int]): return x.repeat(repeats)
def Concat(*xs:Tensor, axis:int): return Tensor.cat(*xs, dim=axis)
def Slice(data:Tensor, starts:list[int], ends:list[int], axes:list[int]|None=None, steps:list[int]|None=None):
axes = axes or list(range(data.ndim))
steps = steps or [1]*data.ndim
slices = [slice(0,x,1) for x in data.shape]
for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
return data[tuple(slices)]
def Split(data:Tensor, split:list[int]|None=None, num_outputs:int=0, axis:int=0):
sz = data.shape[axis]
if split is None: split = [sz // num_outputs + (1 if i < sz % num_outputs else 0) for i in range(num_outputs)]
return data.split(split, axis)
def Pad(x:Tensor, pads:list[int], constant_value:ConstType|None=None, axes:list[int]|None=None,
mode:Literal["constant", "reflect", "edge", "wrap"]="constant", value=0):
value = constant_value or value
axes = axes or list(range(x.ndim))
real_pads = [0] * (x.ndim*2)
for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
return x.pad(padding=_onnx_pads_to_tiny_pads(real_pads), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)
def CenterCropPad(t:Tensor, shape:list[int], axes:list[int]|None=None):
shrink_arg:list[None|tuple[int,int]] = [None] * t.ndim
pad_arg:list[None|tuple[int,int]] = [None] * t.ndim
for s, x in zip(shape, axes or range(t.ndim)):
tx = t.shape[x]
if s < tx: shrink_arg[x] = (tx//2 - (s+1)//2, tx//2 + s//2)
elif s > tx: pad_arg[x] = ((s-tx)//2, (s-tx+1)//2)
return t.shrink(tuple(shrink_arg)).pad(tuple(pad_arg))
# ***** Processing Ops *****
def AveragePool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, count_include_pad:int=0,
dilations:list[int]|int=1, pads:list[int]|int=0, strides:list[int]|int=1):
return X.avg_pool2d(kernel_shape, strides, dilations, _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad),
ceil_mode=ceil_mode, count_include_pad=count_include_pad)
def MaxPool(X: Tensor, kernel_shape:list[int], auto_pad:AUTO_PAD_OPTIONS="NOTSET", ceil_mode:int=0, dilations:list[int]|int=1, pads:list[int]|int=0,
storage_order:int=0, strides:list[int]|int=1):
pads = _resolve_pool_pads(X, pads, kernel_shape, dilations, strides, auto_pad)
ret, idx = X.max_pool2d(kernel_shape, strides, dilations, pads, ceil_mode=ceil_mode, return_indices=True)
return ret, idx.transpose(-2, -1).cast(dtypes.int64) if storage_order else idx.cast(dtypes.int64)
def Conv(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
kernel_shape:list[int]|None=None, pads:list[int]|int=0, strides:list[int]|int=1):
return X.conv2d(W, B, stride=strides, groups=group, dilation=dilations,
padding=_resolve_pool_pads(X, pads, kernel_shape or W.shape[2:], dilations, strides, auto_pad))
def ConvTranspose(X: Tensor, W: Tensor, B:Tensor|None=None, auto_pad:AUTO_PAD_OPTIONS="NOTSET", dilations:list[int]|int=1, group:int=1,
kernel_shape:list[int]|None=None, pads:list[int]|None=None, output_shape:list[int]|None=None, output_padding:list[int]|int=0,
strides:list[int]|int=1):
input_shape, kernel_shape = X.shape[2:], (kernel_shape or W.shape[2:])
strides, dilations, output_padding = (make_tuple(x, len(input_shape)) for x in (strides, dilations, output_padding))
if output_shape is not None: # we pad according to output_shape
pads = _auto_pad([s*(i-1) + op + ((k-1)*d+1) - os for s,i,op,k,d,os in
zip(strides, input_shape, output_padding, kernel_shape, dilations, output_shape)], auto_pad)
if pads is None: # we generate pads
output_shape = output_shape or [X.shape[i+2] * strides[i] for i in range(len(strides))]
pads = [strides[i]*(input_shape[i]-1)+output_padding[i]+((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape) * 2
pads = _onnx_pads_to_tiny_pads(pads)
return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads, output_padding=output_padding)
def MaxUnpool(xT: Tensor, xI: Tensor, outshape: list[int]|None=None, kernel_shape:list[int]=None, pads:list[int]|int=0, strides:list[int]|int=1):
return Tensor.max_unpool2d(xT, xI, kernel_shape, strides, 1, pads, outshape if outshape is None else tuple(outshape))
def GlobalAveragePool(X:Tensor): return X.mean(axis=tuple(range(2, X.ndim)), keepdim=True)
def GlobalMaxPool(X:Tensor): return X.max(axis=tuple(range(2, X.ndim)), keepdim=True)
def Gemm(A:Tensor, B:Tensor, C:Tensor|None=None, alpha:float=1.0, beta:float=1.0, transA:int=0, transB:int=0, broadcast=0):
ret = alpha * (A.transpose(transA) @ B.transpose(transB))
if C is not None: ret = ret + beta * (C if broadcast == 0 else C.reshape([-1 if i < len(C.shape) else 1 for i in range(ret.ndim)][::-1]))
return ret
def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
axis = X._resolve_dim(_resolve_const(axis))
if reverse: X = X.flip(axis)
if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
.shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)
def Trilu(x:Tensor, k:int|list[int]=0, upper:int=1):
k_ = _resolve_const(k)
return x.triu(k_) if upper else x.tril(k_)
def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
extrapolation_value:float=0.0, keep_aspect_ratio_policy:str='stretch', mode:str='nearest', nearest_mode:str='round_prefer_floor'):
def _apply_nearest_mode(index: Tensor, input_dim, mode: str):
if mode == "round_prefer_floor": index = (index - 0.5).ceil()
elif mode == "round_prefer_ceil": index = (index + 0.5).floor()
elif mode in ["floor", "ceil"]: index = getattr(index, mode)()
else: raise ValueError(f"invalid {nearest_mode=}")
return index.cast(dtypes.int32).clip(0, input_dim-1)
def _apply_transformation(index: Tensor, input_dim, scale_dim, mode):
# TODO: needs more testing, not confident in this
# NOTE: their reference implementation differ from the implementation in their reference docs
# https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_resize.py
# https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize
output_dim = scale_dim * input_dim
if mode == "half_pixel": index = (index + 0.5) / scale_dim - 0.5
elif mode == "align_corners": index = index * (input_dim - 1) / (output_dim - 1) if output_dim != 1 else Tensor([0])
elif mode == "asymmetric": index = index / scale_dim
elif mode == "pytorch_half_pixel": index = (index + 0.5) / scale_dim - 0.5 if output_dim != 1 else Tensor([-0.5])
elif mode == "half_pixel_symmetric": index = input_dim / 2 * (1 - int(output_dim) / output_dim) + (index + 0.5) / scale_dim - 0.5
else: raise NotImplementedError(f"invalid {coordinate_transformation_mode=}")
return index.clip(0, input_dim-1)
scales, sizes = (None if scales is None else scales[2-(X.ndim-len(scales)):]), (None if sizes is None else sizes[2-(X.ndim-len(sizes)):])
# we pre permute the axes and permute back after resize
axes, input_shape, = (axes or list(range(X.ndim))), cast(tuple[int, ...], X.shape[2:]),
perm = [a for a in range(len(X.shape)) if a not in axes] + list(axes)
X = X.permute(*perm)
if sizes is not None:
if keep_aspect_ratio_policy in ["not_larger", "not_smaller"]:
scale_fxn = min if keep_aspect_ratio_policy == "not_larger" else max
scales = [scale_fxn([sizes[i] / input_shape[i] for i in range(len(input_shape)) if i+2 in axes])] * 2
sizes = [int((scales[0] * input_shape[i]) + 0.5) if i+2 in axes else input_shape[i] for i in range(X.ndim-2)]
else:
scales = [size / input_shape for size, input_shape in zip(sizes, input_shape)]
else:
sizes = [int(sc*sh) for sc, sh in zip(scales, input_shape)]
# NOTE: this transformation makes it so that we can't just call Tensor.interpolate
# in Tensor.interpolate, we use indexes without any transformation
indexes = []
for shape, size, scale in zip(input_shape, sizes, scales):
indexes.append(_apply_transformation(Tensor.arange(size), shape, scale, coordinate_transformation_mode))
if mode == "nearest":
indexes = [_apply_nearest_mode(index, shape, nearest_mode) for (index, shape) in zip(indexes, input_shape)]
X = X[(..., *Tensor.meshgrid(*indexes))]
if mode == "linear":
expand = list(X.shape)
for i in range(-len(sizes), 0):
reshape, index = [1] * X.ndim, indexes[i]
reshape[i] = expand[i] = sizes[i]
low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
X = X.gather(i, low).lerp(X.gather(i, high), perc)
if mode == "cubic": raise NotImplementedError("cubic interpolation is not implemented")
return X.permute(*argsort(perm)) if perm else X
def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode) # deprecated
def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1): # noqa: A002
val, idx = X.topk(_resolve_const(K), axis, largest, sorted)
return val, idx.cast(dtypes.int64)
# ***** Neural Network Ops *****
def BatchNormalization(X:Tensor, scale:Tensor, B:Tensor, input_mean:Tensor, input_var:Tensor, epsilon:float=1e-05, momentum:float=0.9,
training_mode:int=0, spatial=1, is_test=0):
if training_mode:
x_detached = X.detach()
current_mean = x_detached.mean(axis=(0,2,3))
y = (x_detached - current_mean.reshape(shape=[1, -1, 1, 1]))
current_var = (y*y).mean(axis=(0,2,3))
current_invstd = current_var.add(epsilon).rsqrt()
running_mean = input_mean * momentum + current_mean * (1 - momentum)
running_var = input_var * momentum + current_var * (1 - momentum)
return X.batchnorm(scale, B, current_mean, current_invstd), running_mean, running_var
return X.batchnorm(scale, B, input_mean, (input_var + epsilon).rsqrt())
def GroupNormalization(x:Tensor, scale:Tensor, bias:Tensor, num_groups:int, epsilon:float=1e-05):
x = x.reshape(x.shape[0], num_groups, -1).layernorm(eps=epsilon).reshape(x.shape)
return x * scale.reshape(1, -1, *[1] * (x.ndim-2)) + bias.reshape(1, -1, *[1] * (x.ndim-2))
def InstanceNormalization(x:Tensor, scale:Tensor, bias:Tensor, epsilon:float=1e-05):
return GroupNormalization(x, scale, bias, num_groups=x.shape[1], epsilon=epsilon)
def LayerNormalization(x:Tensor, scale:Tensor, bias:Tensor, axis:int=-1, epsilon:float=1e-05, stash_type:int=1):
assert stash_type == 1, "only float32 is supported"
axes = tuple(i for i in range(axis if axis >= 0 else x.ndim + axis, x.ndim))
mean = x.mean(axis=axes, keepdim=True)
return x.layernorm(axes, epsilon).mul(scale).add(bias), mean, (x.sub(mean)).square().mean(axis=axes, keepdim=True).add(epsilon).rsqrt()
def SkipLayerNormalization(x:Tensor, skip:Tensor, gamma:Tensor, beta:Tensor|None=None, bias:Tensor|None=None, epsilon:float=1e-12):
x = x + skip
if bias is not None: x = x + bias
ret = x.layernorm(eps=epsilon) * gamma
if beta is not None: ret = ret + beta
return ret, None, None, x
def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Tensor, word_embedding:Tensor, position_embedding:Tensor,
segment_embedding:Tensor, gamma=None, beta=None, mask:Tensor|None=None,
position_ids:Tensor|None=None, epsilon=1e-12, mask_index_type=0):
# https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.EmbedLayerNormalization
assert (segment_ids is None) is (segment_embedding is None)
assert mask is None and not mask_index_type, "functionality not supported yet" # TODO
input_shape = input_ids.shape
seq_length = input_shape[1]
compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
vocab_size, max_position_embeddings = word_embedding.shape[0], position_embedding.shape[0]
type_vocab_size = (segment_embedding.shape[0] if compute_seg_emb else None)
def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
# bert embedding layer
if position_ids is None: position_ids = Tensor.arange(seq_length, requires_grad=False).unsqueeze(0).expand(*input_shape)
wrd_embedding_res = embedding(input_ids, vocab_size, word_embedding)
pos_embedding_res = embedding(position_ids, max_position_embeddings, position_embedding)
seg_embedding_res = embedding(segment_ids, type_vocab_size, segment_embedding) if compute_seg_emb else None
embedding_sum = wrd_embedding_res + pos_embedding_res
if seg_embedding_res is not None: embedding_sum = embedding_sum + seg_embedding_res
out = embedding_sum.layernorm(eps=epsilon) * gamma + beta
return out, None, embedding_sum
def MeanVarianceNormalization(x:Tensor, axis:list[int]=[0,2,3]):
return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
def OneHot(indices:Tensor, depth:float|int|list[int|float], values:Tensor, axis:int=-1):
# Scalar or Rank 1 tensor containing exactly one element
depth = int(_resolve_const(depth))
indices = indices.int()
indices = (indices < 0).where(indices+depth, indices)
return indices.unsqueeze(axis)._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
def DepthToSpace(X:Tensor, blocksize:int, mode:str="DCR"):
return X.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)" if mode=="CRD" else "b (h1 w1 c) h w -> b c (h h1) (w w1)", h1=blocksize, w1=blocksize)
def SpaceToDepth(X:Tensor, blocksize:int):
return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)
# Reimplemented here because you need legacy RNG for passing ONNX tests.
def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
import numpy as np
if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
if seed is not None:
rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
else:
rand = data.rand_like(requires_grad=False)
mask = rand >= ratio
return data * mask / (1.0 - ratio), mask
# 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
def dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return dropout_7(data, ratio, training_mode=not is_test)
Dropout = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 7):dropout_7}
def LRN(x:Tensor, size:int, alpha:float=1e-4, beta:float=0.75, bias:float=1.0):
pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
return x / (pooled_x.reshape(x.shape) * alpha + bias).pow(beta)
def NegativeLogLikelihoodLoss(x:Tensor, target:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
return x.nll_loss(target, weight, ignore_index, reduction)
def SoftmaxCrossEntropyLoss(scores:Tensor, labels:Tensor, weights:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean"):
log_probs = scores.log_softmax(1)
return log_probs.nll_loss(labels, weights, ignore_index, reduction), log_probs
def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
N, _, *spatial_dims = size
def generate_grid(steps):
if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)
def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None, do_rotary:int=0, mask_filter_value:float=-10000.0,
num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
assert not do_rotary and not attention_bias, "TODO"
if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
qkv = x.linear(weights, bias)
q, k, v = qkv.split(qkv_hidden_sizes, dim=2)
batch_size, seq_len, _ = x.shape
q_head_size, k_head_size, v_head_size = (sz // num_heads for sz in qkv_hidden_sizes)
q, k, v = (x.reshape(batch_size, seq_len, num_heads, hsz).transpose(1, 2) for x, hsz in zip((q, k, v), (q_head_size, k_head_size, v_head_size)))
present = None
if past is not None:
k, v = past[0].cat(k, dim=2), past[1].cat(v, dim=2)
present = k.stack(v)
if scale is None: scale = 1.0 / math.sqrt(q_head_size)
attn_scores = q @ k.transpose(-1, -2) * scale
if mask_index is not None:
assert 4 >= mask_index.ndim >= 1, f"{mask_index.ndim=}"
if mask_index.ndim != 1: mask = mask_index.bool()
else:
if mask_index.shape[0] == batch_size:
mask = Tensor.arange(attn_scores.shape[-1], requires_grad=False, device=mask_index.device).unsqueeze(0) < mask_index.unsqueeze(1)
elif mask_index.shape[0] == 2*batch_size:
end_positions = mask_index[:batch_size]
start_positions = mask_index[batch_size:]
arange = Tensor.arange(seq_len).unsqueeze(0)
mask = (arange < end_positions.unsqueeze(1)) & (arange >= start_positions.unsqueeze(1))
else: raise NotImplementedError("mask_index with shape (3 * batch_size + 2) is not implemented")
while mask.ndim < 4: mask = mask.unsqueeze(1)
attn_scores = mask.where(attn_scores, mask_filter_value)
if unidirectional:
causal_mask = Tensor.ones((seq_len, seq_len), dtype=dtypes.bool).tril()
attn_scores = causal_mask.where(attn_scores, mask_filter_value)
output = attn_scores.softmax(-1) @ v
output = output.transpose(1, 2).reshape(batch_size, seq_len, -1)
return output, present
def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, past_key:Tensor|None=None, past_value:Tensor|None=None,
is_causal:int=0, kv_num_heads:int|None=None, q_num_heads:int|None=None, qk_matmul_output_mode:int=0, scale:float|None=None,
softcap:float=0.0, softmax_precision:int|None=None):
input_shape_len = Q.ndim
if input_shape_len == 3:
assert q_num_heads is not None and kv_num_heads is not None
Q = Q.reshape(Q.shape[0], q_num_heads, Q.shape[1], -1)
K = K.reshape(K.shape[0], kv_num_heads, K.shape[1], -1)
V = V.reshape(V.shape[0], kv_num_heads, V.shape[1], -1)
if past_key is not None: K = past_key.cat(K, dim=2)
if past_value is not None: V = past_value.cat(V, dim=2)
present_key, present_value = K, V
_q_heads, _kv_heads = q_num_heads or Q.shape[1], kv_num_heads or K.shape[1]
if _q_heads != _kv_heads:
K = K.repeat((1, _q_heads // _kv_heads, 1, 1))
V = V.repeat((1, _q_heads // _kv_heads, 1, 1))
effective_scale = scale if scale is not None else 1.0 / (Q.shape[-1] ** 0.5)
scores = (Q @ K.transpose(-1, -2)) * effective_scale
qk_matmul_return_val = scores
if is_causal:
causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0)
scores = scores.masked_fill(causal_mask.logical_not(), -float("inf"))
if attn_mask is not None:
mask_to_add = attn_mask.where(0, -float("inf")) if attn_mask.dtype == dtypes.bool else attn_mask
scores = scores + mask_to_add
if qk_matmul_output_mode == 1: qk_matmul_return_val = scores
if softcap > 0.0: scores = (scores / softcap).tanh() * softcap
if qk_matmul_output_mode == 2: qk_matmul_return_val = scores
if softmax_precision: scores = scores.cast({1: dtypes.float32, 10: dtypes.float16, 16: dtypes.bfloat16}[softmax_precision])
qk_softmax = scores.softmax(-1).cast(Q.dtype)
if qk_matmul_output_mode == 3: qk_matmul_return_val = qk_softmax
output = (qk_softmax @ V).cast(Q.dtype)
if input_shape_len == 3: output = output.permute(0, 2, 1, 3).reshape(Q.shape[0], Q.shape[2], -1)
return output, present_key, present_value, qk_matmul_return_val
Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}
def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
return X * norm * scale
def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
rotary_embedding_dim:int=0):
original_input_shape = X.shape
if X.ndim == 4: X = X.permute(0, 2, 1, 3)
elif X.ndim == 3:
assert num_heads is not None, "num_heads must be provided for 3D input"
X = X.reshape(*X.shape[:-1], num_heads, X.shape[-1] // num_heads)
head_size = X.shape[-1]
rot_dim = rotary_embedding_dim or head_size
x_rotate, x_pass = X[..., :rot_dim], X[..., rot_dim:]
cos = cos_cache[position_ids] if position_ids is not None else cos_cache[:X.shape[1]]
sin = sin_cache[position_ids] if position_ids is not None else sin_cache[:X.shape[1]]
cos = cos[..., :rot_dim//2].unsqueeze(2)
sin = sin[..., :rot_dim//2].unsqueeze(2)
if interleaved:
x1, x2 = x_rotate[..., ::2], x_rotate[..., 1::2]
real = x1 * cos - x2 * sin
imag = x1 * sin + x2 * cos
x_rotated = Tensor.stack(real, imag, dim=-1).flatten(start_dim=-2)
else:
x1, x2 = x_rotate.chunk(2, dim=-1)
real = x1 * cos - x2 * sin
imag = x1 * sin + x2 * cos
x_rotated = real.cat(imag, dim=-1)
output = x_rotated.cat(x_pass, dim=-1)
return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
# ***** Indexing Ops *****
def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]
def Gather(x:Tensor, indices:Tensor, axis:int=0):
if indices.numel() < 9: # NOTE lessor kernels for smaller indices but kernel number increases depending on size of indices
ret_shape = x.shape[:axis] + indices.shape + x.shape[axis+1:]
if indices.ndim > 1: indices = indices.flatten()
index_consts = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
index_consts = [x.shape[axis]+i if i<0 else i for i in index_consts]
args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x.shape)] for i in index_consts]
return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
# NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
def Scatter(*args, **kwargs): return ScatterElements(*args, **kwargs) # deprecated
def GatherND(x:Tensor, indices:Tensor, batch_dims:int=0):
if batch_dims == 0: return x[tuple(i.squeeze(-1) for i in indices.split(1, -1))]
x_shape, i_shape = x.shape, indices.shape
b = math.prod(x.shape[dim] for dim in range(batch_dims))
# NOTE: each batched dim of both input and indices are equal
x = x.reshape(b, *x.shape[batch_dims:])
indices = indices.reshape(b, *indices.shape[batch_dims:])
b_idx = Tensor.arange(b, device=x.device).reshape(b, *(1,)*(indices.ndim - 2)).expand(*indices.shape[:-1])
ret = x[(b_idx,) + tuple(i.squeeze(-1) for i in indices.split(1, -1))]
return ret.reshape(*x_shape[:batch_dims], *i_shape[batch_dims:-1], *ret.shape[indices.ndim-1:])
def ScatterND(x:Tensor, indices:Tensor, updates:Tensor, reduction:Literal["none", "add", "mul"]='none'):
assert updates.shape == indices.shape[:-1] + x.shape[cast(int, indices.shape[-1]):]
x = x.contiguous()
for index, u in zip(indices.split(1, 0), updates.split(1, 0)):
i = tuple(idx.squeeze(-1) for idx in index.squeeze(0).split(1, -1))
u = u.squeeze(0)
if reduction == "none": x[i] = u
elif reduction == "add": x[i] += u
elif reduction == "mul": x[i] *= u
else: raise NotImplementedError("reduction doesn't support max or min")
return x
def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Literal["none", "add", "mul", "min", "max"]="none"):
indices = (indices < 0).where(x.shape[axis], 0) + indices
if reduction == "none": return x.scatter(axis, indices, updates)
return x.scatter_reduce(axis, indices, updates, {"add": "sum", "mul": "prod", "min": "amin", "max": "amax"}.get(reduction))
def GatherElements(x:Tensor, indices:Tensor, axis:int):
indices = (indices < 0).where(x.shape[axis], 0) + indices
return x.gather(axis, indices)
def Compress(inp:Tensor, condition:list[bool], axis:int|None=None):
if axis is None:
inp = inp.flatten()
axis = 0
if axis < 0: axis += inp.ndim
con = Tensor([i for i,cond in enumerate(condition) if cond]) # compress in python
return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))]
# ***** Quantization Ops *****
def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
elif output_dtype != 0: out_dtype = dtype_parse(output_dtype, "QuantizeLinear op")
else: out_dtype = dtypes.uint8
y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
if out_dtype == dtypes.uchar:
# this appears to work in practice, at least for uchar out_dtype. it folds with the quantize stuff
ret = _clamp_cast((x / y_scale + 0.4999999 + y_zero_point).int(), out_dtype)
else:
ret = _clamp_cast(((x / y_scale).round() + y_zero_point), out_dtype)
return ret.contiguous()
def DynamicQuantizeLinear(x: Tensor):
# only support uint8
qmin, qmax = dtypes.min(dtypes.uint8), dtypes.max(dtypes.uint8)
scale = (x.max().maximum(0) + ((-x).max()).maximum(0)) / (qmax - qmin)
zero_point = _clamp_cast((qmin - x.min() / scale).round(), dtypes.uint8)
y = _clamp_cast((x / scale).round() + zero_point, dtypes.uint8)
return y, scale, zero_point
def DequantizeLinear(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int=0, axis:int=1, block_size:int=0):
x_scale, x_zero_point = _prepare_quantize(x, x_scale, x_zero_point, axis, block_size)
return ((x.int() - x_zero_point) * x_scale).cast(x_scale.dtype)
def QLinearConv(x:Tensor, x_scale:Tensor, x_zero_point:Tensor|int, w:Tensor, w_scale:Tensor, w_zero_point:Tensor|int, y_scale:Tensor,
y_zero_point: Tensor|int, B:Tensor|None=None, **opts):
return _qlinearop_quantized(Conv, [x,w], [x_zero_point,w_zero_point], [x_scale,w_scale], y_scale, y_zero_point, **{"B":B, **opts})
def QLinearMatMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor|int, b:Tensor, b_scale:Tensor, b_zero_point:Tensor|int, y_scale:Tensor,
y_zero_point:Tensor|int) -> Tensor:
return _qlinearop_quantized(Tensor.matmul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], y_scale, y_zero_point)
def QLinearAdd(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
return _qlinearop_float(Tensor.add, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
def QLinearMul(a:Tensor, a_scale:Tensor, a_zero_point:Tensor, b:Tensor, b_scale:Tensor, b_zero_point:Tensor, c_scale:Tensor, c_zero_point:Tensor):
return _qlinearop_quantized(Tensor.mul, [a,b], [a_zero_point,b_zero_point], [a_scale,b_scale], c_scale, c_zero_point)
def QLinearGlobalAveragePool(X:Tensor, x_scale:Tensor, x_zero_point:Tensor, y_scale:Tensor, y_zero_point:Tensor, channels_last:int):
assert channels_last == 0, "TODO NHWC"
return _qlinearop_float(GlobalAveragePool, [X], [x_zero_point], [x_scale], y_scale, y_zero_point)
def ConvInteger(x: Tensor, w: Tensor, x_zero_point: Tensor | int = 0, w_zero_point: Tensor | int = 0, B: Tensor | None = None, **opts) -> Tensor:
return _op_integer(Conv, [x,w], [x_zero_point,w_zero_point], **{"B":B, **opts})
def MatMulInteger(A: Tensor, B: Tensor, a_zero_point: Tensor | int = 0, b_zero_point: Tensor | int = 0) -> Tensor:
return _op_integer(Tensor.matmul, [A,B], [a_zero_point,b_zero_point])
# ***** Training Ops *****
# NOTE: onnx training ops actually don't need the state for optim, all the ops work in a functional way, but we still can reuse optim.py code
@_onnx_training(3)
def Adagrad(R:Tensor, T:int, *inputs:Tensor, decay_factor:float=0.0, epsilon:float=0.0, norm_coefficient:float=0.0):
X, G, H = (i.detach() for i in inputs)
grad = norm_coefficient * X + G
H.assign(H + grad.square())
up = grad / (H.sqrt() + epsilon)
r = R / (1 + T * decay_factor)
X.assign(X.detach() - r * up)
return [X, H]
@_onnx_training(4)
def Adam(R:Tensor, T:int, *inputs:Tensor, alpha:float=0.9, beta:float=0.999, epsilon:float=0.0, norm_coefficient:float=0.0,
norm_coefficient_post:float=0.0):
from tinygrad.nn.optim import Adam as TinyAdam
X, G, V, H = inputs
G, V, H = G.detach(), V.detach(), H.detach()
X.grad = norm_coefficient * X.detach() + G
opt = TinyAdam([X], b1=alpha, b2=beta, eps=epsilon)
opt.m, opt.v, opt.lr = [V], [H], R
# need no-op for m_hat and v_hat if T == 0
if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like()
else:
# `T-1` since it's applied again at the start of `_step`
opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
opt.step()
X = (1 - norm_coefficient_post) * X
return [X, V, H]
@_onnx_training(3)
def Momentum(R:Tensor, T:int, *inputs:Tensor, alpha:float, beta:float, mode:str, norm_coefficient:float):
X, G, V = (i.detach() for i in inputs)
grad = norm_coefficient * X + G
# NOTE: this beta_adjusted term makes it so we can't use SGD for nesterov
beta_adjusted = beta if T > 0 else 1
V.assign(alpha * V + grad * beta_adjusted)
X.assign(X - R * (V if mode == "standard" else (grad + alpha * V)))
return [X, V]
def Gradient(*inputs:Tensor, y:str, intermediate_tensors:dict[str, Tensor], **_):
intermediate_tensors[y].backward()
return tuple([t.grad for t in inputs])
return {
# Tensor ops
**{op: getattr(Tensor, op.lower()) for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan",
"Asin", "Acos", "Atan", "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh",
"Tanh", "Softsign", "Asinh", "Acosh", "Atanh", "Elu", "Celu", "Selu", "Round", "Erf")},
# Implemented ops
**{name:obj for name,obj in locals().items() if isinstance(obj, types.FunctionType) and not name.startswith("_") and name[0].isupper()},
# Version ops
**{name:obj for name,obj in locals().items() if isinstance(obj, dict)},
}
onnx_ops = get_onnx_ops()
tinygrad_repo/extra/onnx_helpers.py
+1 -2
View File
@@ -1,7 +1,6 @@
from tinygrad import Tensor
from tinygrad.tensor import _to_np_dtype
from tinygrad.frontend.onnx import OnnxRunner
from extra.onnx import OnnxValue
from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
import numpy as np
import onnxruntime as ort
tinygrad_repo/extra/onnx_parser.py
-207
View File
@@ -1,207 +0,0 @@
# https://github.com/onnx/onnx/blob/main/onnx/onnx.proto3
import os, pathlib, struct
from io import BufferedReader
from types import SimpleNamespace
from tinygrad.nn.state import TensorIO
from tinygrad.tensor import Tensor
# Protobuf Wire Types
WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
# TensorProto.DataType
class TensorDataType:
UNDEFINED = 0; FLOAT = 1; UINT8 = 2; INT8 = 3; UINT16 = 4; INT16 = 5; INT32 = 6; INT64 = 7 # noqa: E702
STRING = 8; BOOL = 9; FLOAT16 = 10; DOUBLE = 11; UINT32 = 12; UINT64 = 13; COMPLEX64 = 14; COMPLEX128 = 15; BFLOAT16 = 16 # noqa: E702
# AttributeProto.AttributeType
class AttributeType:
UNDEFINED = 0; FLOAT = 1; INT = 2; STRING = 3; TENSOR = 4; GRAPH = 5; SPARSE_TENSOR = 11; TYPE_PROTO = 13; FLOATS = 6; INTS = 7 # noqa: E702
STRINGS = 8; TENSORS = 9; GRAPHS = 10; SPARSE_TENSORS = 12; TYPE_PROTOS = 14 # noqa: E702
class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702
PB_INFOS: dict[str, dict] = {
"OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
"StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
"TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES),
10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
"TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
"TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
"ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
2: ("producer_name", PBType.STRING), 3: ("producer_version", PBType.STRING), 4: ("domain", PBType.STRING), 6: ("doc_string", PBType.STRING),
7: ("graph", PBType.SUB, False, ("GraphProto", lambda: {"node": [], "initializer": [], "input": [], "output": [], "value_info": []})),
8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
"GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
"int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
"NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
"AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
"double_data": None, "uint64_data": None, "raw_data": None}))},
"ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
"TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
"TypeProtoSequence": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
"TypeProtoOptional": {1: ("elem_type", PBType.SUB, False, "TypeProto")},
"TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
}
def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
parser = OnnxParser(fn, load_external_data)
onnx_model = parser.parse()
model = dict_to_namespace(onnx_model)
return model
def gen_result(obj: dict, key_name, val, repeated: bool):
if repeated: obj.setdefault(key_name, []).append(val)
else: obj[key_name] = val
def dict_to_namespace(d):
if isinstance(d, dict): return SimpleNamespace(**{k: dict_to_namespace(v) for k, v in d.items()})
elif isinstance(d, list): return [dict_to_namespace(i) for i in d]
return d
class OnnxParser:
def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
self.file_path: pathlib.Path|None = None
self.load_external_data = load_external_data
if not isinstance(inp, Tensor):
self.file_path = pathlib.Path(inp)
self.tensor = Tensor(self.file_path)
else: self.tensor = inp
self.attr_func_dict = { PBType.BYTES: self._handle_bytes, PBType.SUB: self._handle_sub_message, PBType.FLOATS: self._handle_packed_floats,
PBType.INT: self._handle_int64, PBType.INTS: self._handle_packed_int64s, PBType.STRING: self._handle_string, PBType.FLOAT: self._handle_float}
self.registered_handles = {}
for pb_name in PB_INFOS:
res = {}
for fid, config in PB_INFOS[pb_name].items():
parser_fn, repeated = None, False
if len(config) == 2: name, attr = config
elif len(config) == 3: name, attr, repeated = config
elif len(config) == 4: name, attr, repeated, parser_fn = config
handler_fn = self.attr_func_dict[attr]
def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
res[fid] = _wrapper_handler
self.registered_handles[pb_name] = res
def parse(self):
reader = BufferedReader(TensorIO(self.tensor))
return self._parse_message(reader, "ModelProto", lambda: {"opset_import": [], "domain": None, "graph": None})
def decode_varint(self, reader: BufferedReader) -> int:
result = 0
shift = 0
while True:
data = reader.read(1)
if data == b"": raise EOFError("decode_varint EOF")
result |= (data[0] & 0x7F) << shift
if not (data[0] & 0x80): return result
shift += 7
if shift >= 70: raise ValueError("Varint too long")
def skip_field_value(self, reader: BufferedReader, wire_type):
if wire_type == WIRETYPE_VARINT: self.decode_varint(reader)
elif wire_type == WIRETYPE_FIXED64: reader.seek(8, os.SEEK_CUR)
elif wire_type == WIRETYPE_FIXED32: reader.seek(4, os.SEEK_CUR)
elif wire_type == WIRETYPE_LENGTH_DELIMITED: reader.seek(self.decode_varint(reader), os.SEEK_CUR)
else: raise ValueError(f"Unknown wire type: {wire_type}")
def _parse_message(self, reader, message_field_handlers_name, initial_obj_factory=lambda: {}):
message_field_handlers = self.registered_handles[message_field_handlers_name]
obj = initial_obj_factory()
while True:
try:
tag_val = self.decode_varint(reader)
field_number = tag_val >> 3
wire_type = tag_val & 0x07
if handler := message_field_handlers.get(field_number):
handler(obj, reader, wire_type)
else: self.skip_field_value(reader, wire_type)
except EOFError: break
if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
return obj
def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
str_len = self.decode_varint(reader)
if not use_tensor: return reader.read(str_len)
raw = reader.raw
assert isinstance(raw, TensorIO)
res = raw._tensor[reader.tell():(reader.tell()+str_len)]
reader.seek(str_len, os.SEEK_CUR)
return res
def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
value = self._handle_delimited(reader)
assert isinstance(value, bytes)
gen_result(obj, key_name, value.decode("utf-8"), repeated)
def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for bytes field '{key_name}'")
value = self._handle_delimited(reader, use_tensor=True)
gen_result(obj, key_name, value, repeated)
def _handle_int64(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_VARINT: raise ValueError(f"Expected varint for int64 field '{key_name}'")
val = self.decode_varint(reader)
gen_result(obj, key_name, val - 2**64 if val & (1 << 63) else val, repeated)
def _handle_float(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_FIXED32: raise ValueError(f"Expected fixed32 for float field '{key_name}'")
val, = struct.unpack("<f", reader.read(4))
gen_result(obj, key_name, val, repeated)
def _handle_packed_int64s(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed int64s expected length_delimited")
total_bytes_len = self.decode_varint(reader)
old_pos = reader.tell()
values = []
while reader.tell() < total_bytes_len + old_pos:
val = self.decode_varint(reader) # need copy here because packed ints are varint
values.append(val - 2**64 if val & (1 << 63) else val)
obj[key_name] = values
def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
value = self._handle_delimited(reader, use_tensor=True)
obj[key_name] = value
def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
value = self._handle_delimited(reader, use_tensor=True)
assert isinstance(value, Tensor)
if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
gen_result(obj, key_name, sub_obj, repeated)
def _parse_external_data(self, obj):
if "external_data" not in obj: raise ValueError("no external_data")
location = None
length = None
offset = 0
for kv in obj["external_data"]:
if kv["key"] == "location": location = kv["value"]
if kv["key"] == "offset": offset = int(kv["value"])
if kv["key"] == "length": length = int(kv["value"])
if location is None: raise ValueError("no location in external_data")
if self.file_path is None:
# get onnx file path from Tensor
if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
self.file_path = pathlib.Path(self.tensor.device[5:])
if not (ext_path := self.file_path.parent.joinpath(location)).exists():
raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")
ext_path = self.file_path.parent.joinpath(location)
if not ext_path.exists(): raise Exception(f"external location not exists: {ext_path}")
ext_tensor = Tensor(ext_path)
obj["raw_data"] = ext_tensor[offset:offset+length] if length is not None else ext_tensor[offset:]
obj["data_location"] = 0
tinygrad_repo/extra/optimization/extract_policynet.py
+3 -3
View File
@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
from tinygrad.tensor import Tensor
from tinygrad.nn.optim import Adam
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import actions
from tinygrad.codegen.opt.search import actions
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.helpers import getenv
# stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
INNER = 256
class PolicyNet:
tinygrad_repo/extra/optimization/extract_sa_pairs.py
+3 -3
View File
@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
# more stuff
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import actions
from extra.optimization.helpers import lin_to_feats
from extra.optimization.pretrain_valuenet import ValueNet
from tinygrad.nn.optim import Adam
tinygrad_repo/extra/optimization/get_action_space.py
+3 -3
View File
@@ -1,8 +1,8 @@
import random
from extra.optimization.helpers import load_worlds, ast_str_to_lin
from tinygrad.opt.search import actions
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from tinygrad.helpers import tqdm
tactions = set()
tinygrad_repo/extra/optimization/helpers.py
+3 -3
View File
@@ -1,6 +1,6 @@
# stuff needed to unpack a kernel
from tinygrad import Variable
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.dtype import dtypes, PtrDType
from tinygrad.shape.shapetracker import ShapeTracker
@@ -11,7 +11,7 @@ inf, nan = float('inf'), float('nan')
UOps = Ops
# kernel unpacker
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
def kern_str_to_lin(kern_str:str, opts=None):
@@ -103,7 +103,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
return ret
from tinygrad.device import Device, Buffer
from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float: # noqa: E501
tinygrad_repo/extra/optimization/pretrain_valuenet.py
+2 -2
View File
@@ -1,4 +1,4 @@
from tinygrad.opt.kernel import Kernel
from tinygrad.codegen.opt.kernel import Kernel
from tqdm import tqdm, trange
import math
import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
from tinygrad.shape.view import View
from tinygrad.uop.ops import Variable
inf, nan = float('inf'), float('nan')
from tinygrad.opt.kernel import Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from extra.optimization.helpers import lin_to_feats, MAX_DIMS
tinygrad_repo/extra/optimization/rl.py
+1 -1
View File
@@ -3,7 +3,7 @@ import numpy as np
import math, random
from tinygrad.tensor import Tensor
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import actions, bufs_from_lin, get_kernel_actions
from tinygrad.codegen.opt.search import actions, bufs_from_lin, get_kernel_actions
from tinygrad.nn.optim import Adam
from extra.optimization.extract_policynet import PolicyNet
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
tinygrad_repo/extra/optimization/run_qnet.py
+2 -2
View File
@@ -1,6 +1,6 @@
from typing import List, Tuple
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import get_kernel_actions, actions
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import get_kernel_actions, actions
_net = None
def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:
tinygrad_repo/extra/optimization/search.py
+2 -2
View File
@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
from tinygrad import dtypes
from tinygrad.helpers import BEAM, getenv
from tinygrad.device import Device, Compiled
from tinygrad.opt.kernel import Kernel
from tinygrad.opt.search import beam_search, bufs_from_lin
from tinygrad.codegen.opt.kernel import Kernel
from tinygrad.codegen.opt.search import beam_search, bufs_from_lin
if __name__ == '__main__':
tinygrad_repo/extra/optimization/test_net.py
+2 -2
View File
@@ -6,8 +6,8 @@ from copy import deepcopy
from tinygrad.helpers import getenv, colored
from tinygrad.tensor import Tensor
from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
from tinygrad.opt.search import bufs_from_lin, actions, get_kernel_actions
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.search import bufs_from_lin, actions, get_kernel_actions
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
from extra.optimization.extract_policynet import PolicyNet
from extra.optimization.pretrain_valuenet import ValueNet
tinygrad_repo/extra/optimization/test_time_linearizer.py
+1 -1
View File
@@ -1,5 +1,5 @@
from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
from tinygrad.opt.search import bufs_from_lin, get_kernel_actions
from tinygrad.codegen.opt.search import bufs_from_lin, get_kernel_actions
if __name__ == "__main__":
ast_strs = load_worlds()
tinygrad_repo/extra/perfetto/to_perfetto.py
+7 -3
View File
@@ -1,6 +1,6 @@
import sys, pickle, decimal, json
from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent
from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent, TracingKey
devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
@@ -11,12 +11,14 @@ def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
{"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
{"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
return [{"name": ev.name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
ret = []
for i,e in enumerate(ev.ents):
st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
ret += [{"name": e.name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
for dep in ev.deps[i]:
d = ev.ents[dep]
ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
@@ -24,6 +26,8 @@ def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
return ret
def to_perfetto(profile:list[ProfileEvent]):
# Start json with devices.
profile += [ProfileDeviceEvent("TINY")]
prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
for ev in tqdm(profile, desc="preparing profile"):
if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)
tinygrad_repo/extra/replay_pkl.py
+2 -2
View File
@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
from tinygrad.engine.jit import TinyJit
from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
from tinygrad.renderer import ProgramSpec
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.opt.heuristic import hand_coded_optimizations
from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.heuristic import hand_coded_optimizations
import numpy as np
def move_jit_captured_to_dev(captured, device="DSP"):
tinygrad_repo/extra/torch_backend/backend.py
+14 -3
View File
@@ -128,6 +128,12 @@ def _linalg_eigh(self, UPLO: str = 'U'):
w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
return w.tiny(), v.tiny()
@torch.library.impl("aten::_linalg_det", "privateuseone")
# TODO: move to tinygrad
def _linalg_det(self: torch.Tensor):
result = aten._linalg_det(self.cpu())
return result[0].tiny(), result[1].tiny(), result[2].tiny()
def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()
for i in [
@@ -217,15 +223,18 @@ def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
@torch.library.impl("aten::arange", "privateuseone")
def arange(end, dtype=None, device=None, pin_memory=None):
return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
has_float = isinstance(end, float)
return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::arange.start", "privateuseone")
def arange_start(start, end, dtype=None, device=None, pin_memory=None):
return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
has_float = any(isinstance(x, float) for x in (start, end))
return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::arange.start_step", "privateuseone")
def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or torch.get_default_dtype())))
has_float = any(isinstance(x, float) for x in (start, end, step))
return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::convolution_overrideable", "privateuseone")
def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
@@ -362,6 +371,7 @@ from torch._decomp import get_decompositions
decomps = [
aten.native_batch_norm, aten.native_batch_norm_backward,
aten.native_layer_norm_backward,
aten.linalg_cross,
aten.addmm,
aten.addcmul,
aten.addcdiv,
@@ -371,6 +381,7 @@ decomps = [
aten.elu, # elu has a scale + input_scale param
aten.elu_backward,
aten.softplus,
aten.logaddexp,
aten.threshold,
aten.nll_loss_forward,
aten.nll_loss_backward,
tinygrad_repo/extra/torch_backend/test.py
+12 -1
View File
@@ -135,7 +135,7 @@ class TestTorchBackend(unittest.TestCase):
print(c.cpu())
def test_maxpool2d_backward(self):
x = torch.arange(3*3, device=device).reshape(1, 1, 3, 3).requires_grad_(True)
x = torch.arange(3*3, dtype=torch.float32, device=device).reshape(1, 1, 3, 3).requires_grad_(True)
torch.nn.functional.max_pool2d(x, kernel_size=2, stride=1).sum().backward()
np.testing.assert_equal(x.grad.squeeze().cpu().numpy(), [[0, 0, 0], [0, 1, 1], [0, 1, 1]])
@@ -198,6 +198,17 @@ class TestTorchBackend(unittest.TestCase):
recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)
def test_linalg_det(self):
a = torch.diag(torch.tensor([1,2,3,4,5], dtype = torch.float32, device=device))
b = torch.linalg.det(a)
np.testing.assert_equal(b.cpu().numpy(), 120.0)
def test_linalg_cross(self):
a = torch.tensor([[1, 0, 0], [0, 1, 0]], dtype=torch.float32, device=device)
b = torch.tensor([[0, 0, 1]], dtype=torch.float32, device=device)
cross = torch.linalg.cross(a, b)
np.testing.assert_equal(cross.cpu().numpy(), np.array([[0, -1, 0], [1, 0, 0]], dtype=np.float32))
def test_scalar_assign(self):
a = torch.tensor([1, 2, 3], device=device)
a[1] = 4
tinygrad_repo/extra/torch_muon.py
+75
View File
@@ -0,0 +1,75 @@
import torch
#credit to KellerJordan at https://github.com/KellerJordan/Muon/tree/master
#some changes: classic momentum instead of weighting gradient
#added ns_steps, ns_params, nesterov as hyperparams
def zeropower_via_newtonschulz5(G:torch.tensor, steps:int, params:tuple[int, ...]):
"""
Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a
quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose
of minimizing steps, it turns out to be empirically effective to keep increasing the slope at
zero even beyond the point where the iteration no longer converges all the way to one everywhere
on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T
where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model
performance at all relative to UV^T, where USV^T = G is the SVD.
"""
assert G.ndim >= 2 # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
a, b, c = params
X = G
if G.size(-2) > G.size(-1):
X = X.mT
# Ensure spectral norm is at most 1
X = X / (X.norm(dim=(-2, -1), keepdim=True) + 1e-7)
# Perform the NS iterations
for _ in range(steps):
A = X @ X.mT
B = b * A + c * A @ A # quintic computation strategy adapted from suggestion by @jxbz, @leloykun, and @YouJiacheng
X = a * X + B @ X
if G.size(-2) > G.size(-1):
X = X.mT
return X
def muon_update(grad, momentum, beta=0.95, ns_steps=5, ns_params=(3.4445, -4.7750, 2.0315), nesterov=True):
if beta:
momentum.mul_(beta).add_(grad)
update = grad.add(momentum,alpha=beta) if nesterov else momentum
else: update = grad
if update.ndim == 4: # for the case of conv filters
update = update.view(len(update), -1)
update = zeropower_via_newtonschulz5(update, steps=ns_steps, params=ns_params)
return update
class SingleDeviceMuon(torch.optim.Optimizer):
"""
Muon variant for usage in non-distributed settings.
"""
def __init__(self, params, lr=0.02, weight_decay=0.0, momentum=0.95, ns_steps=5, ns_params=(3.4445, -4.7750, 2.0315), nesterov=True):
defaults = dict(lr=lr, weight_decay=weight_decay, momentum=momentum, ns_steps=ns_steps, ns_params=ns_params, nesterov=nesterov)
super().__init__(params, defaults)
@torch.no_grad()
def step(self, closure=None):
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
for p in group["params"]:
if p.grad is None:
p.grad = torch.zeros_like(p) # Force synchronization
state = self.state[p]
if len(state) == 0:
state["momentum_buffer"] = torch.zeros_like(p)
update = muon_update(p.grad, state["momentum_buffer"], beta=group["momentum"], ns_steps=group["ns_steps"],
ns_params=group["ns_params"], nesterov=group["nesterov"])
p.mul_(1.0 - group["lr"] * group["weight_decay"])
p.add_(update.reshape(p.shape), alpha=-group["lr"])
return loss
tinygrad_repo/pytest.ini
+2
View File
@@ -0,0 +1,2 @@
[pytest]
norecursedirs = extra
tinygrad_repo/ruff.toml
+1
View File
@@ -35,6 +35,7 @@ lint.select = [
line-length = 150
exclude = [
".git/",
"docs/",
"extra/",
"tinygrad/runtime/autogen",
tinygrad_repo/setup.py
+25 -6
View File
@@ -9,7 +9,7 @@ with open(directory / 'README.md', encoding='utf-8') as f:
testing_minimal = [
"numpy",
"torch",
"torch==2.7.1",
"pytest",
"pytest-xdist",
"hypothesis",
@@ -18,16 +18,35 @@ testing_minimal = [
]
setup(name='tinygrad',
version='0.10.3',
version='0.11.0',
description='You like pytorch? You like micrograd? You love tinygrad! <3',
author='George Hotz',
license='MIT',
long_description=long_description,
long_description_content_type='text/markdown',
packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.schedule',
'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
'tinygrad.runtime.support.nv', 'tinygrad.apps'],
packages = [
'tinygrad',
'tinygrad.apps',
'tinygrad.codegen',
'tinygrad.codegen.opt',
'tinygrad.codegen.late',
'tinygrad.engine',
'tinygrad.frontend',
'tinygrad.nn',
'tinygrad.renderer',
'tinygrad.runtime',
'tinygrad.runtime.autogen',
'tinygrad.runtime.autogen.am',
'tinygrad.runtime.autogen.nv',
'tinygrad.runtime.graph',
'tinygrad.runtime.support',
'tinygrad.runtime.support.am',
'tinygrad.runtime.support.nv',
'tinygrad.schedule',
'tinygrad.shape',
'tinygrad.uop',
'tinygrad.viz',
],
package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
classifiers=[
"Programming Language :: Python :: 3",
tinygrad_repo/test/test_amd_llvm.py → tinygrad_repo/test/device/test_amd_llvm.py
View File
tinygrad_repo/test/test_hcq.py → tinygrad_repo/test/device/test_hcq.py
+3 -5
View File
@@ -6,7 +6,7 @@ from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
from tinygrad.runtime.autogen import libc
from tinygrad.runtime.support.system import PCIIfaceBase
from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
from tinygrad.opt.kernel import Kernel, Opt, OptOps
from tinygrad.codegen.opt.kernel import Opt, OptOps
from tinygrad import Variable
MOCKGPU = getenv("MOCKGPU")
@@ -163,10 +163,8 @@ class TestHCQ(unittest.TestCase):
a = Tensor.randint((3, 3, 3), dtype=dtypes.int, device=Device.DEFAULT).realize()
b = a + 1
si = b.schedule()[-1]
k = Kernel(si.ast, opts=TestHCQ.d0.renderer)
for i in range(3): k.apply_opt(Opt(op=OptOps.LOCAL, axis=0, arg=3))
runner = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))
zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -338,7 +336,7 @@ class TestHCQ(unittest.TestCase):
et = float(sig_en.timestamp - sig_st.timestamp)
print(f"exec kernel time: {et:.2f} us")
assert 0.1 <= et <= (15000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
assert 0.1 <= et <= (100000 if MOCKGPU or Device.DEFAULT in {"CPU", "LLVM"} else 100)
def test_speed_copy_bandwidth(self):
if TestHCQ.d0.hw_copy_queue_t is None: self.skipTest("device does not support copy queue")
tinygrad_repo/test/test_hcq_iface.py → tinygrad_repo/test/device/test_hcq_iface.py
View File

Some files were not shown because too many files have changed in this diff Show More