IQ.Pilot Prebuilt Release @ 27f668a

artifacts/package_runtime/tinygrad/codegen/gpudims.py

@@ -0,0 +1,101 @@
+import math
+from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.renderer import Renderer
+
+def _dim_max(d:sint) -> int: return d if isinstance(d, int) else int(d.vmax)
+
+def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):
+  while len(dims) > len(max_sizes) or any(d > m for d,m in zip(dims, max_sizes)):
+    for i,m in enumerate(max_sizes):
+      if i < (len(dims)-1) and _dim_max(dims[i]) * _dim_max(dims[i+1]) <= m:
+        dims = dims[:i] + (dims[i]*dims[i+1],) + dims[i+2:]
+        break
+    else: return None
+  return dims
+
+def _split_dims(dims, max_sizes):
+  if all(d <= m for d,m in zip(dims, max_sizes)): return dims
+  _dims = list(dims) + [1]*(3-len(dims))
+  for i in range(len(_dims)):
+    while _dims[i] > max_sizes[i]:
+      div = next((d for d in range(2, math.ceil(math.sqrt(_dims[i])) + 1) if (_dims[i] % d) == 0), 1)
+      if div == 1: raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+      _dims[i], _dims[(i+1)%len(_dims)] = _dims[i]//div, _dims[(i+1)%len(_dims)]*div
+  return tuple(_dims[:2] if _dims[2] == 1 else _dims)
+
+def get_grouped_dims(prefix, dims:tuple[sint, ...], max_sizes:tuple[int, ...]|None, reverse=False) -> list[UOp]:
+  if reverse: return get_grouped_dims(prefix, dims[::-1], max_sizes)[::-1]
+  if max_sizes is None: limited = dims
+  else:
+    # try to group first: (a, b, c, d) -> (ab, c, d)
+    limited = grouped if (grouped := _group_dims(dims, max_sizes)) else dims
+    # check if grouping failed
+    if len(limited) > len(max_sizes): raise RuntimeError(f"cannot limit dim {dims=}, {max_sizes=}")
+    # try to split up dims: (a,) -> (b, c)
+    if limited == dims: limited = _split_dims(dims, max_sizes)
+  raw_idxs = [UOp.special(s, f"{prefix}{i}") for i,s in enumerate(limited)]
+  flat = sum(idx * math.prod(limited[i+1:]) for i,idx in enumerate(raw_idxs))
+  return [ssimplify(flat // math.prod(dims[i+1:])) if i == 0 else ssimplify((flat // math.prod(dims[i+1:])) % dims[i]) for i in range(len(dims))]
+
+def add_gpudims(ctx:Renderer, s:UOp):
+  if s.arg is None: return None
+  s_topo = list(s.toposort())
+  if any(x.op is Ops.SPECIAL for x in s_topo): return None
+
+  # get ranges
+  all_ranges = {x.arg[0:-1]:x for x in s_topo if x.op is Ops.RANGE}
+
+  # extract global/local dims
+  global_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.GLOBAL, AxisType.THREAD)])
+  local_dims = sorted([x.arg[0:-1] for x in all_ranges.values() if x.arg[-1] in (AxisType.WARP, AxisType.LOCAL, AxisType.GROUP_REDUCE)])
+  if not global_dims and not local_dims: return None
+
+  # get global and local shape
+  global_shape = tuple(ssimplify(all_ranges[r].src[0]) for r in global_dims)
+  local_shape = tuple(ssimplify(all_ranges[r].src[0]) for r in local_dims)
+
+  # get the idxs
+  ki: KernelInfo = s.arg
+  if ctx.has_threads: idxs = [UOp.variable("core_id", 0, int(global_shape[0])-1, dtypes.int).cast(dtypes.weakint)]
+  elif ki.dont_use_locals:
+    assert not local_dims, "can't use locals if there's no local dims"
+    idxs = get_grouped_dims("idx", global_shape, ctx.global_max, reverse=True)
+  else:
+    # define indexes for GPU-like execution
+    local_idxs = get_grouped_dims("lidx", local_shape, ctx.local_max)
+    hw_local = [_dim_max(u.src[0]) for u in local_idxs if u.op is Ops.SPECIAL]
+    global_max = ctx.global_max if ctx.global_prod_max is None else \
+      tuple(min(gm, pm//l) for gm,pm,l in zip(ctx.global_max or ctx.global_prod_max, ctx.global_prod_max, hw_local+[1]*3))
+    idxs = get_grouped_dims("gidx", global_shape, global_max, reverse=True) + local_idxs
+
+  # apply to multiple ranges
+  subs = {}
+  for r in s_topo:
+    # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
+    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
+      missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
+      if len(missing_locals):
+        assert len(idx.src) == 2, "index has 2 sources"
+        mask: UOp = UOp.uprod(*[x.eq(0) for x in missing_locals])
+        subs[idx] = idx.replace(src=(idx.src[0], idx.src[1].valid(mask)))
+    if r.op is not Ops.RANGE: continue
+    try:
+      ii = (global_dims+local_dims).index(r.arg[0:-1])
+      if r.arg[1] == AxisType.REDUCE: continue
+      subs[r] = idxs[ii]
+    except ValueError: continue
+  return s.substitute(subs)
+
+pm_device_to_var = PatternMatcher([
+  # the DEVICE axis is not a program axis, it's bound per device at launch. lower it to the _device_num variable (like SPECIAL for devices)
+  (UPat(Ops.RANGE, name="r"), lambda r: UOp.variable("_device_num", 0, r.vmax, dtype=r.dtype) if r.arg[-1] is AxisType.DEVICE else None),
+  # ENDs that closed a DEVICE range no longer close it
+  (UPat(Ops.END, name="e"), lambda e: e.replace(src=(e.src[0],)+tuple(s for s in e.src[1:] if s.op is not Ops.PARAM))
+   if any(s.op is Ops.PARAM and s.arg.name == '_device_num' for s in e.src[1:]) else None),
+])
+
+pm_add_gpudims = PatternMatcher([
+  # add gpudims must be last
+  (UPat(Ops.SINK, name="s"), add_gpudims),
+])+pm_device_to_var