cleanups

This reverts commit 1e07dff384.

.github/actions/setup-tinygrad/action.yml

@@ -153,7 +153,7 @@ runs:
         fi
         # **** AMD ****
         if [[ "${{ inputs.amd }}" == "true" ]]; then
-          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev"
+          pkgs+=" hsa-rocr comgr hsa-rocr-dev liburing-dev libibverbs-dev libc6-dev"
         fi
         # **** CUDA ****
         if [[ "${{ inputs.cuda }}" == "true" ]]; then

.github/workflows/test.yml

@@ -132,10 +132,13 @@ jobs:
       run: |
         cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
         cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
+        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
         ./autogen_stubs.sh libc
         ./autogen_stubs.sh io_uring
+        ./autogen_stubs.sh ib
         diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
         diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
+        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
     - name: Verify WebGPU autogen
       run: |
         cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak

autogen_stubs.sh

@@ -240,6 +240,21 @@ generate_io_uring() {
   fixup $BASE/io_uring.py
 }
 
+generate_ib() {
+  clang2py -k cdefstum \
+    /usr/include/infiniband/verbs.h \
+    /usr/include/infiniband/verbs_api.h \
+    /usr/include/infiniband/ib_user_ioctl_verbs.h \
+    /usr/include/rdma/ib_user_verbs.h \
+    -o $BASE/ib.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
+  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
+
+  fixup $BASE/ib.py
+}
+
 generate_libc() {
   clang2py -k cdefstum \
     $(dpkg -L libc6-dev | grep sys/mman.h) \
@@ -465,6 +480,7 @@ elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
+elif [ "$1" == "ib" ]; then generate_ib
 elif [ "$1" == "libc" ]; then generate_libc
 elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl

extra/gemm/amd_matmul.py

@@ -19,6 +19,9 @@ if __name__ == "__main__":
   elif getenv("ASM") == -1:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
     prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+  elif getenv("ASM") == -2:
+    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel4_gmem_df.cpp").read_text()
+    prgfast = replace(prg, name="kernel4_gmem_db", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
   else:
     src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
     prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])

extra/gemm/amd_seb/kernel3_registers.cpp

@@ -10,7 +10,8 @@ __attribute__((device)) inline void __syncthreads() {
 }
 
 #define BLOCK_SIZE 256
-extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, float *c)
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel3_registers(float *a, float *b, float *c)
 {
   constexpr int N = 4096;
   constexpr float alpha = 1.0;

extra/gemm/amd_seb/kernel4_gmem_df.cpp

@@ -0,0 +1,172 @@
+typedef long unsigned int size_t;
+extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
+extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
+struct Dim3 { size_t x, y, z; };
+#define __shared__ __attribute__((shared, aligned(16)))
+__attribute__((device)) inline void __syncthreads() {
+  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
+  __builtin_amdgcn_s_barrier();
+  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
+}
+
+#define BLOCK_SIZE 256
+extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
+kernel4_gmem_db(float *a, float *b, float *c)
+{
+  constexpr int N = 4096;
+  constexpr float alpha = 1.0;
+  constexpr float beta = 0.0;
+
+  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
+  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
+
+  // Block Tile size
+  constexpr int BN = 128;
+  constexpr int BM = 128;
+  // Number of Row or column we read per batch
+  constexpr int BK = 8;
+
+  // Thread Tile size
+  constexpr int TN = 4;
+  constexpr int TM = 4;
+
+  constexpr int nbWaves = BLOCK_SIZE / 32;
+  // Wave Tile size
+  constexpr int WN = 64;
+  constexpr int WM = BN * BM / nbWaves / WN;
+
+  // Number of wave on X & Y axis in the Block tile
+  constexpr int nbWaveX = BN / WN;
+  constexpr int nbWaveY = BM / WM;
+
+  const int waveIndex = threadIdx.x / 32;
+  const int waveIdx = waveIndex % nbWaveX;
+  const int waveIdy = waveIndex / nbWaveX;
+  const int indexInWave = threadIdx.x % 32;
+
+  // A wave is a block of 8x4 of the output matrix
+  constexpr int nbThreadXPerWave = 8;
+  constexpr int nbThreadYPerWave = 4;
+
+  // Thread coordinates in Wave
+  const int idxInWave = indexInWave % nbThreadXPerWave;
+  const int idyInWave = indexInWave / nbThreadXPerWave;
+
+  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
+  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
+
+  // Wave Sub-tile size
+  constexpr int SUBWN = WN / nbIterWaveN;
+  constexpr int SUBWM = WM / nbIterWaveM;
+
+  // Thread mapping to read BKxBN block from A
+  int rAIdx = threadIdx.x % BK;
+  int rAIdy = threadIdx.x / BK;
+  // Thread mapping to read BNxBK block from B
+  int rBIdx = threadIdx.x % BN;
+  int rBIdy = threadIdx.x / BN;
+
+  constexpr int strideReadB = BLOCK_SIZE / BN;
+  constexpr int strideReadA = BLOCK_SIZE / BK;
+  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
+  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
+
+  float A_col[nbIterWaveM * TM];
+  float B_row[nbIterWaveN * TN];
+
+  __shared__ float As[BK][BM];
+  __shared__ float Bs[BK][BN];
+
+  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
+
+  for (int i = 0; i < nbReadsB; i++) {
+    int index_x = BN * blockIdx.x + rBIdx;
+    int index_y = rBIdy + i * strideReadB;
+    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
+  }
+
+  for (int i = 0; i < nbReadsA; i++) {
+    int index_x = rAIdx;
+    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
+  }
+
+  __syncthreads();
+  // Iteration over BK blocks.
+  for (int kId = 0; kId < N; kId += BK) {
+    float regA[nbReadsA];
+    float regB[nbReadsB];
+    if (kId < N - BK) {
+      // We populate the Shared Memory with Ks row and columns
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        regB[i] = b[N * index_y + index_x];
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        regA[i] = a[N * index_y + index_x];
+      }
+    }
+
+    for (int k = 0; k < BK; k++) {
+      // we cache A & B for the entire Wave tile
+      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
+        for (int i = 0; i < TN; i++) {
+          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
+          B_row[iterWave * TN + i] = Bs[k][index];
+        }
+      }
+
+      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
+        for (int i = 0; i < TM; i++) {
+          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
+          A_col[iterWave * TM + i] = As[k][index];
+        }
+      }
+
+      // we accumulate to C_regs
+      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+          for (int yt = 0; yt < TM; yt++) {
+            for (int xt = 0; xt < TN; xt++) {
+              const int x = iterWaveN * TN + xt;
+              const int y = iterWaveM * TM + yt;
+              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
+            }
+          }
+        }
+      }
+    }
+    __syncthreads();
+    if (kId < N - BK) {
+      for (int i = 0; i < nbReadsB; i++) {
+        int index_x = BN * blockIdx.x + rBIdx;
+        int index_y = rBIdy + i * strideReadB + kId + BK;
+        Bs[index_y % BK][index_x % BN] = regB[i]; // row
+      }
+
+      for (int i = 0; i < nbReadsA; i++) {
+        int index_x = rAIdx + kId + BK;
+        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
+        As[(index_x % BK)][(index_y % BM)] = regA[i];
+      }
+      __syncthreads();
+    }
+  }
+
+  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
+    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
+      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
+      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
+      for (int yt = 0; yt < TM; yt++) {
+        for (int xt = 0; xt < TN; xt++) {
+          int indexC = N * (yOut + yt) + xOut + xt;
+          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
+        }
+      }
+    }
+  }
+}

extra/gemm/amd_seb/kernel5_lds_optim.cpp

@@ -26,7 +26,7 @@ kernel5_lds_optim(float *a, float *b, float *c)
   // Number of Row or column we read per batch
   constexpr int BK = 8;
 
-  // Thread Tile size . 4x4
+  // Thread Tile size
   constexpr int TN = 4;
   constexpr int TM = 4;
 

extra/gemm/amd_uop_matmul.py

@@ -1,10 +1,13 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
 from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.dtype import AddrSpace
-from tinygrad.schedule.kernelize import merge_views
-from tinygrad.helpers import getenv, colored
-from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.schedule.kernelize import merge_views, view_left
+from tinygrad.helpers import getenv, colored, prod, unwrap
+from tinygrad.shape.shapetracker import ShapeTracker, View, strides_for_shape
+from tinygrad.opt.kernel import axis_colors
+
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
 
 N = 4096
 run_count = 5
@@ -16,21 +19,54 @@ BK = 8
 TN = 4
 TM = 4
 
+# NOTE: this is from testgrad
+# 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"
+
+  # append the reduce shape to each of the views
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]
+
+  # 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), len(view.shape) + len(r.axis_arg)))))
+
+pm = PatternMatcher([
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])
+
+def top_spec_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  sink = c.schedule()[-1].ast
+  L = 16
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
+  sink = graph_rewrite(sink, view_left+pm)
+  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
+  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
+
 def hl_spec_kernel3():
   nbIterWaveM = 2
   nbIterWaveN = 2
 
   # define buffers
   # TODO: remove these views once the defines have a shape
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N))).permute((1,0))
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N)))
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
   As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK*BM,)))
   Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK*BN,)))
   A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
   B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))
-  acc = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM * nbIterWaveN * TN, AddrSpace.REG), arg=2) \
-    .view(ShapeTracker.from_shape((nbIterWaveM * TM * nbIterWaveN * TN,)))
 
   # shape buffers. TODO: permutes
   full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
@@ -42,32 +78,24 @@ def hl_spec_kernel3():
   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)
 
-  acc = acc.reshape((1, nbIterWaveM, 1, TM, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape[:-2]+(1,1))
-
-  #out = a.load() * b.load()
-  #out = As.load(As.store(a.load())) * Bs.load(Bs.store(b.load()))
-  out = A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))
-
-  #out = out.r(Ops.ADD, (8, 9))
-  out = UOp(Ops.REDUCE_INTO, out.dtype, (acc, out), Ops.ADD)
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
+  out = (As.load(As.store(a.load())) * Bs.load(Bs.store(b.load()))).r(Ops.ADD, (8, 9))
+  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))
 
   axis_types = (
     AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
     AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
     AxisType.REDUCE, AxisType.UNROLL)
 
-  from tinygrad.opt.kernel import axis_colors
-  shape = '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
-  print(shape)
-  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+shape, axis_types=axis_types))
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
   sink = graph_rewrite(sink, merge_views)
   return sink
 
-def hand_spec_kernel3():
-  BLOCK_SIZE = 256
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256
 
   nbWaves = BLOCK_SIZE // 32
-  WN = 64
+  WN = 128 if kernel5 else 64
   WM = BN * BM // nbWaves // WN
 
   nbWaveX = BN // WN
@@ -106,14 +134,15 @@ def hand_spec_kernel3():
   blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
   blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
 
-  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)
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
 
   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)
 
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0)
+  BM_As_stride = (BM+4) if kernel5 else BM
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
   Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
 
   c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
@@ -121,51 +150,131 @@ def hand_spec_kernel3():
   i = UOp.range(dtypes.int, c_regs.dtype.size, 16)
   init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
 
-  kId_range = UOp.range(dtypes.int, N//BK, 0)
-  kId = kId_range*BK
+  if kernel4:
+    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
+    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
 
-  # 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)
+    # initial load from globals into locals (0)
+    kId = 0
 
-  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)
+    # load from globals into locals
+    i = UOp.range(dtypes.int, nbReadsB, 0)
+    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)
 
-  barrier = UOp(Ops.BARRIER, src=(As_store, Bs_store))
+    i = UOp.range(dtypes.int, nbReadsA, 1)
+    index_x = rAIdx + kId
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
 
-  k = UOp.range(dtypes.int, BK, 3)
+    # iterate over the middle chunk
+    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
+    kId = kId_range*BK
 
-  # 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)
+    barrier = UOp.barrier(As_store, Bs_store)
 
-  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)
+    # load from globals into registers (next round)
+    i = UOp.range(dtypes.int, nbReadsB, 3)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
 
-  # do the GEMM math
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-  yt = UOp.range(dtypes.int, TM, 9)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 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)
+    i = UOp.range(dtypes.int, nbReadsA, 4)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
+
+    def inner_loop(first_range, inp_dep=()):
+      # inner unroll
+      k = UOp.range(dtypes.int, BK, first_range+0)
+
+      # load from locals into registers
+      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
+      i = UOp.range(dtypes.int, TN, first_range+2)
+      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
+      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
+
+      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
+      i = UOp.range(dtypes.int, TM, first_range+4)
+      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
+      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
+
+      # do the GEMM math
+      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
+      yt = UOp.range(dtypes.int, TM, first_range+6)
+      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
+      xt = UOp.range(dtypes.int, TN, first_range+8)
+      x = iterWaveN * TN + xt
+      y = iterWaveM * TM + yt
+      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+      # sketchy, this should end the kId_range but it doesn't
+      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)
+      return sink
+
+    # TODO: kId_range should endrange after a barrier
+    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
+
+    # load from registers into locals
+    i = UOp.range(dtypes.int, nbReadsB, 14)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
+
+    i = UOp.range(dtypes.int, nbReadsA, 15)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
+
+    # final iteration without the copy
+    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
+  else:
+    kId_range = UOp.range(dtypes.int, N//BK, 0)
+    kId = kId_range*BK
+
+    # 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)
+
+    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_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
+
+    barrier = UOp.barrier(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_As_stride + index].load(barrier), iterWave, i)
+
+    # do the GEMM math
+    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
+    yt = UOp.range(dtypes.int, TM, 9)
+    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
+    xt = UOp.range(dtypes.int, TN, 12)
+    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)
-  yt = UOp.range(dtypes.int, TM, 13)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 14)
-  xt = UOp.range(dtypes.int, TN, 15)
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
+  yt = UOp.range(dtypes.int, TM, 1001)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
+  xt = UOp.range(dtypes.int, TN, 1003)
   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
@@ -175,10 +284,13 @@ def hand_spec_kernel3():
   return sink.sink(arg=KernelInfo(name="tinygemm"))
 
 if __name__ == "__main__":
-  hprg = hl_spec_kernel3() if getenv("HL") else hand_spec_kernel3()
+  HL = getenv("HL")
+  if HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
   prg = get_program(hprg, Device.default.renderer)
   print(prg.src)
-  exit(0)
+  if getenv("SRC"): exit(0)
   hrunner = CompiledRunner(prg)
 
   a = Tensor.randn(N, N).realize()
@@ -190,7 +302,8 @@ if __name__ == "__main__":
     for _ in range(run_count): tc = (a@b).realize()
 
   GlobalCounters.reset()
-  ei = ExecItem(hrunner, [a.uop.buffer, b.uop.buffer, hc.uop.buffer])
+  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
+  ei = ExecItem(hrunner, buffers)
   with Context(DEBUG=2):
     for _ in range(run_count): ei.run(wait=True)
   err = (hc-tc).square().mean().item()

test/test_define_reg.py

@@ -28,15 +28,5 @@ class TestDefineReg(unittest.TestCase):
   @unittest.skipIf(getenv("PTX"), "ptx needs regs to be unrolled")
   def test_simple_loop(self): self.test_simple(AxisType.LOOP)
 
-  def test_reduce_into(self):
-    N = 16
-    b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N*N), arg=0).view(ShapeTracker.from_shape((N,N,1)))
-    a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N*N), arg=1).view(ShapeTracker.from_shape((N,N,N)))
-    a_reg = UOp(Ops.DEFINE_REG, dtypes.float.ptr(N, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((N,N,1), (0,1,0)))
-    ret = UOp(Ops.REDUCE_INTO, dtypes.float, src=(a_reg, a.load()), arg=Ops.ADD)
-    sink = b.store(ret).sink(arg=KernelInfo(name="regcopy"))
-    prg = get_program(sink, Device.default.renderer)
-    print(prg.src)
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/linearize.py

@@ -97,7 +97,7 @@ class BlockContext:
 
 # ***** make blocks *****
 
-DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}
+DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_REG, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST}
 
 def add_blockends(base_block:UOp, new_ctx:tuple[UOp, ...], current_ctx:tuple[UOp, ...], cnt:int=1) -> UOp:
   ends_to_add = [z for z in new_ctx if z not in current_ctx]

tinygrad/codegen/lowerer.py

@@ -2,34 +2,20 @@
 from dataclasses import dataclass
 from typing import cast
 from tinygrad.dtype import dtypes, PtrDType, AddrSpace
-from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType, graph_rewrite
+from tinygrad.uop.ops import KernelInfo, UOp, Ops, PatternMatcher, UPat, sint_to_uop, AxisType
 from tinygrad.helpers import prod, partition, flatten
 
 # ***** indexing *****
 
 @dataclass
 class IndexContext:
-  axis_types: list[AxisType]
   idxs: list[UOp]
-  start: int = 0
-  #ridxs: list[UOp]
-
-def shape_to_idx(s, axis_types, start=0, allow_unroll=False):
-  idxs = []
-  for i, (s, at) in enumerate(zip(s, axis_types)):
-    if at in (AxisType.UPCAST, AxisType.UNROLL) and allow_unroll:
-      assert isinstance(s, int), "needs to be int to upcast/unroll"
-      idxs.append(UOp(Ops.UNROLL, dtypes.int, (UOp.const(dtypes.int.vec(s), tuple(range(s))),), ((i,s),)))
-    else:
-      # all others are RANGES
-      idxs.append(UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), start+i))
-  return idxs
+  ridxs: list[UOp]
 
 def get_index(ast:UOp) -> IndexContext:
   axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else ()
   if len(ast.full_shape) != len(axis_types): axis_types = (AxisType.LOOP,)*len(ast.full_shape)
 
-  """
   # indexes
   idxs = []
   for i, (s, at) in enumerate(zip(ast.full_shape, axis_types)):
@@ -45,85 +31,35 @@ def get_index(ast:UOp) -> IndexContext:
   for i, (s, at) in enumerate(zip(ast.full_shape, axis_types)):
     if at == AxisType.GROUP_REDUCE:
       ridxs[i] = UOp(Ops.RANGE, dtypes.int, (sint_to_uop(s),), 1000+i)
-  """
 
-  return IndexContext(axis_types, []) #idxs, ridxs)
+  return IndexContext(idxs, ridxs)
 
 # ***** lowering (given index) *****
 
-def lower_reduce_into(ctx: IndexContext, x: UOp):
-  # x.src[1] is the big
-  axis_arg = [i for i,(s0,s1) in enumerate(zip(x.src[0].shape, x.src[1].shape)) if s0 != s1]
-  reduce_idxs = shape_to_idx([s for i,s in enumerate(x.src[1].shape) if i in axis_arg],
-                             [s for i,s in enumerate(ctx.axis_types) if i in axis_arg], ctx.start)
-  print("reduce into", [x.arg for x in reduce_idxs])
-  new_idxs = shape_to_idx(x.src[0].shape[:-len(reduce_idxs)], ctx.axis_types[:-len(reduce_idxs)], ctx.start+len(reduce_idxs)) + reduce_idxs
-
-  idx, valid = x.src[0].arg.to_indexed_uops(new_idxs)
-  used_idxs = [x for x in UOp.sink(idx, valid).toposort() if x in new_idxs]
-
-  real_new_idxs = []
-  for i in range(len(x.src[0].shape)):
-    if new_idxs[i] in used_idxs or new_idxs[i] in reduce_idxs or len(ctx.idxs) <= i: real_new_idxs.append(new_idxs[i])
-    else: real_new_idxs.append(ctx.idxs[i])
-  non_replaced = [x for x in real_new_idxs if x not in ctx.idxs]
-
-  acc = x.src[0].load(*reduce_idxs).alu(x.arg, x.src[1])
-
-  lc = IndexContext(ctx.axis_types, tuple(real_new_idxs), ctx.start+len(new_idxs))
-  from tinygrad.codegen.lowerer import pm_lowerer  # TODO: better way to do this?
-  ret = graph_rewrite(acc, pm_lowerer, lc, name="subreduce", bottom_up=True)
-  ctx.start = lc.start
-
-  red = x.src[0].src[0].index(idx, valid).store(ret, *used_idxs, *reduce_idxs)
-  return x.src[0].load(red)
-
 def lower_reduce_axis(ctx: IndexContext, x: UOp):
-  new_idx = shape_to_idx([s for i,s in enumerate(x.src[0].shape) if i in x.axis_arg],
-                         [s for i,s in enumerate(ctx.axis_types) if i in x.axis_arg], ctx.start)
-  full_new_idx = list(ctx.idxs)
-  for i,s in zip(x.axis_arg, new_idx):
-    #assert len(full_new_idx) == i, f"len(full_new_idx) = {len(full_new_idx)}, trying to place {i}"
-    #full_new_idx.append(s)
-    full_new_idx[i] = s
-
-  lc = IndexContext(ctx.axis_types, tuple(full_new_idx), ctx.start+len(new_idx))
-  from tinygrad.codegen.lowerer import pm_lowerer  # TODO: better way to do this?
-  ret = graph_rewrite(x.src[0], pm_lowerer, lc, name="subreduce", bottom_up=True)
-  ctx.start = lc.start
-
   # NOTE: always using ridxs is fine here
-  reduce_range, reduce_expand = partition([full_new_idx[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
+  reduce_range, reduce_expand = partition([ctx.ridxs[i] for i in x.axis_arg], lambda y: y.op is Ops.RANGE)
   assert all(x.op is Ops.UNROLL for x in reduce_expand), f"not all UNROLLS in {reduce_expand} for {x.axis_arg}"
+  ret = x.src[0]
   if len(contract_axis:=flatten(x.arg for x in reduce_expand)):
     ret = UOp(Ops.CONTRACT, x.dtype.vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis))
   # REDUCE supports both "horizontal" reduction and range reduction. the horizontal elements are taken in the nearest group
   return UOp(Ops.REDUCE, x.dtype, (ret,)+tuple(reduce_range), x.arg[0])
 
 def lower_load(ctx: IndexContext, x: UOp, buf: UOp):
-  print("lower_load", [x.arg for x in ctx.idxs])
-  idx, valid = x.st_arg.to_indexed_uops(ctx.idxs)
-  return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + x.src[1:])
-
-  #barrier = tuple([y.barrier() if buf.op is Ops.DEFINE_LOCAL else y for y in x.src[1:]])
+  idx, valid = x.st_arg.to_indexed_uops(ctx.ridxs if buf.op is Ops.DEFINE_LOCAL else ctx.idxs)
+  barrier = tuple([y.barrier() if buf.op is Ops.DEFINE_LOCAL else y for y in x.src[1:]])
+  return UOp(Ops.LOAD, x.dtype, (buf.index(idx, valid),) + barrier)
 
 def lower_store(ctx: IndexContext, x: UOp, buf: UOp):
-  assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}"
-  new_idxs = shape_to_idx(x.src[0].shape, ctx.axis_types, ctx.start)
-  idx, valid = x.st_arg.to_indexed_uops(new_idxs)
-  used_idxs = [x for x in UOp.sink(idx, valid).toposort() if x in new_idxs]
-  real_new_idxs = []
-  for i in range(len(x.src[0].shape)):
-    if new_idxs[i] in used_idxs or len(ctx.idxs) <= i:
-      real_new_idxs.append(new_idxs[i])
-    else:
-      real_new_idxs.append(ctx.idxs[i])
-  print("got", len(real_new_idxs), len(used_idxs))
-  lc = IndexContext(ctx.axis_types, tuple(real_new_idxs), ctx.start+len(new_idxs))
-  from tinygrad.codegen.lowerer import pm_lowerer  # TODO: better way to do this?
-  stored = graph_rewrite(x.src[1], pm_lowerer, lc, name="substore", bottom_up=True)
-  ctx.start = lc.start
-  return buf.index(idx, valid).store(stored, *[x for x in used_idxs if x.op is Ops.RANGE])
+  # TODO: reenable after REDUCE_AXIS is fixed
+  #assert x.src[1].shape == x.src[0].shape, f"shape mismatch on store {x.src[1].shape} != {x.src[0].shape}"
+  idx, valid = x.st_arg.to_indexed_uops(ctx.idxs)
+  if cast(PtrDType, buf.dtype).addrspace == AddrSpace.GLOBAL:
+    # NOTE: only store the local reduceop in the threads that are actually doing the reduce
+    for oidx, ridx in zip(ctx.idxs, ctx.ridxs):
+      if oidx is not ridx: valid = valid * oidx.eq(0)
+  return buf.index(idx, valid).store(x.src[1], *[x for x in UOp.sink(idx, valid).toposort() if x.op is Ops.RANGE])
 
 def lower_const(ctx:IndexContext, view:UOp, c:UOp):
   if all(x.mask is None for x in view.arg.views): return c
@@ -138,7 +74,6 @@ pm_lowerer = PatternMatcher([
   (UPat(Ops.VALID, src=(UPat(Ops.VIEW, name="v"),)).where(UPat.cvar("c"), UPat(Ops.CONST, arg=0)), lambda c,v: c.replace(src=()).view(v.arg)),
 
   # reduce/view_const
-  (UPat(Ops.REDUCE_INTO, name="x"), lower_reduce_into),
   (UPat(Ops.REDUCE_AXIS, name="x"), lower_reduce_axis),
   (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="c"),), name="view"), lower_const),
   # rewrite LOAD/STORE VIEW to LOAD/STORE with indexed

tinygrad/opt/kernel.py

@@ -448,9 +448,7 @@ class Kernel:
       ret = op.replace(src=tuple(fixup_ast(x) for x in op.src)) # noqa: F821
       if op.op in GroupOp.Buffer and op in self.bufs:
         st = self.sts[self.bufs.index(op)]
-        # NOTE: if CONST got masked after applying opts, we create a new VALID
-        if op.op is Ops.CONST and any(v.mask is not None for v in st.views): return op.view(st).valid()
-        # otherwise we just replace the VIEW source
+        # replace the VIEW source
         return ret.replace(src=(ret.src[0].replace(arg=st),)+ret.src[1:])
       if op.op is Ops.SINK:
         # NOTE: should group_for_reduces be added to the local_dims?

tinygrad/renderer/llvmir.py

@@ -48,8 +48,9 @@ def render_wmma_amx(ctx, wmma: UOp) -> str:
 def render_wmma_amd(ctx, wmma: UOp, arch: str) -> str:
   dt_map = {dtypes.half: "f16", dtypes.float: "f32", dtypes.bfloat16: "bf16", dtypes.ushort: "bf16"}
   # https://github.com/llvm/llvm-project/blob/main/clang/test/CodeGenOpenCL/builtins-amdgcn-mfma.cl
-  if arch.split(":")[0] == "gfx942": return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
-    f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
+  if arch.split(":")[0] in {"gfx942", "gfx950"}:
+    return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.mfma.{dt_map[wmma.src[-1].dtype.scalar()]}" + \
+           f".16x16x16{dt_map[wmma.src[0].dtype.scalar()]}(" + ", ".join([f"{ldt(w.dtype)} {ctx[w]}" for w in wmma.src]) + ", i32 0, i32 0, i32 0)"
   # https://github.com/llvm/llvm-project/blob/main/llvm/test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.wmma_32.ll
   # example: %wmma0 = call <8 x float> @llvm.amdgcn.wmma.f32.16x16x16.f16(<16 x half> %v99,<16 x half> %v100,<8 x float> %v101)
   return f"  {ctx[wmma]} = call {ldt(wmma.dtype)} @llvm.amdgcn.wmma.{dt_map[wmma.src[-1].dtype.scalar()]}.16x16x16." + \

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
+from tinygrad.helpers import capstone_flatdump, getenv, from_mv, to_mv, OSX, mv_address, wait_cond, cpu_profile
 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
@@ -92,8 +92,10 @@ class CPUAllocator(HCQAllocatorBase):
     return HCQBuffer(va:=addr, sz:=size, 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))
-  def _copyout(self, dest:memoryview, src): ctypes.memmove(from_mv(dest), src.va_addr, len(dest))
+  def _copyin(self, dest, src:memoryview):
+    with cpu_profile('TINY -> CPU', self.dev.device, is_copy=True): ctypes.memmove(dest.va_addr, from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src):
+    with cpu_profile('CPU -> TINY', self.dev.device, is_copy=True): ctypes.memmove(from_mv(dest), src.va_addr, len(dest))
   def _map(self, buf:HCQBuffer):
     if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
 

tinygrad/runtime/ops_metal.py

@@ -226,6 +226,6 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
   def _as_buffer(self, src:MetalBuffer) -> memoryview:
     self.dev.synchronize()
     return to_mv(cast(int, msg("contents", objc_id)(src.buf).value), src.size + src.offset)[src.offset:]
-  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "CPU -> METAL")
-  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> CPU")
+  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL")
+  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY")
   def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset)

tinygrad/runtime/ops_remote.py

@@ -16,6 +16,7 @@ from tinygrad.helpers import getenv, DEBUG, fromimport, unwrap, LazySeq, Timing
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner, ExecItem, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferXfer
 from tinygrad.device import Compiled, Buffer, Allocator, Compiler, Device, BufferSpec
+from tinygrad.runtime.support.ib import IBCtx, IBConn, SGE
 
 # ***** API *****
 
@@ -35,6 +36,7 @@ class RemoteProperties:
   offset_supported: bool
   graph_supported: bool
   graph_supports_multi: bool
+  ib_gid: bytes|None
 
 @dataclass(frozen=True)
 class GetProperties(RemoteRequest): pass
@@ -45,12 +47,18 @@ class Event(RemoteRequest): event_session: SessionKey; event: int # noqa: E702
 @dataclass(frozen=True)
 class Wait(RemoteRequest): event: int
 
+@dataclass(frozen=True)
+class IBConnect(RemoteRequest): host: str; gid: bytes; qp_num: int # noqa: E702
+
 @dataclass(frozen=True)
 class BufferAlloc(RemoteRequest): buffer_num: int; size: int; options: BufferSpec # noqa: E702
 
 @dataclass(frozen=True)
 class BufferOffset(RemoteRequest): buffer_num: int; size: int; offset: int; sbuffer_num: int # noqa: E702
 
+@dataclass(frozen=True)
+class BufferIOVAS(RemoteRequest): buffer_nums: list[tuple[SessionKey, int]] # noqa: E702
+
 @dataclass(frozen=True)
 class BufferFree(RemoteRequest): buffer_num: int # noqa: E702
 
@@ -111,9 +119,9 @@ class GraphExec(RemoteRequest):
   wait: bool
 
 # for safe deserialization
-eval_globals = {x.__name__:x for x in [SessionKey, SessionFree, RemoteProperties, GetProperties, Event, Wait, BufferAlloc, BufferOffset, BufferFree,
-                                       CopyIn, CopyOut, Transfer, BatchTransfer, ProgramAlloc, ProgramFree, ProgramExec, GraphComputeItem, GraphAlloc,
-                                       GraphFree, GraphExec, BufferSpec, UOp, Ops, dtypes]}
+eval_globals = {x.__name__:x for x in [SessionKey, SessionFree, RemoteProperties, GetProperties, Event, Wait, BufferAlloc, BufferOffset, BufferIOVAS,
+                                       BufferFree, CopyIn, CopyOut, Transfer, BatchTransfer, IBConnect, ProgramAlloc, ProgramFree, ProgramExec,
+                                       GraphComputeItem, GraphAlloc, GraphFree, GraphExec, BufferSpec, UOp, Ops, dtypes]}
 attribute_whitelist: dict[Any, set[str]] = {dtypes: {*DTYPES_DICT.keys(), 'imagef', 'imageh'}, Ops: {x.name for x in Ops}}
 eval_fxns = {ast.Constant: lambda x: x.value, ast.Tuple: lambda x: tuple(map(safe_eval, x.elts)), ast.List: lambda x: list(map(safe_eval, x.elts)),
   ast.Dict: lambda x: {safe_eval(k):safe_eval(v) for k,v in zip(x.keys, x.values)},
@@ -160,6 +168,12 @@ class RemoteHandler:
     self.base_device = base_device
     self.sessions: defaultdict[SessionKey, RemoteSession] = defaultdict(RemoteSession)
 
+    try: self.ib_ctx: IBCtx|None = IBCtx(getenv("IB_DEV", 0))
+    except (IndexError, AttributeError): self.ib_ctx = None
+    self.ib_lock = asyncio.Lock()
+    self.ib_conns: dict[str, IBConn|None] = {}
+    self.iova_cache: dict[tuple[SessionKey, int], tuple[int, int, int]] = {}
+
   async def __call__(self, reader:asyncio.StreamReader, writer:asyncio.StreamWriter):
     while (req_hdr:=(await reader.readline()).decode().strip()):
       req_method, req_path, _ = req_hdr.split(' ')
@@ -171,6 +185,30 @@ class RemoteHandler:
       res_status, res_body = await self.handle(req_method, req_path, req_body)
       writer.write(f"HTTP/1.1 {res_status.value} {res_status.phrase}\r\nContent-Length: {len(res_body)}\r\n\r\n".encode() + res_body)
 
+  async def ib_connect(self, ssession:SessionKey, dsession:SessionKey) -> IBConn|None:
+    if self.ib_ctx is None: return None
+    await self.ib_lock.acquire()
+    conn = RemoteConnection(dsession.host)
+    if dsession.host not in self.ib_conns:
+      props = safe_eval(ast.parse(conn.q(GetProperties(session=dsession), wait=True), mode="eval").body)
+      if props.ib_gid is not None:
+        self.ib_conns[dsession.host] = ib_conn = IBConn(self.ib_ctx)
+        ibxc_ret = conn.q(IBConnect(ssession.host, ib_conn.gid, ib_conn.qp_num, session=dsession), wait=True)
+        ib_conn.connect(*struct.unpack('<16sQ', ibxc_ret))
+      else:
+        self.ib_conns[dsession.host] = None
+    self.ib_lock.release()
+    return self.ib_conns[dsession.host]
+
+  async def get_iovas(self, bufs:list[tuple[SessionKey, int]]) -> list[tuple[int, int, int]]:
+    await self.ib_lock.acquire()
+    if (rbufs:=[buf for buf in bufs if buf not in self.iova_cache]):
+      conn = RemoteConnection(rbufs[0][0].host)
+      resp = await conn.aq(BufferIOVAS(rbufs, session=rbufs[0][0]), wait=True)
+      self.iova_cache.update({rbuf: struct.unpack('<QQQ', resp[i*24:(i+1)*24]) for i,rbuf in enumerate(rbufs)})
+    self.ib_lock.release()
+    return [self.iova_cache[buf] for buf in bufs]
+
   async def handle(self, method:str, path:str, body:bytes) -> tuple[http.HTTPStatus, bytes]:
     status, ret = http.HTTPStatus.OK, b""
     if path == "/batch" and method == "POST":
@@ -188,6 +226,7 @@ class RemoteHandler:
             rp = RemoteProperties(
               real_device=dev.device, renderer=(cls.__module__, cls.__name__, args), offset_supported=hasattr(dev.allocator, '_offset'),
               graph_supported=graph_cls is not None, graph_supports_multi=graph_cls is not None and issubclass(graph_cls, MultiGraphRunner),
+              ib_gid=bytes(self.ib_ctx.gid_attr.raw) if self.ib_ctx is not None else None,
             )
             ret = repr(rp).encode()
           case Event():
@@ -200,9 +239,19 @@ class RemoteHandler:
           case Wait():
             assert await session.events[c.event].wait()
             del session.events[c.event] # do not leak memory
+          case IBConnect():
+            self.ib_conns[c.host] = ibc = IBConn(unwrap(self.ib_ctx))
+            ibc.connect(c.gid, c.qp_num)
+            ret = struct.pack('<16sQ', ibc.gid, ibc.qp_num)
           case BufferAlloc():
             assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already allocated"
             session.buffers[c.buffer_num] = Buffer(dev.device, c.size, dtypes.uint8, options=c.options, preallocate=True)
+          case BufferIOVAS():
+            rets = []
+            for buffer_session,buffer_num in c.buffer_nums:
+              iova, mr = unwrap(self.ib_ctx).reg(buf:=self.sessions[buffer_session].buffers[buffer_num])
+              rets.append(struct.pack("<QQQ", iova, mr.contents.rkey, buf.nbytes))
+            ret = b"".join(rets)
           case BufferOffset():
             assert c.buffer_num not in session.buffers, f"buffer {c.buffer_num} already exists"
             session.buffers[c.buffer_num] = session.buffers[c.sbuffer_num].view(c.size, dtypes.uint8, c.offset).allocate()
@@ -220,16 +269,29 @@ class RemoteHandler:
                 sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
                 dbuf.copyin(data)
             else:
-              conn = RemoteConnection(c.dsession.host)
+              conn, ib_conn = RemoteConnection(c.dsession.host), await self.ib_connect(unwrap(c.session), c.dsession)
               sbuf = session.buffers[c.buffer_num]
-              sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-              conn.q(CopyIn(c.dbuffer_num, conn.req.h(data), session=c.dsession), wait=True)
+              if ib_conn is not None:
+                src_iova, src_mr = unwrap(self.ib_ctx).reg(sbuf)
+                dst_iova, dst_key, dst_size = (await self.get_iovas([(c.dsession, c.dbuffer_num)]))[0]
+                assert sbuf.nbytes == dst_size, f"{sbuf.nbytes} != {dst_size}"
+                for d in Device._opened_devices: Device[d].synchronize()
+                ib_conn.rdma_write([SGE(dst_iova, dst_key, src_iova, src_mr.contents.lkey, dst_size)])
+              else:
+                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
+                await conn.aq(CopyIn(c.dbuffer_num, conn.req.h(data), session=c.dsession), wait=True)
           case BatchTransfer():
-            conn = RemoteConnection(c.dbuffer_nums[0][0].host)
-            for (sbuf_session,sbuf_num),(dbuf_session,dbuf_num) in zip(c.sbuffer_nums, c.dbuffer_nums):
-              sbuf = self.sessions[sbuf_session].buffers[sbuf_num]
-              sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
-              await conn.aq(CopyIn(dbuf_num, conn.req.h(data), session=dbuf_session), wait=True)
+            conn, ib_conn = RemoteConnection(c.dbuffer_nums[0][0].host), await self.ib_connect(c.sbuffer_nums[0][0], c.dbuffer_nums[0][0])
+            if ib_conn is not None:
+              sbufs = [unwrap(self.ib_ctx).reg(self.sessions[s].buffers[bi]) for s,bi in c.sbuffer_nums]
+              dbufs = await self.get_iovas(c.dbuffer_nums)
+              for d in Device._opened_devices: Device[d].synchronize()
+              ib_conn.rdma_write([SGE(di, dk, si, sm.contents.lkey, ds) for (di,dk,ds),(si,sm) in zip(dbufs, sbufs)])
+            else:
+              for (sbuf_session,sbuf_num),(dbuf_session,dbuf_num) in zip(c.sbuffer_nums, c.dbuffer_nums):
+                sbuf = self.sessions[sbuf_session].buffers[sbuf_num]
+                sbuf.copyout(data:=memoryview(bytearray(sbuf.nbytes)))
+                await conn.aq(CopyIn(dbuf_num, conn.req.h(data), session=dbuf_session), wait=True)
           case ProgramAlloc():
             lib = dev.compiler.compile_cached(req._h[c.datahash].decode())
             session.programs[(c.name, c.datahash)] = dev.runtime(c.name, lib)

tinygrad/runtime/support/hcq.py

@@ -496,7 +496,7 @@ class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
 class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
   def _copyin(self, dest:HCQBuffer, src:memoryview):
     assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"CPU -> {self.dev.device}", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE):
       for i in range(0, src.nbytes, self.b[0].size):
         self.b_next = (self.b_next + 1) % len(self.b)
         self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
@@ -528,7 +528,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
     self.dev.synchronize()
 
     assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> CPU", enabled=PROFILE):
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE):
       for i in range(0, dest.nbytes, cp_size:=(self.max_copyout_size or self.b[0].size)):
         self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
                                   .copy(self.b[0].va_addr, src.va_addr+i, lsize:=min(cp_size, dest.nbytes-i)) \

tinygrad/runtime/support/ib.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+import resource, ctypes, weakref, functools, itertools, tinygrad.runtime.autogen.ib as ib
+from typing import Iterator
+from dataclasses import dataclass
+from weakref import WeakKeyDictionary
+from tinygrad.device import Buffer, DMACPURef, DMAFdRef
+from tinygrad.helpers import getenv, round_up, DEBUG
+
+DEFAULT_PORT, DEFAULT_GID = getenv("DEFAULT_PORT", 1), getenv("DEFAULT_GID", 3) # DEFAULT_GID=0 for RXE
+IOVA_ALIGN = resource.getpagesize()
+
+def checkz(x, ret=None):
+  assert x == 0, f'{x} != 0 (errno {ctypes.get_errno()})'
+  return ret
+
+@dataclass(frozen=True)
+class SGE:
+  dst_iova: int
+  dst_key: int
+  src_iova: int
+  src_key: int
+  size: int
+
+class IBCtx:
+  def __init__(self, idx:int):
+    # Open the device (aka Host Channel Adapter in ib-speak)
+    devs = ib.ibv_get_device_list(ctypes.byref(ndevs:=ctypes.c_int32()))
+    if idx >= ndevs.value: raise IndexError(f"{idx} > {ndevs.value}")
+    self.ctx = ib.ibv_open_device(devs[idx])
+    ib.ibv_free_device_list(devs)
+
+    # HACK: remove this (and all usage of `ctx.contents.ops`) when clang2py can deal with `static inline` wrapper-functions
+    self.vctx = ctypes.cast(ctypes.addressof(self.ctx.contents) - ib.struct_verbs_context.context.offset, ctypes.POINTER(ib.struct_verbs_context))
+
+    # Get attributes. Something like port_attr.max_msg_sz sound like it might requre taking the min of host's and remote's attributes if they differ
+    self.device_attr = checkz(ib.ibv_query_device(self.ctx, ctypes.byref(da:=ib.struct_ibv_device_attr())), da)
+    self.port_attr = checkz(self.vctx.contents.query_port(self.ctx, DEFAULT_PORT, ctypes.byref(pa:=ib.struct_ibv_port_attr()), ctypes.sizeof(pa)), pa)
+    self.gid_attr = checkz(ib.ibv_query_gid(self.ctx, DEFAULT_PORT, DEFAULT_GID, ctypes.byref(ga:=ib.union_ibv_gid())), ga)
+
+    # Allocate protection domain
+    self.pd = ib.ibv_alloc_pd(self.ctx)
+    self.next_iova: int = IOVA_ALIGN # don't start at zero (nullptr)
+
+    # weakref(buf) => (iova, mr, mr_dealloc). mr_dealloc is kept here to avoid double freeing mrs that are deallocated in __del__
+    self.mrs: WeakKeyDictionary[Buffer, tuple[int, ctypes._Pointer[ib.struct_ibv_mr], weakref.finalize]] = WeakKeyDictionary()
+
+    # Default soft fd limit is 1024, which is not enough, set soft to hard (maximum allowed by the os)
+    IBCtx.rlimit_fix()
+
+  def __del__(self):
+    # must deallocate all mrs in protection domain before deallocating the protection domain
+    if hasattr(self, "mrs"): [fin() for _,_,fin in self.mrs.values()]
+    if hasattr(self, "pd"): ib.ibv_dealloc_pd(self.pd)
+    if hasattr(self, "ctx"): ib.ibv_close_device(self.ctx)
+
+  @functools.cache # run once
+  @staticmethod
+  def rlimit_fix():
+    soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+    resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
+    if DEBUG>=2: print(f"IB: Increased fd limit from {soft} to {hard}")
+
+  def alloc_iova(self, size:int, required_offset:int):
+    iova = round_up(self.next_iova - required_offset, IOVA_ALIGN) + required_offset
+    self.next_iova = iova + size
+    return iova
+
+  def reg(self, buf:Buffer) -> tuple[int, ctypes._Pointer[ib.struct_ibv_mr]]:
+    buf = buf.base
+    if buf not in self.mrs:
+      if buf.nbytes > self.device_attr.max_mr_size: raise RuntimeError(f"Buffer too big: {buf.nbytes:#x} > {self.device_attr.max_mr_size:#x}")
+      if len(self.mrs) >= self.device_attr.max_mr: raise RuntimeError(f"Out of memory region cap: {len(self.mrs)} >= {self.device_attr.max_mr}")
+      # Local read is implied (but still have to create the memory region, except for short sends/writes with IBV_SEND_INLINE that are inlined by cpu)
+      mr_flags = ib.IBV_ACCESS_LOCAL_WRITE | ib.IBV_ACCESS_REMOTE_READ | ib.IBV_ACCESS_REMOTE_WRITE
+      match (dmaref:=buf.as_dmaref()):
+        case DMACPURef():
+          iova = self.alloc_iova(dmaref.size, dmaref.addr % IOVA_ALIGN)
+          mr = ib.ibv_reg_mr_iova2(self.pd, ctypes.c_void_p(dmaref.addr), dmaref.size, iova, mr_flags)
+        case DMAFdRef():
+          iova = self.alloc_iova(dmaref.size, dmaref.offset % IOVA_ALIGN)
+          mr = ib.ibv_reg_dmabuf_mr(self.pd, dmaref.offset, dmaref.size, iova, dmaref.fd, mr_flags)
+        case _: raise RuntimeError(f"Unknown type of dma ref: {dmaref}")
+      if not mr: raise RuntimeError(f"Couldn't register memory region for {buf} {dmaref} (errno={ctypes.get_errno()})")
+      self.mrs[buf] = (iova, mr, weakref.finalize(buf, ib.ibv_dereg_mr, mr))
+    return self.mrs[buf][0:2]
+
+class IBConn:
+  def __init__(self, ctx:IBCtx):
+    self.ctx = ctx
+
+    # Create Completion Channel. It is a file descriptor that kernel sends notifications through, not a thing in infiniband spec, just linux-ism
+    self.comp_channel = ib.ibv_create_comp_channel(self.ctx.ctx)
+    # Create Completion Queue. When a Work Request with signaled flag is completed a Completion Queue Entry is pushed onto this queue
+    self.cq = ib.ibv_create_cq(self.ctx.ctx, _capacity:=256, _cq_context:=None, self.comp_channel, _comp_vector:=0)
+    self.pending_wrids: set[int] = set()
+    self.wrid_num: Iterator[int] = itertools.count(0) # wc_id is uint64, this will never overflow
+
+    # Create Queue Pair. It's the closest thing to a socket in infiniband with QP num being the closest thing to a port, except it's allocated by hca
+    qp_init_attrs_cap = ib.struct_ibv_qp_cap(max_send_wr=1024, max_recv_wr=64, max_send_sge=8, max_recv_sge=8, max_inline_data=64)
+    qp_init_attrs = ib.struct_ibv_qp_init_attr(send_cq=self.cq, recv_cq=self.cq, cap=qp_init_attrs_cap, qp_type=ib.IBV_QPT_RC) # Reliable Connection
+    self.qp = ib.ibv_create_qp(self.ctx.pd, ctypes.byref(qp_init_attrs))
+    self.qp_cap = qp_init_attrs.cap
+
+    # The most important thing about QPs is their state, when a new QP is created it's in the RESET state, before it can be properly used it has to go
+    # through Init, Ready To Receive, Ready To Send. A good docs on QP state machine: https://www.rdmamojo.com/2012/05/05/qp-state-machine/
+
+    # INIT
+    qp_access_flags = ib.IBV_ACCESS_REMOTE_WRITE | ib.IBV_ACCESS_REMOTE_READ
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_INIT, port_num=DEFAULT_PORT, qp_access_flags=qp_access_flags)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PORT | ib.IBV_QP_ACCESS_FLAGS | ib.IBV_QP_PKEY_INDEX))
+
+    self.gid, self.qp_num = bytes(self.ctx.gid_attr.raw), self.qp.contents.qp_num
+
+  # Exchange GID and QP num with remote. At least in RoCEv2 gid can be guessed from remote's ip, QP num can't.
+
+  def connect(self, remote_gid:bytes, remote_qp_num:int):
+    # RTR
+    qp_ah_attr_grh = ib.struct_ibv_global_route(hop_limit=1, dgid=ib.union_ibv_gid(raw=(ctypes.c_ubyte * 16)(*remote_gid)), sgid_index=DEFAULT_GID)
+    qp_ah_attr = ib.struct_ibv_ah_attr(is_global=1, port_num=DEFAULT_PORT, grh=qp_ah_attr_grh)
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTR, path_mtu=ib.IBV_MTU_4096, dest_qp_num=remote_qp_num, rq_psn=0, max_dest_rd_atomic=1,
+                                min_rnr_timer=12, ah_attr=qp_ah_attr)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_PATH_MTU | ib.IBV_QP_DEST_QPN | ib.IBV_QP_RQ_PSN | \
+                                          ib.IBV_QP_MAX_DEST_RD_ATOMIC | ib.IBV_QP_MIN_RNR_TIMER | ib.IBV_QP_AV))
+
+    # RTS
+    qpa = ib.struct_ibv_qp_attr(qp_state=ib.IBV_QPS_RTS, timeout=14, retry_cnt=7, rnr_retry=7, sq_psn=0, max_rd_atomic=1)
+    checkz(ib.ibv_modify_qp(self.qp, qpa, ib.IBV_QP_STATE | ib.IBV_QP_TIMEOUT | ib.IBV_QP_RETRY_CNT | ib.IBV_QP_RNR_RETRY | ib.IBV_QP_SQ_PSN | \
+                                          ib.IBV_QP_MAX_QP_RD_ATOMIC))
+
+  def __del__(self):
+    self.wait_cq() # need to wait for **everything** to complete before it's safe to dealloc queues and stuff
+    ib.ibv_destroy_qp(self.qp)
+    ib.ibv_destroy_cq(self.cq)
+    ib.ibv_destroy_comp_channel(self.comp_channel)
+
+  def next_wrid(self):
+    self.pending_wrids.add(wrid:=next(self.wrid_num))
+    return wrid
+
+  def wait_cq(self, wr_id: int|None=None):
+    while (wr_id in self.pending_wrids) if wr_id is not None else self.pending_wrids:
+      if self.ctx.ctx.contents.ops.poll_cq(self.cq, _num_entries:=1, ctypes.byref(wc:=ib.struct_ibv_wc())):
+        if wc.status != ib.IBV_WC_SUCCESS:
+          raise RuntimeError(f'Work Request completed with error: wr_id={wc.wr_id} status={ib.ibv_wc_status__enumvalues.get(wc.status, wc.status)}')
+        self.pending_wrids.remove(wc.wr_id)
+
+  def rdma_write(self, sgl:list[SGE]):
+    swr: ctypes._Pointer[ib.struct_ibv_send_wr]|None = None
+    swr_cnt, wr_id = 0, self.next_wrid()
+    def _post():
+      nonlocal swr, swr_cnt, wr_id
+      if swr is not None:
+        # The swr can be freed when this returns, the memory that sge points to can be unmapped after work completion is retrieved from cq
+        checkz(self.ctx.ctx.contents.ops.post_send(self.qp, swr, ctypes.byref(_bad_wr:=ctypes.POINTER(ib.struct_ibv_send_wr)())))
+        # TODO: async
+        self.wait_cq(wr_id)
+        swr, swr_cnt, wr_id = None, 0, self.next_wrid()
+    # Everything is in reverse for elegant chaining
+    for sg in reversed(sgl):
+      # Message size limit (max 2GB per ib spec, 1GB on tinybox mellanoxes) applies to both scatter-gather entries and entire wrs
+      for off in reversed(range(0, sg.size, self.ctx.port_attr.max_msg_sz)):
+        # Scatter-Gather Entry for local memory
+        sge = ctypes.pointer(ib.struct_ibv_sge(addr=sg.src_iova+off, length=min(sg.size-off, self.ctx.port_attr.max_msg_sz), lkey=sg.src_key))
+        # RDMA struct for remote memory
+        wr = ib.union_ibv_send_wr_wr(rdma=ib.struct_ibv_send_wr_1_rdma(remote_addr=sg.dst_iova+off, rkey=sg.dst_key))
+        # Signal (with chosen work request id) if it's the last wr (first in the loop since it's reversed)
+        wid, flags = (wr_id, ib.IBV_SEND_SIGNALED) if swr is None else (0, 0)
+        # Create Send Request
+        swr = ctypes.pointer(ib.struct_ibv_send_wr(opcode=ib.IBV_WR_RDMA_WRITE, sg_list=sge, num_sge=1, wr=wr, wr_id=wid, send_flags=flags, next=swr))
+        # Flush if queue is being overrun
+        if (swr_cnt:=swr_cnt + 1) >= self.qp_cap.max_send_wr: _post()
+    _post()

tinygrad/schedule/kernelize.py

@@ -4,7 +4,7 @@ from tinygrad.uop.ops import track_rewrites, _substitute
 from tinygrad.uop.spec import type_verify, tensor_uop_spec
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import Metadata, all_int, all_same, colored, prod, dedup, unwrap, getenv, pluralize, FUSE_ARANGE, DEBUG, SPLIT_REDUCEOP
-from tinygrad.dtype import ImageDType
+from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View, strides_for_shape, get_contraction_with_reduce
@@ -188,7 +188,7 @@ def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):
 
 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.LOAD, Ops.STORE, Ops.VALID}, name="e"),), name="view"),
+  (UPat(Ops.VIEW, src=(UPat({*GroupOp.ALU, Ops.CAST, Ops.BITCAST, Ops.BIND, Ops.LOAD, Ops.STORE, Ops.VALID, Ops.SINK}, name="e"),), name="view"),
    lambda e,view: e.replace(src=tuple(s.view(view.st) for s in e.src))),
   # if there's ones added after reduce, put this before the reduce
   (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), reduce_push_add_ones),
@@ -258,7 +258,8 @@ add_buffer_ops = PatternMatcher([
   # passthrough ASSIGN
   (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
   # VALID
-  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"), UOp.valid),
+  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
+   lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
 ])
 
 def check_load_st(glbl:UOp, view:UOp):

tinygrad/uop/__init__.py

@@ -40,7 +40,7 @@ class Ops(FastEnum):
   SPECIAL = auto()
 
   # reduce
-  REDUCE_AXIS = auto(); REDUCE_INTO = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702
+  REDUCE_AXIS = auto(); REDUCE = auto(); ALLREDUCE = auto() # noqa: E702
 
   # optimization helper ops
   UNROLL = auto(); CONTRACT = auto(); GEP = auto(); VECTORIZE = auto(); CAT = auto(); PTRCAT = auto() # noqa: E702

tinygrad/uop/ops.py

@@ -138,7 +138,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
   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
-    if self.op is Ops.REDUCE_INTO: return self.src[0].st
     # VIEW and MovementOps define a new ShapeTracker from the arg
     if self.op is Ops.VIEW: return self.arg
     if self.op in GroupOp.Movement: return unwrap(self.src[0].st).mop(self.op, self.arg)
@@ -239,7 +238,7 @@ 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, dtypes.void, (self,)+src, **kwargs)
   def assign(self, x:UOp): return UOp(Ops.ASSIGN, self.dtype, (self, x))
-  def barrier(self): return UOp(Ops.BARRIER, src=(self,))
+  def barrier(self, *src:UOp): return UOp(Ops.BARRIER, src=(self,)+src)
   def alu(self, op, *src:UOp, **kwargs):
     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
@@ -255,7 +254,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if device is not None:
       ret = ret.replace(src=(UOp(Ops.DEVICE, arg=device).view(unwrap(ret.st)),))
     return ret
-  def valid(self): return UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)
   @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, ...]):
@@ -263,13 +261,10 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
     if len(axis) == 0: return self
     # 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:
-      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
+    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)
     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)

tinygrad/uop/spec.py

@@ -198,7 +198,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, Ops.REDUCE_INTO)), lambda: True),
+  (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM, Ops.PRECAST)), lambda: True),
 
   # PTX LOAD/STORE
   (UPat((Ops.LOAD, Ops.STORE), src=(UPat(dtype=dtypes.int64),), allow_any_len=True), lambda: True),

tinygrad/viz/index.html

@@ -104,6 +104,17 @@
   .metadata > * + *, .rewrite-container > * + *, .ctx-list > * + * {
     margin-top: 12px;
   }
+  .stats-list > * + * {
+    margin-top: 8px;
+  }
+  .stats-list > p > * + * {
+    margin-top: 12px;
+  }
+  .stats-list {
+    width: 100%;
+    max-height: 240px;
+    overflow: auto;
+  }
   .ctx-list > ul > * + * {
     margin-top: 4px;
   }
@@ -232,6 +243,32 @@
   .raw-text code {
     max-height: none !important;
   }
+  table {
+    width: 100%;
+    border-collapse: separate;
+    border-spacing: 0;
+    border-radius: 8px;
+    overflow: hidden;
+    background-color: #1a1b26;
+    color: #f0f0f5;
+    font-size: 0.95em;
+  }
+  table td {
+    border-bottom: 1px solid #2c2f40;
+    vertical-align: top;
+  }
+  table tr:last-child > td {
+     border-bottom: none;
+  }
+  tr.main-row:hover {
+    background-color: #2a2d3a;
+  }
+  tr.sub-row {
+    max-width: 150px;
+  }
+  tr.main-row > td, tr.sub-row > td {
+    padding: 8px 12px;
+  }
   </style>
 </head>
 <body>

tinygrad/viz/js/index.js

@@ -162,7 +162,7 @@ async function renderProfiler() {
       else if (ref != null) {
         const start = ref.step>0 ? ref.step+1 : 0;
         const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
-        if (stepIdx !== -1) ref = {ctx:ref.ctx, step:stepIdx};
+        ref = stepIdx === -1 ? null : {ctx:ref.ctx, step:stepIdx};
       }
       const arg = { tooltipText:formatTime(e.dur), ...ref };
       // offset y by depth
@@ -377,6 +377,14 @@ function codeBlock(st, language, { loc, wrap }={}) {
   return ret;
 }
 
+function appendRow(table, name, value, unit, cls) {
+  const tr = table.appendChild(document.createElement("tr"));
+  tr.className = cls;
+  tr.appendChild(document.createElement("td")).innerText = name;
+  tr.appendChild(document.createElement("td")).innerText = unit === "us" ? formatTime(value) : value.toFixed(2)+(unit != null ? " "+unit : "%");
+  return tr;
+}
+
 function setActive(e) {
   if (e == null) return;
   e.classList.add("active");
@@ -519,11 +527,24 @@ async function main() {
   metadata.replaceChildren(codeBlock(step.code_line, "python", { loc:step.loc, wrap:true }), codeBlock(code, lang, { wrap:false }));
   if (ctx.runtime_stats != null) {
     const div = metadata.appendChild(document.createElement("div"));
-    div.style.maxHeight = "200px";
-    div.style.overflow = "auto";
+    div.className = "stats-list";
     for (const [i, s] of ctx.runtime_stats.entries()) {
       const p = div.appendChild(document.createElement("p"));
-      p.innerText = `Run ${i+1} ${formatTime(s.duration)}`;
+      if (ctx.runtime_stats.length > 1) p.innerText = `Run ${i+1}/${ctx.runtime_stats.length}`;
+      const table = div.appendChild(document.createElement("table"));
+      const tbody = table.appendChild(document.createElement("tbody"));
+      for (const { name, value, unit, subunits } of s.data) {
+          const mainRow = appendRow(tbody, name, value, unit, "main-row");
+          if (!subunits?.length) continue;
+          const subunitRow = tbody.appendChild(document.createElement("tr"));
+          subunitRow.style.display = "none";
+          mainRow.onclick = () => subunitRow.style.display = subunitRow.style.display === "none" ? "table-row" : "none";
+          mainRow.style.cursor = "pointer";
+          const td = subunitRow.appendChild(document.createElement("td"));
+          td.colSpan = 2;
+          const table = td.appendChild(document.createElement("table"));
+          for (const u of subunits) appendRow(table, u.name, u.value, unit, "sub-row");
+      }
     }
   }
   // ** rewrite steps

tinygrad/viz/serve.py

@@ -13,7 +13,6 @@ from tinygrad.dtype import dtypes
 
 uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.VCONST: "#e0e0e0", Ops.REDUCE: "#FF5B5B",
                Ops.DEFINE_GLOBAL: "#ffe0b0", Ops.DEFINE_LOCAL: "#ffe0d0", Ops.DEFINE_REG: "#f0ffe0", Ops.REDUCE_AXIS: "#FF6B6B",
-               Ops.REDUCE_INTO: "#FF6B6B",
                Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#909090", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff",
                Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF",
@@ -183,7 +182,7 @@ def get_runtime_stats(key) -> list[dict]:
   ret:list[dict] = []
   for e in profile:
     if isinstance(e, ProfileRangeEvent) and e.en is not None and e.name == key:
-      ret.append({"device":e.device, "duration":float(e.en-e.st)})
+      ret.append({"device":e.device, "data":[{"name":"Duration", "value":float(e.en-e.st), "unit":"us"}]})
   return ret
 
 def get_disassembly(ctx:list[str]):