store

This reverts commit c88ad9d24f.

.github/workflows/test.yml

@@ -335,7 +335,6 @@ jobs:
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
         python -m mypy --strict-equality extra/onnx_parser.py
-        python -m mypy --strict-equality extra/onnx.py
 
   unittest:
     name: Unit Tests
@@ -509,6 +508,10 @@ jobs:
         run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
         run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+      - name: Run REMOTE=1 Test (without process replay)
+        run: |
+          # TODO: re enable process replay, currently remote schedule opens devices
+          CAPTURE_PROCESS_REPLAY=0 REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -532,6 +535,10 @@ jobs:
           opencl: 'true'
       - name: Test ONNX (GPU)
         run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Run REMOTE=1 Test
+        run: |
+          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
       - name: Test Optimization Helpers
         run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
       #- name: Test Action Space
@@ -892,11 +899,12 @@ jobs:
           python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
 
   amdremote:
-    name: Linux (remote)
+    name: Linux (remote amd)
     runs-on: ubuntu-22.04
     timeout-minutes: 20
     env:
       REMOTE: 1
+      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
       PYTHONPATH: ${{ github.workspace }}
     steps:
       - name: Checkout Code
@@ -904,58 +912,38 @@ jobs:
       - name: Setup Environment
         uses: ./.github/actions/setup-tinygrad
         with:
-          key: linux-remote
+          key: linux-remote-amd
           deps: testing_minimal
           amd: 'true'
           llvm: 'true'
-          opencl: 'true'
       - name: Start remote server
         run: |
           start_server() {
             systemd-run --user \
               --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
+              --setenv=REMOTEDEV=AMD \
               --setenv=MOCKGPU=1 \
               --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
+              --setenv=PORT="$2" \
               --working-directory="$(pwd)" \
               python tinygrad/runtime/ops_remote.py
           }
 
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "GPU" 7667
-          start_server "remote-server-cpu" "CPU" 8667
+          start_server "remote-server-1" 6667
+          start_server "remote-server-2" 6668
       - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
           python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+      - name: Run REMOTE=1 Test
         run: |
           python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
-      - name: Run REMOTE=1 Test (GPU)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
       - name: Show remote server logs
         if: always()
         run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
+          journalctl --user -u remote-server-1 --no-pager
+          journalctl --user -u remote-server-2 --no-pager
 
   osxtests:
     strategy:

extra/gemm/amd_matmul.py

@@ -29,7 +29,7 @@ if __name__ == "__main__":
   c = Tensor.zeros(N, N).contiguous().realize()
 
   GlobalCounters.reset()
-  with Context(DEBUG=2):
+  with Context(DEBUG=2, BEAM=4):
     for _ in range(run_count): tc = (a@b).realize()
 
   GlobalCounters.reset()

extra/gemm/amd_seb/kernel3_registers.cpp

@@ -80,8 +80,6 @@ extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, fl
 
   // Iteration over BK blocks.
   for (int kId = 0; kId < N; kId += BK) {
-    __syncthreads();
-
     // We populate the Shared Memory with Ks row and columns
     for (int i = 0; i < nbReadsB; i++) {
       int index_x = BN * blockIdx.x + rBIdx;
@@ -125,6 +123,7 @@ extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, fl
         }
       }
     }
+    __syncthreads();
   }
 
   for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {

extra/gemm/amd_uop_matmul.py

@@ -1,177 +1,168 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite
+from tinygrad.helpers import prod, unwrap
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.opt.kernel import AxisType
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
-from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp, GroupOp
+from tinygrad.shape.shapetracker import ShapeTracker, strides_for_shape
 from tinygrad.schedule.kernelize import merge_views
-from tinygrad.helpers import getenv
+from tinygrad.shape.view import View
+from tinygrad.dtype import AddrSpace
 
 N = 4096
 run_count = 5
 
-BN = 128
-BM = 128
-BK = 8
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
+# src->r->view  -->   src->view->r
+def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
+  if r.tag is not None: return None
+  # confirm the input is in order
+  # TODO: replace this with a UOp that allows for nothing else then remove this
+  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
+  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"
 
-TN = 4
-TM = 4
+  # append the reduce shape to each of the views
+  reduce_count = len(r.axis_arg)
+  prshape = prod(rshape:=src.shape[-reduce_count:])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape[:-reduce_count]+rshape, tuple(x*prshape for x in v.strides[:-reduce_count])+rstrides, v.offset*prshape,
+                    v.mask[:-reduce_count]+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
 
-def hl_spec_kernel3():
+  # no reshape required with shrinking REDUCE_AXIS
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
+             (r.arg[0], tuple(range(len(view.shape)-reduce_count, len(view.shape)))))
+
+early_view_left = merge_views+PatternMatcher([
+  # view before elementwise and buffer ops
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.VALID, Ops.STORE, Ops.LOAD}, name="e"),), name="view"),
+   lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src)) if e.tag is None else None),
+  # push a non contiguous ShapeTracker through reduceop
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])
+
+def hand_spec():
+  # Block Tile size  . 128x128
+  # Thread Tile size . 4x4
+  # Wave Tile size   . 128x32
+  # A wave is        . 8x4
+  # ────── problem size and tiling params (mirror the C kernel) ───────────────────
+  BK = 8             # depth of K-tile
+  BN = BM = 128      # block-tile (output) sizes
+  # the real thread is 16x8 = 128 regs
+  TM = 4
   nbIterWaveM = 2
-  nbIterWaveN = 2
+  TN = 4
+  nbIterWaveN = 4
 
-  # define buffers
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0)
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
+  # ────── shared-memory tile sizes (unchanged) ───────────────────────────────────
+  LDS_A_SZ = BK * BM          # 1024 floats
+  LDS_B_SZ = BK * BN          # 1024 floats
 
-  # shape buffers. TODO: permutes
-  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
-  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
-  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
-  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
-  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
-  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
-  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
-  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)
+  bC = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)   # output C
+  bA = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)   # input A
+  bB = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)   # input B
 
-  out = (A_col.store(As.store(a.load()).load()).load() * B_row.store(Bs.store(b.load()).load()).load()).r(Ops.ADD, (8, 9))
-  sink = c.store(out).sink(arg=KernelInfo(name="tinygemm"))
-  sink = graph_rewrite(sink, merge_views)
-  return sink
+  # TODO: this should not be a string, just a number
+  lAs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(LDS_A_SZ, addrspace=AddrSpace.LOCAL), arg="As")
+  lBs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(LDS_B_SZ, addrspace=AddrSpace.LOCAL), arg="Bs")
 
-def hand_spec_kernel3():
-  BLOCK_SIZE = 256
+  s0 = ShapeTracker.from_shape((N, N, N), (N, 0, 1))
+  s1 = ShapeTracker.from_shape((N, N, N), (0, 1, N))
+  s2 = ShapeTracker.from_shape((N, N, 1), (N, 1, 0))
 
-  nbWaves = BLOCK_SIZE // 32
-  WN = 64
-  WM = BN * BM // nbWaves // WN
+  ls0 = ShapeTracker.from_shape((BM, BK))
+  ls1 = ShapeTracker.from_shape((BN, BK))
 
-  nbWaveX = BN // WN
-  nbWaveY = BM // WM
+  buf_at = [AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.UPCAST, AxisType.UPCAST]
+  buf_bt = [AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.LOCAL, AxisType.UPCAST, AxisType.UPCAST]
+  axis_types = buf_at + buf_bt + [AxisType.REDUCE, AxisType.UNROLL, AxisType.UNROLL, AxisType.UNROLL]
 
-  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
-  waveIndex = threadIdx_x // 32
-  waveIdx = waveIndex % nbWaveX
-  waveIdy = waveIndex // nbWaveX
-  indexInWave = threadIdx_x % 32
+  # 128 x 128 x 8
+  full_shape = (N//BM, 2, 2, 2, 2, 2, 2, 2, N//BN, 2, 2, 2, 2, 2, 2, 2, N//BK, 2, 2, 2)
 
-  nbThreadXPerWave = 8
-  nbThreadYPerWave = 4
+  s0 = s0.reshape(full_shape)
+  s1 = s1.reshape(full_shape)
+  s2 = s2.reshape(full_shape[:-4] + (1,)*4)
 
-  idxInWave = indexInWave % nbThreadXPerWave
-  idyInWave = indexInWave // nbThreadXPerWave
+  ls0 = ls0.reshape((1, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2)).expand(s0.shape)
+  ls1 = ls1.reshape((1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 1, 2, 2, 2)).expand(s1.shape)
+  assert ls0.real_size() == LDS_A_SZ
+  assert ls1.real_size() == LDS_B_SZ
 
-  nbIterWaveN = WN // (nbThreadXPerWave * TN)
-  nbIterWaveM = WM // (nbThreadYPerWave * TM)
+  # BK is a loop of 8
+  #  each loop reads 8 in A, 16 in B
 
-  SUBWN = WN // nbIterWaveN
-  SUBWM = WM // nbIterWaveM
+  print(ls0)
+  print(ls1)
 
-  # Thread mapping to read BKxBN block from A
-  rAIdx = threadIdx_x % BK
-  rAIdy = threadIdx_x // BK
-  # Thread mapping to read BNxBK block from B
-  rBIdx = threadIdx_x % BN
-  rBIdy = threadIdx_x // BN
+  permaxis = []
+  for axis_order in [AxisType.GLOBAL, AxisType.LOCAL, AxisType.LOOP, AxisType.UPCAST, AxisType.GROUP_REDUCE, AxisType.REDUCE, AxisType.UNROLL]:
+    permaxis += [i for i,a in enumerate(axis_types) if a == axis_order]
+  axis_types = [axis_types[x] for x in permaxis]
+  s0, s1, s2, ls0, ls1 = [x.permute(tuple(permaxis)) for x in [s0, s1, s2, ls0, ls1]]
+  print(axis_types)
 
-  strideReadB = BLOCK_SIZE // BN
-  strideReadA = BLOCK_SIZE // BK
-  nbReadsB = BN * BK // BLOCK_SIZE
-  nbReadsA = BM * BK // BLOCK_SIZE
+  lw0, lr0 = ls0, ls0
+  lw1, lr1 = ls1, ls1
 
-  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
-  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
+  # first round of permutes
 
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+  permaxis = (0, 1, 19, 18, 17, 12, 11, 10, 5, 4, 3, 2, 6, 7, 8, 9, 16, 13, 14, 15)
+  s0 = s0.permute(permaxis)
+  lw0 = lw0.permute(permaxis)
 
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
+  permaxis = (0, 1, 15, 14, 9, 8, 7, 6, 13, 19, 18, 17, 5, 4, 3, 2, 16, 12, 11, 10)
+  s1 = s1.permute(permaxis)
+  lw1 = lw1.permute(permaxis)
 
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0)
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
+  # second round of permutes
+  #permaxis = (0, 1, 12, 11, 5, 4, 3, 2, 10, 6, 7, 8, 9, 13, 14, 15, 16, 17, 18, 19)
+  #lw0 = lw0.permute(permaxis)
+  #lr0 = lr0.permute(permaxis)
 
-  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
+  from tinygrad.opt.kernel import axis_colors, colored
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s0.shape, s0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s1.shape, s1.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(s2.shape, s2.views[0].strides, axis_types)]))
+  print("lw")
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lw0.shape, lw0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lw1.shape, lw1.views[0].strides, axis_types)]))
+  print("lr")
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lr0.shape, lr0.views[0].strides, axis_types)]))
+  print('_'.join([colored(f"{s}({st})", axis_colors[x]) for s,st,x in zip(lr1.shape, lr1.views[0].strides, axis_types)]))
 
-  i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
-  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
+  # loads and stores
+  bs0 = bA.view(s0).load()
+  bs1 = bB.view(s1).load()
+  bs0 = lAs.view(lr0).load(lAs.view(lw0).store(bs0))
+  bs1 = lBs.view(lr1).load(lBs.view(lw1).store(bs1))
 
-  kId_range = UOp.range(dtypes.int, N//BK, 0)
-  kId = kId_range*BK
+  mat = (bs0 * bs1).r(Ops.ADD, tuple([i for i,a in enumerate(axis_types) if a in (AxisType.REDUCE, AxisType.UNROLL)]), permute=False)
+  st = bC.view(s2).store(mat)
 
-  # load from globals into locals
-  i = UOp.range(dtypes.int, nbReadsB, 1)
-  index_x = BN * blockIdx_x + rBIdx
-  index_y = rBIdy + i * strideReadB + kId
-  Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
+  ast = st.sink(arg=KernelInfo(axis_types=tuple(axis_types), name="tinygemm"))
+  ast = graph_rewrite(ast, merge_views)
+  prg = get_program(ast, Device.default.renderer)
+  print(prg.src)
+  return prg
 
-  i = UOp.range(dtypes.int, nbReadsA, 2)
-  index_x = rAIdx + kId
-  index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-  As_store = As[(index_x % BK) * BM + index_y % BM].store(a[N * index_y + index_x].load(), i)
-
-  barrier = UOp(Ops.BARRIER, src=(As_store, Bs_store))
-
-  k = UOp.range(dtypes.int, BK, 3)
-
-  # load from locals into registers
-  iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-  i = UOp.range(dtypes.int, TN, 5)
-  index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-  B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
-
-  iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
-  i = UOp.range(dtypes.int, TM, 7)
-  index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-  A_col_store = A_col[iterWave*TM + i].store(As[k*BM + index].load(barrier), iterWave, i)
-
-  # do the GEMM math
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 9)
-  yt = UOp.range(dtypes.int, TM, 10)
-  xt = UOp.range(dtypes.int, TN, 11)
-  x = iterWaveN * TN + xt
-  y = iterWaveM * TM + yt
-  c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-  sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                          iterWaveM, iterWaveN, yt, xt, k, kId_range)
-
-  # store c_regs into c
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 12)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 13)
-  yt = UOp.range(dtypes.int, TM, 14)
-  xt = UOp.range(dtypes.int, TN, 15)
-  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
-  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
-  indexC = N * (yOut + yt) + xOut + xt
-  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink),
-                         iterWaveM, iterWaveN, yt, xt)
-
-  return sink.sink(arg=KernelInfo(name="tinygemm"))
 
 if __name__ == "__main__":
-  hprg = hl_spec_kernel3() if getenv("HL") else hand_spec_kernel3()
-  prg = get_program(hprg, Device.default.renderer)
-  print(prg.src)
-  hrunner = CompiledRunner(prg)
+  hprg = hand_spec()
+  hrunner = CompiledRunner(hprg)
 
   a = Tensor.randn(N, N).realize()
   b = Tensor.randn(N, N).realize()
   hc = Tensor.zeros(N, N).contiguous().realize()
 
   GlobalCounters.reset()
-  with Context(DEBUG=2):
+  with Context(DEBUG=2, BEAM=4):
     for _ in range(run_count): tc = (a@b).realize()
 
   GlobalCounters.reset()
-  ei = ExecItem(hrunner, [a.uop.buffer, b.uop.buffer, hc.uop.buffer])
+  ei = ExecItem(hrunner, [hc.uop.buffer, a.uop.buffer, b.uop.buffer])
   with Context(DEBUG=2):
     for _ in range(run_count): ei.run(wait=True)
   err = (hc-tc).square().mean().item()
   print(f"hrunner {err}")
-  if err > 1e-06: raise RuntimeError("matmul is wrong!")
+  assert err < 1e-06

extra/onnx.py

@@ -1,4 +1,3 @@
-# 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
@@ -84,7 +83,7 @@ class OnnxValue:
   is_optional: bool
   is_sequence: bool
 
-class Domain(enum.Enum):
+class Domain(enum.StrEnum):
   ONNX = "ai.onnx"
   ONNX_ML = "ai.onnx.ml"
   AI_ONNX_TRAINING = "ai.onnx.training"
@@ -798,11 +797,12 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   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:]
+      x_sh = list(x.shape)
+      ret_shape = x_sh[:axis] + list(indices.shape) + x_sh[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]
+      indices = [_cached_to_python_const(indices)] if indices.shape == () else _cached_to_python_const(indices)
+      indices = [x_sh[axis]+x if x<0 else x for x in indices]
+      args = [[(0,x) if j != axis else (i,i+1) for j, x in enumerate(x_sh)] for i in indices] # type: ignore
       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)])]

test/external/external_test_am.py

@@ -21,12 +21,12 @@ class FakeAM:
   def __init__(self):
     self.is_booting, self.smi_dev = True, False
     self.pcidev = FakePCIDev()
-    self.vram_size = (512 << 20)
+    self.vram_size = (4 << 30)
     self.vram_mv = memoryview(bytearray(self.vram_size))
     self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
     self.gmc = FakeGMC(self)
     self.mm = AMMemoryManager(self, self.vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39], va_bits=48,
-      first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base,
+      first_lv=am.AMDGPU_VM_PDB1, va_base=AMMemoryManager.va_allocator.base,
       palloc_ranges=[(1 << (i + 12), 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
     self.is_booting = False
     self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
@@ -56,8 +56,6 @@ def helper_read_entry_components(entry_val):
           "read": (entry_val >> 5) & 0x1, "write": (entry_val >> 6) & 0x1, "exec": (entry_val >> 4) & 0x1,
           "mtype": (entry_val >> 48) & 0x7, "T": (entry_val >> 51) & 0x1, "L": (entry_val >> 55) & 0x1, "F": (entry_val >> 56) & 0x1}
 
-def helper_va(va:int): return va + AMMemoryManager.va_allocator.base
-
 class TestAMPageTable(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
@@ -68,9 +66,10 @@ class TestAMPageTable(unittest.TestCase):
 
     for va,sz in [(0x10000, 0x3000), (0x11000, 0x300000), (0x10000, 0x2000), (0x11000, 0x5000),
                   (0x2000000, 0x2000), (0x4000000, 0x4000000), (0x38000, 0x303000), (0x8000, 0x1000)]:
-      mm.map_range(vaddr=helper_va(va), size=sz, paddrs=[(va, sz)])
+      exteranl_va = va + AMMemoryManager.va_allocator.base
+      mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
-      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, helper_va(va))
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
       results = list(ctx.next(sz))
 
       total_covered = 0
@@ -87,7 +86,7 @@ class TestAMPageTable(unittest.TestCase):
           assert pte['paddr'] == va + _offset + i * _pte_covers, f"Expected paddr {pte['paddr']:#x} to be {va + _offset + i * _pte_covers:#x}"
           assert pte['valid'] == 1
 
-      mm.unmap_range(helper_va(va), sz)
+      mm.unmap_range(va, sz)
 
       for tup in results:
         _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
@@ -100,16 +99,18 @@ class TestAMPageTable(unittest.TestCase):
     mm0 = self.d[0].mm
 
     for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      mm0.map_range(vaddr=helper_va(va1), size=sz1, paddrs=[(va1, sz1)])
-      mm0.map_range(vaddr=helper_va(va2), size=sz2, paddrs=[(va2, sz2)])
-      mm0.unmap_range(helper_va(va2), sz2)
-      mm0.unmap_range(helper_va(va1), sz1)
+      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
+      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
+      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
+      mm0.unmap_range(va2, sz2)
+      mm0.unmap_range(va1, sz1)
 
   def test_double_map(self):
     mm0 = self.d[0].mm
 
     for va,sz in [(0x10000, 0x3000), (0x1000000, 0x1000000), (0x12000, 0x4000)]:
-      exteranl_va = helper_va(va)
+      exteranl_va = va + AMMemoryManager.va_allocator.base
       mm0.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
       with self.assertRaises(AssertionError):
@@ -143,36 +144,36 @@ class TestAMPageTable(unittest.TestCase):
     mm0 = self.d[0].mm
 
     with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(helper_va(0x10000), 0x3000)
+    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(0x10000, 0x3000)
 
     with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)
 
-    mm0.map_range(helper_va(0x10000), 0x3000, paddrs=[(0x10000, 0x3000)])
-    mm0.unmap_range(helper_va(0x10000), 0x3000)
+    mm0.map_range(0x10000, 0x3000, paddrs=[(0x10000, 0x3000)])
+    mm0.unmap_range(0x10000, 0x3000)
 
     with self.assertRaises(AssertionError):
-      mm0.unmap_range(helper_va(0x10000), 0x3000)
+      mm0.unmap_range(0x10000, 0x3000)
 
   def test_free_pt(self):
     mm0 = self.d[0].mm
 
     # offset from start
     for off in [0, 0x3000, 0x10000]:
-      mm0.map_range(helper_va(0x1000000) + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(helper_va(0x1000000) + off, (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+      mm0.map_range(0x1000000 + off, (2 << 20)  - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(0x1000000 + off, (2 << 20) - off)
+      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(0x1000000, 2 << 20)
 
     # offset from end
     for off in [0x1000, 0x20000]:
-      mm0.map_range(helper_va(0x1000000), (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
-      mm0.unmap_range(helper_va(0x1000000), (2 << 20) - off)
-      mm0.map_range(helper_va(0x1000000), 2 << 20, paddrs=[(0x10000, 2 << 20)])
-      mm0.unmap_range(helper_va(0x1000000), 2 << 20)
+      mm0.map_range(0x1000000, (2 << 20) - off, paddrs=[(0x10000, 0x1000)] * (512 - off // 0x1000))
+      mm0.unmap_range(0x1000000, (2 << 20) - off)
+      mm0.map_range(0x1000000, 2 << 20, paddrs=[(0x10000, 2 << 20)])
+      mm0.unmap_range(0x1000000, 2 << 20)
 
   def test_frag_size(self):
     mm0 = self.d[0].mm

test/external/external_test_onnx_backend.py

@@ -184,11 +184,6 @@ backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad d
 backend_test.exclude('test_scatternd_min_cpu') # min not yet supported
 backend_test.exclude('test_scatternd_max_cpu') # max not yet supported
 
-# regression from removing StrEnum in Domain
-backend_test.exclude('test_adam_cpu')
-backend_test.exclude('test_gradient_of_add_and_mul_cpu')
-backend_test.exclude('test_gradient_of_add_cpu')
-
 if Device.DEFAULT in ['GPU', 'METAL']:
   backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')
   backend_test.exclude('test_resize_upsample_sizes_nearest_axes_3_2_cpu')

test/external/external_test_onnx_ops.py

@@ -251,7 +251,6 @@ class TestTrainingOnnxOps(TestOnnxOps):
         outputs = ["X_out", "V_out"]
         self._validate_training("Momentum", onnx_fxn, inputs, attributes, outputs)
 
-  @unittest.expectedFailure  # TODO: regression from removing StrEnum in Domain
   def test_adam_t_greater_than_zero(self):
     from onnx.backend.test.case.node.adam import apply_adam
     for t in [1, 3, 100]:

test/external/fuzz_fast_idiv.py

@@ -2,7 +2,7 @@ import random
 from z3 import Int, Solver, sat
 from tinygrad import dtypes, Device
 from tinygrad.uop.ops import UOp, Ops, UPat, graph_rewrite, PatternMatcher
-from tinygrad.codegen.optional import fast_idiv
+from tinygrad.codegen.devectorizer import fast_idiv
 random.seed(42)
 
 z3_renderer = PatternMatcher([

test/test_copy_speed.py

@@ -1,9 +1,10 @@
-import unittest, numpy as np
-from tinygrad import Tensor, Device, TinyJit
+import unittest
+from tinygrad import Tensor
+from tinygrad import Device
 from tinygrad.helpers import Timing, CI, OSX
 import multiprocessing.shared_memory as shared_memory
 
-N = 256 if CI else 4096
+N = 4096
 class TestCopySpeed(unittest.TestCase):
   @classmethod
   def setUpClass(cls): Device[Device.DEFAULT].synchronize()
@@ -48,32 +49,6 @@ class TestCopySpeed(unittest.TestCase):
       with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
         t.to('CPU').realize()
 
-  def testCopyDefaulttoCPUJit(self):
-    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
-
-    @TinyJit
-    def _do_copy(t): return t.to('CPU').realize()
-
-    t = Tensor.randn(N, N, 4).contiguous().realize()
-    for _ in range(5):
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
-        x = _do_copy(t)
-        Device[Device.DEFAULT].synchronize()
-      np.testing.assert_equal(t.numpy(), x.numpy())
-
-  def testCopytoCPUtoDefaultJit(self):
-    if Device.DEFAULT == "CPU": return unittest.skip("CPU to CPU copy is a no-op")
-
-    @TinyJit
-    def _do_copy(x): return t.to(Device.DEFAULT).realize()
-
-    for _ in range(5):
-      t = Tensor.randn(N, N, 4, device="CPU").contiguous().realize()
-      with Timing("sync:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns:.2f} GB/s"):
-        x = _do_copy(t)
-        Device[Device.DEFAULT].synchronize()
-      np.testing.assert_equal(t.numpy(), x.numpy())
-
   @unittest.skipIf(CI, "CI doesn't have 6 GPUs")
   @unittest.skipIf(Device.DEFAULT != "GPU", "only test this on GPU")
   def testCopyCPUto6GPUs(self):

test/test_hcq.py

@@ -1,6 +1,6 @@
 import unittest, ctypes, struct, os, random, numpy as np
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import getenv, CI, mv_address, DEBUG
+from tinygrad.helpers import getenv, CI, mv_address
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
@@ -513,31 +513,6 @@ class TestHCQ(unittest.TestCase):
 
       assert buf2.as_buffer()[0] == i
 
-  def test_map_cpu_buffer_to_device(self):
-    if Device[Device.DEFAULT].hw_copy_queue_t is None: self.skipTest("skip device without copy queue")
-
-    sz = 0x2000
-    cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
-    cpu_buffer._buf.cpu_view().view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
-
-    for devid in range(6):
-      if DEBUG >= 2: print(f"Testing map to device {Device.DEFAULT}:{devid}")
-
-      try: d = Device[f"{Device.DEFAULT}:{devid}"]
-      except Exception: break
-
-      local_buf = Buffer(f"{Device.DEFAULT}:{devid}", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
-
-      d.allocator.map(cpu_buffer._buf)
-
-      d.hw_copy_queue_t().wait(d.timeline_signal, d.timeline_value - 1) \
-                         .copy(local_buf._buf.va_addr, cpu_buffer._buf.va_addr, sz) \
-                         .signal(d.timeline_signal, d.timeline_value).submit(d)
-      d.timeline_signal.wait(d.timeline_value)
-      d.timeline_value += 1
-
-      np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
-
   @unittest.skipUnless(MOCKGPU, "Emulate this on MOCKGPU to check the path in CI")
   def test_on_device_hang(self):
     if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")

test/test_linearizer.py

@@ -267,7 +267,7 @@ class TestLinearizer(unittest.TestCase):
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
     assert len(stores) == 1
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[-1].dtype == dtypes.float.vec(4)
 
   # NOTE: can reenable, it does work. it just makes BEAM slow
   @unittest.expectedFailure
@@ -294,10 +294,10 @@ class TestLinearizer(unittest.TestCase):
     stores = [u for u in program.uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
 
     # the first store is to lds and can be upcasted
-    assert stores[0].src[1].dtype == dtypes.float.vec(4)
+    assert stores[0].src[-1].dtype == dtypes.float.vec(4)
     assert any(x.op is Ops.DEFINE_LOCAL for x in stores[0].toposort())
     # the second store is to gds with no upcasts
-    assert stores[1].src[1].dtype == dtypes.float
+    assert stores[1].src[-1].dtype == dtypes.float
     assert any(x.op is Ops.DEFINE_GLOBAL for x in stores[1].toposort())
 
   def test_zero_fold(self):
@@ -648,7 +648,7 @@ class TestLinearizer(unittest.TestCase):
     k = helper_linearizer_opt(out)[-1]
     uops = get_program(k.get_optimized_ast(), k.opts).uops
     # check that the float4 cast collapses
-    store_vals = [u.src[1] for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
+    store_vals = [u.src[-1] for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
     for val in store_vals:
       assert val.dtype == dtypes.float.vec(4) # and val.op is not Ops.VECTORIZE
 
@@ -671,7 +671,7 @@ class TestLinearizer(unittest.TestCase):
     x = Tensor.randn((4,3,6,6)).realize()
     out = x.flip((0,1)).contiguous()
     k = helper_linearizer_opt(out)[-1]
-    store_val = [u.src[1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
+    store_val = [u.src[-1] for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
     assert store_val.dtype == dtypes.float.vec(4) and store_val.op is not Ops.VECTORIZE
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -690,7 +690,7 @@ class TestLinearizer(unittest.TestCase):
     barrier = [u for u in uops if u.op is Ops.BARRIER][0]
     # check that the float4 cast collapses for all stores
     for store in local_stores+global_stores:
-      assert store.src[1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
+      assert store.src[-1].dtype.count > 1 # and store.src[2].op is not Ops.VECTORIZE
     # # check the children's vins
     # TODO: src ALU are not the same, should it?
     # assert barrier.src == tuple(local_stores)
@@ -707,11 +707,11 @@ class TestLinearizer(unittest.TestCase):
     stores = [u for u in uops if u.op is Ops.STORE and u.dtype.addrspace != AddrSpace.REG]
 
     # the float4 value stores directly in lds and we skip upcast
-    self.assertEqual(stores[0].src[1].dtype, dtypes.float.vec(4))
+    self.assertEqual(stores[0].src[-1].dtype, dtypes.float.vec(4))
     #assert stores[0].src[-1].op is not Ops.VECTORIZE
 
     # the global store doesn't change
-    assert stores[1].src[1].dtype == dtypes.float
+    assert stores[1].src[-1].dtype == dtypes.float
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -730,7 +730,7 @@ class TestLinearizer(unittest.TestCase):
     ]
     k = helper_linearizer_ast(ast, [Tensor.randn(240*40).realize()], opts=[opt])[-1]
     out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
-    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype == dtypes.float.vec(4)
+    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype == dtypes.float.vec(4)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "test requires float4")
@@ -748,18 +748,18 @@ class TestLinearizer(unittest.TestCase):
             Opt(op=OptOps.UPCAST, axis=1, arg=0), Opt(op=OptOps.UPCAST, axis=0, arg=2)]
     k = helper_linearizer_ast(ast, [Tensor.randn(8*32).realize()], opts=[opt])[-1]
     out = [u for u in get_program(k.get_optimized_ast(), k.opts).uops if u.op is Ops.STORE][0]
-    assert out.src[1].op is Ops.VECTORIZE and out.src[1].dtype.count != 1
+    assert out.src[-1].op is Ops.VECTORIZE and out.src[-1].dtype.count != 1
 
 @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need backends that support float4")
 class TestFloat4(unittest.TestCase):
   @staticmethod
   def count_float4(uops: list[UOp], n=4):
     return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.float.vec(n)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.float.vec(n)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.float.vec(n)]))
   @staticmethod
   def count_half4(uops: list[UOp]):
     return (len([uop for uop in uops if uop.op is Ops.LOAD and uop.dtype == dtypes.half.vec(4)]),
-            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[1].dtype == dtypes.half.vec(4)]))
+            len([uop for uop in uops if uop.op is Ops.STORE and uop.src[-1].dtype == dtypes.half.vec(4)]))
 
   def test_float4_basic(self):
     a = Tensor.empty(2, 8).realize()

test/test_tensor.py

@@ -512,24 +512,6 @@ class TestTinygrad(unittest.TestCase):
     subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
                     shell=True, check=True)
 
-    if Device.DEFAULT != "CPU":
-      # setting multiple devices fail
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
-      # setting device via DEV
-      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                      shell=True, check=True)
-
-      with self.assertRaises(subprocess.CalledProcessError):
-        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
-                        shell=True, check=True)
-
   def test_no_attributeerror_after_apply_uop_exception(self):
     try:
       Tensor.arange(4).reshape(3,2)

test/unit/test_keccak.py

@@ -2,21 +2,6 @@ from typing_extensions import Callable
 import hashlib, random, unittest
 from tinygrad import Tensor, Device, getenv, dtypes
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import CI
-
-@unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
-@unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
-class TestHashing(unittest.TestCase):
-  def _python_hash_1mb(self, data:bytes):
-    chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
-    chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
-    return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
-
-  @unittest.skipIf(CI, "very slow")
-  def test_abc(self):
-    expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
-    out = Tensor(b"abc").hash()
-    self.assertEqual(bytes(out.data()), expected)
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uint8) and is_dtype_supported(dtypes.uint64), "Device must support uint8 and uint64")
 @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT == "NV", "crashes in NV CI")
@@ -48,25 +33,19 @@ class TestKeccak(unittest.TestCase):
       self.assertEqual(ha_ref, Tensor(a).keccak(name).data())
       self.assertEqual(hb_ref, hb)
 
-  def test_referenced(self):
+  def test_abc(self):
     # https://www.di-mgt.com.au/sha_testvectors.html
-    self.assertEqual(bytes(Tensor(b"abc").keccak().tolist()),
-                     bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
-    self.assertEqual(bytes(Tensor(b"").keccak().tolist()),
-                     bytearray.fromhex("a7ffc6f8bf1ed766 51c14756a061d662 f580ff4de43b49fa 82d80a4b80f8434a"))
-    t = Tensor(b"abcdefghbcdefghicdefghijdefghijkefghijklfghijklmghijklmnhijklmnoijklmnopjklmnopqklmnopqrlmnopqrsmnopqrstnopqrstu").keccak()
-    self.assertEqual(bytes(t.tolist()),
-                     bytearray.fromhex("916f6061fe879741 ca6469b43971dfdb 28b1a32dc36cb325 4e812be27aad1d18"))
-    # TODO: this does not run or very slow
-    # self.assertEqual(bytes(Tensor(b"a" * 1000000).keccak().tolist()),
-    #                  bytearray.fromhex("5c8875ae474a3634 ba4fd55ec85bffd6 61f32aca75c6d699 d0cdcb6c115891c1"))
+    out = Tensor(b"abc").keccak()
+    self.assertEqual(bytes(out.tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
 
   def test_long(self):
     data = b"\x00" * 4
     self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
-    data = b"\x00" * (1000 if CI else 4096)
-    self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
+    data = b"\x00" * 4096
+    with self.assertRaises(RecursionError):
+      # TODO: fix
+      self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
 
 if __name__ == "__main__":
   unittest.main()

test/unit/test_uop_symbolic.py

@@ -14,7 +14,7 @@ def render(self) -> tuple[str, ConstType, ConstType]:
   # NOTE: we need STORE so the ALU op has children
   glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
   uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), self)).sink())
-  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
+  rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]
   return rewritten_uop.render(simplify=False), rewritten_uop.vmin, rewritten_uop.vmax
 
 def uconst(val): return UOp.const(dtypes.int, val)
@@ -642,7 +642,7 @@ class TestSymbolic(unittest.TestCase):
     # TODO: copied from render, render does not support cast
     glbl = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(), arg=0)
     uops = full_rewrite(UOp(Ops.STORE, dtypes.void, (glbl.index(UOp.const(dtypes.int, 0)), expr)).sink())
-    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[1]
+    rewritten_uop = [uop for uop in uops if uop.op is Ops.STORE][0].src[-1]
 
     self.assertEqual(rewritten_uop, cond.where(a.cast(dtypes.half), b.cast(dtypes.half)))
 

tinygrad/codegen/__init__.py

@@ -12,9 +12,8 @@ from tinygrad.codegen.quantize import pm_quant
 from tinygrad.codegen.gpudims import pm_add_gpudims
 from tinygrad.uop.symbolic import sym, symbolic_simple, gep_pushing
 from tinygrad.codegen.expander import migrate_indexing, expander
-from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, pm_reduce, \
-  ReduceContext, correct_load_store, pm_render
-from tinygrad.codegen.optional import get_late_rewrite_patterns
+from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexing, devectorize, \
+  pm_reduce, ReduceContext, correct_load_store, pm_render, get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
 
 @dataclass

tinygrad/codegen/devectorizer.py

@@ -1,11 +1,13 @@
-from typing import Any, cast
+from typing import Any, Callable, cast
 import functools, operator, itertools
 from collections import defaultdict
 from dataclasses import dataclass
-from tinygrad.dtype import dtypes, ImageDType, PtrDType, DType, AddrSpace
+from tinygrad.device import is_dtype_supported
+from tinygrad.dtype import dtypes, ImageDType, PtrDType, promo_lattice, DType, AddrSpace
 from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher, graph_rewrite, GroupOp, identity_element
 from tinygrad.uop.symbolic import split_uop, uop_given_valid, parse_valid, simplify_valid, sym, symbolic_flat
 from tinygrad.helpers import getenv, flatten, AMX, prod, partition, all_same
+from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
 from tinygrad.renderer import Renderer
 
 # ***** image load valid simplification *****
@@ -45,10 +47,10 @@ def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
   new_valid = functools.reduce(operator.and_, ss) if (ss:=[s for s in split_uop(valid, Ops.AND) if s not in drop_stmt]) else None
   return buf.index(idx, new_valid)
 
-def delete_redundant_gates(store:UOp, buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
+def delete_redundant_gates(buf:UOp, idx:UOp, val:UOp, store_gate:UOp, cast:UOp|None=None) -> UOp|None:
   if store_gate not in [gate.src[0] for gate in val.toposort() if gate.op is Ops.IF]: return None
   # remove the gate from the index
-  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val, *store.src[2:])
+  return UOp.store(buf.index(idx).cast(cast.dtype) if cast is not None else buf.index(idx), val)
 
 load_store_indexing = PatternMatcher([
   # simplify valid
@@ -59,7 +61,7 @@ load_store_indexing = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("start_idx"), UPat(Ops.CONST, arg=True))), lambda buf,start_idx: buf.index(start_idx)),
   # delete_redundant_gates (after expand)
   (UPat(Ops.STORE, src=(UPat.any(stidx:=UPat.var("buf").index(UPat.var("idx"), UPat.var("store_gate")), stidx.cast().named("cast")),
-                                  UPat.var("val")), name="store", allow_any_len=True), delete_redundant_gates),
+                                  UPat.var("val"))), delete_redundant_gates),
 ])
 
 # ***** load/store grouping *****
@@ -141,6 +143,55 @@ load_store_folding = PatternMatcher([
   (UPat(Ops.STORE, src=(UPat(Ops.PTRCAT, name="cat"), UPat(name="data")), allow_any_len=True, name="sto"), cat_after_store),
 ])
 
+# ***** optional patterns *****
+
+@functools.lru_cache(None)
+def magicgu(vmax:int, d:int) -> tuple[int,int]:
+  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
+  nc = (vmax+1)//(d) * d - 1
+  nbits = vmax.bit_length()
+  for s in range(0, 2*nbits + 1):
+    if 2**s > nc*(d - 1 - (2**s - 1) % d):
+      m = (2**s + d - 1 - (2**s - 1) % d)//d
+      return m, s
+  assert False
+
+def fast_idiv(ctx: Renderer|None, x: UOp, d: int) -> UOp|None:
+  # idiv is truncated division, but arithmetic shift is floored division, so can only do non-negative numbers!
+  if x.vmin<0: return None
+  sign = 1 if d > 0 else -1
+  m,s = magicgu(vmax := min(x.vmax, dtypes.max(x.dtype)), abs(d))
+  if m * vmax <= dtypes.max(x.dtype): return sign * ((x*m) >> s)
+  # promo_lattice needs to return an unsigned type
+  if ctx is not None and dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, ctx.device):
+    if m * vmax <= dtypes.max(next_dtype): return sign * ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
+  return None
+
+powers_of_two = {2**i:i for i in range(64)}
+@functools.cache
+def get_late_rewrite_patterns(ops, force_transcendental=False):
+  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
+           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
+  # rewrite SQRT to xpow 0.5
+  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
+  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
+  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
+  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
+  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+  if Ops.SHR in ops:
+    # no reason to check x<0 for uints
+    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
+    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
+      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
+    if not getenv("DISABLE_FAST_IDIV"):
+      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
+      pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]
+  if Ops.NEG in ops:
+    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
+    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
+  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
+  return PatternMatcher(pat)
+
 # *** correct load/store ***
 
 def split_load_store(ctx:Renderer|None, ls:UOp, idx:UOp):
@@ -259,9 +310,9 @@ pm_render = PatternMatcher([
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat(), UPat(), UPat())).or_casted(),), allow_any_len=True, name="x"),
    lambda x: x.replace(src=(x.src[0], x.const_like(0))+x.src[1:]) if len(x.src) == 1 or x.src[1].op is Ops.CUSTOM else None),
   # gate any stores that aren't gated with ifs
-  (UPat(Ops.STORE, src=(UPat(src=(UPat(), UPat(), UPat(dtype=dtypes.bool)), name="idx").or_casted(), UPat()), name="store", allow_any_len=True),
-    lambda store,idx: UOp(Ops.STORE, dtype=store.dtype, src=store.src[:2]+(UOp(Ops.IF, src=(idx.src[2],)),)+store.src[2:]) if \
-      len(store.src) <= 2 or store.src[2].op != Ops.IF else None),
+  (UPat(Ops.STORE, src=(UPat(src=(UPat.var("buf"), UPat.var("idx"), UPat(name="gate", dtype=dtypes.bool))).or_casted(), UPat.var("val")),
+        name="store", allow_any_len=True),
+        lambda gate,store,buf,idx,val: UOp(Ops.STORE, dtype=store.dtype, src=(buf.index(idx), val, UOp(Ops.IF, src=(gate,)),)+store.src[2:])),
 ])
 
 # *** Ops.REDUCE -> Ops.DEFINE_ACC ***
@@ -284,10 +335,9 @@ def reduce_to_acc(ctx:ReduceContext, red:UOp):
   assert all(x.dtype == red.dtype for x in lst), f"horizontal reduction mismatch {lst[0].dtype} != {red.dtype}"
   # if we have a range
   if len(reduce_range) != 0:
-    input_ranges = tuple([x for x in inp.toposort(gate=lambda x: x.op is not Ops.STORE) if x.op is Ops.RANGE and x not in reduce_range])
-    identity = red.const_like(identity_element(red.arg, red.dtype.scalar()))
-    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG), arg=(ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
-    lst = [acc.store(identity, UOp(Ops.NOOP, src=input_ranges)).load(*reduce_range)] + lst  # put acc as the first element
+    acc = UOp(Ops.DEFINE_REG, red.dtype.ptr(size=1, addrspace=AddrSpace.REG),
+              (red.const_like(identity_element(red.arg, red.dtype.scalar())),) + tuple(reduce_range), (ctx.acc_num,)).index(UOp.const(dtypes.int, 0))
+    lst = [acc.load()] + lst  # put acc as the first element
     ctx.acc_num += 1
   ret = functools.reduce(lambda x,y: x.alu(red.arg, y), lst)
   return acc.store(ret, *reduce_range).load() if len(reduce_range) != 0 else ret

tinygrad/codegen/linearize.py

@@ -86,13 +86,22 @@ class BlockContext:
         ctx.child_count[s] += 1
         this_block_ctx += ctx.last_ctx(s)
 
-      # save the block ctx. SINK never has anything
-      ctx.block_ctxs[u] = _sort_ctx(this_block_ctx) if u.op is not Ops.SINK else ()
+      # save the block ctx
+      ctx.block_ctxs[u] = _sort_ctx(this_block_ctx)
 
       # RANGE/IF add to the next ctx
       # STORE/ASSIGN subtract from the next ctx
       if u.op in {Ops.RANGE, Ops.IF}: ctx.child_ctxs[u] = _sort_ctx(ctx.block_ctxs[u] + (u,))
-      elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([y for y in ctx.block_ctxs[u] if y not in u.src])
+      elif u.op is Ops.STORE:
+        if len(definereg:=[x for x in u.src[0].toposort() if x.op is Ops.DEFINE_REG]):
+          # old assign logic
+          ctx.child_ctxs[u] = tuple([y for y in ctx.last_ctx(u.src[1]) if y not in definereg[0].src[1:]])
+        elif any(x.op is Ops.DEFINE_LOCAL for x in u.src[0].toposort()):
+          # deal with non-reduce locals. probably wrong
+          idx_context, store_context = ctx.last_ctx(u.src[0]), ctx.last_ctx(u.src[1])
+          ctx.child_ctxs[u] = tuple([y for y in store_context if y not in idx_context and y.op is Ops.RANGE])
+        else: ctx.child_ctxs[u] = ()
+      #elif u.op is Ops.STORE: ctx.child_ctxs[u] = tuple([x for x in ctx.block_ctxs[u] if x not in u.src])
     return ctx
 
 # ***** make blocks *****

tinygrad/codegen/optional.py

@@ -1,58 +0,0 @@
-# should this merge with transcendental?
-from typing import Callable
-import functools
-from tinygrad.device import is_dtype_supported
-from tinygrad.dtype import dtypes, promo_lattice
-from tinygrad.uop.ops import UOp, Ops, UPat, PatternMatcher
-from tinygrad.helpers import getenv
-from tinygrad.uop.transcendental import xexp2, xlog2, xsin, xpow, TRANSCENDENTAL_SUPPORTED_DTYPES
-from tinygrad.renderer import Renderer
-
-# ***** optional patterns *****
-
-@functools.lru_cache(None)
-def magicgu(vmax:int, d:int) -> tuple[int,int]:
-  # calculate m,s such that x//d == (x*m) >> s for all 0 <= x <= vmax, d>0; adapted from Hacker's Delight, Chapter 10
-  nc = (vmax+1)//(d) * d - 1
-  nbits = vmax.bit_length()
-  for s in range(0, 2*nbits + 1):
-    if 2**s > nc*(d - 1 - (2**s - 1) % d):
-      m = (2**s + d - 1 - (2**s - 1) % d)//d
-      return m, s
-  assert False
-
-def fast_idiv(ctx: Renderer|None, x: UOp, d: int) -> UOp|None:
-  # idiv is truncated division, but arithmetic shift is floored division, so can only do non-negative numbers!
-  if x.vmin<0: return None
-  sign = 1 if d > 0 else -1
-  m,s = magicgu(vmax := min(x.vmax, dtypes.max(x.dtype)), abs(d))
-  if m * vmax <= dtypes.max(x.dtype): return sign * ((x*m) >> s)
-  # promo_lattice needs to return an unsigned type
-  if ctx is not None and dtypes.is_int(next_dtype := promo_lattice[x.dtype][-1]) and is_dtype_supported(next_dtype, ctx.device):
-    if m * vmax <= dtypes.max(next_dtype): return sign * ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
-  return None
-
-powers_of_two = {2**i:i for i in range(64)}
-@functools.cache
-def get_late_rewrite_patterns(ops, force_transcendental=False):
-  pat: list[tuple[UPat, Callable]] = [(UPat(op, dtype=TRANSCENDENTAL_SUPPORTED_DTYPES, src=(UPat.var("d"),)), f) for op,f in \
-           ((Ops.EXP2, xexp2), (Ops.LOG2, xlog2), (Ops.SIN, xsin)) if op not in ops or force_transcendental]
-  # rewrite SQRT to xpow 0.5
-  if Ops.SQRT not in ops: pat.append((UPat(Ops.SQRT, src=UPat.var("d")), lambda d: xpow(d, d.const_like(0.5))))
-  # rewrite MOD to AND (which should always be supported, but not for generic in tests): x % (2**y) -> x & (2**y-1)
-  if Ops.AND in ops: pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("c"), lambda x,c: x & (c.arg-1) if c.arg in powers_of_two else None)]
-  # rewrite MUL/IDIV to SHL+SHR: x*(2**y) -> shl(x,y) and x//(2**y) -> shr(x,y)
-  if Ops.SHL in ops: pat += [(UPat.var("x", dtypes.ints)*UPat.cvar("c"), lambda c,x: x << v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-  if Ops.SHR in ops:
-    # no reason to check x<0 for uints
-    pat += [(UPat.var("x", dtypes.uints)//UPat.cvar("c"), lambda x,c: x >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]
-    pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("c"), lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(
-      c-1, 0)) >> v if (v:=powers_of_two.get(c.arg, 0)) else None)]  # (x+(x<0).where(c-1, 0)) >> v
-    if not getenv("DISABLE_FAST_IDIV"):
-      pat += [(UPat.var("x", dtypes.ints)//UPat.cvar("d"), lambda ctx, x, d: fast_idiv(ctx, x, d.arg))]
-      pat += [(UPat.var("x", dtypes.ints)%UPat.cvar("d"), lambda ctx, x, d: x - d*f if (f:=fast_idiv(ctx, x, d.arg)) is not None else None)]
-  if Ops.NEG in ops:
-    pat += [(UPat.var('x')*-1, lambda x: x.alu(Ops.NEG))]
-    if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda x,y: x.alu(Ops.SUB, y))]
-  if Ops.MULACC in ops: pat += [(UPat.var('a')*UPat.var('b')+UPat.var('c'), lambda a,b,c: a.alu(Ops.MULACC, b, c))]
-  return PatternMatcher(pat)

tinygrad/device.py

@@ -4,7 +4,7 @@ from collections import defaultdict
 from typing import Any, Generic, TypeVar, Iterator
 import importlib, inspect, functools, pathlib, os, platform, contextlib, sys, re, atexit, pickle, decimal, time
 from tinygrad.helpers import CI, OSX, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, flat_mv, PROFILE, temp, \
-                             colored, Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent, dedup
+                             colored, Context, DISABLE_COMPILER_CACHE, ALLOW_DEVICE_USAGE, cpu_events, ProfileEvent
 from tinygrad.dtype import DType, ImageDType, PtrDType, dtypes, _to_np_dtype
 from tinygrad.renderer import Renderer
 
@@ -37,8 +37,7 @@ class _Device:
       with contextlib.suppress(Exception): yield self[device].device
   @functools.cached_property
   def DEFAULT(self) -> str:
-    dev = [dev] if (dev:=getenv("DEV", "").upper()) else []
-    from_env = dedup(dev + [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1])
+    from_env = [d for d in self._devices if d not in ["DISK", "NPY"] and getenv(d) == 1]
     assert len(from_env) < 2, f"multiple devices set in env: {from_env}"
     if len(from_env) == 1: return from_env[0]
     try:

tinygrad/engine/jit.py

@@ -42,13 +42,17 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]
 
   for ji in jit_cache:
     match ji.prg:
-      case CompiledRunner(): ji_graph_dev = ji.prg.dev
-      case BufferXfer(): ji_graph_dev = Device[unwrap(ji.bufs[0]).device]
-      case BufferCopy(): ji_graph_dev = next((Device[unwrap(b).device] for b in ji.bufs if unwrap(b).device not in {"CPU", "LLVM"}), None)
+      case CompiledRunner():
+        ji_graph_dev = ji.prg.dev
+        # All GraphRunners can graph CompiledRunners
+        can_be_graphed = ji_graph_dev.graph is not None
+      case BufferXfer():
+        ji_graph_dev = Device[unwrap(ji.bufs[0]).device]
+        # All *Multi*GraphRunner support graphing BufferXfers
+        can_be_graphed = ji_graph_dev.graph is not None and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
       case ViewOp(): continue # ViewOps are just ignored
-      case _: ji_graph_dev = None # Everything else is not graphed and flushes existing graph if it's being constructed
+      case _: can_be_graphed = False # Everything else is not graphed and flushes existing graph if it's being constructed
 
-    can_be_graphed = ji_graph_dev is not None and ji_graph_dev.graph is not None and graph_class(ji_graph_dev).supports_exec_item(ji_graph_dev, ji)
     is_multigraph = can_be_graphed and issubclass(graph_class(ji_graph_dev), MultiGraphRunner)
     can_share_graph = can_be_graphed and (type(ji_graph_dev) is type(current_device) if is_multigraph else ji_graph_dev == current_device)
     can_extend_graph_batch = can_share_graph and (max_batch_size == 0 or len(current_batch) < max_batch_size)
@@ -126,13 +130,8 @@ class GraphRunner(Runner):
 
     return list({id(x):x for x in wait_nodes}.values())
 
-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, CompiledRunner)
-
 # a marker for your graph supporting multiple devices of the same type
-class MultiGraphRunner(GraphRunner):
-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool: return isinstance(ei.prg, (CompiledRunner, BufferXfer))
+class MultiGraphRunner(GraphRunner): pass
 
 def get_out_buffers_for_ei(ei:ExecItem) -> list[Buffer]:
   if isinstance(ei.prg, CompiledRunner): return [cast(Buffer, ei.bufs[out]) for out in ei.prg.p.outs if out not in ei.prg.p.ins]

tinygrad/opt/kernel.py

@@ -5,7 +5,7 @@ from collections import defaultdict
 from typing import cast, Final, Callable, Sequence
 from enum import Enum, auto
 
-from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, AxisType
+from tinygrad.uop.ops import GroupOp, KernelInfo, UOp, Ops, can_pad, resolve, Variable, sint, graph_rewrite, smax, AxisType
 from tinygrad.uop.spec import type_verify, ast_spec
 from tinygrad.device import Device
 from tinygrad.opt.tc import TensorCore
@@ -73,8 +73,7 @@ class Kernel:
       self.sts.append(unwrap(x.src[0].st))
 
     # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
-    self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
+    self.sts.append(ShapeTracker.from_shape(tuple([smax(*s) for s in zip(*[x.shape for x in self.sts])]), (0,)*len(self.sts[0].shape)))
 
     # parameters for optimization
     self.tensor_core: TensorCore|None = None

tinygrad/renderer/cstyle.py

@@ -9,7 +9,7 @@ from tinygrad.renderer import Renderer
 from tinygrad.codegen.devectorizer import no_vectorized_alu
 
 base_rewrite = PatternMatcher([
-  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
+  (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x: f"{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}] = {{{ctx[x.src[0]]}}};"),
   (UPat(Ops.IF, name="x"), lambda ctx,x: f"if ({ctx[x.src[0]]}) {{"),
   (UPat((Ops.ENDIF, Ops.ENDRANGE)), lambda ctx: "}"),
   (UPat(Ops.WMMA, name="x"), lambda ctx,x: f"__{x.arg[0]}({ctx[x.src[0]]}, {ctx[x.src[1]]}, {ctx[x.src[2]]})"),
@@ -25,7 +25,7 @@ base_rewrite = PatternMatcher([
   (UPat(Ops.BITCAST, name="x"), lambda ctx,x: f"(*(({ctx.buffer_prefix}{ctx.render_dtype(x.dtype)}*)&{ctx[x.src[0]]}))"),
   (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x: f"{ctx.smem_align}{ctx.smem_prefix}{ctx.render_dtype(x.dtype.base)} {ctx[x]}[{x.dtype.size}];"),
   (UPat(Ops.BARRIER), lambda ctx: ctx.barrier),
-  (UPat(Ops.PRECAST, name="x"), lambda ctx,x: ctx[x.src[0]]),
+  (UPat(Ops.NOOP, name="x"), lambda ctx,x: ctx[x.src[0]]),
   (UPat(Ops.SPECIAL, name="x"), lambda ctx,x: f"{ctx.code_for_workitem[x.arg[0][0]](x.arg[0][-1])}; /* {x.arg[1]} */"),
   # const
   (UPat(Ops.CONST, arg=math.inf, name="x"), lambda ctx, x: f"({ctx.render_cast(x.dtype, ctx.infinity)})"),
@@ -60,9 +60,9 @@ base_rewrite = PatternMatcher([
 ])
 
 extra_pm = PatternMatcher([
-  # insert a PRECAST before BITCAST to force it to be rendered. not needed on all backends?
-  (UPat(Ops.BITCAST, name="x"), lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.PRECAST, x.src[0].dtype, x.src),))
-   if x.src[0].op not in {Ops.PRECAST, Ops.LOAD, Ops.CUSTOM} else None),
+  # insert a NOOP before BITCAST to force it to be rendered. not needed on all backends?
+  (UPat(Ops.BITCAST, name="x"),
+   lambda x: UOp(Ops.BITCAST, x.dtype, (UOp(Ops.NOOP, x.src[0].dtype, x.src),)) if x.src[0].op not in {Ops.NOOP, Ops.LOAD, Ops.CUSTOM} else None),
   # rewrite MAX to CMPLT + WHERE (max function is annoying on many cstyle backends)
   (UPat(Ops.MAX, name="m"), lambda m: (m.src[0] < m.src[1]).where(m.src[1], m.src[0])),
   # devectorize any bools
@@ -135,7 +135,6 @@ class CStyleLanguage(Renderer):
     c: defaultdict[str, int] = defaultdict(int)
     name = "test"
     for u in uops:
-      if u.op is Ops.NOOP: continue
       if u.op is Ops.SINK:
         if u.arg is not None: name = u.arg.function_name
         continue
@@ -155,7 +154,7 @@ class CStyleLanguage(Renderer):
       elif u.op is Ops.RANGE: r[u] = f"ridx{u.arg}"
       else:
         prefix = {Ops.WMMA: "wmma", Ops.DEFINE_LOCAL: "temp", Ops.CONST: "const",
-                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.PRECAST: "precast",
+                  Ops.CAST: "cast", Ops.BITCAST: "cast", Ops.GEP: "gep", Ops.VECTORIZE: "cast", Ops.NOOP: "precast",
                   Ops.INDEX: "bidx", Ops.DEFINE_REG: "acc", Ops.LOAD: "val"}.get(u.op, "alu")
         r[u] = f"{prefix}{c[prefix]}"
 

tinygrad/renderer/llvmir.py

@@ -160,7 +160,6 @@ class LLVMRenderer(Renderer):
 
     name = "test"
     for u in uops:
-      if u.op is Ops.NOOP: continue
       if u.op is Ops.SINK:
         if u.arg is not None: name = u.arg.function_name
         continue
@@ -171,7 +170,13 @@ class LLVMRenderer(Renderer):
         r[u] = f"%{'local' if u.op is Ops.DEFINE_LOCAL else 'reg'}_{str(u.arg).replace('(', '').replace(')', '').replace(',', '_').replace(' ', '')}"
         assert isinstance(u.dtype, PtrDType)
         if self.device == "LLVM" or u.op is Ops.DEFINE_REG:
-          kernel.append(f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]")
+          # put alloca in the beginning of the function always
+          kernel = [f"  {r[u]} = alloca [{u.dtype.size} x {ldt(u.dtype.base)}]"] + kernel
+          if u.op is Ops.DEFINE_REG:
+            # store the const here. TODO: this should be INDEX and STORE and shouldn't be handcoded here
+            for i in range(u.dtype.size):
+              kernel.append(f"  {r[u]}_idx_{i} = getelementptr inbounds {ldt(u.dtype.base)}, {ldt(u.dtype)} {r[u]}, i32 {i}")
+              kernel.append(f"  store {ldt(u.src[0].dtype)} {r[u.src[0]]}, {ldt(u.dtype)} {r[u]}_idx_{i}")
         else:
           local_args.append(f"@{r[u][1:]} = internal unnamed_addr addrspace(3) global [{u.dtype.size} x {ldt(u.dtype)}] undef, align 16")
           kernel.append(f"  {r[u]} = addrspacecast [{u.dtype.size} x {ldt(u.dtype)}] addrspace(3)* @{r[u][1:]} to [{u.dtype.size} x {ldt(u.dtype)}]*")

tinygrad/renderer/ptx.py

@@ -54,7 +54,7 @@ ptx_matcher = PatternMatcher([
   # move mask from INDEX to the load/store to enable pointer arithmetic
   (UPat(Ops.LOAD, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat.var("gate"))), UPat.var("alt"))),
    lambda buf,idx,gate,alt: UOp(Ops.LOAD, alt.dtype, (buf.index(idx), alt, gate))),
-  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat())), UPat.var("val"), UPat.var("gate")), allow_any_len=True),
+  (UPat(Ops.STORE, src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx"), UPat())), UPat.var("val"), UPat.var("gate"))),
    lambda buf,idx,val,gate: UOp.store(buf.index(idx), val, gate)),
   # ptx shr and shl instructions require y to be uint
   (UPat.var("x") << UPat.var("y"), lambda x,y: UOp(Ops.SHL, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
@@ -110,7 +110,10 @@ string_rewrite = PatternMatcher([
   (UPat(Ops.LOAD, name="x", src=(UPat.var('loc'),), allow_any_len=True),
    lambda ctx, x, loc: f"ld.{mem_type(x)}.v{x.dtype.count}.{ctx.mem_types[x.dtype.scalar()]} {{{', '.join(ctx.r[x])}}}, [{ctx.r[loc]}+0];" \
      if x.dtype.count > 1 else f"ld.{mem_type(x)}.{ctx.mem_types[x.dtype]} {ctx.r[x]}, [{ctx.r[loc]}+0];"),
-  (UPat(Ops.DEFINE_REG, src=()), lambda ctx: []),
+  (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred", dtype=dtypes.bool),), allow_any_len=True), lambda ctx, x, pred: [
+    f"setp.ne.s16 {ctx.r[pred]}, {render_val(pred.arg, pred.dtype)}, 0;", f"mov.pred {ctx.r[x]}, {ctx.r[pred]};"]),
+  (UPat(Ops.DEFINE_REG, name="x", src=(UPat.cvar("pred"),), allow_any_len=True),
+   lambda ctx, x, pred: f"mov.b{ctx.types[x.dtype.base][1:]} {ctx.r[x]}, {render_val(pred.arg, x.dtype.base)};"),
   (UPat(Ops.RANGE, name="x"), lambda ctx, x: [f"mov.u32 {ctx.r[x]}, 0;", "LOOP_" + f"{ctx.r[x][1:]}:"]),
   (UPat(Ops.ENDRANGE, name="x", src=(UPat.var("src0"),)), lambda ctx, x, src0: [
     ctx.code_for_op[Ops.ADD](ctx.r[src0], ctx.r[src0], "1", dtypes.int, ctx.types[dtypes.int]),
@@ -173,7 +176,6 @@ class PTXRenderer(Renderer):
 
     name = "test"
     for u in uops:
-      if u.op is Ops.NOOP: continue
       if u.op is Ops.SINK:
         if u.arg is not None: name = u.arg.function_name
         continue

tinygrad/renderer/wgsl.py

@@ -40,6 +40,11 @@ wgsl_matcher = PatternMatcher([
   (UPat.var("x") >> UPat.var("y"), lambda x,y: UOp(Ops.SHR, x.dtype, (x,y.cast(dtypes.uint))) if y.dtype != dtypes.uint else None),
   ]) + extra_pm
 
+def webgpu_define_reg(ctx, x):
+  ret = [f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"]
+  for i in range(x.dtype.size): ret.append(f"{ctx[x]}[{i}] = {ctx[x.src[0]]};")
+  return ' '.join(ret)
+
 class WGSLRenderer(CStyleLanguage):
   device = "WEBGPU"
   global_max = (65535, 65535, 65535)
@@ -59,8 +64,7 @@ class WGSLRenderer(CStyleLanguage):
      lambda x: f"bitcast<u32>({x.arg})" if x.arg < 0 else f"{x.arg&0xFFFFFFFF}u"),
     (UPat(Ops.DEFINE_LOCAL, name="x"), lambda ctx,x:
       f"var<workgroup> {ctx[x]}: array<{ctx.buf_map(x.dtype.base)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"),
-    (UPat(Ops.DEFINE_REG, name="x"), lambda ctx,x:
-     f"var {ctx[x]}: array<{ctx.buf_map(x.dtype)},{x.dtype.size//(4//x.dtype.itemsize) if is_packed(x.dtype) else x.dtype.size}>;"),
+    (UPat(Ops.DEFINE_REG, name="x"), webgpu_define_reg),
     (UPat(Ops.BITCAST, dtype=dtypes.half, name="x", src=(UPat(dtype=(dtypes.short, dtypes.ushort, dtypes.uint32),),)),
      lambda ctx,x: f"bitcast<vec2<f16>>({ctx[x.src[0]]})[0]"),
     (UPat(Ops.BITCAST, dtype=(dtypes.char, dtypes.uchar), name="x"), lambda ctx,x: f"bitcast<{ctx.type_map[x.dtype]}>({ctx[x.src[0]]}&0xFF)"),

tinygrad/runtime/graph/hcq.py

@@ -1,11 +1,11 @@
 import collections, time
 from typing import Any, cast
-from tinygrad.helpers import round_up, PROFILE, merge_dicts, getenv
+from tinygrad.helpers import round_up, PROFILE, merge_dicts
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Variable
-from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner, BufferCopy
+from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
 from tinygrad.engine.jit import MultiGraphRunner
 
 class HCQGraph(MultiGraphRunner):
@@ -13,9 +13,6 @@ class HCQGraph(MultiGraphRunner):
     super().__init__(jit_cache, input_rawbuffers, var_vals)
     self.devices = list(set(cast(HCQCompiled, d) for ji in jit_cache for d in [Device[cast(Buffer, x).device] for x in ji.bufs]))
 
-    # CPU Device is always last
-    self.devices = sorted(self.devices, key=lambda x: 1 if x._is_cpu() else 0)
-
     # Replace input buffers with variables.
     self.hcq_bufs = [[cast(Buffer, x)._buf for x in ji.bufs] for ji in jit_cache]
     self.input_replace_to_var: dict[tuple[int, int], Variable] = {}
@@ -51,8 +48,7 @@ class HCQGraph(MultiGraphRunner):
     self.comp_queues: dict[HCQCompiled, HWQueue] = {dev: dev.hw_compute_queue_t() for dev in self.devices}
     self.copy_queues: dict[HCQCompiled, HWQueue] = {} # lazy allocation
 
-    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices if dev.device != "CPU"},
-      **{"KICK": self.devices[0].new_signal(value=0)}, **{dev: self.devices[0].new_signal(value=0) for dev in self.devices if dev.device == "CPU"}}
+    self.signals: dict[Any, HCQSignal] = {**{dev: dev.new_signal(value=0) for dev in self.devices}, **{"KICK": self.devices[0].new_signal(value=0)}}
     self.kickoff_value: int = 0
     self.kickoff_var = UOp.variable("kickoff_var", 0, 0xffffffff, dtype=dtypes.uint32)
 
@@ -68,15 +64,10 @@ class HCQGraph(MultiGraphRunner):
 
     for dev, queue in self.comp_queues.items(): dev_access[queue].add(dev)
 
-    self.input_replace_map: dict[HCQCompiled, set[int]] = collections.defaultdict(set)
     self.fixedvars: dict[HCQCompiled, dict[Variable, int]] = {}
 
     for j,ji in enumerate(jit_cache):
-      if is_exec_prg:=isinstance(ji.prg, CompiledRunner): enqueue_dev: HCQCompiled = ji.prg.dev
-      else:
-        # For copy ops prioritize enqeueuing on the dest device, so reverse the buffers.
-        for b in cast(list[Buffer], ji.bufs[::-1]):
-          if (enqueue_dev:=cast(HCQCompiled, Device[b.device])).hw_copy_queue_t is not None: break
+      enqueue_dev: HCQCompiled = ji.prg.dev if (is_exec_prg:=isinstance(ji.prg, CompiledRunner)) else Device[ji.bufs[1].device] #type:ignore
 
       # set any fixedvars on the device
       self.fixedvars[enqueue_dev] = merge_dicts([self.fixedvars.get(enqueue_dev, {}), ji.fixedvars])
@@ -157,11 +148,10 @@ class HCQGraph(MultiGraphRunner):
       # Encode main commands based on ji type.
       if isinstance(ji.prg, CompiledRunner):
         enqueue_queue.exec(ji.prg._prg, self.ji_args[j], tuple(ji.prg.p.global_size or (1,1,1)), tuple(ji.prg.p.local_size or (1,1,1)))
-      elif isinstance(ji.prg, (BufferXfer, BufferCopy)):
+      elif isinstance(ji.prg, BufferXfer):
         dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
-        for bufid, src in enumerate(cast(list[Buffer], ji.bufs)):
-          if (inprep_idx:=self.input_replace.get((j, bufid))) is not None: self.input_replace_map[enqueue_dev].add(inprep_idx)
-          else: cast(HCQAllocator, enqueue_dev.allocator).map(self.hcq_bufs[j][bufid])
+        cast(HCQAllocator, Device[src.device].allocator).map(dest._buf)
+
         enqueue_queue.copy(self.hcq_bufs[j][0].va_addr, self.hcq_bufs[j][1].va_addr, dest.nbytes)
         self.copy_to_devs[cast(HCQCompiled, Device[dest.device])].add(cast(HCQCompiled, Device[src.device]))
 
@@ -187,9 +177,6 @@ class HCQGraph(MultiGraphRunner):
     for sig in self.queue_signals_to_reset: sig.value = 0
     self.signals['KICK'].value = self.kickoff_value
 
-    for dev in self.devices:
-      for idx_to_map in self.input_replace_map[dev]: cast(HCQAllocator, dev.allocator).map(input_rawbuffers[idx_to_map]._buf)
-
     if PROFILE and self.kickoff_value > 1: self.collect_timestamps()
 
     hcq_var_vals = {self.kickoff_var: self.kickoff_value, **var_vals,
@@ -223,9 +210,3 @@ class HCQGraph(MultiGraphRunner):
     if PROFILE and self.kickoff_value >= 1: self.collect_timestamps()
 
     for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True))
-
-  @staticmethod
-  def supports_exec_item(dev, ei:ExecItem) -> bool:
-    # MOCKGPU is not supported, since it can't execute commands in parallel
-    copy = (isinstance(ei.prg, BufferCopy) and cast(HCQCompiled, dev).hw_copy_queue_t is not None) and not getenv("MOCKGPU")
-    return all(issubclass(type(Device[b.device]), HCQCompiled) for b in ei.bufs if b) and (isinstance(ei.prg, (CompiledRunner, BufferXfer)) or copy)

tinygrad/runtime/ops_amd.py

@@ -479,7 +479,7 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
       self.dev.iface.free(opaque)
     except AttributeError: pass
 
-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): self.dev.iface.map(buf._base if buf._base is not None else buf)
 
 @dataclass(frozen=True)
 class ProfileSQTTEvent(ProfileEvent): device:str; se:int; blob:bytes; itrace:bool # noqa: E702
@@ -563,7 +563,7 @@ class KFDIface:
     self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
     self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False) -> HCQBuffer:
     flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
 
     if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
@@ -572,7 +572,7 @@ class KFDIface:
     if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC
 
     if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
-      buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
+      buf = addr = FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
     else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0)
 
     try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, base=addr, length=size, gpu_id=self.gpu_id,
@@ -606,8 +606,6 @@ class KFDIface:
     return dmaref
 
   def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-
     c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
     stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
     assert stm.n_success == 1

tinygrad/runtime/ops_cpu.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
-import platform, subprocess, sys, ctypes, functools, time, mmap
-from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, wait_cond
+import platform, subprocess, sys, ctypes, functools, time
+from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, round_up, wait_cond
 from tinygrad.device import Compiler, BufferSpec, DMACPURef
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQBuffer, HWQueue, HCQArgsState, HCQSignal, HCQProgram, MMIOInterface
 from tinygrad.runtime.support.elf import jit_loader
@@ -62,9 +62,10 @@ class CPUProgram(HCQProgram):
       ctypes.windll.kernel32.FlushInstructionCache(ctypes.c_void_p(proc), ctypes.c_void_p(self.mem), ctypes.c_size_t(len(lib)))
       self.fxn = ctypes.CFUNCTYPE(None)(self.mem)
     else:
+      from mmap import mmap, PROT_READ, PROT_WRITE, PROT_EXEC, MAP_ANON, MAP_PRIVATE
       # On apple silicon with SPRR enabled (it always is in macos) RWX pages are unrepresentable: https://blog.svenpeter.dev/posts/m1_sprr_gxf/
       # MAP_JIT allows us to easily flip pages from RW- to R-X and vice versa. It is a noop on intel cpus. (man pthread_jit_write_protect_np)
-      self.mem = mmap.mmap(-1, len(lib), mmap.MAP_ANON|mmap.MAP_PRIVATE|(MAP_JIT if OSX else 0), mmap.PROT_READ|mmap.PROT_WRITE|mmap.PROT_EXEC)
+      self.mem = mmap(-1, len(lib), MAP_ANON | MAP_PRIVATE | (MAP_JIT if OSX else 0), PROT_READ | PROT_WRITE | PROT_EXEC)
 
       if OSX: CPUProgram.rt_lib.pthread_jit_write_protect_np(False)
       self.mem.write(lib)
@@ -81,15 +82,15 @@ class CPUProgram(HCQProgram):
     super().__init__(HCQArgsState, dev, name, kernargs_alloc_size=0)
 
   def __del__(self):
-    if getattr(sys, 'is_finalizing', lambda: True)(): return
     if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.mem), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE
 
 class CPUAllocator(HCQAllocatorBase):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
-    if options.external_ptr: addr, buf = options.external_ptr, None
-    elif sys.platform == "win32": addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
-    else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_PRIVATE, mmap.PROT_READ | mmap.PROT_WRITE))
-    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
+    if options.external_ptr: buf = (ctypes.c_uint8 * size).from_address(options.external_ptr)
+    else:
+      offset = round_up(ctypes.addressof(tmpbuf:=(ctypes.c_uint8 * (size + 0x1000))()), 0x1000) - ctypes.addressof(tmpbuf)
+      buf = (ctypes.c_uint8 * size).from_buffer(tmpbuf, offset)
+    return HCQBuffer(va:=ctypes.addressof(buf), sz:=ctypes.sizeof(buf), meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
   def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
   def _as_dmaref(self, buf): return DMACPURef(buf.va_addr, buf.size)
   def _copyin(self, dest, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), len(src))

tinygrad/runtime/ops_nv.py

@@ -282,7 +282,7 @@ class NVAllocator(HCQAllocator['NVDevice']):
       self.dev.iface.free(opaque)
     except AttributeError: pass
 
-  def _map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+  def _map(self, buf:HCQBuffer): self.dev.iface.map(buf._base if buf._base is not None else buf)
 
 @dataclass
 class GPFifo:
@@ -382,14 +382,14 @@ class NVKIface:
     if made.params.status != 0: raise RuntimeError(f"_gpu_map_to_cpu returned {get_error_str(made.params.status)}")
     return fd_dev.mmap(target, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if target is not None else 0), 0)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0) -> HCQBuffer:
     # Uncached memory is "system". Use huge pages only for gpu memory.
     page_size = (4 << (12 if OSX else 10)) if uncached or host else ((2 << 20) if size >= (8 << 20) else (4 << (12 if OSX else 10)))
     size = round_up(size, page_size)
     va_addr = self._alloc_gpu_vaddr(size, alignment=page_size, force_low=cpu_access)
 
     if host:
-      va_addr = cpu_addr or FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ|mmap.PROT_WRITE, MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, 0)
+      va_addr = FileIOInterface.anon_mmap(va_addr, size, mmap.PROT_READ | mmap.PROT_WRITE, MAP_FIXED | mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
 
       flags = (nv_gpu.NVOS02_FLAGS_PHYSICALITY_NONCONTIGUOUS << 4) | (nv_gpu.NVOS02_FLAGS_COHERENCY_CACHED << 12) \
             | (nv_gpu.NVOS02_FLAGS_MAPPING_NO_MAP << 30)
@@ -438,11 +438,7 @@ class NVKIface:
       hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid],
       has_cpu_mapping=has_cpu_mapping), view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None, owner=self.dev)
 
-  def map(self, mem:HCQBuffer):
-    if mem.owner is not None and mem.owner._is_cpu():
-      if not any(x.device.startswith("NV") for x in mem.mapped_devs): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mem = mem.mappings[next(x for x in mem.mapped_devs if x.device.startswith("NV"))]
-    self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False)
+  def map(self, mem:HCQBuffer): self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False)
 
   def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):
     return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment)

tinygrad/runtime/ops_python.py

@@ -40,7 +40,8 @@ class PythonProgram:
       loop_ends: dict[int, int] = {}
       while i < len(self.uops):
         uop, dtype, idp, arg = self.uops[i]
-        void_ops = {Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP}
+        void_ops = {Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK}
+        if uop is Ops.DEFINE_REG: idp = [idp[0]]
         inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops]
         dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops]
         if getenv("TRACE"): print(i, uop, dtype, arg, inp, dtp)
@@ -48,13 +49,14 @@ class PythonProgram:
           loop_ends[idp[0]] = i
           i = idp[0]
           continue
-        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK, Ops.NOOP):
+        if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK):
           # in the python emulator, the warp is always in sync
           i += 1
           continue
         assert dtype is not None, f"{uop} is missing a dtype"
         dl[i] = dtype
         if uop is Ops.STORE:
+          #assert len(inp) == 2, "expected store is ([(memory, offset, gate)], [value])"
           for j,val in enumerate(inp[1] if dtp[1].count > 1 else [inp[1]]):
             for (m,o,g),v in zip(inp[0], val):
               if g: _store(m, o+j, v)
@@ -67,6 +69,8 @@ class PythonProgram:
           if uop is Ops.DEFINE_REG:
             # REGs are per thread
             ul[i] = [memoryview(bytearray(dtype.size*dtype.itemsize)).cast(dtype.fmt) for _ in range(warp_size)]
+            for buf, val in zip(ul[i], inp[0]):
+              for x in range(dtype.size): buf[x] = val
           else:
             buf = memoryview(bytearray(dtype.size*dtype.itemsize)) if uop is not Ops.DEFINE_GLOBAL else pbufs.pop(0)
             ul[i] = [buf.cast(dtype.fmt)] * warp_size

tinygrad/runtime/support/hcq.py

@@ -406,8 +406,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)
 
   def _at_profile_finalize(self):
-    self.synchronize() # Expect device to be synchronizes
-
     def _sync(d:HCQCompiled, q_t:Callable[[], HWQueue]):
       q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d)
       st = time.perf_counter_ns()
@@ -439,8 +437,6 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       except Exception: errs += f"\n{iface_t.__name__}: {traceback.format_exc()}"
     raise RuntimeError(f"Cannot find a usable interface for {type(self).__name__[:-6]}:{self.device_id}:\n{errs}")
 
-  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] in ("CPU", "LLVM")
-
   def finalize(self):
     try: self.synchronize() # Try to finalize device in any case.
     except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
@@ -452,8 +448,7 @@ class HCQBuffer:
   def __init__(self, va_addr:sint, size:int, texture_info:Any=None, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None,
                owner:HCQCompiled|None=None):
     self.va_addr, self.size, self.texture_info, self.meta, self._base, self.view = va_addr, size, texture_info, meta, _base, view
-    self._devs, self.owner = ([owner] if owner is not None else []), owner
-    self._mappings:dict[HCQCompiled, HCQBuffer] = {} # mapping to the other devices
+    self.devs, self.owner = ([owner] if owner is not None else []), owner
 
   def offset(self, offset:int=0, size:int|None=None) -> HCQBuffer:
     return HCQBuffer(self.va_addr+offset, size or (self.size - offset), owner=self.owner, texture_info=self.texture_info, meta=self.meta,
@@ -464,10 +459,7 @@ class HCQBuffer:
     return self.view
 
   @property
-  def mappings(self): return self._mappings if self._base is None else self._base._mappings
-
-  @property
-  def mapped_devs(self): return self._devs if self._base is None else self._base._devs
+  def mapped_devs(self): return self.devs if self._base is None else self._base.devs
 
 class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
   """
@@ -485,10 +477,7 @@ class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
     if self.dev in buf.mapped_devs: return
     if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
     if not hasattr(self, '_map'): raise NotImplementedError("map failed: no method implemented")
-
-    # Since it's unified memory space, any buffer mapping is valid for all devices after successful map.
-    # Devices can save mappings and internal metadata as a new buffer.
-    if (mb:=self._map(buf)) is not None: buf.mappings[self.dev] = mb
+    self._map(buf)
     buf.mapped_devs.append(self.dev)
 
   def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)

tinygrad/runtime/support/system.py

@@ -12,20 +12,16 @@ class _System:
 
   def memory_barrier(self): lib.atomic_thread_fence(__ATOMIC_SEQ_CST:=5) if (lib:=self.atomic_lib()) is not None else None
 
-  def lock_memory(self, addr:int, size:int):
-    if libc.mlock(ctypes.c_void_p(addr), size): raise RuntimeError(f"Failed to lock memory at {addr:#x} with size {size:#x}")
-
-  def system_paddrs(self, vaddr:int, size:int) -> list[int]:
-    self.pagemap().seek(vaddr // mmap.PAGESIZE * 8)
-    return [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))]
-
   def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False, data:bytes|None=None) -> tuple[int, list[int]]:
     assert not contiguous or size <= (2 << 20), "Contiguous allocation is only supported for sizes up to 2MB"
     flags = (libc.MAP_HUGETLB if contiguous and (size:=round_up(size, mmap.PAGESIZE)) > 0x1000 else 0) | (MAP_FIXED if vaddr else 0)
     va = FileIOInterface.anon_mmap(vaddr, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS|MAP_POPULATE|MAP_LOCKED|flags, 0)
 
     if data is not None: to_mv(va, len(data))[:] = data
-    return va, self.system_paddrs(va, size)
+
+    # Read pagemap to get the physical address of each page. The pages are locked.
+    self.pagemap().seek(va // mmap.PAGESIZE * 8)
+    return va, [(x & ((1<<55) - 1)) * mmap.PAGESIZE for x in array.array('Q', self.pagemap().read(size//mmap.PAGESIZE*8, binary=True))]
 
   def pci_reset(self, gpu): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{gpu}/reset'")
   def pci_scan_bus(self, target_vendor:int, target_devices:list[int]) -> list[str]:
@@ -159,12 +155,6 @@ class PCIIfaceBase:
     if b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
 
   def map(self, b:HCQBuffer):
-    if b.owner is not None and b.owner._is_cpu():
-      System.lock_memory(cast(int, b.va_addr), b.size)
-      paddrs, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(cast(int, b.va_addr), round_up(b.size, 0x1000))], True, False
-    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
-      paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
-      snooped, uncached = b.meta.mapping.snooped, b.meta.mapping.uncached
-    else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
-
-    self.dev_impl.mm.map_range(cast(int, b.va_addr), round_up(b.size, 0x1000), paddrs, system=True, snooped=snooped, uncached=uncached)
+    if (ifa:=getattr(b.owner, "iface", None)) is None or not isinstance(ifa, PCIIfaceBase): raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
+    paddrs = [(paddr if b.meta.mapping.system else (paddr + ifa.p2p_base_addr), size) for paddr,size in b.meta.mapping.paddrs]
+    self.dev_impl.mm.map_range(cast(int, b.va_addr), b.size, paddrs, system=True, snooped=b.meta.mapping.snooped, uncached=b.meta.mapping.uncached)

tinygrad/tensor.py

@@ -1976,14 +1976,12 @@ class Tensor(MathTrait):
 
     # https://keccak.team/keccak_specs_summary.html
 
-    def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64):
-      # TODO: contiguous is here for compile speed
-      return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l)).contiguous()
+    def ctensor(l: Sequence[ConstType], dtype: DType = dtypes.uint64): return Tensor.stack(*(Tensor(v, dtype=dtype, device=self.device) for v in l))
     rot_offsets = [44, 43, 21, 14, 28, 20, 3, 45, 61, 1, 6, 25, 8, 18, 27, 36, 10, 15, 56, 62, 55, 39, 41, 2]
     rot_offsets_v0, rot_offsets_v1 =  ctensor([0] + [1 << v for v in rot_offsets]), ctensor([1] + [1 << (64 - v) for v in rot_offsets])
 
     # calculated from π step
-    reorder_indexes = ctensor([0,6,12,18,24,3,9,10,16,22,1,7,13,19,20,4,5,11,17,23,2,8,14,15,21], dtype=dtypes.int32)
+    reorder_indexes = ctensor([0,6,12,18,24,3,9,10,16,22,1,7,13,19,20,4,5,11,17,23,2,8,14,15,21])
     rnd_const_masks = [ctensor([v]).pad((0, 24)) for v in (1, 0x8082, 0x800000000000808a, 0x8000000080008000, 0x808b, 0x80000001, 0x8000000080008081,
     0x8000000000008009, 0x8a, 0x88, 0x80008009, 0x8000000a, 0x8000808b, 0x800000000000008b, 0x8000000000008089, 0x8000000000008003,
     0x8000000000008002, 0x8000000000000080, 0x800a, 0x800000008000000a, 0x8000000080008081, 0x8000000000008080, 0x80000001, 0x8000000080008008)]
@@ -1994,9 +1992,9 @@ class Tensor(MathTrait):
     data = data.pad((None, (0, data_pad))).reshape(bs := data.shape[0], -1, rate).pad((None, None, (0, 200 - rate)))
 
     # create pad mask
-    lbe = prod(data.shape[1:]) + rate - data_pad - 200
-    if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (200 - rate, 0)]
-    else: mb = [(lbe, 0), (1, dsbyte), (data_pad - 2, 0), (1, 0x80), (200 - rate, 0)]
+    lbe = (blen := prod(data.shape[1:])) + rate - data_pad - 200
+    if data_pad == 1: mb = [(lbe, 0), (1, dsbyte ^ 0x80), (blen - lbe - 1, 0)]
+    else: mb = [(lbe, 0), (1, dsbyte), (blen + rate - lbe - 202, 0), (1, 0x80), (200 - rate, 0)]
     pad_mask = Tensor.cat(*(Tensor(v, dtype=dtypes.uint8, device=data.device).expand(l) for l, v in mb if l > 0)).unsqueeze(0)
 
     data = (data.flatten(1) ^ pad_mask).reshape(*data.shape[:2], 200).bitcast(dtypes.uint64)
@@ -2015,42 +2013,8 @@ class Tensor(MathTrait):
         # χ and ι step
         state = state.bitwise_xor(~state.roll(shifts=-1, dims=2) & state.roll(shifts=-2, dims=2))
         state = state.flatten(1) ^ rnd_const_masks[i]
-      # NOTE: kernelize here to prevent internal stack from growing propotional to data size
-      state = state.kernelize()
     return state.bitcast(dtypes.uint8)[:,:(obytes:=(200 - rate) // 2)].reshape(*self.shape[:-1], obytes)
 
-  def _hash_1mb(self) -> Tensor:
-    assert self.dtype == dtypes.uint8, "only support uint8 tensors for hashing"
-    assert self.ndim == 2, "only support batched 1d tensors"
-    assert self.shape[1] == 1024 * 1024, "only support messages of 1mb"
-
-    blocks = self.shape[0] * self.shape[1] // 4096
-    data = self.reshape(blocks, 4096)
-    block_hashes = data.keccak("shake_128").reshape(self.shape[0], 4096)
-    return block_hashes.keccak("shake_128").reshape(self.shape[0], 16)
-
-  def hash(self) -> Tensor:
-    """
-    Calculates a 16-byte hash of the tensor.
-    ```python exec="false source="above" session="tensor" result="python"
-    t = Tensor(b"Hello World!").hash()
-    print(t.data().hex())
-    ```
-    """
-
-    data = self.flatten().bitcast(dtypes.uint8)
-    if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-    base_chunks = ceildiv(data.shape[0], 2**20)
-    tree_depth = math.ceil(math.log(base_chunks, 65536)) if base_chunks > 1 else 0
-
-    level_chunks = base_chunks
-    for _ in range(tree_depth + 1):
-      data = data.reshape(level_chunks, 2**20)._hash_1mb().flatten()
-      if (tsize := data.shape[0]) % 2**20 != 0: data = data.pad((0, 2**20 - tsize % 2**20))
-      level_chunks = ceildiv(data.shape[0], 2**20)
-
-    return data[:16]
-
   def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Tensor, Tensor, Tensor]:
     m = self - self.max(axis=axis, keepdim=True).detach()
     if dtype is not None: m = m.cast(dtype)

tinygrad/uop/__init__.py

@@ -9,7 +9,7 @@ class FastEnum(IntEnum):
 # the order of these Ops controls the order of the toposort
 class Ops(FastEnum):
   # uops that aren't rendered
-  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto(); PRECAST = auto()  # noqa: E702
+  NOOP = auto(); SINK = auto(); UNIQUE = auto(); DEVICE = auto(); KERNEL = auto()  # noqa: E702
 
   # buffer ops
   COPY = auto(); BUFFER = auto(); BUFFER_VIEW = auto(); MSELECT = auto(); MSTACK = auto() # noqa: E702

tinygrad/uop/mathtraits.py

@@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes
 
 class MathTrait:
   # required to implement
-  def alu(self:T, op:Ops, *src) -> T: raise NotImplementedError
+  def alu(self:T, arg:Ops, *src) -> T: raise NotImplementedError
   def const_like(self:T, b) -> T: raise NotImplementedError
 
   # great functions you get!

tinygrad/uop/ops.py

@@ -136,7 +136,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
 
   @functools.cached_property
   def st(self) -> ShapeTracker|None:
-    if self.op in GroupOp.Block or self.op is Ops.INDEX: return None
     from tinygrad.shape.shapetracker import ShapeTracker
     # VIEW and MovementOps define a new ShapeTracker from the arg
     if self.op is Ops.VIEW: return self.arg
@@ -144,13 +143,12 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     # CONST with a DEVICE has a shape of ()
     if self.op is Ops.CONST and len(self.src) and self.src[0].op is Ops.DEVICE: return ShapeTracker.from_shape(())
     # BufferOps and ASSIGN flow ShapeTracker from a direct edge
-    if self.op in {Ops.STORE, Ops.ASSIGN, Ops.LOAD}: return self.src[0].st
     if self.op in GroupOp.Buffer: return views[0] if (views:=[x.st for x in self.src if x.op is Ops.VIEW]) else None
+    if self.op is Ops.ASSIGN: return self.src[0].st
 
     # BUFFER/BUFFER_VIEW and KERNEL only have a size
     if self.op in {Ops.BUFFER, Ops.BUFFER_VIEW}: return ShapeTracker.from_shape((self.size,))
     if self.op is Ops.KERNEL: return ShapeTracker.from_shape((self.arg.ast.size,))
-    #if self.op in {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG}: return ShapeTracker.from_shape((self.dtype.size,))
 
     # otherwise we get the shape from sources
     if not (src_sts := [x.st for x in self.src if x.st is not None]): return None
@@ -169,7 +167,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if self.op is Ops.VIEW: return self.shape
     # NOTE: if a parent doesn't have st its full_shape is empty
     parent_shapes = [x.full_shape for x in self.src]
-    return tuple(smax(x) for x in itertools.zip_longest(*parent_shapes, fillvalue=1))
+    return tuple(smax(x) for x in zip(*[x for x in parent_shapes if x != ()]))
   @property
   def shape(self) -> tuple[sint, ...]: return unwrap(self.st).shape
   @property
@@ -213,7 +211,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def sink(self, *srcs:UOp|None, **kwargs): return UOp(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def detach(self): return UOp(Ops.DETACH, self.dtype, (self,))
   def index(self, idx:UOp, valid:UOp|None=None): return UOp(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx))
-  def __getitem__(self, idx): return self.index(idx)
   def const_like(self, b:ConstLike):
     # constants can optionally have a DEVICE source
     return UOp.const(self.dtype, b, device=self._device, shape=self.shape if self.st is not None else None)
@@ -238,10 +235,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   def load(self, *src:UOp, **kwargs): return UOp(Ops.LOAD, dtype=kwargs.pop("dtype", self.dtype.base), src=(self,)+src, **kwargs)
   def store(self, *src:UOp, **kwargs): return UOp(Ops.STORE, self.dtype, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
-  def alu(self, op, *src:UOp, **kwargs):
+  def alu(self, arg, *src:UOp):
     out_dtype = (self, *src)[-1].dtype
-    if op in {Ops.CMPLT, Ops.CMPNE}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
-    return UOp(op, out_dtype, (self,)+src, **kwargs)
+    if arg in {Ops.CMPLT, Ops.CMPNE}: out_dtype = dtypes.bool.vec(out_dtype.count) if out_dtype.count > 1 else dtypes.bool
+    return UOp(arg, out_dtype, (self,)+src)
   @staticmethod
   def const(dtype:DType, b:ConstLike, device:str|tuple[str, ...]|None=None, shape:tuple[sint, ...]|None=None):
     if isinstance(b, UOp): return b.unbind()[0] if b.op is Ops.BIND else b
@@ -257,12 +254,18 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   @staticmethod
   def range(dtype:DType, end:sint, idx:int): return UOp(Ops.RANGE, dtype=dtype, src=(sint_to_uop(end),), arg=idx)
   def r(self, op:Ops, axis:tuple[int, ...], permute=True):
-    assert permute
     axis = tuple(sorted([x for x in axis if resolve(self.shape[x] != 1)]))
     if len(axis) == 0: return self
-    ret = self.permute(tuple([i for i in range(len(self.shape)) if i not in axis])+axis)
-    ret = ret.reshape(tuple([s for s in ret.shape if resolve(s != 1)]))
-    ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, tuple(range(len(ret.shape)-len(axis), len(ret.shape)))))
+    # move any non reduce axis before the first reduce axis
+    move_early, rest = partition(range(axis[0], len(self.shape)), lambda i: i not in axis and resolve(self.shape[i] != 1))
+    if move_early and permute:
+      permaxis = tuple(range(axis[0])) + tuple(move_early) + tuple(rest)
+      ret = self.permute(permaxis)
+      new_axis = tuple([x for x in range(axis[0]+len(move_early), len(self.shape)) if resolve(ret.shape[x] != 1)])
+      assert len(axis) == len(new_axis)
+    else:
+      ret, new_axis = self, axis
+    ret = UOp(Ops.REDUCE_AXIS, self.dtype, (ret,), (op, new_axis))
     return ret.reshape(tuple([x if i not in axis else 1 for i,x in enumerate(self.shape)]))
   def reduce(self, *src:UOp, **kwargs): return UOp(Ops.REDUCE, kwargs.pop('dtype', self.dtype), src=(self,)+src, **kwargs)
   def contiguous(self): return self.alu(Ops.CONTIGUOUS)
@@ -333,7 +336,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     return self
   def view(self, new_st:ShapeTracker) -> UOp: return UOp(Ops.VIEW, self.dtype, (self,), new_st)
 
-  def _mop(self, op:Ops, arg) -> UOp:
+  def _mop(self, op:Ops, arg):
     ret = UOp(op, self.dtype, (self,), arg)
     if self.st == ret.st: return self  # ignore NOOPs, also check ret.st
     return ret
@@ -366,7 +369,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       return self.src[0].device[self.arg]
     if self.op is Ops.MSTACK: return tuple(cast(str, x.device) for x in self.src)
     if self.op in {Ops.COPY, Ops.BUFFER, Ops.ALLREDUCE}: return self.src[1].device
-    return next((x._device for x in self.src if x._device is not None), None)
+    return dsrcs[0]._device if len(dsrcs:=[x for x in self.src if x._device is not None]) != 0 else None
   @property
   def buf_uop(self) -> UOp:
     if self.op is Ops.BUFFER: return self
@@ -853,7 +856,7 @@ if TRACK_MATCH_STATS or PROFILE:
         print(f"rewrote {len(tracked_ctxs)} graphs and matched {sum(len(r.matches) for x in tracked_ctxs for r in x)} times, saved to {fn}")
         pickle.dump((tracked_keys, tracked_ctxs, uop_fields), f)
     if VIZ: launch_viz("VIZ", temp("rewrites.pkl", append_user=True))
-    if getenv("PRINT_MATCH_STATS", TRACK_MATCH_STATS.value):
+    if getenv("PRINT_MATCH_STATS", 1):
       ret = [0,0,0.0,0.0]
       for k,v in sorted(list(match_stats.items()), key=lambda x: x[1][2]+x[1][3]):
         loc_str = f"{k.location[0].split('/')[-1]}:{k.location[1]}"

tinygrad/uop/spec.py

@@ -130,7 +130,8 @@ index_pat = UPat(Ops.INDEX, name="idx").or_casted()
 spec = PatternMatcher([
   (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: isinstance(x.dtype, (PtrDType, ImageDType)) and x.dtype.addrspace == AddrSpace.GLOBAL),
   (UPat(Ops.DEFINE_LOCAL, name="x"), lambda x: isinstance(x.dtype, PtrDType) and x.dtype.addrspace == AddrSpace.LOCAL),
-  (UPat(Ops.DEFINE_REG, src=()), lambda: True),
+  (UPat(Ops.DEFINE_REG, src=(UPat.var("c"),), name="x", allow_any_len=True),
+   lambda x,c: all(y.op is Ops.RANGE for y in x.src[1:]) and c.dtype.base == x.dtype.base),
   (UPat(Ops.DEFINE_VAR, name="x"), lambda x: isinstance(x.arg[1], int) and isinstance(x.arg[2], int)),
 
   (UPat(Ops.RANGE, src=(UPat.var("x"),), name="rng"), lambda rng,x: rng.dtype == x.dtype and isinstance(rng.arg, int)),
@@ -158,11 +159,13 @@ spec = PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat((Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG)), UPat(), UPat(dtype=dtypes.bool))), lambda: True),
 
   # LOAD on STORE
-  (UPat(Ops.LOAD, src=(UPat(Ops.STORE),), allow_any_len=True), lambda: True),
+  (UPat(Ops.LOAD, src=(UPat(Ops.STORE),)), lambda: True),
 
   # LOAD takes a <bufidx, alt?, barrier?>
-  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.IF, name="cond")), allow_any_len=True), lambda idx,cond: validate_index(idx,cond.src[0])),
-  (UPat(Ops.LOAD, src=(index_pat,), allow_any_len=True), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat,)), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.BARRIER))), validate_index),
+  (UPat(Ops.LOAD, src=(index_pat, UPat(Ops.IF, name="cond"))), lambda idx,cond: validate_index(idx,cond.src[0])),
+  (UPat(Ops.LOAD, src=(index_pat, UPat.var("alt")), name="ld"), lambda ld,alt,idx: ld.dtype == alt.dtype and validate_index(idx)),
 
   # STORE takes a <bufidx, val, gate?>
   (UPat(Ops.STORE, src=(index_pat, UPat(name="val"), UPat(Ops.IF, name="gate")), allow_any_len=True), validate_store),
@@ -198,7 +201,7 @@ spec = PatternMatcher([
   # NOTE: for testing, we let sinks be anything
   #(UPat(Ops.SINK, src=UPat(Ops.STORE)), lambda: True),
   (UPat(Ops.SINK, dtypes.void), lambda: True),
-  (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM, Ops.PRECAST)), lambda: True),
+  (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM)), lambda: True),
 
   # PTX LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat(dtype=dtypes.int64),), allow_any_len=True), lambda: True),
@@ -208,7 +211,7 @@ spec = PatternMatcher([
 
 def verify_sink_dims(sink:UOp):
   if not all_same([s.shape for s in sink.src]): return False
-  for dims in zip(*[x.shape for x in sink.toposort() if x.op is Ops.VIEW]):
+  for dims in zip(*[x.shape for x in sink.toposort() if x.st is not None]):
     if len(n_dims:={s for s in dims if resolve(s!=1)}) > 1:
       print(f"# INVALID KERNEL DIMS: can only have 1 or n in each dimension: {n_dims}")
       return False

tinygrad/uop/symbolic.py

@@ -437,6 +437,7 @@ sym = symbolic_flat+PatternMatcher([
   ((UPat.var('x', dtypes.uint64)&(UPat.var('y').where(UPat.const(dtypes.uint64, 0xFFFFFFFF), UPat.const(dtypes.uint64, 0)))).cast(dtypes.uint32),
    lambda x,y: y.where(x.cast(dtypes.uint32), UOp.const(dtypes.uint32, 0))),
   # ** self folding **
+  (UPat(Ops.DEFINE_REG, src=(UPat.var("x"),)), lambda x: x),            # a DEFINE_ACC without ranges is a CONST
   # x!=0 -> (bool)x
   (UPat.var("x")!=0, lambda x: x.cast(dtypes.bool.vec(x.dtype.count))),
   # ** where **
@@ -450,9 +451,8 @@ sym = symbolic_flat+PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat.var("b"), UPat.var("idx"), UPat.const(dtypes.bool, True))), lambda b, idx: b.index(idx)),
   # ** load/store folding **
   (UPat.store(UPat(Ops.INDEX, name="index"), UPat.load(UPat(Ops.INDEX, name="index"))), lambda index: UOp(Ops.NOOP)),
-  (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"),
-                                                                    UPat.load(UPat(Ops.INDEX, name="index"))), allow_any_len=True, name="store"),
-   lambda index, gate, alt, store: UOp.store(index.src[0].index(index.src[1], gate), alt, *store.src[2:])),
+  (UPat.store(UPat(Ops.INDEX, name="index"), UPat.var("gate").where(UPat.var("alt"), UPat.load(UPat(Ops.INDEX, name="index")))),
+   lambda index, gate, alt: UOp.store(index.src[0].index(index.src[1], gate), alt)),
   # fold gated LOAD/STORE
   (UPat().index(UPat(), UPat.const(dtypes.bool, True)).named("idx"), lambda idx: idx.replace(src=idx.src[0:2])), # remove True
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat().index(UPat(), UPat.const(dtypes.bool, False)).or_casted(),), allow_any_len=True, name="x"),

tinygrad/viz/index.html

@@ -215,10 +215,11 @@
   #device-list > div {
     min-height: 32px;
     max-width: 100px;
-    overflow-x: auto;
-    overflow-y: hidden;
+    overflow: hidden;
+    text-overflow: ellipsis;
     white-space: nowrap;
     display: flex;
+    cursor: pointer;
   }
   #device-list > div:hover {
     background-color: rgba(20, 23, 35, 0.3);
@@ -237,10 +238,10 @@
         </svg>
       </button>
     </div>
-    <div class="progress-message"></div>
     <div class="container ctx-list-parent"><div class="ctx-list"></div></div>
     <div class="view profiler"></div>
     <div class="view graph">
+      <div class="progress-message">Rendering new layout...</div>
       <svg id="graph-svg" preserveAspectRatio="xMidYMid meet">
         <g id="render">
           <g id="edges"></g>

tinygrad/viz/js/index.js

@@ -25,10 +25,6 @@ function intersectRect(r1, r2) {
 let [workerUrl, worker, timeout] = [null, null, null];
 async function renderDag(graph, additions, recenter=false) {
   // start calculating the new layout (non-blocking)
-  const progressMessage = document.querySelector(".progress-message");
-  progressMessage.innerText = "Rendering new graph...";
-  if (timeout != null) clearTimeout(timeout);
-  timeout = setTimeout(() => {progressMessage.style.display = "block"}, 2000);
   if (worker == null) {
     const resp = await Promise.all(["/assets/dagrejs.github.io/project/dagre/latest/dagre.min.js","/js/worker.js"].map(u => fetch(u)));
     workerUrl = URL.createObjectURL(new Blob([(await Promise.all(resp.map((r) => r.text()))).join("\n")], { type: "application/javascript" }));
@@ -37,6 +33,9 @@ async function renderDag(graph, additions, recenter=false) {
     worker.terminate();
     worker = new Worker(workerUrl);
   }
+  if (timeout != null) clearTimeout(timeout);
+  const progressMessage = document.querySelector(".progress-message");
+  timeout = setTimeout(() => {progressMessage.style.display = "block"}, 2000);
   worker.postMessage({graph, additions, ctxs});
   worker.onmessage = (e) => {
     displayGraph("graph");
@@ -172,7 +171,6 @@ async function renderProfiler() {
     const startY = offsetY+(levelHeight*timeline.maxDepth)+padding/2;
     let area = mem.shapes.length === 0 ? 0 : areaScale(mem.peak);
     if (area === 0) div.style.pointerEvents = "none";
-    else div.style.cursor = "pointer";
     if (k === focusedDevice) {
       // expand memory graph for the focused device
       area = maxArea*4;
@@ -230,7 +228,6 @@ async function renderProfiler() {
       ctx.fillRect(x, e.y, width, e.height);
       rectLst.push({ y0:e.y, y1:e.y+e.height, x0:x, x1:x+width, arg:e.arg });
       // add label
-      if (e.label == null) continue;
       ctx.textAlign = "left";
       ctx.textBaseline = "middle";
       let [labelX, labelWidth] = [x+2, 0];