needed for mock

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

@@ -49,10 +49,6 @@ inputs:
     description: "Install ninja?"
     required: false
     default: 'false'
-  autogen:
-    description: "Install autogen support packages?"
-    required: false
-    default: 'false'
 runs:
   using: "composite"
   steps:
@@ -137,7 +133,7 @@ runs:
 
     # ******************* apt *******************
     - name: Setup apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       shell: bash
       run: |
         sudo mkdir -p /var/cache/apt/archives
@@ -158,14 +154,14 @@ runs:
         echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
 
     - name: Add LLVM Repo (Linux)
-      if: (inputs.llvm == 'true' || inputs.autogen == 'true') && runner.os == 'Linux'
+      if: inputs.llvm == 'true' && runner.os == 'Linux'
       shell: bash
       run: |
         wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
         echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
 
     - name: Compute Package List + Hash
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true')
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
       id: apt-pkgs
       shell: bash
       run: |
@@ -194,10 +190,6 @@ runs:
         if [[ "${{ inputs.ninja }}" == "true" ]]; then
           pkgs+=" ninja-build"
         fi
-        # **** autogen ****
-        if [[ "${{ inputs.autogen }}" == "true" ]]; then
-          pkgs+=" libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev"
-        fi
 
         echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
         echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
@@ -211,20 +203,20 @@ runs:
         echo "installed=$installed" >> "$GITHUB_OUTPUT"
 
     - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
       uses: actions/cache/restore@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
     - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
       uses: actions/cache@v5
       with:
         path: /var/cache/apt/archives/
         key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
 
     - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true' || inputs.autogen == 'true') && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
       shell: bash
       run: |
         sudo apt -qq update || true

.github/workflows/autogen.yml

@@ -35,8 +35,9 @@ jobs:
         key: 'autogen'
         amd: 'true'
         llvm: 'true'
-        autogen: 'true'
         deps: 'autogen'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
     - name: Regenerate autogen files
       run: |
         find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete

.github/workflows/benchmark.yml

@@ -431,8 +431,6 @@ jobs:
       run: |
         ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
         ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
-    - name: reset chestnut
-      run: python3 extra/usbgpu/debug.py -rnw
     - name: UsbGPU boot time
       run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
     - name: UsbGPU tiny tests

.github/workflows/docs.yml

@@ -11,14 +11,13 @@ jobs:
     runs-on: ubuntu-24.04
     steps:
       - uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          deps: docs
       - name: Configure Git Credentials
         run: |
           git config user.name github-actions[bot]
           git config user.email 41898282+github-actions[bot]@users.noreply.github.com
+      - uses: actions/setup-python@v6
+        with:
+          python-version: 3.x
       - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV
       - uses: actions/cache@v5
         with:
@@ -26,5 +25,6 @@ jobs:
           path: .cache
           restore-keys: |
             mkdocs-material-
+      - run: pip install -e .[docs]
       - run: mkdocs build --strict
       - run: mkdocs gh-deploy --force

.github/workflows/test.yml

@@ -578,7 +578,7 @@ jobs:
           DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run MXFP4 Llama training on NULL backend
         if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
-        run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
+        run: HCQ2=0 PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
       - name: Run pytest (amd)
         run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq2.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
       - name: Run opt tests
@@ -621,7 +621,7 @@ jobs:
         # skip multitensor because it's slow
         run: python -m pytest -n=auto test/backend --ignore test/backend/test_multitensor.py --durations=20
       - name: Run TestOps.test_add with PMA
-        run: VIZ=-1 PMA=1 PMA_BUFFER_SIZE=8 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
+        run: VIZ=-1 PMA=1 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 

examples/compile_tensorflow.py

@@ -10,6 +10,7 @@ import tensorflow as tf
 import tf2onnx
 from tinygrad.nn.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
+from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model
 
 def get_uncompiled_model2(dataset_size=32, output_size=4):
@@ -46,8 +47,8 @@ def compile_onnx_model(onnx_model):
   cprog.append("void initialize(float *weights) {")
   weights = bytes()
   for name,cl in bufs_to_save.items():
-    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl.nbytes});")
-    weights += cl.as_memoryview()
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
+    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
   cprog.append("}")
 
   # write the weights to disk

examples/mlperf/model_train.py

@@ -1750,7 +1750,7 @@ def train_gptoss():
   def _scale_key(n):
     if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_scale.{b[1]}") in model_state: return c
     return f"{n}_scale"
-  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE and not getattr(t, '_prestore_wT', False)}
+  fp8_scale_names = {n: _scale_key(n) for n, t in model_state.items() if t.dtype == FP8_DTYPE}
   fp8_inv_scales = [model_state[sname] for sname in fp8_scale_names.values()]
   for wname, sname in fp8_scale_names.items():
     w, scale = model_state[wname], model_state[sname]
@@ -1761,23 +1761,11 @@ def train_gptoss():
       bs = _mx_block_scale(inv.reshape(-1, inv.shape[-1])).reshape(w.shape)
       master.assign((master * bs).contiguous())
 
-  fp8_wT_tensors = []
-  if getenv("PRESTORE_WT", 0):
-    def _wt_key(n, suffix):
-      if "." in n and (c:=f"{(b:=n.rsplit('.',1))[0]}_{suffix}.{b[1]}") in model_state: return c
-      return f"{n}_{suffix}"
-    for wname in fp8_scale_names:
-      wtq_name, wte_name = _wt_key(wname, "wT"), _wt_key(wname, "wT_scale")
-      if wtq_name in model_state and wte_name in model_state:
-        w = model_state[wname]
-        w._wT_q, w._wT_e8 = model_state[wtq_name], model_state[wte_name]
-        fp8_wT_tensors += [w._wT_q, w._wT_e8]
-
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
 
   if optim.master_params:
     for m in optim.master_params: m.realize()
-  Tensor.realize(*optim.params, *fp8_inv_scales, *fp8_wT_tensors)
+  Tensor.realize(*optim.params, *fp8_inv_scales)
 
   @TinyJit
   @Context(TRAINING=1)

examples/mlperf/models/gpt_oss.py

@@ -15,13 +15,12 @@ from extra.llama_kernels.rmsnorm import rmsnorm
 from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine, router_mfma
-from extra.gptoss_kernels.embedding import GPTOSSEmbedding
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_MAX = 448.0
 INIT_STD = 0.02
 ASM_GEMM = getenv("ASM_GEMM", 0)
-PRESTORE_WT = getenv("PRESTORE_WT", 0)
+
 
 def _quant_dequant_fwd(x:Tensor) -> Tensor:
   # x (2d bf16) -> bf16 value after an mxfp8 round-trip (1x32 block scaling on the last axis)
@@ -149,16 +148,10 @@ class GPTOSS:
     self.w_gate_up_bias = Tensor.zeros(n_layers, n_experts, intermediate_size * 2, dtype=dtypes.bfloat16).contiguous()
     self.w_down, self.w_down_scale = self._quant_weight(n_layers, n_experts, dim, intermediate_size, std=scaled_std, moe=True)
     self.w_down_bias = Tensor.zeros(n_layers, n_experts, dim, dtype=dtypes.bfloat16).contiguous()
-    if PRESTORE_WT:
-      self.w_gate_up_wT, self.w_gate_up_wT_scale = self._make_wT(self.w_gate_up, self.w_gate_up_scale)
-      self.w_down_wT, self.w_down_wT_scale = self._make_wT(self.w_down, self.w_down_scale)
 
     # output
     self.norm = nn.RMSNorm(dim, norm_eps)
-    if getenv("GPTOSS_EMBEDDING", 0):
-      self.tok_embeddings = GPTOSSEmbedding(vocab_size, dim)
-    else:
-      self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    self.tok_embeddings = nn.Embedding(vocab_size, dim)
     self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)
@@ -174,15 +167,6 @@ class GPTOSS:
       return [q[0] for q in qs], [q[1] for q in qs]
     return _one(*shape)
 
-  def _make_wT(self, weights:list[Tensor], scales:list[Tensor]):
-    wtq, wte = [], []
-    for w_q, w_e8 in zip(weights, scales):
-      wq, we, _ = quantize_mxfp8((w_q.cast(dtypes.bfloat16) * _mx_block_scale_3d(w_e8).cast(dtypes.bfloat16)).transpose(1, 2))
-      wq, we = wq.is_param_(False), we.is_param_(False)
-      wq._prestore_wT = we._prestore_wT = True
-      wtq.append(wq); wte.append(we)
-    return wtq, wte
-
   def _attn_mask(self, seqlen:int, dtype) -> Tensor:
     i, j = Tensor.arange(seqlen).reshape(seqlen, 1), Tensor.arange(seqlen).reshape(1, seqlen)
     return (j <= i).where(0.0, -1e30).cast(dtype).contiguous()

examples/mlperf/optim.py

@@ -10,7 +10,6 @@ ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
 FP8_AMAX_MARGIN = getenv("FP8_AMAX_MARGIN", 1.1)
 IMMEDIATE_SCALE = getenv("IMMEDIATE_SCALE", 0)
 MXFP8 = getenv("MXFP8", 0)
-PRESTORE_WT = getenv("PRESTORE_WT", 0)
 
 def stochastic_round_bf16(x:Tensor) -> Tensor:
   bits = x.bitcast(dtypes.uint32)
@@ -105,12 +104,6 @@ class GradAccClipAdamW(Optimizer):
         new_e8 = w_e8.reshape(t._inv_scale.shape)
         t._inv_scale.assign(new_e8.shard_like(t._inv_scale) if offloaded else new_e8)
         ret = w_q.reshape(t.shape)
-        if PRESTORE_WT and hasattr(t, '_wT_q'):
-          from extra.gemm.cdna_asm_gemm import _mx_block_scale_3d
-          w_phys = ret.cast(dtypes.bfloat16) * _mx_block_scale_3d(new_e8).cast(dtypes.bfloat16)
-          wT_q, wT_e8, _ = quantize_mxfp8(w_phys.transpose(1, 2))
-          t._wT_q.assign(wT_q.shard_like(t._wT_q) if offloaded else wT_q)
-          t._wT_e8.assign(wT_e8.shard_like(t._wT_e8) if offloaded else wT_e8)
         return ret.shard_like(t) if offloaded else ret
       from examples.mlperf.models.flat_llama import FP8_MAX
       if IMMEDIATE_SCALE:

extra/dsp/invoke_bug.py

@@ -315,7 +315,7 @@ return 0; }
 if __name__ == "__main__":
   dev = DSPDevice()
 
-  bufs = [dev.allocator.alloc(0x60000).buf for _ in range(4)]
+  bufs = [dev.allocator.alloc(0x60000)[0][0] for _ in range(4)]
 
   only_entry = dev.compiler.compile(entry)
   app1 = dev.runtime("test", only_entry)

extra/dsp/invoke_bug_2.py

@@ -268,7 +268,7 @@ return HAP_perf_get_time_us() == 1 ? 4 : 0;
 if __name__ == "__main__":
   dev = DSPDevice()
 
-  bufs = [dev.allocator.alloc(0x60000).buf for _ in range(4)]
+  bufs = [dev.allocator.alloc(0x60000)[0][0] for _ in range(4)]
 
   only_entry = dev.compiler.compile(entry)
   app1 = dev.runtime("test", only_entry)

extra/export_model.py

@@ -4,7 +4,7 @@ from tinygrad.tensor import Tensor
 from tinygrad.device import Device, Buffer
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, prod
+from tinygrad.helpers import Context, to_mv, prod
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.codegen import to_program
 import json
@@ -69,7 +69,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
 
   if not wasm:
     for name,cl in bufs_to_save.items():
-      weight = ''.join(["\\x%02X"%x for x in cl.as_memoryview()])
+      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
       cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
     cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]

extra/gemm/hip_matmul.py

@@ -34,9 +34,9 @@ num_threads = prod(local_size)
 # Can AMDAllocator initialized as device=0 by default?
 device = AMDDevice()
 hipallocator = AMDAllocator(device)
-a = hipallocator.alloc(N*N*4).buf
-b = hipallocator.alloc(N*N*2).buf
-c = hipallocator.alloc(N*N*2).buf
+a = hipallocator.alloc(N*N*4)[0][0]
+b = hipallocator.alloc(N*N*2)[0][0]
+c = hipallocator.alloc(N*N*2)[0][0]
 na = np.empty(N*N, np.float32)
 nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
 nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)

extra/gptoss_kernels/embedding/__init__.py

@@ -1,90 +0,0 @@
-from __future__ import annotations
-import functools, pathlib
-from tinygrad import Tensor, dtypes, nn
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, sint
-from tinygrad.renderer import Estimates
-from extra.llama_kernels import alloc_like, compile_hip
-
-VOCAB, EMBED = 128256, 2880
-FWD_THREADS, ROWS_PER_WG = 512, 8
-
-@functools.cache
-def _custom_embedding_fwd(out:UOp, idx:UOp, weight:UOp) -> UOp:
-  tokens = idx.numel()
-  threads, workgroups = UOp.special(FWD_THREADS, "lidx0"), UOp.special(tokens // ROWS_PER_WG, "gidx0")
-  sink = UOp.sink(out.base, idx.base, weight.base, threads, workgroups,
-                  arg=KernelInfo(f"gptoss_embedding_fwd_{tokens}_{VOCAB}_{EMBED}_v16_t{FWD_THREADS}_nt",
-                                 estimates=Estimates(mem=tokens*4 + tokens*EMBED*4)))
-  src = (pathlib.Path(__file__).parent/"embedding_fwd.cpp").read_text()
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, []))))
-
-def gptoss_embedding_fwd(weight:Tensor, idx:Tensor) -> Tensor:
-  out_shape = idx.shape + (EMBED,)
-  out = alloc_like(out_shape, dtypes.bfloat16, idx.device, idx.uop.axis).contiguous()
-  out, *_ = Tensor.custom_kernel(out, idx.reshape(-1), weight, fxn=_custom_embedding_fwd)
-  return out
-
-THREADS = 256
-
-@functools.cache
-def _custom_init_heads(head:UOp) -> UOp:
-  vocab = head.numel()
-  threads, workgroups = UOp.special(THREADS, "lidx0"), UOp.special((vocab+THREADS-1)//THREADS, "gidx0")
-  sink = UOp.sink(head.base, threads, workgroups,
-                  arg=KernelInfo(f"embedding_bwd_init_heads_{vocab}", estimates=Estimates(mem=vocab*4)))
-  src = (pathlib.Path(__file__).parent/"embedding_bwd.cpp").read_text()
-  defines = [f"-DVOCAB={vocab}", f"-DTHREADS={THREADS}", "-DINIT_HEADS=1"]
-  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-@functools.cache
-def _custom_build_links(next_idx:UOp, head:UOp, idx:UOp) -> UOp:
-  tokens, vocab = idx.numel(), head.numel()
-  threads, workgroups = UOp.special(THREADS, "lidx0"), UOp.special((tokens+THREADS-1)//THREADS, "gidx0")
-  sink = UOp.sink(next_idx.base, head.base, idx.base, threads, workgroups,
-                  arg=KernelInfo(f"embedding_bwd_build_links_{tokens}_{vocab}", estimates=Estimates(ops=tokens, mem=3*tokens*4)))
-  src = (pathlib.Path(__file__).parent/"embedding_bwd.cpp").read_text()
-  defines = [f"-DTOKENS={tokens}", f"-DVOCAB={vocab}", f"-DTHREADS={THREADS}", "-DBUILD_LINKS=1"]
-  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-@functools.cache
-def _custom_reduce(out:UOp, grad_emb:UOp, head:UOp, next_idx:UOp) -> UOp:
-  vocab, embed = out.shape
-  tokens = next_idx.numel()
-  threads = UOp.special(THREADS, "lidx0")
-  workgroups = UOp.special(vocab*((embed+THREADS-1)//THREADS), "gidx0")
-  sink = UOp.sink(out.base, grad_emb.base, head.base, next_idx.base, threads, workgroups,
-                  arg=KernelInfo(f"embedding_bwd_owner_reduce_{tokens}_{vocab}_{embed}",
-                                 estimates=Estimates(ops=tokens*embed, mem=tokens*embed*2+vocab*embed*2)))
-  src = (pathlib.Path(__file__).parent/"embedding_bwd.cpp").read_text()
-  defines = [f"-DTOKENS={tokens}", f"-DVOCAB={vocab}", f"-DEMBED={embed}", f"-DTHREADS={THREADS}"]
-  return UOp(Ops.PROGRAM,
-             src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=compile_hip(src, defines))))
-
-def embedding_bwd_owner(grad_emb:Tensor, idx:Tensor, vocab:sint) -> Tensor:
-  grad_emb = grad_emb.reshape(idx.numel(), grad_emb.shape[-1])
-  device = grad_emb.device
-  head = alloc_like((vocab,), dtypes.int32, device)
-  next_idx = alloc_like((idx.numel(),), dtypes.int32, device)
-  out = alloc_like((vocab, grad_emb.shape[-1]), dtypes.bfloat16, device)
-  head, *_ = Tensor.custom_kernel(head, fxn=_custom_init_heads)
-  next_idx, *_ = Tensor.custom_kernel(next_idx, head, idx.reshape(-1), fxn=_custom_build_links)
-  out, *_ = Tensor.custom_kernel(out, grad_emb, head, next_idx, fxn=_custom_reduce)
-  return out
-
-@functools.cache
-def _embedding_fwd_fxn(wp:UOp, ip:UOp, device:str|tuple[str, ...]) -> Tensor:
-  return gptoss_embedding_fwd(Tensor(wp, device=device), Tensor(ip, device=device))
-
-def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple:
-  weight, idx = call.src[1:]
-  device = Tensor(weight).device
-  if isinstance(device, tuple):
-    grad_emb, idx = grad_emb.copy_to_device(device), idx.copy_to_device(device)
-  return embedding_bwd_owner(Tensor(grad_emb, device=device), Tensor(idx, device=device), weight.shape[0]).uop, None
-
-class GPTOSSEmbedding(nn.Embedding):
-  def __call__(self, idx:Tensor) -> Tensor:
-    fxn = _embedding_fwd_fxn(self.weight.as_param(0).uop, idx.as_param(1).uop, self.weight.device)
-    return Tensor.call(self.weight, idx, fxn=fxn, grad_fxn=_embedding_bwd)

extra/gptoss_kernels/embedding/embedding_bwd.cpp

@@ -1,46 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef TOKENS
-#define TOKENS 131072
-#endif
-#ifndef VOCAB
-#define VOCAB 128256
-#endif
-#ifndef EMBED
-#define EMBED 2880
-#endif
-#ifndef THREADS
-#define THREADS 256
-#endif
-
-#ifdef INIT_HEADS
-extern "C" __global__ __launch_bounds__(THREADS) void embedding_bwd_init_heads(int *__restrict__ head) {
-  const int v=blockIdx.x*THREADS+threadIdx.x;
-  if(v<VOCAB) head[v]=-1;
-}
-#elif defined(BUILD_LINKS)
-// atomicExch is one operation per token, not one operation per embedding element.
-extern "C" __global__ __launch_bounds__(THREADS) void embedding_bwd_build_links(
-    int *__restrict__ next_idx, int *__restrict__ head, const int *__restrict__ idx) {
-  const int i=blockIdx.x*THREADS+threadIdx.x;
-  if(i<TOKENS) {
-    const int token=max(0,min(VOCAB-1,idx[i]));
-    next_idx[i]=atomicExch(&head[token],i);
-  }
-}
-#else
-// One workgroup owns a 256-column block of a vocabulary row. It traverses that token's occurrence list,
-// accumulates in FP32, and rounds once to BF16. Empty rows naturally write zero, replacing the old fill.
-extern "C" __global__ __launch_bounds__(THREADS) void embedding_bwd_owner_reduce(
-    __hip_bfloat16 *__restrict__ out, const __hip_bfloat16 *__restrict__ grad_emb,
-    const int *__restrict__ head, const int *__restrict__ next_idx) {
-  constexpr int D_BLOCKS=(EMBED+THREADS-1)/THREADS;
-  const int token=blockIdx.x/D_BLOCKS, d=(blockIdx.x%D_BLOCKS)*THREADS+threadIdx.x;
-  if(d<EMBED) {
-    float sum=0.0f;
-    for (int i=head[token]; i>=0; i=next_idx[i]) sum+=(float)grad_emb[(long long)i*EMBED+d];
-    out[(long long)token*EMBED+d]=(__hip_bfloat16)sum;
-  }
-}
-#endif

extra/gptoss_kernels/embedding/embedding_fwd.cpp

@@ -1,37 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-// GPT-OSS production keeps two 8,192-token sequences on each DP rank. One wave
-// owns one token row, so its 16-byte weight reads and output writes are fully
-// contiguous instead of interleaving four unrelated vocabulary rows per wave.
-using uint4v = unsigned int __attribute__((ext_vector_type(4)));
-
-constexpr int TOKENS = 16384;
-constexpr int VOCAB = 128256;
-constexpr int EMBED = 2880;
-constexpr int VECS_PER_ROW = EMBED / 8;
-constexpr int THREADS = 512;
-constexpr int WAVE = 64;
-constexpr int ROWS_PER_WG = THREADS / WAVE;
-
-static_assert(EMBED % 8 == 0 && TOKENS % ROWS_PER_WG == 0);
-
-extern "C" __global__ __launch_bounds__(THREADS) void gptoss_embedding_fwd(
-    __hip_bfloat16 *__restrict__ out, const int *__restrict__ idx,
-    const __hip_bfloat16 *__restrict__ weight) {
-  const int lid = threadIdx.x;
-  const int lane = lid & (WAVE - 1);
-  const int row = blockIdx.x * ROWS_PER_WG + (lid / WAVE);
-  const int token = idx[row];
-  const uint4v zero = {0u, 0u, 0u, 0u};
-
-#pragma unroll
-  for (int j = lane; j < VECS_PER_ROW; j += WAVE) {
-    // The embedding is a one-use, random-row source. Bypass temporal caching so
-    // it does not displace the freshly written activation consumed by RMSNorm.
-    const uint4v value = static_cast<unsigned>(token) < VOCAB ?
-      __builtin_nontemporal_load(reinterpret_cast<const uint4v *>(weight) +
-                                 static_cast<long long>(token) * VECS_PER_ROW + j) : zero;
-    reinterpret_cast<uint4v *>(out)[static_cast<long long>(row) * VECS_PER_ROW + j] = value;
-  }
-}

extra/hcq1/graph.py

@@ -1,9 +1,9 @@
 import collections, time
 from typing import Any, cast
 from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, suppress_finalizing, TracingKey, unwrap
-from extra.hcq1.hcq import HCQBuffer, HCQCompiled, HCQAllocator, HCQSignal, HWQueue, HCQArgsState
-from tinygrad.runtime.support.hcq import BumpAllocator, MMIOInterface
-from tinygrad.device import BufferStorage, Buffer, BufferSpec, Compiled, Device, MultiBuffer, ProfileGraphEntry, ProfileGraphEvent
+from extra.hcq1.hcq import HCQCompiled, HCQAllocator, HCQSignal, HWQueue, HCQArgsState
+from tinygrad.runtime.support.hcq import HCQBuffer, BumpAllocator, MMIOInterface
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
 from tinygrad.engine.jit import GraphRunner, MultiGraphRunner
@@ -30,7 +30,7 @@ class HCQGraph(MultiGraphRunner):
     for runtime in self.runtimes:
       if runtime is None: continue
       kernargs_size[runtime.dev] += round_up(runtime.kernargs_alloc_size, 16)
-    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True)).buf for d,sz in kernargs_size.items()}
+    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True))[0][0] for d,sz in kernargs_size.items()}
 
     # Fill initial arguments.
     self.ji_args: dict[int, HCQArgsState] = {}
@@ -314,7 +314,7 @@ class HCQGraph(MultiGraphRunner):
 
     if PROFILE and self.kickoff_value >= 1: self.collect_timestamps()
 
-    for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(BufferStorage(buf, buf.meta, buf.view), BufferSpec(cpu_access=True))
+    for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(buf, BufferSpec(cpu_access=True))
 
   @staticmethod
   def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:

extra/hcq1/hcq.py

@@ -3,39 +3,17 @@ from typing import cast, Callable, Type, TypeVar, Generic, Any
 import contextlib, decimal, statistics, time, ctypes, array, collections, itertools
 from tinygrad.helpers import PROFILE, getenv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
 from tinygrad.helpers import suppress_finalizing, TracingKey
-from tinygrad.device import BufferStorage, Buffer, BufferSpec, Compiled, Allocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
+from tinygrad.device import BufferSpec, Compiled, Allocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
 from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
 from tinygrad.renderer import Renderer
+from tinygrad.runtime.support.hcq import HCQBuffer
 
 SignalType = TypeVar('SignalType', bound='HCQSignal')
 HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled')
 ProgramType = TypeVar('ProgramType', bound='HCQProgram')
 ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState')
 
-class HCQBuffer:
-  def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None, owner:Any=None):
-    self.va_addr, self.size, self.meta, self._base, self.view = va_addr, size, meta, _base, view
-    self._devs, self.owner = ([owner] if owner is not None else []), owner
-    self._mappings:dict[Compiled, HCQBuffer] = {} # mapping to the other devices
-
-  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, meta=self.meta,
-      _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
-
-  def cpu_view(self) -> MMIOInterface:
-    assert self.view is not None, "buffer has no cpu_view"
-    return self.view
-
-  @property
-  def base(self) -> HCQBuffer: return self._base or self
-
-  @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
-
 class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
   """
   A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
@@ -304,7 +282,7 @@ class HCQProgram(Program[HCQDeviceType]):
     if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter, self.profile_key)]
 
   @staticmethod
-  def _fini(dev, buf, spec): dev.allocator.free(BufferStorage(buf, buf.meta, buf.view), buf.size, spec)
+  def _fini(dev, buf, spec): dev.allocator.free(((buf, buf.meta), buf.view), buf.size, spec)
 
   def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int|None, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
     """
@@ -381,7 +359,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
       self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
 
     if comp_queue_t is not None:
-      self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True)).buf
+      self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))[0][0]
       self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
 
     self.can_recover = can_recover # Whether the device can recover from faults or timeouts
@@ -415,7 +393,7 @@ class HCQCompiled(Compiled, Generic[SignalType]):
   def new_signal(self, **kwargs) -> SignalType:
     assert self.signal_t is not None, "Device does not support signals"
     if not HCQCompiled.signal_pool[pg:=self.peer_group]:
-      HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True)).buf)
+      HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True))[0][0])
       HCQCompiled.signal_pool[pg] += [alc.offset(offset=off, size=16) for off in range(0, alc.size, 16)]
       for dev in HCQCompiled.peer_groups[pg]: cast(HCQAllocator, dev.allocator)._map(alc)
     return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)
@@ -447,11 +425,11 @@ class HCQCompiled(Compiled, Generic[SignalType]):
     cast(HCQAllocatorBase, self.allocator).b_timeline = [0] * len(cast(HCQAllocatorBase, self.allocator).b)
 
   def _realloc(self, oldbuf:HCQBuffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQBuffer, bool]:
-    if oldbuf is not None: self.allocator.free(BufferStorage(oldbuf, oldbuf.meta, oldbuf.view), oldbuf.size, options=options)
-    try: buf, realloced = self.allocator.alloc(new_size, options=options).buf, True
+    if oldbuf is not None: self.allocator.free(((oldbuf, oldbuf.meta), oldbuf.view), oldbuf.size, options=options)
+    try: buf, realloced = self.allocator.alloc(new_size, options=options)[0][0], True
     except MemoryError:
       if force: raise
-      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options).buf, False
+      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options)[0][0], False
     return buf, realloced
 
   def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
@@ -477,11 +455,9 @@ class HCQAllocatorBase(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
 
   def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, **kwargs):
     super().__init__(dev, **kwargs)
-    self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)).buf for _ in range(batch_cnt)]
+    self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True))[0][0] for _ in range(batch_cnt)]
     self.b_timeline, self.b_next = [0] * len(self.b), 0
 
-  def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(*self._map(buf.ensure_allocated()._buf))
-
   def _map(self, buf:HCQBuffer) -> tuple:
     if self.dev not in buf.mapped_devs:
       if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
@@ -492,10 +468,10 @@ class HCQAllocatorBase(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
     return mapped, mapped.meta
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options:BufferSpec|None=None):
-    for dev in storage.buf.mapped_devs: dev.synchronize()
-    for d, mb in storage.buf.mappings.items(): d.allocator._do_unmap(mb)
-    if hasattr(self, '_do_free'): self._do_free(storage.buf, options)
+  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
+    for dev in buf.mapped_devs: dev.synchronize()
+    for d, mb in buf.mappings.items(): d.allocator._do_unmap(mb)
+    if hasattr(self, '_do_free'): self._do_free(buf, options)
 
   def _do_unmap(self, mb): self.dev.iface.free(mb)
 

extra/hcq1/ops_amd_old.py

@@ -3,11 +3,11 @@ from typing import cast
 import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit, time
 assert sys.platform != 'win32'
 from dataclasses import dataclass
-from extra.hcq1.hcq import HCQBuffer, HCQCompiled, HCQAllocator, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, hcq_profile
-from tinygrad.runtime.support.hcq import FileIOInterface
+from extra.hcq1.hcq import HCQCompiled, HCQAllocator, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, hcq_profile
+from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
-from tinygrad.device import BufferStorage, Compiled, BufferSpec, TinyELF
+from tinygrad.device import Compiled, BufferSpec, TinyELF
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, TracingKey
 from tinygrad.helpers import ceildiv, unwrap, pluralize
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
@@ -55,7 +55,7 @@ class AMDComputeQueue(HWQueue):
 
   def __del__(self):
     if self.binded_device is not None:
-      self.binded_device.allocator.free(BufferStorage(self.hw_page, self.hw_page.meta, self.hw_page.view), self.hw_page.size,
+      self.binded_device.allocator.free(((self.hw_page, self.hw_page.meta), self.hw_page.view), self.hw_page.size,
                                        BufferSpec(cpu_access=True, nolru=True, uncached=True))
 
   def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, len(vals) - 1), *vals)
@@ -396,7 +396,7 @@ class AMDComputeQueue(HWQueue):
 
   def bind(self, dev:AMDDevice):
     self.binded_device = dev
-    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
+    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True))[0][0]
     hw_view = self.hw_page.cpu_view().view(fmt='I')
     for i, value in enumerate(self._q): hw_view[i] = value
 
@@ -447,7 +447,7 @@ class AMDComputeAQLQueue(AMDComputeQueue):
 
   def bind(self, dev:AMDDevice):
     self.binded_device = dev
-    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
+    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True))[0][0]
     self._cmds = self._prep_aql(self._q, self.hw_page)
     self._q = self.hw_page.cpu_view().view(fmt='I')
     return self
@@ -512,7 +512,7 @@ class AMDCopyQueue(HWQueue):
     if not getenv("AMD_SDMA_BIND", 0) or not dev.is_am(): return
 
     self.binded_device = dev
-    self.hw_page = dev.allocator.alloc((qsz:=round_up(len(self._q), 8)) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
+    self.hw_page = dev.allocator.alloc((qsz:=round_up(len(self._q), 8)) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True))[0][0]
     hw_view = self.hw_page.cpu_view().view(fmt='I')
     for i in range(qsz): hw_view[i] = self._q[i] if i < len(self._q) else 0
 
@@ -572,7 +572,7 @@ class AMDProgram(HCQProgram['AMDDevice']):
       if typ == 5: image[apply_image_offset:apply_image_offset+8] = struct.pack('<q', rel_sym_offset - apply_image_offset + addent) # R_AMDGPU_REL64
       else: raise RuntimeError(f"unknown AMD reloc {typ}")
 
-    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True)).buf
+    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True))[0][0]
     self.dev.allocator._copyin(self.lib_gpu, image)
     self.dev.synchronize()
 
@@ -643,9 +643,9 @@ class AMDAllocator(HCQAllocator['AMDDevice']):
     super().__init__(dev, copy_bufs=getattr(dev.iface, 'copy_bufs', None),
                      supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
 
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     opaque = self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
-    return BufferStorage(opaque, opaque.meta, opaque.view)
+    return (opaque, opaque.meta), opaque.view
 
   def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
 
@@ -981,19 +981,6 @@ class PCIIface(PCIIfaceBase):
 
   def device_fini(self): self.dev_impl.fini()
 
-  def alloc(self, size:int, **kwargs) -> HCQBuffer:
-    storage = super().alloc(size, **kwargs)
-    return HCQBuffer(storage.buf, storage.meta.mapping.size, meta=storage.meta, view=storage.host, owner=self.dev)
-
-  def free(self, buf:HCQBuffer):
-    if buf.owner == self.dev: super().free(BufferStorage(buf.va_addr, buf.meta, buf.view))
-    else: self.dev_impl.mm.unmap_range(buf.va_addr, round_up(buf.size, 0x1000))
-
-  def map(self, buf:HCQBuffer):
-    from types import SimpleNamespace
-    super().map(SimpleNamespace(device=buf.owner.device, _buf=buf.va_addr, nbytes=buf.size, meta=buf.meta))
-    return HCQBuffer(buf.va_addr, buf.size, meta=buf.meta, owner=buf.owner)
-
 class USBIface(PCIIface):
   def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
     if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
@@ -1099,7 +1086,7 @@ class AMDDevice(HCQCompiled):
 
       with (q:=cast(AMDComputeQueue, unwrap(self.hw_compute_queue_t)())).pred_exec((1 << self.xccs) - 1):
         q.pmc_start([(k, *self.pmc_counters[k]) for k in PMC_COUNTERS]).submit(self)
-      self.pmc_buffer = self.allocator.alloc(self.pmc_sched[-1].off + self.pmc_sched[-1].size, BufferSpec(nolru=True, uncached=True)).buf
+      self.pmc_buffer = self.allocator.alloc(self.pmc_sched[-1].off + self.pmc_sched[-1].size, BufferSpec(nolru=True, uncached=True))[0][0]
       self.allocator._copyin(self.pmc_buffer, memoryview(bytearray(self.pmc_buffer.size))) # zero pmc buffers, some counters have only lo part.
 
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
@@ -1108,8 +1095,8 @@ class AMDDevice(HCQCompiled):
       self.iface.require_profile_mode()
 
       SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)).buf for _ in range(self.se_cnt * self.xccs)]
-      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True)).buf
+      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True))[0][0] for _ in range(self.se_cnt * self.xccs)]
+      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))[0][0]
       self.sqtt_next_cmd_id = itertools.count(0)
 
     if self.is_am():

extra/hcq1/ops_rdma.py

@@ -2,8 +2,8 @@ from __future__ import annotations
 import mmap, struct, functools, atexit
 from typing import cast
 from tinygrad.uop.ops import sint
-from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, HCQBuffer
-from tinygrad.runtime.support.hcq import FileIOInterface
+from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue
+from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace
 from tinygrad.runtime.support.mlx.mlxdev import MLXDev, MLXQP

extra/hcq1/test_hcq.py

@@ -3,7 +3,8 @@ from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import mv_address, DEBUG, DEV
 from test.helpers import slow, replace_opts
 from tinygrad.device import Buffer, BufferSpec
-from extra.hcq1.hcq import HCQCompiled, HCQBuffer
+from extra.hcq1.hcq import HCQCompiled
+from tinygrad.runtime.support.hcq import HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runtime

extra/mlx_driver/loopback.py

@@ -2,8 +2,7 @@
 # GMMU=0 MLX_PCI=0000:41:00.0 PYTHONPATH=. python3 extra/mlx_driver/loopback.py
 import struct
 from tinygrad.helpers import getenv, round_up
-from tinygrad.device import Device, BufferSpec, Buffer
-from tinygrad.dtype import dtypes
+from tinygrad.device import Device, BufferSpec
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.memory import AddrSpace
 from tinygrad.runtime.ops_amd import AMDComputeQueue
@@ -31,8 +30,8 @@ print(f"[init] loopback connect QP 0x{qp.qp_info['qpn']:x}")
 qp.connect(qp.qp_info['qpn'], dev.mac, int.from_bytes(dev.local_gid, 'big'))
 
 # allocate src/dst via AMD GPU allocator
-buf_src = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
-buf_dst = Buffer(gpu.device, BUF_SIZE, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
+buf_src = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
+buf_dst = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
 
 bar_base = gpu.iface.pci_dev.bar_info(gpu.iface.vram_bar)[0]
 src_paddr = buf_src.meta.mapping.paddrs[0][0] + bar_base
@@ -41,8 +40,8 @@ print(f"src paddr=0x{src_paddr:x} dst paddr=0x{dst_paddr:x}")
 
 # fill src, zero dst
 test_msg = b"Hello from loopback send/recv!"
-buf_src.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00')))))
-buf_dst.copy_from(Buffer("PYTHON", BUF_SIZE, dtypes.uint8, opaque=memoryview(bytearray(BUF_SIZE))))
+gpu.allocator._copyin(buf_src, memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00'))))
+gpu.allocator._copyin(buf_dst, memoryview(bytearray(BUF_SIZE)))
 gpu.synchronize()
 
 # post recv WQE on RQ from CPU (scatter entry: byte_count, lkey, addr)
@@ -91,7 +90,7 @@ qp.poll_cq()
 
 # read back
 result = bytearray(BUF_SIZE)
-result[:] = buf_dst.as_memoryview()
+gpu.allocator._copyout(memoryview(result), buf_dst)
 gpu.synchronize()
 
 got = bytes(result[:len(test_msg)])

extra/usbgpu/debug.py

@@ -57,34 +57,18 @@ class USBGPUDebug:
     self.provisioned = True
     print("Provisioning complete.")
 
-  def reset(self, bootloader=False, wait=False):
+  def reset(self, bootloader=False):
     if not self.provisioned:
       raise RuntimeError("Device not provisioned for usbgpu debugging. Use --provision to provision it.")
 
     self.ftdi.set_cbus_gpio(self.CBUS_RESET | (self.CBUS_BOOTLOADER if bootloader else 0))
     time.sleep(0.5)
     self.ftdi.set_cbus_gpio(self.CBUS_BOOTLOADER if bootloader else 0)
-    if bootloader or wait: self._wait_for(bootloader)
-    if bootloader: self.ftdi.set_cbus_gpio(0)
+    if bootloader:
+      time.sleep(1)
+      self.ftdi.set_cbus_gpio(0)
     print("Device reset complete.")
 
-  def _wait_for(self, bootloader: bool, timeout=10.0):
-    """Wait for the ASM2464 bootloader (or the device itself) to enumerate on USB."""
-    import usb.core
-    SUPPORTED_CONTROLLERS = [
-      (0x174C, 0x2464),
-      (0x174C, 0x2463),
-      (0x3801, 0x0001),
-    ] if bootloader else [(0x3801, 0x0001)]
-    start = time.time()
-    while time.time() - start < timeout:
-      for vendor, device in SUPPORTED_CONTROLLERS:
-        dev = usb.core.find(idVendor=vendor, idProduct=device)
-        if dev is not None:
-          return  # Found it!
-      time.sleep(0.1)
-    raise RuntimeError(f"{'Bootloader' if bootloader else 'Device'} did not enumerate within {timeout}s")
-
   def read(self) -> bytes:
     return self.ftdi.read_data(256).decode('utf-8', errors='replace')
 
@@ -95,9 +79,7 @@ if __name__ == "__main__":
   args.add_argument('--provision', '-p', action='store_true', default=False, help="Provision the connected FTDI for usbgpu debugging")
   args.add_argument('--reset', '-r', action='store_true', default=False, help="Reset the device")
   args.add_argument('--bootloader', '-b', action='store_true', default=False, help="Reset to bootloader")
-  args.add_argument('--wait', '-w', action='store_true', default=False, help="Wait for the device to enumerate after reset")
   args.add_argument('--no-read', '-n', action='store_true', default=False, help="Do not read debug output")
-  args.add_argument('--timeout', '-t', type=float, default=None, help="Timeout in seconds for reading")
 
   args = args.parse_args()
 
@@ -106,17 +88,15 @@ if __name__ == "__main__":
       dbg.provision()
 
     if args.reset:
-      dbg.reset(bootloader=False, wait=args.wait)
+      dbg.reset(bootloader=False)
 
     if args.bootloader:
-      dbg.reset(bootloader=True, wait=args.wait)
+      dbg.reset(bootloader=True)
 
     if not args.no_read:
       print("Starting debug output. Press Ctrl-C to exit.\n------")
-      start_time = time.perf_counter()
       while True:
         sys.stdout.write(dbg.read())
         sys.stdout.flush()
-        if args.timeout is not None and (time.perf_counter() - start_time) >= args.timeout:
-          break
         time.sleep(0.001)
+

spec/tinyspec.tex

@@ -109,14 +109,14 @@ A value \op{Call} is void: its \op{Sink} body stores to output \op{Param}s bound
 \end{tabular}
 
 %% ============================================================
-\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace, anonymous store}
+\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
 
 \begin{tabular}{@{}l l l l@{}}
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Load} & (buf, alt?, gate?) & --- & Read from buffer into AddrSpace.ALU. \\
-\op{Copy} & (buf,) & device & Copy buf to device arg. \\
+\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\
+  & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\
 \bottomrule
 \end{tabular}
 

test/amd/hw/helpers.py

@@ -170,10 +170,9 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
 
 def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
   """Run instructions on real AMD hardware via HIPCompiler and the AMD runtime."""
-  from tinygrad.device import Device, TinyELF, Buffer
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import Target
-  from tinygrad.dtype import dtypes
+  from tinygrad.helpers import Target, flat_mv
 
   dev = Device["AMD"]
   compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -226,11 +225,12 @@ amdhsa.kernels:
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
 
   buf_sz = _out_bytes(n_lanes)
-  out_gpu = Buffer(dev.device, buf_sz, dtypes.uint8, preallocate=True)
-  assert out_gpu._buf % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu._buf:x}"
-  prg(out_gpu._buf, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
+  out_gpu = dev.allocator.alloc(buf_sz)[0][0]
+  assert out_gpu.va_addr % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu.va_addr:x}"
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
 
-  out_buf = out_gpu.as_memoryview()
+  out_buf = bytearray(buf_sz)
+  dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
 
   return parse_output(bytes(out_buf), n_lanes)
 

test/amd/hw/test_cdna_vop3.py

@@ -5,7 +5,7 @@ gfx950 hardware when USE_HW=1.
 """
 import ctypes, struct, unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import Target
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -42,14 +42,13 @@ def _run_emu(instructions: list, out_reg: int = 2) -> int:
   return out_buf[0]
 
 def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.dtype import dtypes
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
   if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
-  out_gpu = Buffer(dev.device, LANES * 4, dtypes.uint8, preallocate=True)
-  code = _code(instructions, out_reg, out_gpu._buf)
+  out_gpu = dev.allocator.alloc(LANES * 4)[0][0]
+  code = _code(instructions, out_reg, out_gpu.va_addr)
   byte_str = ", ".join(f"0x{b:02x}" for b in code)
   asm_src = f""".text
 .globl test
@@ -88,7 +87,8 @@ amdhsa.kernels:
 """
   prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
   prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
-  out = out_gpu.as_memoryview()
+  out = bytearray(LANES * 4)
+  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
   return struct.unpack("<I", out)[0]
 
 def run_cdna(instructions: list, out_reg: int = 2) -> int:

test/amd/hw/test_dpp.py

@@ -6,7 +6,7 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import Target
+from tinygrad.helpers import Target, flat_mv
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
@@ -37,8 +37,7 @@ def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
   return list(out_buf)
 
 def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.dtype import dtypes
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device["AMD"]
@@ -85,9 +84,10 @@ amdhsa.kernels:
 """
   lib = compiler.compile(asm_src)
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  out_gpu = Buffer(dev.device, WAVE64 * 4, dtypes.uint8, preallocate=True)
-  prg(out_gpu._buf, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
-  out = out_gpu.as_memoryview()
+  out_gpu = dev.allocator.alloc(WAVE64 * 4)[0][0]
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
+  out = bytearray(WAVE64 * 4)
+  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
   return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(WAVE64)]
 
 def run_wave64(instructions: list, out_reg: int = 1) -> list[int]:

test/amd/hw/test_rdna4_permlane_var.py

@@ -5,7 +5,7 @@ real RDNA4 hardware when USE_HW=1.
 """
 import ctypes, unittest
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
-from tinygrad.helpers import Target
+from tinygrad.helpers import Target, flat_mv
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
@@ -36,8 +36,7 @@ def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
   return list(out_buf)
 
 def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.dtype import dtypes
+  from tinygrad.device import Device, TinyELF
   from tinygrad.runtime.support.compiler_amd import HIPCompiler
 
   dev = Device['AMD']
@@ -86,9 +85,10 @@ amdhsa.kernels:
 """
   lib = compiler.compile(asm_src)
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  out_gpu = Buffer(dev.device, LANES * 4, dtypes.uint8, preallocate=True)
-  prg(out_gpu._buf, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
-  out = out_gpu.as_memoryview()
+  out_gpu = dev.allocator.alloc(LANES * 4)[0][0]
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
+  out = bytearray(LANES * 4)
+  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
   return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(LANES)]
 
 def run_rdna4(instructions: list, out_reg: int = 2) -> list[int]:

test/amd/test_roundtrip.py

@@ -79,7 +79,7 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
             buf_pool[dst_id] = dst_buf.nbytes
           # Get source data if it's from numpy/CPU
           if hasattr(src_buf, 'base') and src_buf.base is not None and src_buf.base.is_allocated():
-            src_data = bytes(src_buf.base.as_memoryview())
+            src_data = bytes(src_buf.base._buf)
             buf_data[dst_id] = src_data
       elif ast.op is Ops.PROGRAM:
         info = ast.arg

test/amd/test_sqtt_profiler.py

@@ -1,4 +1,4 @@
-import unittest, contextlib, functools
+import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit, dtypes
 from tinygrad.dtype import AddrSpace
 from test.helpers import is_hcq2_device
@@ -10,7 +10,6 @@ from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
 from tinygrad.renderer.amd.sqtt import decode, print_packets
 from tinygrad.renderer.amd.dsl import s, v
-from tinygrad.helpers import getenv
 
 @contextlib.contextmanager
 def save_sqtt():
@@ -21,19 +20,6 @@ def save_sqtt():
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
   data[:] = [e for e in Compiled.profile_events[:profile_start] if isinstance(e, ProfileProgramEvent)]+Compiled.profile_events[profile_start:]
-  if getenv("PRINT_PKTS"):
-    sqtt_kernels = set()
-    for event in data:
-      if not isinstance(event, ProfileSQTTEvent) or not event.itrace: continue
-      print(f"\n=== SE {event.se} ===")
-      print_packets(decode(event.blob))
-      sqtt_kernels.add(event.kern)
-    for event in data:
-      if not isinstance(event, ProfileProgramEvent) or event.tag not in sqtt_kernels: continue
-      from test.null.test_viz import write_files, run_cli
-      with write_files(profile=data) as files:
-        out = run_cli(*files, "-s", f"{event.name} SQTT SE:0 PKTS", json_fmt=False)[0]["out"]
-      print(out)
 
 def map_sqtt(profile:list) -> list[dict]:
   load_amd_counters(data:=VizData(), profile)
@@ -113,37 +99,16 @@ class TestSQTTProfiler(unittest.TestCase):
 
   def test_asm(self):
     t = Tensor.empty(1)
-    with save_sqtt():
+    with save_sqtt() as data:
       t.custom_kernel(fxn=custom_asm_cdna if self.arch == "gfx950" else custom_asm_rdna)[0].realize()
-
-  def test_setprio(self):
-    if self.arch == "gfx950":
-      from tinygrad.runtime.autogen.amd.cdna import ins as isa
-      hw_id, wave_size, add = isa.HWREG.HW_REG_HW_ID.value, 64, isa.s_add_u32
-      barrier = [isa.s_barrier()]
-    elif self.arch.startswith("gfx12"):
-      from tinygrad.runtime.autogen.amd.rdna4 import ins as isa
-      hw_id, wave_size, add = isa.HWREG.HW_REG_WAVE_HW_ID1.value, 32, isa.s_add_co_u32
-      barrier = [isa.s_barrier_signal(ssrc0=-1), isa.s_barrier_wait(simm16=-1)]
-    else: self.skipTest("tested on CDNA4 and RDNA4")
-    def setprio_kernel(A, priority=0):
-      insts = [
-        isa.s_getreg_b32(s[0], hw_id),
-        isa.s_mov_b32(s[1], 0),
-        isa.s_setprio(0),
-        isa.s_cmp_eq_u32(s[0], 0),
-        isa.s_cbranch_scc1(1),
-        isa.s_setprio(priority),
-        *barrier,
-      ]
-      # eight waves contend for scalar issue slots
-      insts += [add(s[1], s[1], 1) for _ in range(64)]
-      insts += [isa.s_setprio(0), *barrier, isa.s_endpgm()]
-      return custom_asm(A, insts, wave_size*8, (96 if self.arch == "gfx950" else 64)*1024)
-
-    with Context(SQTT_LIMIT_SE=1), save_sqtt():
-      Tensor.empty(1).custom_kernel(fxn=functools.partial(setprio_kernel, priority=3))[0].realize()
-      Tensor.empty(1).custom_kernel(fxn=functools.partial(setprio_kernel, priority=0))[0].realize()
+    for event in data:
+      if not isinstance(event, ProfileSQTTEvent) or not event.itrace: continue
+      print(f"\n=== SE {event.se} ===")
+      print_packets(decode(event.blob))
+    from test.null.test_viz import write_files, run_cli
+    with write_files(profile=data) as files:
+      out = run_cli(*files, "-s", "asm SQTT SE:0 PKTS", json_fmt=False)[0]["out"]
+    print(out)
 
   def test_multiple_runs(self):
     t = Tensor.empty(1) + 1

test/backend/test_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException, is_hcq2_device
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -153,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(2)
+      assert_kernel_count(3 if is_hcq2_device() else 2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu
+from test.helpers import assert_kernel_count, needs_second_gpu, is_hcq2_device
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -43,7 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -619,7 +619,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(3)  # TODO: first copy is dead, could be 2
+    assert_kernel_count(5 if is_hcq2_device() else 3)  # TODO: first copy is dead, could be 2
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -629,7 +629,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_assign_temporary_copy_reshape(self):
@@ -637,7 +637,7 @@ class TestAssign(unittest.TestCase):
     c = Tensor.empty(2, 2).assign(a.to(None))
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
     self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
 
 class TestAssignOrdering(unittest.TestCase):

test/backend/test_call.py

@@ -1,8 +1,7 @@
 import unittest, ctypes
 from tinygrad import Tensor, UOp
 from tinygrad.device import Device
-from tinygrad.dtype import dtypes, AddrSpace
-from tinygrad.codegen import to_program
+from tinygrad.dtype import dtypes
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.uop.ops import KernelInfo
 
@@ -38,10 +37,4 @@ class TestCall(unittest.TestCase):
     c.realize()
     self.assertEqual(c.item(), 44)
 
-  def test_call_stack_pointer(self):
-    slot = UOp.placeholder((1,), dtypes.uint32, addrspace=AddrSpace.REG)
-    call = UOp.custom_function("callback", UOp.const(0, dtypes.uint64)).call(slot[0], ret_dtype=dtypes.void)
-    prg = to_program(call.sink(arg=KernelInfo("call_stack")), Device["CPU"].renderer)
-    self.assertIn("(unsigned int*)((buf", prg.src[2].arg)
-
 if __name__ == "__main__": unittest.main()

test/backend/test_profiler.py

@@ -3,7 +3,7 @@ from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from extra.hcq1.hcq import HCQCompiled
-from tinygrad.runtime.support.hcq2 import HCQ_DEVS
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 
@@ -35,7 +35,7 @@ def helper_profile_filter_device(profile, device:str):
   assert len(dev_events) == 1, "only one device registration event is expected"
   return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
 
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT], HCQCompiled) or Device.DEFAULT in HCQ_DEVS | {"CPU", "METAL"}, "Dev not supported")
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT], (HCQCompiled, HCQ2Compiled)) or Device.DEFAULT == "METAL", "Dev not supported")
 class TestSimpleProfiler(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT == "CPU", "fails in CPU")
   def test_profiler(self):

test/backend/test_uops.py

@@ -359,7 +359,8 @@ class TestUOpPrograms(unittest.TestCase):
     prog = C.end(i, j, k)
 
     # run program
-    self._run(prog.sink(arg=KernelInfo()), a, b, c)
+    # TODO: make this work with opts_to_apply
+    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
 
     with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)
 

test/backend/test_wait_loop.py

@@ -141,7 +141,7 @@ class TestWaitLoop(unittest.TestCase):
 class TestVolatileLoops(unittest.TestCase):
   def test_async_wait_ext(self):
     sig_buf = Buffer(Device.DEFAULT, 1, dtypes.int, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
-    try: sig_view = sig_buf.host.view(fmt='i')
+    try: sig_view = sig_buf.as_memoryview(force_zero_copy=True).cast('i')
     except (AssertionError, NotImplementedError): self.skipTest(f"{Device.DEFAULT} does not support host-visible buffers")
     sig_view[0] = 0
 

test/device/test_hcq2.py

@@ -1,7 +1,7 @@
-import unittest, contextlib, ctypes, gc, struct, numpy as np
+import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
-from tinygrad.device import Buffer, Compiled
+from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context, dedup, partition, unwrap
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
@@ -11,16 +11,16 @@ from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.c import init_c_struct_t
 import tinygrad.runtime.support.hcq2 as hcq2
-from tinygrad.runtime.support.hcq2 import HCQ_DEVS, all_devices_in, hcq_compile_cache, link_linear_cache
+from tinygrad.runtime.support.hcq2 import HCQ_DEVS, HCQ2Compiled, all_devices_in, hcq_compile_cache, link_linear_cache
 from test.helpers import call_is_hcq
 
 @contextlib.contextmanager
-def rt_buffers():
-  calls, orig = [], Compiled.rt_buffer
+def rt_views():
+  calls, orig = [], HCQ2Compiled.rt_view
   def track(dev, *args, **kwargs):
     calls.append(dev)
     return orig(dev, *args, **kwargs)
-  with patch.object(Compiled, "rt_buffer", track): yield calls
+  with patch.object(HCQ2Compiled, "rt_view", track): yield calls
 
 def chain(x:Tensor, n:int) -> Tensor:
   for _ in range(n): x = (x + 1).contiguous()
@@ -81,19 +81,19 @@ class TestHCQ2Schedule(unittest.TestCase):
       f(x)
       return f(x), f.captured._linear, [x.uop.base]
     out = chain(x, n)
-    return out, compile_linear(out.schedule_linear(), input_uops=inputs, cache=True), inputs
+    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
 
   def test_jit_has_no_rt_buffers(self):
     dev = Device[Device.DEFAULT]
     rings = [dev.rt_buffer(True, host) for host in (False, True)]
-    ranges = [(b._buf, b._buf + b.nbytes) for b in rings]
+    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
     for n in (1, 65):
       with self.subTest(kernels=n):
         x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
         for _ in range(2): f(x)
         for u in f.captured.linear.toposort():
           if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
-            addr = buf._buf
+            addr = buf._buf.va_addr
             self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
 
   def test_small_eager_cached(self):
@@ -113,6 +113,21 @@ class TestHCQ2Schedule(unittest.TestCase):
     self.assertEqual(device, Device.DEFAULT)
     self.assertEqual(call.src[1 + index].buffer.dtype, dtypes.uint64)
 
+  def test_host_copies(self):
+    dev = Device[Device.DEFAULT]
+    if not dev.has_copy_queue: self.skipTest("copy queue required")
+    for host_device in ("CPU", "NPY", "DISK"):
+      for direct in (False, True):
+        for upload in (False, True):
+          with self.subTest(host_device=host_device, direct=direct, upload=upload):
+            host, gpu = UOp.new_buffer(host_device, 4, dtypes.uint8), UOp.new_buffer(dev.device, 4, dtypes.uint8)
+            src, dst = (host, gpu) if upload else (gpu, host)
+            linear = UOp(Ops.LINEAR, src=(src.copy_to_device(dst.device).call(dst, src),))
+            with patch.object(dev, "host_devs", frozenset({"CPU", host_device}) if direct else frozenset({"CPU"})):
+              compiled = compile_linear(linear, profile=False)
+            self.assertEqual(len(compiled.src), 1 if direct or host_device == "CPU" else 2)
+            self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
+
   def test_large_eager_not_cached(self):
     _, compiled, inputs = self.compiled(65)
     linked = link_linear(compiled, input_uops=inputs)
@@ -126,9 +141,9 @@ class TestHCQ2Schedule(unittest.TestCase):
           out, compiled, inputs = self.compiled(n, jit=jit)
           linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
           before = tuple(inputs)
-          with rt_buffers() as borrowed:
+          with rt_views() as borrowed:
             for linear in (compiled, linked):
-              self.assertIs(compile_linear(linear, input_uops=inputs, cache=not jit), linear)
+              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
           self.assertEqual(tuple(inputs), before)
           self.assertFalse(borrowed)
           run_linear(linked, input_uops=inputs, jit=True, wait=True)
@@ -140,7 +155,7 @@ class TestHCQ2Schedule(unittest.TestCase):
         with self.subTest(kernels=n, jit=jit):
           out, compiled, inputs = self.compiled(n, jit=jit)
           linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
-          with rt_buffers() as borrowed:
+          with rt_views() as borrowed:
             again = link_linear(linked, input_uops=inputs, allow_cache=not jit)
           self.assertIs(again, linked)
           self.assertFalse(borrowed)
@@ -168,9 +183,20 @@ class TestHCQ2Schedule(unittest.TestCase):
   def test_map_cpu_buffer_preserves_contents(self):
     src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
     data = bytes(range(16))
-    src.host[:] = data
+    src.as_memoryview(force_zero_copy=True)[:] = data
     src.get_buf(Device.DEFAULT)
-    self.assertEqual(bytes(src.as_memoryview()), data)
+    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
+
+  def test_staged_copy_roundtrip(self):
+    # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
+    stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
+    for npdt in (np.uint8, np.float32):
+      with self.subTest(dtype=npdt.__name__):
+        n = (size // 2 // np.dtype(npdt).itemsize) * 9 + 7 # nine rotations of a two slot ring, plus a short tail
+        data = np.arange(n, dtype=np.int64).astype(npdt)
+        with patch.object(hcq2, "STAGING_SIZE", size), patch.object(hcq2, "STAGING_SLOTS", 2), patch.object(hcq2, "_staging", lambda: stage):
+          out = Tensor(data).to(Device.DEFAULT).contiguous().realize()
+          np.testing.assert_equal(out.numpy(), data)
 
   def test_rt_patches_are_inputs_and_vars_only(self):
     x = Tensor.rand(17, 33).contiguous().realize()
@@ -204,12 +230,11 @@ class TestHCQ2Schedule(unittest.TestCase):
   def test_caches_hold_no_buffers(self):
     # an eager template caches without its buffers and the jit's linear compiles once uncached: freeing the tensors frees the device memory
     def step(i):
-      buf = Buffer("NPY", 1024, dtypes.float32, initial_value=struct.pack("f", i) * 1024)
-      x = Tensor(UOp.from_buffer(buf)).to(Device.DEFAULT).realize()
+      x = Tensor(np.full(1024, i, np.float32)).to(Device.DEFAULT).realize()
       @TinyJit
       def f(a): return (a * 2 + 1).contiguous().realize()
       for _ in range(3): out = f(x)
-      self.assertEqual(out.to("CPU").tolist(), [2.0 * i + 1] * 1024)
+      self.assertEqual(out.tolist(), [2.0 * i + 1] * 1024)
     step(1) # warms the programs, templates and rings
     gc.collect()
     used = GlobalCounters.mem_used
@@ -261,7 +286,7 @@ class TestHCQ2Schedule(unittest.TestCase):
       linked = hcq2.hcq_link(UOp(Ops.LINEAR, src=(call,)), allow_cache=False).src[0]
       inner_buf, outer_buf = linked.src[1].buffer, linked.without_after.src[1].buffer
       self.assertEqual(inner_buf.host.view(fmt='I')[1], 42)
-      self.assertEqual(outer_buf.host.view(fmt='Q')[0], inner_buf._buf + 4)
+      self.assertEqual(outer_buf.host.view(fmt='Q')[0], inner_buf._buf.va_addr + 4)
 
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):

test/device/test_metal.py

@@ -3,6 +3,7 @@ from tinygrad.device import CompileError, Device, BufferSpec, TinyELF
 from tinygrad.helpers import Target
 if Device.DEFAULT=="METAL":
   from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler
+  from tinygrad.runtime.autogen import metal
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
 class TestMetal(unittest.TestCase):
   def test_alloc_oom(self):
@@ -10,6 +11,13 @@ class TestMetal(unittest.TestCase):
     with self.assertRaises(MemoryError):
       device.allocator.alloc(10000000000000000000)
 
+  def test_import_private_buffer(self):
+    device = Device['METAL']
+    private = device.sysdevice.newBufferWithLength_options(4, metal.MTLResourceStorageModePrivate)
+    (buf, _), host = device.allocator.alloc(4, BufferSpec(external_ptr=private.value))
+    self.assertEqual(buf.buf.value, private.value)
+    self.assertIsNone(host)
+
   def test_compile_error(self):
     compiler = MetalCompiler()
     with self.assertRaises(CompileError):

test/external/external_test_am_fault_recovery.py

@@ -19,7 +19,7 @@ def _run_asm(asm_src:str) -> subprocess.CompletedProcess:
   return _run('from tinygrad.device import Device, TinyELF; from tinygrad.helpers import Target; '
               'from tinygrad.runtime.support.compiler_amd import HIPCompiler; dev = Device["AMD"]; '
               f'dev.runtime(TinyELF(HIPCompiler(dev.arch).compile("""{asm_src}"""), "test", Target("AMD", arch=dev.arch), ()))('
-              'dev.allocator.alloc(64).buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)')
+              'dev.allocator.alloc(64)[0][0], global_size=(1,1,1), local_size=(1,1,1), wait=True)')
 
 def _verify_recovery() -> subprocess.CompletedProcess:
   return _run('from tinygrad import Tensor; t = Tensor([1.0, 2.0], device="AMD").realize(); assert (t + 1).numpy().tolist() == [2.0, 3.0]')

test/external/external_test_hive_reset.py

@@ -20,7 +20,7 @@ extern "C" __attribute__((global)) void broken(int* dummy) {
 '''
 broken_lib = compile_hip(broken_src, dev.arch)
 broken_prg = AMDProgram(dev, "broken", broken_lib)
-buf = dev.allocator.alloc(64).buf
+buf = dev.allocator.alloc(64)[0][0]
 try:
   broken_prg(buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)
   print("  ERROR: Kernel did not fault!")

test/external/external_test_speed_llama.py

@@ -4,7 +4,7 @@ from examples.llama import Transformer, MODEL_PARAMS
 from tinygrad.tensor import Tensor
 from tinygrad import Device
 from tinygrad.nn.state import get_state_dict
-from tinygrad.device import BufferStorage, Allocator, Compiled
+from tinygrad.device import Allocator, Compiled
 from tinygrad.codegen import to_program_cache
 from tinygrad.helpers import Profiling
 
@@ -13,7 +13,7 @@ class FakeProgram:
   def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
 
 class FakeAllocator(Allocator[Compiled]):
-  def _alloc(self, sz, options): return BufferStorage(None)
+  def _alloc(self, sz, options): return (None, None), None
   def _copyin(self, dest, src:memoryview): pass
 
 class TestLLaMASpeed(unittest.TestCase):

test/external/fuzz_kfd.py

@@ -9,7 +9,7 @@ if __name__ == "__main__":
   dev: List[AMDDevice] = [Device[f"KFD:{i}"] for i in range(6)]
   print(f"got {len(dev)} devices")
 
-  buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000)).buf) for i in range(100)]
+  buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000))[0][0]) for i in range(100)]
 
   for _ in trange(100000):
     d1, b1 = random.choice(buffers)

test/helpers.py

@@ -3,7 +3,6 @@ from dataclasses import replace
 from typing import Any, Callable
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.device import Buffer
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen import to_program
@@ -126,14 +125,13 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   allocator = dev.allocator
   bufs = []
   for buf_dt, data in inputs or []:
-    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
-    allocator._copyin(buf.buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
+    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)[0][0])
+    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
-  out_buf = Buffer("PYTHON", 1, uop.dtype, preallocate=True)
-  prog(out_buf._buf, *[b.buf for b in bufs], vals=vals)
-  return out_buf.as_memoryview().cast(uop.dtype.fmt or "").tolist()[0]
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize)[0][0], *bufs, vals=vals)
+  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
   sink = UOp.group(*u)

test/mockgpu/amd/emu.py

@@ -1883,14 +1883,14 @@ class WaveState:
     # CDNA (wave64) has separate ACCVGPR file; RDNA shares with VGPR
     if wave_size == 64:
       self.accvgpr_buf = Buffer('CPU', vgpr_size, dtypes.uint32).ensure_allocated()
-      ctypes.memset(self.accvgpr_buf._buf, 0, vgpr_size * 4)
+      ctypes.memset(self.accvgpr_buf._buf.va_addr, 0, vgpr_size * 4)
     else:
       self.accvgpr_buf = self.vgpr_buf
-    self._vgpr_mv = self.vgpr_buf.host.view(fmt='I').mv
-    self._sgpr_mv = self.sgpr_buf.host.view(fmt='I').mv
+    self._vgpr_mv = self.vgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
+    self._sgpr_mv = self.sgpr_buf.as_memoryview(force_zero_copy=True, no_sync=True).cast('I')
     # Zero memory using ctypes memset (much faster than Python loops)
-    ctypes.memset(self.vgpr_buf._buf, 0, vgpr_size * 4)
-    ctypes.memset(self.sgpr_buf._buf, 0, SGPR_COUNT * 4)
+    ctypes.memset(self.vgpr_buf._buf.va_addr, 0, vgpr_size * 4)
+    ctypes.memset(self.sgpr_buf._buf.va_addr, 0, SGPR_COUNT * 4)
     # Pre-populate inline constants at indices 128-255
     for i in range(65): self._write_sgpr(128 + i, i)  # 128-192: integers 0-64
     for i in range(16): self._write_sgpr(193 + i, (-(i + 1)) & MASK32)  # 193-208: -1 to -16
@@ -1992,10 +1992,11 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
     waves: list[tuple[WaveState, list]] = []
     for wave_start in range(0, total_threads, wave_size):
       st = _init_wave(lib, wave_start, total_threads, lx, ly, lz, args_ptr, rsrc2, scratch_size, arch, gidx, gidy, gidz, user_data, wave_size)
-      scratch_base = scratch_buf._buf + (wave_start // wave_size) * scratch_size * wave_size if scratch_buf else 0
-      waves.append((st, [ctypes.c_uint64(st.sgpr_buf._buf), ctypes.c_uint64(st.vgpr_buf._buf),
-                         ctypes.c_uint64(vmem_buf._buf), ctypes.c_uint64(lds_buf._buf),
-                         ctypes.c_uint64(scratch_base if scratch_buf else 0), ctypes.c_uint64(st.accvgpr_buf._buf)]))
+      scratch_base = scratch_buf._buf.va_addr + (wave_start // wave_size) * scratch_size * wave_size if scratch_buf else 0
+      waves.append((st, [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
+                         ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
+                         ctypes.c_uint64(scratch_base if scratch_buf else 0),
+                         ctypes.c_uint64(st.accvgpr_buf._buf.va_addr)]))
     done = [False] * len(waves)
     for _ in range(10_000_000):
       if all(done): return
@@ -2026,7 +2027,7 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
     for gidz, gidy, gidx in itertools.product(range(gz), range(gy), range(gx)):
       _run_workgroup(gidx, gidy, gidz, tracing)
       tracing = False  # only trace the first workgroup
-      if lds_size > 0: ctypes.memset(lds_buf._buf, 0, max(lds_size, 4))  # reset LDS for next workgroup
+      if lds_size > 0: ctypes.memset(lds_buf._buf.va_addr, 0, max(lds_size, 4))  # reset LDS for next workgroup
 
   if PROFILE: sqtt_traces.append(sqtt_finalize())
   return 0

test/mockgpu/mockgpu.py

@@ -29,6 +29,8 @@ class TrackedMemoryView:
     self.mv = self.mv.cast('B').cast(new_type, **kwargs)
     return self
 
+  @property
+  def obj(self): return self.mv.obj
   @property
   def nbytes(self): return self.mv.nbytes
   def __len__(self): return len(self.mv)

test/mockgpu/nv/nvdriver.py

@@ -268,8 +268,8 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_REGISTER_CHANNEL: pass
     elif nr == nv_gpu.UVM_FREE:
       st = nv_gpu.UVM_FREE_PARAMS.from_address(argp)
-      if st.base not in self.host_ranges: libc.munmap(st.base, st.length) # registered host memory belongs to its allocator
       self.host_ranges.discard(st.base)
+      libc.munmap(st.base, st.length)
     else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
     return 0
 

test/opt/test_kernel_opts.py

@@ -249,32 +249,6 @@ class TestKernelOpts(unittest.TestCase):
     # both reduce axes padded: the outer clause lifts out, the inner clause is the inner reduce's identity
     helper_linearizer_opt(a.max(1).sum(0), [[Opt(OptOps.PADTO, 0, 4), Opt(OptOps.PADTO, 1, 4)]], wanna_output=[[3+6]])
 
-  def test_padto_unindexed_reduce(self):
-    # TODO: broken now. a repeat/cat along the reduced axis leaves a reduce range no buffer index uses, and the padded iterations are reduced too
-    a = Tensor.arange(7*5, dtype=dtypes.float).reshape(7, 5).clone().realize()
-    with self.assertRaises(AssertionError):  # sums 32 not 16
-      helper_linearizer_opt(a.repeat((3, 16)).sum(1), [[Opt(OptOps.PADTO, 2, 32)]])
-    with self.assertRaises(AssertionError):  # sums 6 not 4
-      helper_linearizer_opt(a.repeat((1, 16)).sum(1), [[Opt(OptOps.SPLIT, 1, (4, AxisType.UNROLL)), Opt(OptOps.PADTO, 1, 3)]])
-    with self.assertRaises(AssertionError):  # sums 4 not 2
-      helper_linearizer_opt(a.cat(a, dim=1).sum(1), [[Opt(OptOps.PADTO, 1, 4)]])
-    a = Tensor.full((7, 5), 2.0).clone().realize()
-    with self.assertRaises(AssertionError):  # 2**30 not 2**20
-      helper_linearizer_opt(a.repeat((1, 4)).prod(1), [[Opt(OptOps.PADTO, 1, 3)]])
-
-  def test_padto_masked_reduce(self):
-    # TODO: broken now. a where with a defined false arm gives the padded iterations a value, and it is reduced too
-    a = Tensor.arange(7*17, dtype=dtypes.float).reshape(7, 17).clone().realize()
-    m = (Tensor.arange(7).reshape(7, 1) % 2 == 0).expand(7, 17)
-    with self.assertRaises(AssertionError):  # sums 7 extra 1.0s
-      helper_linearizer_opt(m.where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 8)]])
-    with self.assertRaises(AssertionError):  # sums 15 extra 1.0s
-      helper_linearizer_opt(m.where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 32)]])
-    with self.assertRaises(AssertionError):  # sums 7 extra 1.0s
-      helper_linearizer_opt((Tensor.arange(17).reshape(1, 17) < 5).expand(7, 17).where(a, 1.0).sum(1), [[Opt(OptOps.PADTO, 1, 8)]])
-    with self.assertRaises(AssertionError):  # 2**16 not 2**11
-      helper_linearizer_opt(m[:, :11].where(Tensor.ones(7, 11), 2.0).prod(1), [[Opt(OptOps.PADTO, 1, 8)]])
-
   def test_padto_unrolled_prod(self):
     a = (Tensor.arange(4*17, dtype=dtypes.float).reshape(4, 17) / 100 + 1).clone().realize()
     helper_linearizer_opt(a.prod(1), [[Opt(OptOps.PADTO, 1, 32), Opt(OptOps.SPLIT, 1, (0, AxisType.UNROLL)),

test/opt/test_tensor_cores.py

@@ -208,24 +208,6 @@ class TestTensorCores(unittest.TestCase):
                                                             Opt(OptOps.SPLIT, axis, (0, AxisType.UNROLL))]],
                           check_default_opt=False, atol=3e-2, rtol=1e-3)
 
-  @Context(ALLOW_TF32=1)
-  @unittest.skipUnless(any(tc.dtype_in in (dtypes.half, dtypes.float) for tc in Device[Device.DEFAULT].renderer.tensor_cores),
-                       "test requires half or float tensor cores")
-  def test_tensor_cores_padto_masked_operand(self):
-    # tc_opt=2 pads K. an ALU between the load and the multiply is fine, a where with a defined false arm is not
-    tc = next(tc for tc in Device[Device.DEFAULT].renderer.tensor_cores if tc.dtype_in in (dtypes.half, dtypes.float))
-    Tensor.manual_seed(3)
-    a = Tensor.rand(tc.dims[1]*2+1, tc.dims[2]*3-1, dtype=tc.dtype_in).realize()
-    b = Tensor.rand(tc.dims[2]*3-1, tc.dims[0]*2+1, dtype=tc.dtype_in).realize()
-    tc_opt = Opt(OptOps.TC, 0, (-1, 2, 1))
-    helper_linearizer_opt((a+1).matmul(b+1, dtype=tc.dtype_out), [[tc_opt]], check_default_opt=False, atol=3e-2, rtol=1e-3)
-    one = Tensor(1, dtype=tc.dtype_in)
-    ma = (Tensor.rand(a.shape[0], 1) > 0.5).expand(a.shape).where(a, one)
-    mb = (Tensor.rand(1, b.shape[1]) > 0.5).expand(b.shape).where(b, one)
-    # TODO: broken now, the padded K lanes multiply 1.0*1.0
-    with self.assertRaises(AssertionError):
-      helper_linearizer_opt(ma.matmul(mb, dtype=tc.dtype_out), [[tc_opt]], check_default_opt=False, atol=3e-2, rtol=1e-3)
-
   @Context(ALLOW_TF32=1)
   @unittest.skipIf(Device.DEFAULT == "PYTHON", "not generated on EMULATED device")
   @slow

test/unit/test_after.py

@@ -42,17 +42,19 @@ class TestAfterCounterexamples(unittest.TestCase):
     # y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
 
+  @unittest.expectedFailure
   def test_partial_store_gradient(self):
     x = Tensor([2., 3.]).realize()
     y = Tensor(x.uop.after(x[:1].uop.store(4)))
-    # y = [4, x[1]]; only the untouched element depends on x.
+    # y = [4, x[1]]. Currently returns [0., 0.].
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
+  @unittest.expectedFailure
   def test_partial_store_source_gradient(self):
     x = Tensor([4.])
     y = Tensor([2., 3.]).realize()
     z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
-    # x contributes once, not twice.
+    # x contributes once, not twice. Currently returns [2.].
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
   def test_unrelated_store_gradient(self):
@@ -62,26 +64,14 @@ class TestAfterCounterexamples(unittest.TestCase):
     # Zeroing y does not change x.
     self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
 
+  @unittest.expectedFailure
   def test_after_dependency_gradient(self):
     x = Tensor([2., 3.])
     y = x.clone()
     y[:1].assign(0)
-    # View assign is an AFTER on a partial STORE; only the untouched element depends on x.
+    # View assign creates a nested AFTER; currently raises in backward.
     self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
 
-  def test_view_assign_gradient(self):
-    for view, expected in ((lambda t: t.reshape(3, 2)[1:], [[1., 1., 0.], [0., 0., 0.]]),
-                           (lambda t: t.permute(1, 0)[1:], [[1., 0., 0.], [1., 0., 0.]]),
-                           (lambda t: t.flip((0, 1))[:1], [[1., 1., 1.], [0., 0., 0.]])):
-      with self.subTest(expected=expected):
-        x = Tensor([[1., 2., 3.], [4., 5., 6.]])
-        y = x.clone()
-        v = Tensor.full(view(y).shape, 7.)
-        view(y).assign(v)
-        gx, gv = y.sum().gradient(x, v)
-        self.assertEqual(gx.tolist(), expected)
-        self.assertEqual(gv.tolist(), Tensor.ones(v.shape).tolist())
-
   @unittest.expectedFailure
   def test_unordered_overlapping_stores_rejected(self):
     x = Tensor([0.]).realize().uop

test/unit/test_buffer.py

@@ -10,31 +10,30 @@ class TestBuffer(unittest.TestCase):
     host = v.host
     host.view(fmt='H')[0] = 0x1234
     self.assertEqual(b.host.view(fmt='H')[2], 0x1234)
-    self.assertEqual(v._buf, b._buf + 4)
+    self.assertEqual(v._buf.va_addr, b._buf.va_addr + 4)
     self.assertIs(v.host, host)
     self.assertIs(v.meta, b.meta)
 
   def test_mapping(self):
     b = Buffer("CPU", 8, dtypes.uint8, initial_value=b"abcdefgh")
-    self.assertEqual(b.get_buf("PYTHON"), b._buf)
+    self.assertIs(b.get_storage("PYTHON")[0][1], b.get_buf("PYTHON"))
     v = b.view(4, dtypes.uint8, 2)
     mapped = v.get_storage("PYTHON")
-    self.assertEqual(mapped.buf, b._buf + 2)
-    self.assertEqual(bytes(mapped.host.mv), b"cdef")
-    self.assertIs(mapped.host, v.host)
-    self.assertIsNone(mapped.meta)
-    self.assertIs(v.get_storage("PYTHON"), mapped)
+    self.assertEqual(bytes(mapped[0][0]), b"cdef")
+    self.assertIs(mapped[1], v.host)
+    self.assertIsNone(mapped[0][1])
+    self.assertIs(v.get_storage("PYTHON")[0], mapped[0])
 
   def test_view_reallocation(self):
     b = Buffer("CPU", 8, dtypes.uint8)
     v = b.view(4, dtypes.uint8, 2)
-    old = v.get_storage("PYTHON")
+    old = v.get_storage("PYTHON")[0]
     b.deallocate()
     b.allocate()
     self.assertFalse(v.is_allocated())
     v.host[:] = b"test"
-    self.assertIsNot(v.get_storage("PYTHON"), old)
-    self.assertEqual(bytes(v.get_storage("PYTHON").host.mv), b"test")
+    self.assertIsNot(v.get_storage("PYTHON")[0], old)
+    self.assertEqual(bytes(v.get_buf("PYTHON")), b"test")
 
   def test_cache_owned_storage_only(self):
     for opaque in (None, memoryview(bytearray(8))):

test/unit/test_setitem_schedule.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor, Device, dtypes, GlobalCounters
+from tinygrad import Tensor, dtypes, GlobalCounters
 from test.helpers import assert_kernel_count
 
 class TestSetitemInto(unittest.TestCase):
@@ -139,16 +139,5 @@ class TestSetitemInto(unittest.TestCase):
     assert_kernel_count(1)
     self.assertEqual(GlobalCounters.global_mem, 100*4)  # full buffer written
 
-  @unittest.skipUnless(Device.DEFAULT != "CPU", "source must be on another device")
-  def test_setitem_slice_assign_from_other_device(self):
-    # NOTE: this is 2 kernels, the cross-device copy should fuse with the assign into one
-    a = Tensor.ones(20, device="CPU")
-    b = Tensor.arange(20).float().clone()
-    Tensor.realize(a, b)
-    GlobalCounters.reset()
-    a[10:12].assign(b[13:15].to(a.device)).realize()
-    assert_kernel_count(2)
-    self.assertListEqual(a.tolist(), [1.0]*10 + [13.0, 14.0] + [1.0]*8)
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/opt/heuristic.py

@@ -25,7 +25,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   TC_MIN_GLOBALS -- do not upcast N when it would drop the specified global count
   """
   # NOTE: unless TC_OPT is > 0, we only trigger tensor cores if there's only one reduce axis
-  if USE_TC > 0 and (len(k.reduce_axes) == 1 or (TC_OPT.value >= 1)):
+  if USE_TC > 0 and (len(k.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
     for axis in range(3):
       tk = k.copy()
       # check TC first and apply hand-coded opts if successful
@@ -79,7 +79,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
             return k
 
   # are we grouping? (requires local shape support)
-  if resolve(prod(k.full_shape[i] for i in k.upcastable_dims) <= (240 if k.ren.target.device == "QCOM" else 2048), False):
+  if resolve(prod(k.output_shape[i] for i in k.upcastable_dims) <= (240 if k.ren.target.device == "QCOM" else 2048), False):
     for axis, sz in itertools.product(k.axes_of(AxisType.REDUCE)[:3], (16,)):
       try:
         k.apply_opt(Opt(OptOps.SPLIT, axis, (sz, AxisType.GROUP_REDUCE, True)))
@@ -112,7 +112,7 @@ def hand_coded_optimizations(k:Scheduler) -> Scheduler:
   # potentially do more upcasts of non reduce axes based on a heuristic
   is_dsp = k.ren is not None and k.ren.target.device == "DSP"
   upcasted_axis: set[int] = set()
-  while resolve(prod(k.full_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
+  while resolve(prod(k.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
     xb_choices = []
     # consider all upcastable axes with 3 or 4 upcast (128 on the DSP)
     for axis, upcast_amount in itertools.product(k.upcastable_dims, ([128] if not len(upcasted_axis) else []) if is_dsp else [3,4]):

tinygrad/codegen/opt/postrange.py

@@ -99,19 +99,14 @@ class Scheduler:
   def ranges_of(self, *axis_type:AxisType) -> list[UOp]: return [r for r in self.rngs if r.arg[-1] in axis_type]
   def axes_of(self, *axis_type:AxisType) -> list[int]: return [i for i,t in enumerate(self.axis_types) if t in axis_type]
 
-  @property
-  def reduce_axes(self) -> list[int]:
-    red = {r for u in self.ast.backward_slice if u.op is Ops.REDUCE for s in u.src[1:] for r in s.ranges}
-    return [i for i,r in enumerate(self.rngs) if r in red]
-
   def upcast_size(self): return prod(self.full_shape[a] for a in self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
 
   @property
   def upcastable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GLOBAL, AxisType.LOCAL, AxisType.WEAK) \
                                                   if isinstance(s:=self.full_shape[i], int) and s > 1]
   @property
-  def unrollable_dims(self) -> list[int]: return [i for i in self.reduce_axes if self.axis_types[i] in (AxisType.GROUP_REDUCE, AxisType.REDUCE) \
-                                                  and isinstance(s:=self.full_shape[i], int) and s > 1]
+  def unrollable_dims(self) -> list[int]: return [i for i in self.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
+                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
 
   def real_axis(self, op:OptOps, axis:int|None) -> int:
     if axis is None or op is OptOps.TC: return -1
@@ -296,6 +291,9 @@ class Scheduler:
   @property
   def bufs(self) -> list[UOp]: return [x for x in self.ast.toposort() if x.op is Ops.INDEX][::-1]
   @property
+  def output_shape(self):
+    return [s if at not in {AxisType.REDUCE, AxisType.UNROLL, AxisType.GROUP_REDUCE} else 1 for s,at in zip(self.full_shape, self.axis_types)]
+  @property
   def upcasted(self) -> int: return len(self.axes_of(AxisType.UPCAST, AxisType.UNROLL))
   @property
   def group_for_reduces(self) -> int: return len(self.axes_of(AxisType.GROUP_REDUCE))

tinygrad/codegen/opt/search.py

@@ -106,9 +106,9 @@ def get_kernel_actions(s:Scheduler, include_0=True, max_up:int|None=None) -> dic
   return acted
 
 BEAM_DEBUG = getenv("BEAM_DEBUG")
-def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True):
+def beam_search(s:Scheduler, rawbufs:list[Buffer], var_vals:dict[str,int], amt:int, allow_test_size=True, disable_cache=IGNORE_BEAM_CACHE.value):
   key = {"ast": s.ast.key, "amt": amt, "allow_test_size": allow_test_size, "device": s.ren.target.device, "suffix": s.ren.suffix}
-  if not IGNORE_BEAM_CACHE and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
+  if not disable_cache and CACHELEVEL >= 1 and (val:=diskcache_get("beam_search", key)) is not None:
     ret = s.copy()
     for o in val[len(s.applied_opts):]: ret.apply_opt(o)
     return ret

tinygrad/device.py

@@ -1,17 +1,14 @@
 from __future__ import annotations
-from dataclasses import dataclass, replace, field
+from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
-import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct, mmap, time, statistics
-from tinygrad.helpers import WIN, mv_address, to_mv, LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
+import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
+from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, GlobalCounters, PROFILE, temp, colored
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
-from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up, is_numpy_ndarray
-from tinygrad.helpers import cpu_profile, perf_counter_us
-from tinygrad.dtype import dtypes, DType, _to_np_dtype
-from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
-if TYPE_CHECKING:
-  from tinygrad.renderer import Renderer
-  from tinygrad.uop.ops import UOp
+from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
+from tinygrad.dtype import DType, _to_np_dtype
+from tinygrad.runtime.support.memory import MMIOInterface
+if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
 # **************** Device ****************
 
@@ -100,9 +97,6 @@ class MultiBuffer:
   def is_allocated(self): return all(x.is_allocated() for x in self.bufs)
   def __repr__(self): return f"<multibuf real:{self.is_allocated()} device:{tuple(x.device for x in self.bufs)} size:{self.size} dtype:{self.dtype}>"
 
-@dataclass(frozen=True)
-class BufferStorage: buf:Any; meta:Any=None; host:MMIOInterface|None=None; maps:dict[str, BufferStorage]=field(default_factory=dict) # noqa: E702
-
 class Buffer:
   profile_events:list[ProfileEvent] = []
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
@@ -110,14 +104,14 @@ class Buffer:
     assert isinstance(dtype, DType)
     self.device, self.size, self.dtype, self.offset, self.allocated_views, self._base = Device.canonicalize(device), size, dtype, offset, 0, base
     self.options = options if options is not None else BufferSpec()
-    self._storage:BufferStorage|None = None
+    self._storage:tuple|None = None
+    self._maps:dict[str, tuple] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
       if opaque is not None: self.allocate(opaque)
       if initial_value is not None:
         self.allocate()
-        if (host:=self.get_storage().host) is not None: host[:] = memoryview(initial_value).cast('B')
-        else: self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
+        self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
         if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
@@ -136,25 +130,24 @@ class Buffer:
   @functools.cached_property
   def allocator(self) -> Allocator: return self.base.allocator if self._base is not None else Device[self.device].allocator
   @property
-  def _buf(self) -> Any: return self.get_storage().buf
+  def _buf(self) -> Any: return self.get_storage()[0][0]
   @property
-  def host(self) -> MMIOInterface: return unwrap(self.get_storage().host)
+  def host(self) -> MMIOInterface: return unwrap(self.get_storage()[1])
   @property
-  def meta(self) -> Any: return self.get_storage().meta
+  def meta(self) -> Any: return self.get_storage()[0][1]
   @property
   def nbytes(self): return self.size * self.dtype.itemsize
 
-  def get_storage(self, device:str|None=None) -> BufferStorage:
+  def get_storage(self, device:str|None=None) -> tuple:
     storage = unwrap(self.ensure_allocated()._storage)
     device = Device.canonicalize(device) if device is not None else self.device
     if device == self.device: return storage
-    if device not in storage.maps:
-      alloc = Device[device].allocator
-      storage.maps[device] = BufferStorage(alloc._offset(self.base.get_buf(device), self.nbytes, self.offset)) if self._base else alloc.map(self)
-    if storage.maps[device].host is not storage.host: storage.maps[device] = replace(storage.maps[device], host=storage.host)
-    return storage.maps[device]
+    if device not in self._maps:
+      allocator = Device[device].allocator
+      self._maps[device] = (allocator._offset(self.base.get_buf(device), self.nbytes, self.offset), None) if self._base else allocator.map(self)
+    return self._maps[device], storage[1]
 
-  def get_buf(self, device:str) -> Any: return self.get_storage(device).buf
+  def get_buf(self, device:str) -> Any: return self.get_storage(device)[0][0]
 
   def is_allocated(self) -> bool: return self._storage is not None and (self._base is None or self._base_storage is self.base._storage)
   def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self
@@ -165,24 +158,19 @@ class Buffer:
       raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
     if external_ptr is not None: self.options = replace(self.options, external_ptr=external_ptr)
     if self._base is not None:
-      storage = replace(self.base.get_storage(), buf=self.allocator._offset(self.base._buf, self.nbytes, self.offset), maps={})
-    elif opaque is not None:
-      self.options = replace(self.options, nolru=True)
-      if is_numpy_ndarray(opaque):
-        if not opaque.flags.c_contiguous: opaque = opaque.copy(order='C')
-        opaque = BufferStorage(addr:=opaque.ctypes.data, memoryview(opaque), MMIOInterface(addr, self.nbytes))
-      elif isinstance(opaque, memoryview):
-        opaque = BufferStorage(addr:=mv_address(opaque) if self.nbytes else 0, opaque, MMIOInterface(addr, self.nbytes))
-      storage = opaque if isinstance(opaque, BufferStorage) else BufferStorage(opaque)
-    else: storage = self.allocator.alloc(self.nbytes, self.options)
-    storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None)
+      (buf, meta), host = self.base.get_storage()
+      mapping = self.allocator._offset(buf, self.nbytes, self.offset), meta
+    else:
+      if opaque is not None: self.options = replace(self.options, nolru=True)
+      mapping, host = ((opaque, None), None) if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
+    storage = mapping, host.view(self.offset, self.nbytes, fmt='B') if host is not None else None
     if self._base is None:
       if not self.device.startswith("DISK") and self.options.external_ptr is None:
         GlobalCounters.mem_used += self.nbytes
         GlobalCounters.mem_used_per_device[self.device] += self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
     elif self._storage is None: self.base.allocated_views += 1
-    self._storage, self._base_storage = storage, self.base._storage if self._base else None
+    self._storage, self._maps, self._base_storage = storage, {}, self.base._storage if self._base else None
     return self
 
   def deallocate(self):
@@ -193,18 +181,16 @@ class Buffer:
         GlobalCounters.mem_used -= self.nbytes
         GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
+      for dev, mb in self._maps.items(): Device[dev].allocator._unmap(mb[0])
       self.allocator.free(self._storage, self.nbytes, self.options)
     else: self.base.allocated_views -= 1
-    self._storage, self._base_storage = None, None
+    self._storage, self._maps, self._base_storage = None, {}, None
 
   def __reduce_ex__(self, protocol):
     buf:bytearray|pickle.PickleBuffer|None = None
     if self._base is not None:
       return self.__class__, (self.device, self.size, self.dtype, None, None, None, self.base, self.offset, self.is_allocated())
-    if self.device == "NPY":
-      import numpy as np
-      arr = np.frombuffer(self.meta, _to_np_dtype(self.dtype)) # over the storage itself, so an out-of-band pickle buffer keeps it alive
-      return self.__class__, (self.device, self.size, self.dtype, arr, self.options, None)
+    if self.device == "NPY": return self.__class__, (self.device, self.size, self.dtype, self._buf, self.options, None)
     if self.is_allocated():
       buf = pickle.PickleBuffer(self.as_memoryview()) if protocol >= 5 else bytearray(self.as_memoryview())
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf)
@@ -215,15 +201,19 @@ class Buffer:
     return self._trace_num
 
   def _host_mv(self) -> memoryview|None:
-    if self.is_allocated() and hasattr(host:=self.get_storage().host, 'mv'): return unwrap(host).view(fmt='B').mv
+    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'):
+      mv = unwrap(host).view(fmt='B').mv
+      mv.obj._buffer = self  # raw ctypes views do not own their memory; keep the allocation alive for asynchronous copies
+      return mv
     if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
     return None
 
-  def as_memoryview(self, allow_zero_copy=False) -> memoryview:
-    if (mv:=self._host_mv()) is not None:
-      self.allocator.dev.synchronize()
-      if allow_zero_copy: return mv
-      with cpu_profile(f"{self.device} -> TINY", f"{self.device}:COPY"): return memoryview(bytearray(mv))
+  def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False, no_sync=False) -> memoryview:
+    # zero copy with as_memoryview (disabled by default due to use after free)
+    if (force_zero_copy or allow_zero_copy) and (mv:=self._host_mv()) is not None:
+      if not no_sync: self.allocator.dev.synchronize()
+      return mv
+    assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv
 
@@ -254,10 +244,10 @@ class Allocator(Generic[DeviceType]):
   def __init__(self, dev:DeviceType, supports_copy_from_disk:bool=True, supports_transfer:bool=True):
     self.dev: DeviceType = dev
     self.default_buffer_spec: BufferSpec = BufferSpec()
-    self.cache:dict[tuple[int, BufferSpec|None], list[BufferStorage]] = defaultdict(list)
+    self.cache:dict[tuple[int, BufferSpec|None], list[tuple]] = defaultdict(list)
     self.supports_copy_from_disk, self.supports_transfer = supports_copy_from_disk, supports_transfer
 
-  def alloc(self, size:int, options:BufferSpec|None=None) -> BufferStorage:
+  def alloc(self, size:int, options:BufferSpec|None=None) -> tuple:
     assert size > 0, f"alloc size must be positive, getting {size}"
     if len(c:=self.cache[(size, options)]): return c.pop()
     spec = options if options is not None else self.default_buffer_spec
@@ -267,51 +257,29 @@ class Allocator(Generic[DeviceType]):
     except (RuntimeError, MemoryError) as e: raise MemoryError(f"Allocation of {size_to_str(size)} failed on {self.dev.device}. "
                                                             f"Used: {size_to_str(GlobalCounters.mem_used_per_device[self.dev.device])}") from e
 
-  def free(self, storage:BufferStorage, size:int, options:BufferSpec|None=None):
+  def free(self, storage:tuple, size:int, options:BufferSpec|None=None):
     spec = options if options is not None else self.default_buffer_spec
     if LRU and self.lru and not (spec.nolru or spec.zero) and spec.external_ptr is None: self.cache[(size, options)].append(storage)
-    else: self.do_free(storage, spec)
+    else: self._free(storage[0][0], spec)
 
   def free_cache(self):
     for (_, options), storages in self.cache.items():
-      for storage in storages: self.do_free(storage, options if options is not None else self.default_buffer_spec)
+      for storage in storages: self._free(storage[0][0], options if options is not None else self.default_buffer_spec)
       storages.clear()
 
-  def do_free(self, storage:BufferStorage, options:BufferSpec):
-    for dev in storage.maps: Device[dev].synchronize()
-    for dev, mb in storage.maps.items(): Device[dev].allocator._unmap(mb)
-    if options.external_ptr is None: self._free(storage, options)
-
-  def map(self, buf:Buffer) -> BufferStorage: return self._map(buf.ensure_allocated())
+  def map(self, buf:Buffer) -> tuple: return self._map(buf.ensure_allocated()._buf)
 
   # implemented by the runtime
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: raise NotImplementedError("need alloc")
-  def _free(self, storage:BufferStorage, options:BufferSpec): pass  # if opaque is a Python object, you don't need a free
+  def _alloc(self, size:int, options:BufferSpec) -> tuple: raise NotImplementedError("need alloc")
+  def _free(self, opaque, options:BufferSpec): pass  # if opaque is a Python object, you don't need a free
   def _copyin(self, dest, src:memoryview): raise NotImplementedError("need copyin")
   def _copyout(self, dest:memoryview, src): raise NotImplementedError("need copyout")
-  def _map(self, buf) -> BufferStorage: raise NotImplementedError("need map")
+  def _map(self, buf) -> tuple: raise NotImplementedError("need map")
   def _unmap(self, mb): pass  # default no-op; override if _map allocates iface-side state
   def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
   # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
   def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional
 
-class HostAllocator(Allocator):
-  def __init__(self, dev): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    if options.external_ptr is not None: addr, buf = options.external_ptr, None
-    elif WIN: 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_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
-    return BufferStorage(addr, buf, MMIOInterface(addr, size, fmt='B'))
-
-  def _copyin(self, dest:int, src:memoryview):
-    self.dev.synchronize()
-    with cpu_profile(f"TINY -> {self.dev.device}", f"{self.dev.device}:COPY"): to_mv(dest, src.nbytes)[:] = src.cast('B')
-  def _copyout(self, dest:memoryview, src:int):
-    self.dev.synchronize()
-    with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): dest[:] = to_mv(src, dest.nbytes)[:]
-  def _map(self, buf:Buffer) -> BufferStorage: return BufferStorage(buf.host.addr)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
-
 class DepsTracker:
   def __init__(self):
     # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.
@@ -391,42 +359,18 @@ class Compiled:
   ifaces:list[Callable] = []
   profile_events:list[ProfileEvent] = [ProfileDeviceEvent("CPU")] # NOTE: CPU is the default device.
 
-  timestamp_divider: float = 1000.0
-  wait_timeout_ms: float = 30000.0
-  sleep_timeout_ms: int|None = None
-  can_recover:bool = False
-  rtalloc_size:int = 64<<20 # the pool every per-linear buffer is carved out of
-  var_vals: dict[str, int] = {}
+  has_copy_queue:bool = True
 
-  # hcq2
   pm_batch:Any = None
   pm_encode:Any = None
   pm_lower:Any = None
+  pm_bufferize:Any = None
 
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):
     from tinygrad.renderer import Renderer
-    from tinygrad.uop.ops import Ops, UPat, PatternMatcher
-    from tinygrad.runtime.support.hcq2 import cfunc_buf
-
     self.device, self.allocator, self.runtime_t, self.graph, self.renderers = device, allocator, runtime, graph, renderers or [Renderer]
     self.device_id, self.arch = (int(idx) if ":" in device and (idx:=device.split(":")[1]).isdigit() else 0), arch
     self.cached_renderer:dict[Any, Renderer] = {}
-    self.pending:dict[str, int] = {} # timeline values of the devices that touched our memory
-
-    # hcq2
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag="timeline"), lambda ctx: ctx.timeline),
-      (UPat(Ops.PARAM, tag="program", name="b"),
-       lambda ctx, b: ctx.prog_bufs.setdefault(b, Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)))),
-      (UPat(Ops.PARAM, name="b"), lambda b, cfunc_buf=cfunc_buf: cfunc_buf(*b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == "cfunc" else None),
-    ])
-
-    # profiling
-    self.prog_bufs:dict[UOp, Buffer] = {} # cache bufferized for programs
-    self.prof_ents:dict[tuple[Buffer, int], ProfileGraphEntry] = {} # (a batch's timestamps, start slot) -> entry, read at synchronize
-
-  @property
-  def has_copy_queue(self) -> bool: return True
 
   @property
   def renderer(self) -> Renderer: return self._select_renderer()
@@ -438,58 +382,6 @@ class Compiled:
 
   def runtime(self, obj:TinyELF) -> Program[Self]: return unwrap(self.runtime_t)(self, obj)
 
-  @functools.cache
-  def rt_allocator(self, uncached:bool=True, host:bool=False) -> BumpAllocator: return BumpAllocator(self.rtalloc_size)
-
-  @functools.cache
-  def rt_buffer(self, uncached:bool=True, host:bool=False) -> Buffer:
-    spec = BufferSpec(host=host, uncached=uncached, cpu_access=True)
-    return Buffer(self.device, self.rt_allocator(uncached, host).size, dtypes.uint8, options=spec, preallocate=True)
-
-  @functools.cached_property
-  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]
-    return Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), initial_value=bytes(16))
-
-  def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
-    timeout = timeout if timeout is not None and self.can_recover else None
-    st, done = time.perf_counter(), sig[0]
-    while done < value:
-      if done != (done:=sig[0]): st = time.perf_counter()
-      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: raise RuntimeError(f"{self.device} signal wait timed out")
-      elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
-
-  def synchronize(self, timeout:int|None=None):
-    try:
-      self._wait_signal(tl:=self.timeline.host.view(fmt='Q'), tl[1], timeout)
-      for d, v in self.pending.items(): Device[d]._wait_signal(Device[d].timeline.host.view(fmt='Q'), v, timeout)
-    except RuntimeError:
-      self.on_device_hang()
-      raise
-    if self.prof_ents: self.collect_prof()
-
-  def count(self) -> int:
-    """
-    Returns the number of physical accelerators available to the runtime.
-    """
-    return self.iface.count if hasattr(self, 'iface') else 1
-
-  def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
-
-  def on_sleep(self):
-    if (iface:=getattr(self, "iface", None)) is not None and hasattr(iface, "sleep"): iface.sleep(self.sleep_timeout_ms)
-
-  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
-
-  def finalize(self):
-    """
-    Called at the end of process lifetime to allow the device to finalize.
-    """
-    try: self.synchronize() # try to finalize the device in any case
-    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
-
-  # helpers
-
   def _renderer_name(self, r:type[Renderer]) -> str:
     return r.__name__.upper().removesuffix("RENDERER").removeprefix(devname:=self.device.split(':')[0].upper()) or devname
 
@@ -510,28 +402,29 @@ class Compiled:
     return select_first_inited([functools.partial(iface, self, self.device_id) for iface in filtered],
                                f"No interface for {dev}:{self.device_id} is available")
 
-  # profiling
+  def count(self) -> int:
+    """
+    Returns the number of physical accelerators available to the runtime.
+    """
+    return self.iface.count if hasattr(self, 'iface') else 1
 
-  def collect_prof(self):
-    if PROFILE:
-      es = list(self.prof_ents.items())
-      sigs = [buf.host.view(fmt='Q')[i]/decimal.Decimal(self.timestamp_divider) for (buf, _), e in es for i in (e.st_id, e.en_id)]
-      Compiled.profile_events.append(ProfileGraphEvent([replace(e, st_id=2*i, en_id=2*i+1) for i,(_, e) in enumerate(es)], [], sigs))
-    self.prof_ents.clear()
+  def synchronize(self):
+    """
+    Synchronize all pending operations on the device.
 
+    This method ensures that all previously queued operations on the device have been completed before proceeding.
+    """
+    # override this in your device implementation
   def _at_profile_finalize(self):
-    if self.pm_encode is None: return
-    from tinygrad.tensor import Tensor
-    tdiffs = []
-    for _ in range(5):
-      with Context(DEBUG=0, BEAM=0, TRACK_MATCH_STATS=0): Tensor.ones(1, device=self.device).contiguous().realize()
-      if not (ents:=list(self.prof_ents.items())): return
-      self.prof_ents.clear()
-      st = perf_counter_us()
-      self.synchronize()
-      gpu = max(buf.host.view(fmt='Q')[e.en_id] for (buf, _), e in ents)/decimal.Decimal(self.timestamp_divider)
-      tdiffs.append((st+perf_counter_us())/2 - gpu)
-    Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))
+    """
+    Called at the end of profiling to allow the device to finalize any profiling.
+    """
+    # override this in your device implementation
+  def finalize(self):
+    """
+    Called at the end of process lifetime to allow the device to finalize.
+    """
+    if hasattr(self, 'iface') and hasattr(self.iface, 'device_fini'): self.iface.device_fini()
 
 if PROFILE:
   @atexit.register

tinygrad/engine/jit.py

@@ -69,7 +69,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
   # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
   linear = linear.substitute({u: UOp.param(i, u.dtype, u.max_numel(), u.device) for i,u in enumerate(input_uops)}, walk=True)
   linear = memory_plan_rewrite(linear, held_bufs)
-  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), input_uops=input_uops, cache=False)
+  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
   if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
   if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View graphed linear")
   return linear

tinygrad/engine/realize.py

@@ -3,7 +3,7 @@ from typing import cast, Iterator, Any, Sequence
 import weakref, decimal, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, prod, flatten, Context, to_tuple, tqdm, dedup
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us, cpu_profile
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates, Renderer
@@ -156,10 +156,7 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
     elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
          and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
       dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif dest.get_storage().host is not None and src.get_storage().host is not None:
-      for b in (dest, src): b.allocator.dev.synchronize()
-      with cpu_profile(f"{src.device} -> {dest.device}", f"{src.device}:COPY"): dest.host[:] = src.host[:]
-    elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(allow_zero_copy=True), src._buf)
+    elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
     else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
   return []
 
@@ -197,11 +194,10 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   if (info:=call.arg.aux).inputs:
-    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev) + off for u, dev, off in info.inputs]
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr + off for u, dev, off in info.inputs]
     cast(Buffer, call.src[1 + info.table].buffer).host.view(fmt='Q')[:] = array.array('Q', addrs)
   ctx = replace(ctx, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}})
   ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,))
-  for host, dev in info.host_deps: Device[host].pending[dev] = Device[dev].timeline.host.view(fmt='Q')[1]
   if not (ctx.wait or PROFILE): return ets
 
   slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots}
@@ -283,11 +279,11 @@ pm_exec = PatternMatcher([
 
 from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV, HCQInfo # noqa: E402 # down here, hcq2 imports realize
 
-def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None, cache=False) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)
   if (beam_val:=BEAM.value if beam is None else beam) >= 1: linear = graph_rewrite(linear, pm_beam, ctx=beam_val, walk=True)
   linear = lower_and_compile(linear)
-  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile, cache=cache)
+  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
 def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
@@ -295,13 +291,13 @@ def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) ->
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, cache=True), input_uops=inputs)
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), input_uops=inputs)
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
   for call in linear.src: track_stats(ctx, call.without_after, perf_counter_us(), pm_exec.rewrite(call.without_after, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True, cache=False), allow_cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), allow_cache=ctx.cache)
   while True:
     if clear_l2:
       if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()

tinygrad/helpers.py

@@ -5,8 +5,7 @@ import os, functools, re, contextlib, operator, hashlib, pickle, sqlite3, tempfi
 from collections import defaultdict
 import shutil, math, types, copyreg, inspect, importlib, decimal, itertools, difflib
 from dataclasses import dataclass, field, replace
-from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload, TYPE_CHECKING
-if TYPE_CHECKING: import numpy
+from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator, cast, overload
 
 T = TypeVar("T")
 U = TypeVar("U")
@@ -36,7 +35,6 @@ def get_shape(x) -> tuple[int, ...]:
   return (len(subs),) + (subs[0] if subs else ())
 def is_image_shape(shape): return shape is not None and len(shape) == 3 and shape[-1] == 4
 def all_int(t: Sequence[Any]) -> TypeGuard[tuple[int, ...]]: return all(isinstance(s, int) for s in t)
-def is_numpy_ndarray(x) -> TypeGuard[numpy.ndarray]: return str(type(x)) == "<class 'numpy.ndarray'>"
 def colored(st, color:str|None, background=False): # replace the termcolor library
   if NO_COLOR: return st
   colors = ['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white']

tinygrad/mixin/gradient.py

@@ -2,7 +2,7 @@ from typing import cast
 import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
-from tinygrad.dtype import dtypes, sum_acc_dtype
+from tinygrad.dtype import sum_acc_dtype
 from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
@@ -67,19 +67,6 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   ret_set = set(ret_pos)
   return (None,) + tuple(None if i in ret_set else (bwd_outs[gb_map[i]] if i in gb_map else None) for i in range(len(args)))
 
-def partial_store_gradient(ctx:UOp, dest:UOp, view:UOp):
-  # A write through a non-overlapping view replaces only that region of the returned state.
-  path, base = [], view
-  while base is not dest and base.op in {Ops.RESHAPE, Ops.SHRINK, Ops.PERMUTE, Ops.FLIP}:
-    path.append(base)
-    base = base.src[0]
-  if base is not dest: return None
-  grad = ctx
-  for mop in reversed(path): grad = mop.replace(src=(grad,)+mop.src[1:])
-  mask = grad.const_like(1)
-  for mop in path: mask = pm_gradient.rewrite(mop, ctx=mask)[0]
-  return mask.cast(dtypes.bool).where(0, ctx), grad
-
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -117,7 +104,6 @@ pm_gradient = PatternMatcher([
    lambda ctx, dest, t: (ctx, None) if t.buf_uop is not dest.buf_uop else None),
   # clone/assign gradient passes through to val
   (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="dest"), UPat())))), lambda ctx,dest: (None, ctx)),
-  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="view"), UPat())))), partial_store_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/renderer/amd/sqtt.py

@@ -648,14 +648,6 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
   from tinygrad.viz.serve import amd_decode
   pc_map = amd_decode(lib, target)
   wave_pc:dict[tuple[int, int], int] = {}
-  cdna_imm_queue:dict[tuple[int, int], list[CDNA_ISSUE|None]] = {}
-  def cdna_imm_dequeue(key:tuple[int, int]) -> Iterator[tuple[PacketType, InstructionInfo]]:
-    pending = cdna_imm_queue[key]
-    while pending and (p:=pending[0]) is not None:
-      pending.pop(0)
-      if (inst:=pc_map[pc:=wave_pc[key]]).op_name not in {'S_NOP', 'S_WAITCNT', 'S_SETPRIO'}: continue
-      wave_pc[key] += inst.size()
-      yield (p, InstructionInfo(pc, key[1], inst))
   # RDNA selects one SIMD for instruction tracing, CDNA traces multiple SIMDs
   simd:int = 0
   for p in decode(data):
@@ -676,11 +668,12 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
           yield (p, InstructionInfo(pc, wave, inst))
     elif isinstance(p, CDNA_ISSUE):
       for wave in range(10):
-        if (status:=(p.inst >> (wave * 2)) & 3) in {2, 3}:
-          cdna_imm_queue.setdefault(key:=(p.simd, wave), []).append(p if status == 3 else None)
-          yield from cdna_imm_dequeue(key)
+        if (p.inst >> (wave * 2)) & 3 == 3:
+          inst = pc_map[pc:=wave_pc[(p.simd, wave)]]
+          if getattr(inst, 'op_name', '') not in {'S_NOP', 'S_WAITCNT'}: continue
+          wave_pc[(p.simd, wave)] += inst.size()
+          yield (p, InstructionInfo(pc, wave, inst))
     elif isinstance(p, CDNA_INST):
-      cdna_imm_queue[(p.simd, p.wave)].pop(0)
       inst = pc_map[pc:=wave_pc[(p.simd, p.wave)]]
       if p.op == InstOpCDNA.JUMP:
         x = getattr(inst, 'simm16') & 0xffff
@@ -688,7 +681,6 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
       else:
         wave_pc[(p.simd, p.wave)] += inst.size()
       yield (p, InstructionInfo(pc, p.wave, inst))
-      yield from cdna_imm_dequeue((p.simd, p.wave))
     # map INST events on this SIMD to the program counter, we know the waves
     elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):
       inst = pc_map[pc:=wave_pc[(simd, p.wave)]]

tinygrad/renderer/cstyle.py

@@ -187,8 +187,7 @@ class CStyleLanguage(Renderer):
       return prefix + self.type_map.get(dtype, dtype.name).replace(" ", "_") + str(sz) + suffix
     return prefix + self.type_map.get(dtype, dtype.name) + suffix
 
-  def render_type(self, u:UOp):
-    return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape, override_ptr=u.op is Ops.INDEX and u.addrspace is AddrSpace.REG)
+  def render_type(self, u:UOp): return self._render_dtype(u.dtype, u.max_numel(), u.addrspace, shape=u._shape)
   def render_ptr(self, u:UOp):
     # the address of an access, vector-cast if the access reads/writes more lanes than the pointer's scalar type
     if u.max_numel() > 1 or u.dtype != u.src[0].dtype:

tinygrad/renderer/llvmir.py

@@ -152,7 +152,7 @@ class LLVMRenderer(Renderer):
 
   extra_matcher = create_non_native_float_pats((dtypes.bfloat16,)) + pm_manual_bf16_cast
   def _render_fn(self, name:str, args:list[tuple[str,UOp]], kernel:list[str], prefix:list[str]|None=None) -> str:
-    # NOTE: HostAllocator promises 0x20 alignment
+    # NOTE: CPUAllocator promises 0x20 alignment
     sargs = ", ".join([f"{ldt(u.dtype, ptr=u.addrspace == AddrSpace.GLOBAL)}{' noalias align 32' if u.addrspace == AddrSpace.GLOBAL else ''} " + \
       name for name,u in args])
     return "\n".join((prefix or []) + [f"define{' ' + self.abi if self.abi else ''} void @{name}({sargs}) #0", "{"] + kernel + ["  ret void\n}"])

tinygrad/runtime/graph/metal.py

@@ -26,8 +26,8 @@ class MetalGraph(GraphRunner):
 
     self.var_bind_data = []
     if len(self.vars):
-      storage = self.dev.allocator.alloc(sum(dt.itemsize for r in self.runtimes for (_,_,dt,s) in unwrap(r).signature if s == ()))
-      self.var_buf, self.var_buf_view, var_buf_offset = storage.buf, unwrap(storage.host).mv, 0
+      (self.var_buf, _), host = self.dev.allocator.alloc(sum(dt.itemsize for r in self.runtimes for (_,_,dt,s) in unwrap(r).signature if s == ()))
+      self.var_buf_view, var_buf_offset = unwrap(host).mv, 0
 
     all_pipelines, all_resources = [], [self.var_buf.buf] if len(self.vars) else []
     for j, ((_, ast, bufs, _), runtime, replace) in enumerate(zip(self.calls, self.runtimes, self.uop_replace)):

tinygrad/runtime/ops_amd.py

@@ -3,21 +3,21 @@ from typing import cast, Any
 import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, hashlib, itertools, collections, atexit
 assert sys.platform != 'win32'
 from dataclasses import dataclass, replace
-from tinygrad.runtime.support.hcq2 import HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
 from tinygrad.uop.ops import sint, UOp, ProgramInfo
-from tinygrad.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, Compiled, ProfileProgramEvent
+from tinygrad.device import BufferSpec, Buffer, Device, Compiled, ProfileProgramEvent
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, prod, colored
-from tinygrad.helpers import ceildiv, unwrap, pluralize, HCQ2, ContextVar, VIZ
+from tinygrad.helpers import ceildiv, unwrap, pluralize, HCQ2, mv_address, ContextVar, VIZ
 from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, hcq_filter_visible_devices
+from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
-from tinygrad.runtime.support.system import PCIIfaceBase, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.usb import USB3, pm_usb_batch, pm_usb_lower, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
@@ -476,7 +476,7 @@ class AMDSDMAQueue(HWQueue):
     sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
     hdr = self.sdma.SDMA_OP_COPY | self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR)
     for off in range(0, sz, self.max_copy_size):
-      self.q(hdr, min(sz-off, self.max_copy_size)-1, 0,
+      self.q(hdr, self.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, self.max_copy_size)-1), 0,
              *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(self.devs), call.src[1].getaddr(self.devs))))
 
   def wait(self, signal:UOp, value:UOp, eq:bool=False):
@@ -555,16 +555,19 @@ def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]:
     enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
   return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
 
-class AMDAllocator(Allocator['AMDDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
+class AMDAllocator(HCQAllocator['AMDDevice']):
+  def __init__(self, dev:AMDDevice):
+    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
 
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(storage)
-  def _map(self, buf:Buffer) -> BufferStorage: return self.dev.iface.map(buf)
-  def _unmap(self, mapping:BufferStorage): self.dev.iface.unmap(mapping)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    opaque = self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
+    return (opaque, opaque.meta), opaque.view
+
+  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
+
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
 
 @dataclass
 class AMDQueueDesc:
@@ -573,7 +576,7 @@ class AMDQueueDesc:
 
 class KFDIface:
   kfd:FileIOInterface|None = None
-  event_page:Buffer
+  event_page:HCQBuffer|None = None
   gpus:list[FileIOInterface] = []
   count:int = 0
 
@@ -608,7 +611,25 @@ class KFDIface:
     kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
     if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> BufferStorage:
+    # Set these for our device.
+    if KFDIface.event_page is None:
+      KFDIface.event_page = self.alloc(0x8000, uncached=True)
+      kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
+    else: self.map(KFDIface.event_page)
+
+    # Event to wait for queues completion
+    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
+    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
+
+    # OS events to collect memory and hardware faults
+    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)
+
+    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
+      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
+    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> 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
@@ -634,54 +655,43 @@ class KFDIface:
       buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset)
       assert addr == buf == mem.va_addr
 
-    self._map_handle(mem.handle)
-    return BufferStorage(mem.va_addr, mem, MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None)
+    view = MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None
+    self.map(hcqbuf:=HCQBuffer(mem.va_addr, mem.size, meta=mem, view=view, owner=self.dev))
+    return hcqbuf
 
-  def free(self, storage:BufferStorage):
-    self._unmap_handle(storage.meta.handle)
-    if storage.buf: FileIOInterface.munmap(storage.buf, storage.meta.size)
-    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=storage.meta.handle)
+  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
-  def unmap(self, mapping:BufferStorage):
-    handle, owned = mapping.meta
-    self._unmap_handle(handle)
-    if owned: kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=handle)
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
-  def map(self, buf:Buffer) -> BufferStorage:
-    if buf.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if buf._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
-      return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta.handle, True))
-    if buf.device.split(":")[0] != "AMD": raise RuntimeError(f"Cannot map {buf.device} on {self.dev.device}")
-    self._map_handle(buf.meta.handle)
-    return BufferStorage(buf._buf, (buf.meta.handle, False))
-
-  def _map_handle(self, handle):
+  def _unmap(self, mem):
     gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
     assert stm.n_success == 1
 
-  def _unmap_handle(self, handle):
-    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+  def map(self, mem):
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      cast(Any, mapped)._owns_kfd_handle = True
+      return mapped
+
+    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
+    return HCQBuffer(mem.va_addr, mem.size, meta=mem.meta, owner=mem.owner)
 
   def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
                    xcc_id=0, idx=0):
-    if not hasattr(self, 'queue_event_arr'):
-      if not hasattr(KFDIface, 'event_page'):
-        KFDIface.event_page = Buffer(self.dev.device, 0x8000, dtypes.uint8, options=BufferSpec(uncached=True), preallocate=True)
-        kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
-
-      KFDIface.event_page.get_buf(self.dev.device)
-      self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(*[kfd.struct_kfd_event_data(event_id=kfd.AMDKFD_IOC_CREATE_EVENT(
-        KFDIface.kfd, event_type=t, auto_reset=int(t == kfd.KFD_IOC_EVENT_SIGNAL)).event_id)
-        for t in (kfd.KFD_IOC_EVENT_SIGNAL, kfd.KFD_IOC_EVENT_MEMORY, kfd.KFD_IOC_EVENT_HW_EXCEPTION)])
-
-    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf, ring_size=ring.nbytes, gpu_id=self.gpu_id,
+    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
       queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
-      eop_buffer_address=eop_buffer._buf if eop_buffer else 0, eop_buffer_size=eop_buffer.nbytes if eop_buffer else 0,
-      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
-      write_pointer_address=gart._buf+wptr, read_pointer_address=gart._buf+rptr+8*xcc_id)
+      eop_buffer_address=eop_buffer._buf.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer._buf.size if eop_buffer else 0,
+      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf.va_addr if cwsr_buffer else 0,
+      ctx_save_restore_size=ctx_save_restore_size,
+      write_pointer_address=gart._buf.va_addr+wptr, read_pointer_address=gart._buf.va_addr+rptr+8*xcc_id)
 
     if not hasattr(self, 'doorbells'):
       self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
@@ -694,7 +704,7 @@ class KFDIface:
       write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer)
 
   def sleep(self, tm:int):
-    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=ctypes.addressof(self.queue_event_arr), num_events=3, wait_for_all=0, timeout=tm)
+    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
     if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang()
 
   def on_device_hang(self):
@@ -740,6 +750,7 @@ class PCIIface(PCIIfaceBase):
 
   def require_profile_mode(self): return True
   def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)
 
   def _compute_props(self):
     self.ip_versions = self.dev_impl.ip_ver
@@ -764,10 +775,11 @@ class PCIIface(PCIIfaceBase):
 
     rcvr_params: tuple
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr, gart._buf+wptr, idx)))
+      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, idx)))
     else:
-      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr,
-        gart._buf+wptr, eop_buffer._buf, eop_buffer.nbytes, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
+      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, eop_buffer._buf.va_addr, eop_buffer._buf.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
 
     (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True)).host.view(fmt='Q')[0] = 0
     doorbell = Buffer("CPU", 1, dtypes.uint64, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True)
@@ -799,8 +811,9 @@ class PCIIface(PCIIfaceBase):
   def device_fini(self): self.dev_impl.fini()
 
 class USBAllocator(AMDAllocator): # the host program reads another device's memory in place: its bytes are the mapping
-  def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(buf.host.addr, buf.host.mv)
-  def _unmap(self, mapping:BufferStorage): pass
+  def map(self, buf:Buffer) -> tuple:
+    mv = buf.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True)
+    return HCQBuffer(addr:=mv_address(mv), mv.nbytes, meta=mv, view=MMIOInterface(addr, mv.nbytes, fmt='B'), owner=self.dev), mv
 
 class USBIface(PCIIface):
   def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
@@ -809,19 +822,15 @@ class USBIface(PCIIface):
     self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
     self.dev_impl = AMDev(self.pci_dev)
     self._compute_props()
-
-  @functools.cached_property
-  def ctrl(self) -> Buffer:
     # the controller's memory the queue and the host share, one range (usb.py slices it): the sys page at 0, the cq page at 0x1000, the sram at
     # 0x5000. the host's view starts at the sys page's controller address 0xa000, which puts the sram on its scsi window 0xf000
     vaddr, pieces = self.dev_impl.mm.alloc_vaddr(size=0x85000), [(0x0, 0x820000, 0x1000), (0x1000, 0x822000, 0x1000), (0x5000, 0x200000, 0x80000)]
-    for off, paddr, n in pieces: self.dev_impl.mm.map_range(vaddr + off, n, [(paddr, n)], aspace=AddrSpace.SYS, uncached=True)
-    view = self.pci_dev.dma_view(0xa000, 0x85000)
-    for off, n in ((0x800, 4), (0x5000, 0x80000)): view.view(off, n)[:] = bytes(n) # no stale fence or sentinel
-    return Buffer(self.dev.device, 0x85000, dtypes.uint8, options=BufferSpec(external_ptr=vaddr), opaque=BufferStorage(vaddr, host=view))
+    maps = [self.dev_impl.mm.map_range(vaddr + off, n, [(sys, n)], aspace=AddrSpace.SYS, uncached=True) for off, sys, n in pieces]
+    self.ctrl = HCQBuffer(vaddr, 0x85000, meta=PCIAllocationMeta(maps[0], has_cpu_mapping=False), view=self.pci_dev.dma_view(0xa000, 0x85000),
+                          owner=self.dev)
+    for off, n in ((0x800, 4), (0x5000, 0x80000)): unwrap(self.ctrl.view).view(off, n)[:] = bytes(n) # no stale fence or sentinel
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False,
-            **kwargs) -> BufferStorage:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
     # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
     return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
 
@@ -829,7 +838,7 @@ class USBIface(PCIIface):
 
 def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
 
-class AMDDevice(Compiled):
+class AMDDevice(HCQ2Compiled):
   timestamp_divider = 100.0  # AMD GPU clock: ticks/us
   sleep_timeout_ms = 200
   max_scratch_psize = 0
@@ -847,7 +856,7 @@ class AMDDevice(Compiled):
   def __init__(self, device:str=""):
     self.iface = self._select_iface(device)
     self.is_usb = isinstance(self.iface, USBIface)
-    self.can_recover, self.rtalloc_size = self.is_am(), (4 if self.is_usb else 64)<<20
+    if self.is_usb: self.rt_nbytes = 4 << 20
 
     self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
     self.arch = "gfx%d%x%x" % self.target
@@ -871,11 +880,12 @@ class AMDDevice(Compiled):
                       bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
 
     self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
-    self.max_copy_size = 0x40000000 if (4, 4, 2) <= (v:=self.iface.ip_versions[am.SDMA0_HWIP]) < (5, 0, 0) or v >= (5, 2, 0) else 0x400000
+    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
     self.sdma_queues:dict = {}
+    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
 
     allocator = USBAllocator(self) if self.is_usb else AMDAllocator(self)
-    super().__init__(device, allocator, [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, arch=self.arch)
+    super().__init__(device, allocator, [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
 
     # Scratch setup
     self.max_private_segment_size = 0
@@ -887,6 +897,7 @@ class AMDDevice(Compiled):
     if self.is_usb: # the submits write the rings over the link, the copies go through the controller's sram (usb.py)
       self.pm_batch, self.pm_lower = pm_usb_batch, pm_usb_lower
       self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
+      self.host_devs = frozenset({"CPU", "NPY", "DISK"}) # the host program streams numpy and files in place
 
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.pmc_enabled, self.sqtt_enabled = PROFILE > 0 and PMC > 0, PROFILE > 0 and SQTT > 0
@@ -946,11 +957,6 @@ class AMDDevice(Compiled):
       ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
       debug_memory_size=round_up(self.wave_cnt * 32, 64))
 
-  @functools.cached_property
-  def has_copy_queue(self) -> bool:
-    self.has_copy_queue = False # queue setup can allocate buffers that check this property
-    return self.sdma_queue(0) is not None
-
   def sdma_queue(self, idx:int):
     if getenv("AMD_DISABLE_SDMA"): return None
     if idx in self.sdma_queues: return self.sdma_queues[idx]
@@ -993,7 +999,7 @@ class AMDDevice(Compiled):
     rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
     rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
 
-    base = self.scratch._buf
+    base = self.scratch._buf.va_addr
     self.aql_desc.scratch_backing_memory_location = base
     self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size
     self.aql_desc.scratch_resource_descriptor[:] = [lo32(base), int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(base), SWIZZLE_ENABLE=1), 'little'),
@@ -1022,13 +1028,13 @@ class AMDDevice(Compiled):
       buf = self.prog_bufs[b] = Buffer(self.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
       if PROFILE:
         name, lib, key = _amd_program_prof[b]
-        Compiled.profile_events.append(ProfileProgramEvent(self.device, name, lib, buf._buf, b.arg.slot, key))
+        Compiled.profile_events.append(ProfileProgramEvent(self.device, name, lib, buf._buf.va_addr, b.arg.slot, key))
     return self.prog_bufs[b]
 
   def sqtt_trace(self, slot:int, se:int) -> bytes:
     off = (se * self.prof_slots + slot) * self.sqtt_win
     wptr = (self.sqtt_wptrs.host.view(fmt='I')[slot * self.sqtt_ses + se] & 0x1FFFFFFF) * 32
-    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf + off) // 32) & 0x1FFFFFFF) * 32
+    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf.va_addr + off) // 32) & 0x1FFFFFFF) * 32
     assert 0 <= wptr <= self.sqtt_win, f"{wptr} > {self.sqtt_win}, should never happen"
     if wptr >= self.sqtt_win - 32: # the wptr stops at the last dword when the window overflows
       print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {se})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))

tinygrad/runtime/ops_cl.py

@@ -3,9 +3,9 @@ from typing import cast
 import ctypes, hashlib
 from tinygrad.runtime.autogen import opencl as cl
 from tinygrad.runtime.support import c
-from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, suppress_finalizing, unwrap, round_up, is_image_shape
+from tinygrad.helpers import to_char_p_p, from_mv, OSX, DEBUG, mv_address, suppress_finalizing, unwrap, round_up, is_image_shape
 from tinygrad.renderer.cstyle import OpenCLRenderer
-from tinygrad.device import BufferStorage, BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
+from tinygrad.device import BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
 
 CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
 BP_CB = c.CFUNCTYPE[None, [cl.cl_program, c.POINTER[None]]]
@@ -76,14 +76,15 @@ class CLProgram(Program['CLDevice']):
     return None
 
 class CLAllocator(Allocator['CLDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return BufferStorage(checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status))
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), None), None
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options:BufferSpec): check(cl.clReleaseMemObject(storage.buf))
+  def _free(self, opaque:cl.cl_mem, options:BufferSpec): check(cl.clReleaseMemObject(opaque))
   def _copyin(self, dest:cl.cl_mem, src:memoryview):
-    self.dev.pending_copyin.append(src:=memoryview(bytearray(src))) # NOTE: these can't be freed until the GPU actually executes this command
+    if mv_address(src) % 16: src = memoryview(bytearray(src))
     check(cl.clEnqueueWriteBuffer(self.dev.queue, dest, False, 0, len(src)*src.itemsize, from_mv(src), 0, None, None))
+    self.dev.pending_copyin.append(src)    # NOTE: these can't be freed until the GPU actually executes this command
   def _copyout(self, dest:memoryview, src:cl.cl_mem):
     check(cl.clEnqueueReadBuffer(self.dev.queue, src, False, 0, len(dest)*dest.itemsize, from_mv(dest), 0, None, None))
     self.dev.synchronize()
@@ -125,6 +126,6 @@ class CLDevice(Compiled):
 
   def count(self) -> int: return len(unwrap(self.device_ids))
 
-  def synchronize(self, timeout:int|None=None):
+  def synchronize(self):
     check(cl.clFinish(self.queue))
     self.pending_copyin.clear()

tinygrad/runtime/ops_cpu.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 import platform, sys, ctypes, mmap, struct, time
 from typing import cast
-from tinygrad.helpers import OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
-from tinygrad.device import Compiled, TinyELF, Program, HostAllocator
+from tinygrad.helpers import to_mv, from_mv, OSX, WIN, mv_address, suppress_finalizing, unwrap, data64_le
+from tinygrad.device import BufferSpec, TinyELF, Program, Device
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator
 from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
@@ -54,17 +56,17 @@ class CPUProgram(Program['CPUDevice']):
 
       self.fxn = ctypes.CFUNCTYPE(None, ctypes.c_void_p)(self.addr) if self.lvp else ctypes.CFUNCTYPE(None)(self.addr)
 
-  def __call__(self, *bufs:int, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
                vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
     st = time.perf_counter()
     if self.lvp:
       lvp_args = bytearray(12 + (len(bufs) + len(vals)) * 8)
       addr = mv_address(lvp_args)
-      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *bufs)
+      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *[b.va_addr for b in bufs])
       for v,(off,dt) in zip(vals, TinyELF.iter_sig(self.signature[-len(vals):], len(bufs)*8)): struct.pack_into(f'<{dt.fmt}', lvp_args, 12+off, v)
       self.fxn(addr)
     else:
-      args = [*bufs, *cast(tuple[int, ...], vals)]
+      args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
       self.fxn(*[ctypes.c_uint64(x) for x in args])
     return time.perf_counter() - st if wait else None
 
@@ -72,12 +74,33 @@ class CPUProgram(Program['CPUDevice']):
   def __del__(self):
     if sys.platform == 'win32': ctypes.windll.kernel32.VirtualFree(ctypes.c_void_p(self.addr), ctypes.c_size_t(0), 0x8000) #0x8000 - MEM_RELEASE
 
-class CPUDevice(Compiled):
-  wait_timeout_ms = 30000
+class CPUAllocator(HCQAllocator['CPUDevice']):
+  def __init__(self, dev:CPUDevice): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    if options.external_ptr is not None: addr, buf = options.external_ptr, None
+    elif WIN: 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_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
+    return (opaque:=HCQBuffer(addr, size, meta=buf, view=MMIOInterface(addr, size, fmt='B'), owner=self.dev), opaque.meta), opaque.view
 
-  @property
-  def has_copy_queue(self) -> bool: return False
+  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    self.dev.synchronize()
+    ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    self.dev.synchronize()
+    dest[:] = to_mv(int(src.va_addr), dest.nbytes)[:]
+  def _do_map(self, buf:HCQBuffer):
+    if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
+    return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)
+  def _do_unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release
+
+class CPUDevice(HCQ2Compiled):
+  wait_timeout_ms, has_copy_queue = 30000, False
 
   def __init__(self, device:str=""):
-    super().__init__(device, HostAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
+    super().__init__(device, CPUAllocator(self), [ClangRenderer, CPULLVMRenderer, LVPRenderer, X86Renderer], CPUProgram,
       arch={'amd64':'x86_64', 'aarch64':'arm64'}.get(m:=platform.machine().lower(), m)+",native")
+
+  def synchronize(self, timeout:int|None=None): # a host read is safe once every device timeline caught up
+    for dev in [Device[d] for d in Device._opened_devices if not d.startswith("CPU")]:
+      if isinstance(dev, HCQ2Compiled): dev.synchronize(timeout)

tinygrad/runtime/ops_cuda.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes
 from tinygrad.helpers import DEBUG, DEV, getenv, mv_address, suppress_finalizing
-from tinygrad.device import BufferStorage, MMIOInterface, Compiled, BufferSpec, Allocator, Program, TinyELF
+from tinygrad.device import MMIOInterface, Compiled, BufferSpec, Allocator, Program, TinyELF
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
@@ -65,23 +65,24 @@ class CUDAProgram(Program['CUDADevice']):
     return cu_time_execution(lambda: check(cuda.cuLaunchKernel(self.prg, *global_size, *local_size, self.smem, None, None, self.vargs)), enable=wait)
 
 class CUDAAllocator(Allocator['CUDADevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     check(cuda.cuCtxSetCurrent(self.dev.context))
     if options.external_ptr: opaque = cuda.CUdeviceptr_v2(options.external_ptr)
     elif options.host: opaque = init_c_var(ctypes.c_void_p, lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
     else: opaque = init_c_var(cuda.CUdeviceptr, lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
-    return BufferStorage(opaque, None, MMIOInterface(opaque.value, size) if options.host else None)
+    return (opaque, None), MMIOInterface(opaque.value, size) if options.host else None
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    if options.host: check(cuda.cuMemFreeHost(storage.buf))
-    else: check(cuda.cuMemFree_v2(storage.buf))
+  def _free(self, opaque, options:BufferSpec):
+    if options.external_ptr: return
+    if options.host: check(cuda.cuMemFreeHost(opaque))
+    else: check(cuda.cuMemFree_v2(opaque))
   def _copyin(self, dest, src:memoryview):
     check(cuda.cuCtxSetCurrent(self.dev.context))
     host_mem = self.alloc(len(src), BufferSpec(host=True))
     self.dev.pending_copyin.append((host_mem, len(src), BufferSpec(host=True)))
-    ctypes.memmove(host_mem.buf, mv_address(src), len(src))
-    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem.buf, len(src), None))
+    ctypes.memmove(host_mem[0][0], mv_address(src), len(src))
+    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem[0][0], len(src), None))
   def _copyout(self, dest:memoryview, src):
     CUDADevice.synchronize_system()
     check(cuda.cuCtxSetCurrent(self.dev.context))
@@ -115,7 +116,7 @@ class CUDADevice(Compiled):
       check(cuda.cuCtxEnablePeerAccess(dev.context, 0))
       CUDADevice.peer_access = True
 
-    self.pending_copyin: list[tuple[BufferStorage, int, BufferSpec|None]] = []
+    self.pending_copyin: list[tuple[tuple, int, BufferSpec|None]] = []
     CUDADevice.devices.append(self)
 
     from tinygrad.runtime.graph.cuda import CUDAGraph
@@ -124,7 +125,7 @@ class CUDADevice(Compiled):
 
   def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(cuda.cuDeviceGetCount(ctypes.byref(x)))).value
 
-  def synchronize(self, timeout:int|None=None):
+  def synchronize(self):
     check(cuda.cuCtxSetCurrent(self.context))
     check(cuda.cuCtxSynchronize())
     for opaque,sz,options in self.pending_copyin: self.allocator.free(opaque, sz, options)

tinygrad/runtime/ops_disk.py

@@ -1,7 +1,7 @@
 import os, sys, mmap, io, ctypes, contextlib, pathlib
 from typing import Generator, Callable
 from tinygrad.helpers import OSX, round_up, mv_address
-from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator
+from tinygrad.device import MMIOInterface, Compiled, Allocator
 with contextlib.suppress(ImportError):
   import _posixshmem
   from tinygrad.runtime.autogen import io_uring, libc
@@ -9,8 +9,6 @@ with contextlib.suppress(ImportError):
 class DiskDevice(Compiled):
   _tried_io_uring_init = False
 
-  def synchronize(self, timeout:int|None=None): pass
-
   def __init__(self, device:str):
     if not DiskDevice._tried_io_uring_init: self._iouring_setup()
 
@@ -81,11 +79,11 @@ class DiskBuffer:
 MAP_LOCKED, MAP_POPULATE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000)
 class DiskAllocator(Allocator):
   lru = False
-  def _alloc(self, size:int, options) -> BufferStorage:
+  def _alloc(self, size:int, options) -> tuple:
     self.dev._might_open(size)
-    return BufferStorage(opaque:=DiskBuffer(self.dev, size), None, MMIOInterface(mv_address(opaque._buf()), size))
+    return (opaque:=DiskBuffer(self.dev, size), None), MMIOInterface(mv_address(opaque._buf()), size)
 
-  def _free(self, storage:BufferStorage, options): self.dev._might_close()
+  def _free(self, opaque, options): self.dev._might_close()
   def _as_buffer(self, src:DiskBuffer): return src._buf()
   def _copyin(self, dest:DiskBuffer, src:memoryview): dest._buf()[:] = src
   def _copyout(self, dest:memoryview, src:DiskBuffer):

tinygrad/runtime/ops_dsp.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import ctypes, os, mmap, tempfile, pathlib, array, threading, contextlib, sys, subprocess, struct
 assert sys.platform != 'win32'
-from tinygrad.device import BufferStorage, MMIOInterface, BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
+from tinygrad.device import MMIOInterface, BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
 from tinygrad.dtype import dtypes, AddrSpace
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, round_up, mv_address, to_mv, cpu_objdump, system, DEBUG, suppress_finalizing, Target, unwrap
@@ -75,20 +75,20 @@ class DSPBuffer:
     self.va_addr, self.size, self.share_info, self.offset = va_addr, size, share_info, offset
 
 class DSPAllocator(Allocator['DSPDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     if getenv("MOCKDSP"): fd, share_info, flags = -1, None, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS
     else:
       b = qcom_dsp.ION_IOC_ALLOC(self.dev.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
       fd, flags = (share_info:=qcom_dsp.ION_IOC_SHARE(self.dev.ion_fd, handle=b.handle)).fd, mmap.MAP_SHARED
     opaque = DSPBuffer(libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, flags, fd, 0), size, share_info, offset=0)
-    return BufferStorage(opaque, opaque.share_info, MMIOInterface(opaque.va_addr, size))
+    return (opaque, opaque.share_info), MMIOInterface(opaque.va_addr, size)
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    libc.munmap(storage.buf.va_addr, storage.buf.size)
-    if storage.buf.share_info is not None:
-      os.close(storage.buf.share_info.fd)
-      qcom_dsp.ION_IOC_FREE(self.dev.ion_fd, handle=storage.buf.share_info.handle)
+  def _free(self, opaque:DSPBuffer, options:BufferSpec):
+    libc.munmap(opaque.va_addr, opaque.size)
+    if opaque.share_info is not None:
+      os.close(opaque.share_info.fd)
+      qcom_dsp.ION_IOC_FREE(self.dev.ion_fd, handle=opaque.share_info.handle)
 
   def _as_buffer(self, src:DSPBuffer) -> memoryview: return to_mv(src.va_addr, src.size)
   def _copyin(self, dest:DSPBuffer, src:memoryview): ctypes.memmove(dest.va_addr, mv_address(src), src.nbytes)
@@ -132,7 +132,7 @@ class DSPDevice(Compiled):
       self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
       super().__init__(device, DSPAllocator(self), [DSPRenderer], DSPProgram)
       fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
-      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True)).buf
+      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))[0][0]
       ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
 
       self.init_dsp()

tinygrad/runtime/ops_hip.py

@@ -1,6 +1,6 @@
 import ctypes
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import BufferStorage, Compiled, Allocator, BufferSpec, Program, TinyELF
+from tinygrad.device import Compiled, Allocator, BufferSpec, Program, TinyELF
 from tinygrad.runtime.autogen import hip
 from tinygrad.renderer.cstyle import HIPRenderer
 from tinygrad.runtime.support.c import init_c_var, init_c_struct_t
@@ -19,7 +19,7 @@ class HIPDevice(Compiled):
 
   def count(self) -> int: return init_c_var(ctypes.c_int, lambda x: check(hip.hipGetDeviceCount(x))).value
 
-  def synchronize(self, timeout:int|None=None):
+  def synchronize(self):
     check(hip.hipSetDevice(self.device_id))
     check(hip.hipDeviceSynchronize())
 
@@ -57,11 +57,11 @@ class HIPProgram(Program[HIPDevice]):
       return ret.value * 1e-3
 
 class HIPAllocator(Allocator[HIPDevice]):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     check(hip.hipSetDevice(self.dev.device_id))
-    return BufferStorage(init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size))))
+    return (init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size))), None), None
 
-  def _free(self, storage:BufferStorage, options:BufferSpec): check(hip.hipFree(storage.buf))
+  def _free(self, opaque, options:BufferSpec): check(hip.hipFree(opaque))
   def _copyin(self, dest, src: memoryview):
     check(hip.hipSetDevice(self.dev.device_id))
     check(hip.hipMemcpy(dest, mv_address(src), len(src), hip.hipMemcpyHostToDevice))

tinygrad/runtime/ops_metal.py

@@ -1,7 +1,7 @@
 import subprocess, pathlib, struct, ctypes, tempfile, functools, decimal, platform
 from tinygrad.helpers import prod, to_mv, round_up, cache_dir, PROFILE, ProfileRangeEvent, cpu_profile, unwrap, suppress_finalizing
 import tinygrad.runtime.support.objc as objc
-from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Compiler, CompileError, Program, TinyELF, Allocator, ProfileDeviceEvent
+from tinygrad.device import MMIOInterface, Compiled, Compiler, CompileError, Program, TinyELF, Allocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support.c import DLL
@@ -50,7 +50,7 @@ class MetalDevice(Compiled):
                      MetalGraph if 'virtual' not in from_ns_str(self.sysdevice.name()).lower() else None,
                      arch=metal.enum_MTLGPUFamily[check_family("Apple") or check_family("Mac")][12:])
 
-  def synchronize(self, timeout:int|None=None):
+  def synchronize(self):
     for cbuf in self.mtl_buffers_in_flight:
       wait_check(cbuf)
       st, en = decimal.Decimal(cbuf.GPUStartTime()) * 1000000, decimal.Decimal(cbuf.GPUEndTime()) * 1000000
@@ -155,15 +155,16 @@ class MetalBuffer:
   def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
 
 class MetalAllocator(Allocator[MetalDevice]):
-  def _alloc(self, size:int, options) -> BufferStorage:
+  def _alloc(self, size:int, options) -> tuple:
     ret = metal.MTLBuffer(options.external_ptr) if options.external_ptr else \
       self.dev.sysdevice.newBufferWithLength_options(size, metal.MTLResourceStorageModeShared)
     setattr(ret, "retain", False) # Buffer is explicitly released in _free()
     if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
-    return BufferStorage(MetalBuffer(ret, size), None, MMIOInterface(addr, size) if (addr:=ret.contents()) is not None else None)
+    return (MetalBuffer(ret, size), None), MMIOInterface(addr, size) if (addr:=ret.contents()) is not None else None
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options): storage.buf.buf.release()
+  def _free(self, opaque:MetalBuffer, options):
+    if not options.external_ptr: opaque.buf.release()
   def _transfer(self, dest:MetalBuffer, src:MetalBuffer, sz:int, src_dev:MetalDevice, dest_dev:MetalDevice):
     dest_dev.synchronize()
     src_command_buffer = src_dev.mtl_queue.commandBuffer().retained()

tinygrad/runtime/ops_npy.py

@@ -1,4 +1,14 @@
-from tinygrad.device import Compiled, HostAllocator
+import numpy as np
+from tinygrad.helpers import flat_mv
+from tinygrad.device import MMIOInterface, Compiled, Allocator
+
+class NpyAllocator(Allocator['NpyDevice']):
+  def _alloc(self, size:int, options) -> tuple: return (arr:=np.empty(size, dtype=np.uint8), arr), MMIOInterface(arr.ctypes.data, size)
+
+  def _as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data)
+  def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src)
+  def _offset(self, buf:np.ndarray, size:int, offset:int) -> np.ndarray:
+    return np.require(buf, requirements='C').reshape(-1).view(np.uint8)[offset:offset+size]
 
 class NpyDevice(Compiled):
-  def __init__(self, device:str): super().__init__(device, HostAllocator(self), [], None)
+  def __init__(self, device:str): super().__init__(device, NpyAllocator(self), [], None)

tinygrad/runtime/ops_null.py

@@ -1,5 +1,5 @@
 import inspect, math
-from tinygrad.device import BufferStorage, Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
+from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl
 from tinygrad.renderer.cstyle import CStyleLanguage
@@ -22,7 +22,7 @@ class NullProgram(Program['NullDevice']):
     with cpu_profile(self.name, self.device, profile_key=self.profile_key): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):
-  def _alloc(self, size:int, options) -> BufferStorage: return BufferStorage(None)
+  def _alloc(self, size:int, options) -> tuple: return (None, None), None
 
   def _copyin(self, dest, src:memoryview): pass
   def _copyout(self, dest:memoryview, src):
@@ -51,8 +51,6 @@ class NullGraph(MultiGraphRunner):
     return 1e-1
 
 class NullDevice(Compiled):
-  def synchronize(self, timeout:int|None=None): pass
-
   def __init__(self, device:str):
     assert (emu:=getenv("EMULATE", "")) == "", \
       "EMULATE is deprecated, use DEV=NULL:HIP:"+{"AMD":"gfx1100", "AMD_RDNA4":"gfx1201", "AMD_CDNA4":"gfx950"}.get(emu, "<arch>")

tinygrad/runtime/ops_nv.py

@@ -2,21 +2,21 @@ from __future__ import annotations
 import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, itertools
 assert sys.platform != 'win32'
 from typing import Any
-from dataclasses import dataclass, replace
-from tinygrad.runtime.support.hcq2 import HWQueue, encode_submit, patch, to_name, unwrap_view, make_submit, timeline, HCQInfo, lower_call, hcq_link
-from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, BumpAllocator, hcq_filter_visible_devices
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
-from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops, lower_and_compile, run_linear
-from tinygrad.device import BufferStorage, Buffer, BufferSpec, Allocator, Compiled, Device, TinyELF
+from dataclasses import dataclass
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, patch, to_name, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface, FileIOInterface, BumpAllocator, hcq_filter_visible_devices
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, TinyELF
 from tinygrad.dtype import dtypes, DType
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, PROFILE, ContextVar, VIZ
-from tinygrad.helpers import ProfileEvent, unwrap
+from tinygrad.helpers import ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.runtime.autogen import nv_570, nv_580, nv_610, mesa
 from tinygrad.runtime.support.elf import elf_loader
 from tinygrad.runtime.support.nv.nvdev import NVDev, NVMemoryManager
-from tinygrad.runtime.support.system import PCIIfaceBase, MAP_FIXED
+from tinygrad.runtime.support.system import System, PCIIfaceBase, MAP_FIXED
 from tinygrad.renderer.nir import NAKRenderer
 if getenv("IOCTL"): import extra.nv_gpu_driver.nv_ioctl # noqa: F401 # pylint: disable=unused-import
 
@@ -95,9 +95,6 @@ class QMD:
 
 class NVQueue(HWQueue):
   dev:NVDevice
-  q_rewrite = HWQueue.q_rewrite + PatternMatcher([
-    (UPat(Ops.INS, arg=("nv", dtypes.void), name="u"), lambda ctx, u: ctx.q(*u.src)),
-  ])
 
   def nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
@@ -293,19 +290,16 @@ def nv_build_program(dev:NVDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[NVPro
     cached = _nv_program_cache[key] = (data, patch(buf, rows, data.image))
   return cached
 
-class NVAllocator(Allocator['NVDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero)
+class NVAllocator(HCQAllocator['NVDevice']):
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    return (opaque:=self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero), opaque.meta), opaque.view
 
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(storage)
-  def _map(self, buf:Buffer) -> BufferStorage: return self.dev.iface.map(buf)
-  def _unmap(self, mapping:BufferStorage): self.dev.iface.unmap(mapping)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _encode_decode(self, bufout:int, bufin:int, desc_buf:int, hist:list[int], shape:tuple[int,...], frame_pos:int):
-    assert all(h % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _encode_decode(self, bufout:HCQBuffer, bufin:HCQBuffer, desc_buf:HCQBuffer, hist:list[HCQBuffer], shape:tuple[int,...], frame_pos:int):
+    assert all(h.va_addr % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
 
     h, w = ((2 * shape[0]) // 3 if shape[0] % 3 == 0 else (2 * shape[0] - 1) // 3), shape[1]
     dev, chroma_off = self.dev, round_up(w, 64) * round_up(h, 64)
@@ -314,25 +308,25 @@ class NVAllocator(Allocator['NVDevice']):
     cmds = nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
     cmds += nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
                 err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin.va_addr >> 8)
     for pos, buf in zip([(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)] + [frame_pos], hist + [bufout]):
-      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf >> 8)
-      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, (buf + chroma_off) >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, (desc_buf + 0x200) >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf) >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, desc_buf.offset(0x200).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf.va_addr) >> 8)
     cmds += nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_addr + dev.intra_top_off) >> 8)
     if dev.intra_unk_off is not None: cmds += nvm(4, 0x4dc, (filter_addr + dev.intra_unk_off) >> 8)
     cmds += nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
-    dev._submit_cmds("NVDEC:0", *cmds)
+    dev._submit_cmds(dev.fifos["NVDEC:0"], *cmds)
 
 # *****************
 # device
 
 @dataclass
-class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; notifier: Buffer; entries: int; token: int # noqa: E702
+class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; entries: int; token: int # noqa: E702
 
 class NVKIface:
   root = None
@@ -442,7 +436,7 @@ 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, **kwargs) -> BufferStorage:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
     # Uncached memory is "system". Use huge pages only for gpu memory.
     page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (mmap.PAGESIZE if isinstance(self, MOCKIface) else
                                                                                              4 << 10))
@@ -482,20 +476,17 @@ class NVKIface:
 
     return self._gpu_uvm_map(va_addr, size, mem_handle, has_cpu_mapping=cpu_access or host)
 
-  def free(self, storage:BufferStorage):
-    mem = storage.meta
-    if mem.hMemory > NVKIface.host_object_enumerator: # not a host object, clear phys mem.
-      made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, hObjectOld=mem.hMemory)
+  def free(self, mem:HCQBuffer):
+    if mem.owner != self.dev: return
+    if mem.meta.hMemory > NVKIface.host_object_enumerator: # not a host object, clear phys mem.
+      made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, hObjectOld=mem.meta.hMemory)
       nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_FREE, made)
       if made.status != 0: raise RuntimeError(f"_gpu_free returned {get_error_str(made.status)}")
-    self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=storage.buf, length=mem.length))
-    if storage.host is not None: FileIOInterface.munmap(storage.buf, mem.length)
 
-  def unmap(self, mapping:BufferStorage):
-    mem, owns_range = mapping.meta
-    if owns_range: self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=mapping.buf, length=mem.length))
+    self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=int(mem.va_addr), length=mem.size))
+    if mem.view is not None: FileIOInterface.munmap(int(mem.va_addr), mem.size)
 
-  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False) -> BufferStorage:
+  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False, owner=None) -> HCQBuffer:
     if create_range:
       self.uvm(nv_gpu.UVM_CREATE_EXTERNAL_RANGE, nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS(base=va_base, length=size))
       made = nv_gpu.NVOS46_PARAMETERS(hClient=self.root, hDevice=self.dev.nvdevice, hDma=self.dev.virtmem, hMemory=mem_handle, length=size,
@@ -509,16 +500,14 @@ class NVKIface:
 
     self.uvm(nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION, uvm_map:=nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS(base=va_base, length=size,
       rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid]))
-    return BufferStorage(va_base, uvm_map, MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None)
+    return HCQBuffer(va_base, size, meta=uvm_map, view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None,
+                     owner=self.dev if owner is None else owner)
 
-  def map(self, buf:Buffer) -> BufferStorage:
-    mem = buf.meta
-    if buf.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if buf._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
-      if (mem:=next((m.meta[0] for d, m in buf.get_storage().maps.items() if d.startswith("NV")), None)) is None:
-        return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta, True))
-    elif buf.device.split(":")[0] != "NV": raise RuntimeError(f"Cannot map {buf.device} on {self.dev.device}")
-    return replace(mapping:=self._gpu_uvm_map(buf._buf, mem.length, mem.hMemory, create_range=False), meta=(mapping.meta, False))
+  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"))]
+    return self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False, owner=mem.owner)
 
   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)
@@ -554,13 +543,12 @@ class PCIIface(PCIIfaceBase):
 
 class MOCKIface(NVKIface): count = 1
 
-class NVDevice(Compiled):
+class NVDevice(HCQ2Compiled):
   ifaces = [NVKIface, PCIIface, MOCKIface]
   sleep_timeout_ms = 200
   pm_encode = PatternMatcher([
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_compute", name="submit"), lambda ctx, submit: encode_submit(NVComputeQueue(ctx, submit))),
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_copy", name="submit"), lambda ctx, submit: encode_submit(NVCopyQueue(ctx, submit))),
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_raw", name="submit"), lambda ctx, submit: encode_submit(NVQueue(ctx, submit))),
   ])
 
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)
@@ -588,6 +576,9 @@ class NVDevice(Compiled):
     channel_params = nv_gpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS(engineType=nv_gpu.NV2080_ENGINE_TYPE_GRAPHICS)
     self.channel_group = self.iface.rm_alloc(self.nvdevice, nv_gpu.KEPLER_CHANNEL_GROUP_A, channel_params)
 
+    self.gpfifo_mem = self.iface.alloc(0x300000, contiguous=True, cpu_access=True, force_devmem=True,
+      map_flags=(nv_gpu.NVOS33_FLAGS_CACHING_TYPE_WRITECOMBINED<<23))
+
     self.ctxshare = self.iface.rm_alloc(self.channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A,
       nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC))
 
@@ -609,28 +600,26 @@ class NVDevice(Compiled):
 
   @functools.cached_property
   def fifos(self) -> dict[str, GPFifo]:
-    mem = self.iface.alloc(3<<20, contiguous=True, cpu_access=True, force_devmem=True, map_flags=nv_gpu.NVOS33_FLAGS_CACHING_TYPE_WRITECOMBINED<<23)
-    self.gpfifo_buf = Buffer(self.device, 3<<20, dtypes.uint8, opaque=mem)
-
+    self.gpfifo_buf = Buffer(self.device, self.gpfifo_mem.size, dtypes.uint8, options=BufferSpec(external_ptr=self.gpfifo_mem.va_addr, nolru=True)) \
+                        .allocate(opaque=self.gpfifo_mem)
     compute = self._new_gpu_fifo("COMPUTE:0", self.ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
     copy = self._new_gpu_fifo("COPY:0", self.ctxshare, self.channel_group, offset=0x100000, entries=0x10000)
     self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
-    self.fifos = {"COMPUTE:0": compute, "COPY:0": copy}
 
-    self._submit_cmds("COMPUTE:0", *nvm(1, nv_gpu.NVC6C0_SET_OBJECT, self.iface.compute_class),
+    self._submit_cmds(compute, *nvm(1, nv_gpu.NVC6C0_SET_OBJECT, self.iface.compute_class),
                        *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_WINDOW_A, *data64(self.local_mem_window)),
                        *nvm(1, nv_gpu.NVC6C0_SET_SHADER_SHARED_MEMORY_WINDOW_A, *data64(self.shared_mem_window)))
-    self._submit_cmds("COPY:0", *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.dma_class))
+    self._submit_cmds(copy, *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.dma_class))
 
     if self.pma_enabled: self._prof_init() # the sampler binds to the channel group, so it only comes up once the channels do
-    return self.fifos
+    return {"COMPUTE:0": compute, "COPY:0": copy}
 
   def _new_gpu_fifo(self, name:str, ctxshare, channel_group, offset=0, entries=0x400, compute=False, video=False) -> GPFifo:
-    notifier = Buffer(self.device, size:=48 << 20, dtypes.uint8, opaque=self.iface.alloc(size, uncached=True))
-    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_buf._buf+offset, gpFifoEntries=entries,
-      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_buf.meta.hMemory,
-      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_buf.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
-      engineType=19 if video else 0, hContextShare=ctxshare,
+    notifier = self.iface.alloc(48 << 20, uncached=True)
+    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_mem.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
+      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_mem.meta.hMemory,
+      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_mem.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
+      engineType=19 if video else 0,
       hVASpace=self.vaspace if video and self.is_nvd() else 0) # gsp has no default vaspace, rm maps the decoder ctx into its own
     gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params)
 
@@ -653,7 +642,7 @@ class NVDevice(Compiled):
     fifo = GPFifo(ring=self.gpfifo_buf.view(entries, dtypes.uint64, offset).ensure_allocated(),
       gpput=self.gpfifo_buf.view(1, dtypes.uint32, gpput_off).ensure_allocated(),
       doorbell=Buffer("CPU", 1, dtypes.uint32, options=BufferSpec(external_ptr=self.gpu_mmio.addr + 0x90), preallocate=True),
-      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), notifier=notifier, entries=entries, token=ws_token_params.workSubmitToken)
+      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), entries=entries, token=ws_token_params.workSubmitToken)
     self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag=to_name(n, name)), lambda ctx, b=getattr(fifo, n): b)
                                         for n in ("ring", "gpput", "doorbell", "put_value")]) + self.pm_bufferize
     return fifo
@@ -671,16 +660,24 @@ class NVDevice(Compiled):
       nv_gpu.NV2080_CTRL_GR_GET_INFO_PARAMS(grInfoListSize=len(infos), grInfoList=ctypes.addressof(infos)))
     return [x.data for x in infos]
 
-  def _submit_cmds(self, queue:str, *cmds:int): # channel setup and video decode use the same runtime submit as kernels
-    tl = timeline(devs:=(self.device,))
-    value = tl.index(1).load()
-    submit = make_submit(
-      UOp(Ops.INS, arg=("wait", dtypes.void), src=(tl, value)),
-      UOp(Ops.INS, arg=("nv", dtypes.void), src=(UOp(Ops.BINARY, arg=array.array('I', cmds).tobytes()),)),
-      UOp(Ops.INS, arg=("store", dtypes.void), src=(tl, value + 1)), devs=devs, queue=queue).replace(arg="submit_nv_raw")
-    call = UOp.sink(tl.after(submit).index(1).store(value + 1), arg=KernelInfo("nv_submit")).call(aux=HCQInfo(devs))
-    linear = lower_and_compile(UOp(Ops.LINEAR, src=(unwrap(lower_call(call)),)))
-    run_linear(hcq_link(linear, allow_cache=True), jit=True, update_stats=False, wait=True)
+  def _push(self, fifo:GPFifo, cmds:list[int]): # a pushbuffer built in python: channel setup and video decode
+    (buf:=self.rt_view(len(cmds) * 4)).host.view(fmt='I')[:] = array.array('I', cmds)
+
+    put = fifo.put_value.host.view(fmt='Q')
+    fifo.ring._buf.cpu_view().view(fmt='Q')[put[0] % fifo.entries] = buf._buf.va_addr | (len(cmds) << 42) | (1 << 41)
+    fifo.gpput._buf.cpu_view().view(fmt='I')[0] = (put[0] + 1) % fifo.entries
+
+    System.memory_barrier()
+    self.gpu_mmio[0x90 // 4] = fifo.token
+    put[0] += 1
+
+  def _submit_cmds(self, fifo:GPFifo, *cmds:int): # runs cmds once everything already submitted is done, then bumps the timeline
+    tl, addr = self.timeline.host.view(fmt='Q'), self.timeline._buf.va_addr
+    self._push(fifo, nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1]),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit")) + list(cmds) +
+                     nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1] + 1),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit")))
+    tl[1] += 1
 
   def _ensure_has_local_memory(self, required):
     if self.slm_per_thread >= required: return
@@ -690,7 +687,7 @@ class NVDevice(Compiled):
     self.shader_local_mem = Buffer(self.device, round_up(bytes_per_tpc*self.num_tpc_per_gpc*self.num_gpcs, 0x20000), dtypes.uint8,
                                    options=BufferSpec(nolru=True), preallocate=True)
 
-    self._submit_cmds("COMPUTE:0", *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf)),
+    self._submit_cmds(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf.va_addr)),
                        *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(bytes_per_tpc), 0xff))
 
   def _ensure_has_vid_hw(self, w, h):
@@ -706,7 +703,7 @@ class NVDevice(Compiled):
     if "NVDEC:0" not in self.fifos:
       self.fifos["NVDEC:0"] = self._new_gpu_fifo("NVDEC:0", 0, self.nvdevice, offset=0x200000, entries=2048, video=True)
       self.vid_coloc_buf, self.vid_filter_buf, self.vid_stat_buf = _vid_buf(coloc_sz), _vid_buf(filter_sz), _vid_buf(0x1000)
-      self._submit_cmds("NVDEC:0", *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.viddec_class))
+      self._submit_cmds(self.fifos["NVDEC:0"], *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.viddec_class))
     else:
       if coloc_sz > self.vid_coloc_buf.nbytes: self.vid_coloc_buf = _vid_buf(coloc_sz)
       if filter_sz > self.vid_filter_buf.nbytes: self.vid_filter_buf = _vid_buf(filter_sz)
@@ -753,14 +750,12 @@ class NVDevice(Compiled):
       (nv_gpu.NVB0CC_CTRL_POWER_FEATURE_MASK_IDLE_SLOWDOWN_DISABLE << 8) | (nv_gpu.NVB0CC_CTRL_POWER_FEATURE_MASK_VAT_DISABLE << 10))
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_POWER_REQUEST_FEATURES, power_params)
 
-    self.pma_buf = Buffer(self.device, size:=getenv("PMA_BUFFER_SIZE", 512) << 20, dtypes.uint8,
-                          opaque=self.iface.alloc(size, uncached=True, cpu_cached=True, cpu_access=True))
-    self.pma_bytes = Buffer(self.device, size:=0x1000, dtypes.uint8,
-                            opaque=self.iface.alloc(size, uncached=True, cpu_cached=True, cpu_access=self.is_nvd(), read_only=True))
+    self.pma_buf = self.iface.alloc(getenv("PMA_BUFFER_SIZE", 512) << 20, uncached=True, cpu_cached=True, cpu_access=True)
+    self.pma_bytes = self.iface.alloc(0x1000, uncached=True, cpu_cached=True, cpu_access=self.is_nvd(), read_only=True)
     self.pma_rptr = 0
 
     pma_stream = nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS(hMemPmaBuffer=self.pma_buf.meta.hMemory,
-      pmaBufferSize=self.pma_buf.nbytes, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf._buf)
+      pmaBufferSize=self.pma_buf.size, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf.va_addr)
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM, pma_stream, extra=(self.pma_buf, self.pma_bytes))
 
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_RESERVE_HWPM_LEGACY, nv_gpu.struct_NVB0CC_CTRL_RESERVE_HWPM_LEGACY_PARAMS(ctxsw=0))
@@ -827,9 +822,9 @@ class NVDevice(Compiled):
     if params.bOverflowStatus: raise RuntimeError("PMA profiler: buffer overflow detected")
     if params.bytesAvailable == 0: return None
 
-    start, end, view = self.pma_rptr, self.pma_rptr + params.bytesAvailable, self.pma_buf.host
-    pma_data = bytes(view[start:min(end, self.pma_buf.nbytes)]) + bytes(view[:max(0, end - self.pma_buf.nbytes)])
-    self.pma_rptr = end % self.pma_buf.nbytes
+    start, end = self.pma_rptr, self.pma_rptr + params.bytesAvailable
+    pma_data = bytes(self.pma_buf.cpu_view()[start:min(end, self.pma_buf.size)]) + bytes(self.pma_buf.cpu_view()[:max(0, end - self.pma_buf.size)])
+    self.pma_rptr = end % self.pma_buf.size
 
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT,
       nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS(bytesConsumed=params.bytesAvailable))

tinygrad/runtime/ops_python.py

@@ -6,8 +6,8 @@ from typing import Any, TYPE_CHECKING
 import pickle, base64, itertools, time, sys, functools, ctypes
 from dataclasses import replace
 from tinygrad.dtype import bitcast, DType, dtypes, AddrSpace, truncate, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
-from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, to_mv, mv_address
-from tinygrad.device import HostAllocator, Compiled, Compiler, Program, TinyELF
+from tinygrad.helpers import all_same, getenv, flatten, Target, IMAGE, is_image_shape, cpu_profile, mv_address
+from tinygrad.device import MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.renderer import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -55,7 +55,7 @@ class PythonProgram(Program['PythonDevice']):
     warp_size = len(warp)
     for idxs in itertools.product(*[range(x) for x in global_size[::-1]]):
       values: dict[UOp, Any] = {}
-      pbufs: list[int] = list(bufs)
+      pbufs: list[memoryview] = list(bufs)
       pvals: list[int] = list(vals)
       exec_masks = [[True] * warp_size]
       i = 0
@@ -101,8 +101,7 @@ class PythonProgram(Program['PythonDevice']):
             # REGs are per thread
             values[u] = [memoryview(bytearray(u.max_numel()*u.dtype.itemsize)).cast(storage_fmt) for _ in range(warp_size)]
           else:
-            size = u.max_numel() * u.dtype.itemsize
-            buf = memoryview(bytearray(size)) if u.op is not Ops.PARAM else to_mv(pbufs.pop(0), size)
+            buf = memoryview(bytearray(u.max_numel()*u.dtype.itemsize)) if u.op is not Ops.PARAM else pbufs.pop(0)
             values[u] = [buf.cast(storage_fmt)] * warp_size
         elif u.op is Ops.SPECIAL:
           if u.arg[0] == 'g': values[u] = [idxs[2-int(u.arg[-1])]] * warp_size
@@ -237,6 +236,17 @@ class PythonRenderer(Renderer):
 
   def supported_dtypes(self): return {d for d in super().supported_dtypes() if d != dtypes.half or sys.version_info >= (3, 12)}
 
+class PythonAllocator(Allocator['PythonDevice']):
+  def _alloc(self, size:int, options) -> tuple: return (buf:=memoryview(bytearray(size)), buf), MMIOInterface(mv_address(buf), size)
+
+  def _as_buffer(self, src) -> memoryview: return src
+  def _copyin(self, dest, src:memoryview):
+    with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
+  def _copyout(self, dest:memoryview, src):
+    with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
+  def map(self, buf:Buffer) -> tuple: return (mv:=buf.as_memoryview(force_zero_copy=True)), mv
+  def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
+
 class PythonDevice(Compiled):
   def __init__(self, device:str):
-    super().__init__(device, HostAllocator(self), [PythonRenderer], PythonProgram)
+    super().__init__(device, PythonAllocator(self), [PythonRenderer], PythonProgram)

tinygrad/runtime/ops_qcom.py

@@ -2,9 +2,9 @@ from __future__ import annotations
 import os, ctypes, functools, mmap, struct, array, math, sys, contextlib
 assert sys.platform != 'win32'
 from typing import Any
-from tinygrad.device import Compiled, BufferStorage, BufferSpec, Buffer, Device, Allocator, TinyELF
-from tinygrad.runtime.support.hcq2 import HWQueue, HCQ_RUNTIME_DEV, encode_submit, ccall, cstruct, patch, unwrap_view
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
+from tinygrad.device import BufferSpec, Buffer, Device, TinyELF
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, HCQ_RUNTIME_DEV, encode_submit, ccall, cstruct, patch, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa, libc
 from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import IR3Renderer
@@ -301,26 +301,21 @@ def qcom_build_program(dev:QCOMDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[Q
     cached = _qcom_program_cache[key] = (data, patch(buf, [], image))
   return cached
 
-class QCOMAllocator(Allocator['QCOMDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return self.dev._gpu_map(options.external_ptr, size) if options.external_ptr else self.dev._gpu_alloc(size)
+class QCOMAllocator(HCQAllocator['QCOMDevice']):
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    return (opaque:=self.dev._gpu_map(options.external_ptr, size) if options.external_ptr else self.dev._gpu_alloc(size), opaque.meta), opaque.view
 
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev._gpu_free(storage)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _do_free(self, opaque, options:BufferSpec): self.dev._gpu_free(opaque)
 
 def flag(nm, val): return (val << getattr(kgsl, f"{nm}_SHIFT")) & getattr(kgsl, f"{nm}_MASK")
 
-class QCOMDevice(Compiled):
+class QCOMDevice(HCQ2Compiled):
   timestamp_divider = 19.2
+  has_copy_queue = False
   pm_encode = PatternMatcher([
     (UPat(Ops.CUSTOM_FUNCTION, arg="submit_qcom_compute", name="submit"), lambda ctx, submit: encode_submit(QCOMComputeQueue(ctx, submit))),
   ])
 
-  @property
-  def has_copy_queue(self) -> bool: return False
-
   def __init__(self, device:str=""):
     self.fd = FileIOInterface('/dev/kgsl-3d0', os.O_RDWR)
 
@@ -359,9 +354,11 @@ class QCOMDevice(Compiled):
 
   @functools.cached_property
   def border_color(self) -> Buffer: # zeros: the samplers clamp to a black border
-    return Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), initial_value=bytes(0x1000))
+    (b:=Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)) \
+      .as_memoryview(force_zero_copy=True)[:] = bytes(0x1000)
+    return b
 
-  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> BufferStorage:
+  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
     flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
     if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)
 
@@ -369,24 +366,24 @@ class QCOMDevice(Compiled):
     va_addr = self.fd.mmap(0, bosz, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, alloc.id * 0x1000)
 
     if fill_zeroes: ctypes.memset(va_addr, 0, size)
-    return BufferStorage(va_addr, (alloc, True), MMIOInterface(va_addr, size, fmt='B'))
+    return HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self)
 
-  def _gpu_map(self, ptr:int, size:int) -> BufferStorage:
+  def _gpu_map(self, ptr:int, size:int) -> HCQBuffer:
     ptr_aligned, size_aligned = (ptr & ~0xfff), round_up(size + (ptr & 0xfff), 0x1000)
     dcache_flush().fxn(ctypes.c_uint64(ptr_line_aligned:=ptr & ~63), ceildiv(ptr + size - ptr_line_aligned, 64))
     try:
       mi = kgsl.IOCTL_KGSL_MAP_USER_MEM(self.fd, hostptr=ptr_aligned, len=size_aligned, memtype=kgsl.KGSL_USER_MEM_TYPE_ADDR)
-      return BufferStorage(mi.gpuaddr + (ptr - ptr_aligned), (mi, False), MMIOInterface(ptr, size, fmt='B'))
+      return HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
     except OSError as e:
-      if e.errno == 14: return BufferStorage(ptr, (None, False), MMIOInterface(ptr, size, fmt='B'))
+      if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
       raise RuntimeError("Failed to map external pointer to GPU memory") from e
 
-  def _gpu_free(self, storage:BufferStorage):
-    if storage.meta[0] is None: return # external (gpu) ptr
-    if not storage.meta[1]: kgsl.IOCTL_KGSL_SHAREDMEM_FREE(self.fd, gpuaddr=storage.meta[0].gpuaddr) # external (cpu) ptr
+  def _gpu_free(self, mem:HCQBuffer):
+    if mem.meta[0] is None: return # external (gpu) ptr
+    if not mem.meta[1]: kgsl.IOCTL_KGSL_SHAREDMEM_FREE(self.fd, gpuaddr=mem.meta[0].gpuaddr) # external (cpu) ptr
     else:
-      kgsl.IOCTL_KGSL_GPUOBJ_FREE(self.fd, id=storage.meta[0].id)
-      FileIOInterface.munmap(storage.buf, storage.meta[0].mmapsize)
+      kgsl.IOCTL_KGSL_GPUOBJ_FREE(self.fd, id=mem.meta[0].id)
+      FileIOInterface.munmap(mem.va_addr, mem.meta[0].mmapsize)
 
   def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
     if sig[0] < value:

tinygrad/runtime/ops_webgpu.py

@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import BufferStorage, Compiled, Allocator, BufferSpec, Program, TinyELF
+from tinygrad.device import Compiled, Allocator, BufferSpec, Program, TinyELF
 from tinygrad.renderer.wgsl import WGSLRenderer
 from tinygrad.helpers import round_up, suppress_finalizing, getenv, to_mv
 from tinygrad.runtime.autogen import webgpu
@@ -147,10 +147,10 @@ class WebGPUProgram(Program['WebGpuDevice']):
     return None
 
 class WebGpuAllocator(Allocator['WebGpuDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     # WebGPU buffers have to be 4-byte aligned
-    return BufferStorage(webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
-      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc)))
+    return (webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
+      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc)), None), None
   def _copyin(self, dest:webgpu.WGPUBuffer, src:memoryview):
     if src.nbytes % 4:
       padded_src = bytearray(round_up(src.nbytes, 4))
@@ -160,7 +160,7 @@ class WebGpuAllocator(Allocator['WebGpuDevice']):
     dest[:] = buf_to_mv(tmp_buf:=self.dev._readable_buffer(src))[:dest.nbytes]
     self.dev.free(tmp_buf)
 
-  def _free(self, storage:BufferStorage, options:BufferSpec): self.dev.free(storage.buf)
+  def _free(self, opaque:webgpu.WGPUBuffer, options:BufferSpec): self.dev.free(opaque)
 
 class WebGpuDevice(Compiled):
   def __init__(self, device:str):
@@ -189,7 +189,7 @@ class WebGpuDevice(Compiled):
     super().__init__(device, WebGpuAllocator(self), [WGSLRenderer], WebGPUProgram,
                      arch="shader-f16" * (webgpu.WGPUFeatureName_ShaderF16 in self.features))
 
-  def synchronize(self, timeout:int|None=None): QueueOnSubmittedWorkDone(self.queue)
+  def synchronize(self): QueueOnSubmittedWorkDone(self.queue)
 
   @suppress_finalizing
   def free(self, buf:webgpu.WGPUBuffer):

tinygrad/runtime/support/hcq.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
+from typing import Any
 import ctypes, os
 try: import fcntl # windows misses that
 except ImportError: fcntl = None #type:ignore[assignment]
 from tinygrad.helpers import DEV, getenv, pluralize
+from tinygrad.device import Compiled
+from tinygrad.uop.ops import sint
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.memory import MMIOInterface as MMIOInterface, BumpAllocator as BumpAllocator
 
@@ -52,3 +55,26 @@ def hcq_filter_visible_devices(devs, device):
   else: ids = [int(x) for x in idstr.split(',') if x.strip()]
   assert all(x < len(devs) for x in ids), f"invalid visibility filter: {ids} ({pluralize('device', len(devs))} available)"
   return [devs[x] for x in ids] if ids else devs
+
+class HCQBuffer:
+  def __init__(self, va_addr:sint, size:int, meta:Any=None, _base:HCQBuffer|None=None, view:MMIOInterface|None=None, owner:Any=None):
+    self.va_addr, self.size, self.meta, self._base, self.view = va_addr, size, meta, _base, view
+    self._devs, self.owner = ([owner] if owner is not None else []), owner
+    self._mappings:dict[Compiled, HCQBuffer] = {} # mapping to the other devices
+
+  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, meta=self.meta,
+      _base=self._base or self, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
+
+  def cpu_view(self) -> MMIOInterface:
+    assert self.view is not None, "buffer has no cpu_view"
+    return self.view
+
+  @property
+  def base(self) -> HCQBuffer: return self._base or self
+
+  @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

tinygrad/runtime/support/hcq2.py

@@ -1,17 +1,24 @@
 from __future__ import annotations
-from typing import cast, Any
-import functools, itertools, weakref, ctypes, importlib
+from typing import cast, TypeVar, Generic, Any, TYPE_CHECKING
+import functools, time, itertools, decimal, weakref, statistics, ctypes, importlib
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import dedup, pluralize, unwrap, VIZ, HCQ2, to_tuple, ContextVar, Context, panic, partition, DEV, ALL2ALL, getenv
-from tinygrad.device import Device, Buffer, BufferSpec, DepsTracker
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
+from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2, cpu_profile, mv_address
+from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, perf_counter_us, DEV
+from tinygrad.device import Device, Buffer, BufferSpec, Compiled, Allocator, DepsTracker
+from tinygrad.device import ProfileGraphEntry, ProfileGraphEvent, ProfileDeviceEvent
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
 from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
-from tinygrad.renderer import Estimates
-from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear, lower_and_compile, _resolve
+from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
+from tinygrad.renderer import Renderer, Estimates
+from tinygrad.engine.realize import get_call_arg_uops, get_call_name, get_call_outs_ins, estimate_uop, pm_flatten_linear
+from tinygrad.engine.realize import lower_and_compile
+
+if TYPE_CHECKING: from tinygrad.runtime.support.hcq import HCQBuffer # TODO: remove that
 
 # *****************
 # 0. helpers
 
+HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQ2Compiled')
 HCQ_RUNTIME_DEV = ContextVar("HCQ_RUNTIME_DEV", "PYTHON" if DEV.interface.startswith("MOCK") else "CPU")
 HCQ_CACHE_THRESH = ContextVar("HCQ_CACHE_THRESH", 64)
 HCQ_DEVS = frozenset(("NV", "QCOM")) | (frozenset(("AMD",)) if HCQ2 else frozenset())
@@ -27,7 +34,6 @@ class HCQInfo:
   table:int = -1
   inputs:tuple[tuple[UOp, str, int], ...] = ()
   slots:tuple[tuple[str, int], ...] = () # per device, the position of its batch slots in the args
-  host_deps:tuple[tuple[str, str], ...] = () # (memory owner, accessing device)
 
 def all_devices_in(d:Any, c:frozenset[str]) -> bool: return {x.split(":")[0] for x in to_tuple(d)} <= c
 
@@ -37,8 +43,10 @@ def get_enqueue_devs(call:UOp) -> Any|None:
   if call.src[0].op is Ops.COPY: bufs = bufs[::-1] # copies push from the src device: p2p writes are faster than reads
   devs = min(bufs, key=lambda b: not all_devices_in(b.device, HCQ_DEVS)).device
   if not all_devices_in(devs, HCQ_DEVS): return None
-  if call.src[0].op is Ops.COPY and to_tuple(devs)[0].startswith("QCOM"): return None # QCOM is unified memory and uses host copies
-  return devs
+  dev = cast(HCQ2Compiled, Device[to_tuple(devs)[0]])
+  if not all(all_devices_in(b.device, HCQ_DEVS | dev.host_devs) for b in bufs): return None
+  # a device without a copy queue leaves copies to its allocator
+  return devs if call.src[0].op is not Ops.COPY or dev.has_copy_queue else None
 
 def unwrap_view(v:UOp) -> tuple[UOp, int]: # look through views to (base, byte offset)
   if v.op in (Ops.BITCAST, Ops.AFTER): return unwrap_view(v.src[0])
@@ -98,7 +106,34 @@ def replace_buffer(ctx:tuple[bool, list[UOp], dict[UOp, int]], b:UOp) -> UOp:
 pm_replace_buffers = PatternMatcher([(UPat(Ops.BUFFER, name="b"), replace_buffer)])
 
 # *****************
-# 1.1. prep: unwrap multi
+# 1.1. prep: staging copies
+
+STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
+
+@functools.cache
+def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
+
+def _need_staging(a, b):
+  if not all_devices_in(a.device, HCQ_DEVS): return False
+  dev = cast(HCQ2Compiled, Device[to_tuple(a.device)[0]])
+  return not all_devices_in(b.device, HCQ_DEVS | dev.host_devs) and dev.has_copy_queue
+
+def stage_copy(dst:UOp, src:UOp) -> UOp|None:
+  if not (_need_staging(src, dst) or _need_staging(dst, src)): return None
+
+  base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
+  chunk = (STAGING_SIZE // STAGING_SLOTS) // it
+  for i, off in enumerate(range(0, src.max_numel(), chunk)):
+    stage = base[(so:=(i % STAGING_SLOTS) * chunk * it):so + (n:=min(chunk, src.max_numel() - off)) * it]
+    copies += [src[off:off+n].copy_to_device("CPU").call(stage, src[off:off+n]), stage.copy_to_device(dst.device).call(dst[off:off+n], stage)]
+  return UOp(Ops.LINEAR, src=tuple(copies))
+
+pm_insert_copy_staging = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src"))), stage_copy),
+])
+
+# *****************
+# 1.2. prep: one call per device: the args pick their lane, the DEVICE axis binds to it
 
 def unwrap_call(call:UOp) -> UOp|None:
   if get_enqueue_devs(call) is None or (n:=max(len(to_tuple(a.device)) for a in get_call_arg_uops(call))) == 1: return None
@@ -107,36 +142,6 @@ def unwrap_call(call:UOp) -> UOp|None:
                                    for i in range(n)))
 pm_unwrap_multi = PatternMatcher([(UPat(Ops.CALL, name="call"), unwrap_call)])
 
-# *****************
-# 1.2. prep: staging copies
-
-STAGING_SIZE, STAGING_SLOTS = (4 if DEV.interface.startswith("MOCK") else 128) << 20, 2
-
-@functools.cache
-def _staging() -> Buffer: return Buffer("CPU", STAGING_SIZE, dtypes.uint8, preallocate=True)
-
-def stage_copy(ctx:tuple[UOp, ...], call:UOp, dst:UOp, src:UOp) -> UOp|None:
-  if (device:=get_enqueue_devs(call)) is None: return None
-  try:
-    for b in (dst, src): cast(Buffer, _resolve(b, ctx).buffer).get_buf(device)
-  except (RuntimeError, OSError):
-    _staging().get_buf(device)
-    base, it, copies = UOp.from_buffer(_staging()), src.dtype.itemsize, []
-    chunk = (STAGING_SIZE // STAGING_SLOTS) // it
-    for i, off in enumerate(range(0, src.max_numel(), chunk)):
-      stage = base[(so:=(i % STAGING_SLOTS) * chunk * it):so + (n:=min(chunk, src.max_numel() - off)) * it]
-      copies += [src[off:off+n].copy_to_device("CPU").call(stage, src[off:off+n]), stage.copy_to_device(dst.device).call(dst[off:off+n], stage)]
-    return UOp(Ops.LINEAR, src=tuple(copies))
-
-  if Device[device].has_copy_queue: return None
-  out, inp = (UOp.param(i, dtypes.uint8, b.nbytes(), device=device) for i, b in enumerate((dst, src)))
-  ast = out.index(r:=UOp.range(src.nbytes(), 0)).store(inp.index(r).load()).end(r).sink(arg=KernelInfo())
-  return lower_and_compile(call.replace(src=(ast, *call.src[1:])))
-
-pm_insert_copy_staging = PatternMatcher([
-  (UPat(Ops.CALL, src=(UPat(Ops.COPY), UPat(name="dst"), UPat(name="src")), name="call", allow_any_len=True), stage_copy),
-])
-
 # *****************
 # 2. deps
 
@@ -235,24 +240,12 @@ def _finalize_batch(ctx:BatchCtx) -> UOp:
   sink = UOp.sink(*merged, arg=KernelInfo("hcq_submit"), tag=1)
   for pm in [Device[d].pm_batch for d in ctx.queues if Device[d].pm_batch is not None]: # a device adds its own work to the batch
     if (r:=pm.rewrite(sink)) is not None: sink = r
-  host_deps = tuple(dedup((host, devs[0]) for call, devs, _ in ctx.batch for buf in get_call_arg_uops(call)
-                         for host in to_tuple(buf.device) if host not in ctx.queues))
-  return sink.call(*(ctx.slots.values() if ctx.profile else ()),
-                   aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns), estimates=estimates, host_deps=host_deps))
+  return sink.call(*(ctx.slots.values() if ctx.profile else ()), aux=HCQInfo(tuple(ctx.queues), kernels=tuple(kerns), estimates=estimates))
 
 @rewrite_group(new_ctx=False)
 def sched_batches(l:UOp, profile:bool) -> UOp:
   devs = [() if (d:=get_enqueue_devs(c)) is None else tuple(Device.canonicalize(x) for x in to_tuple(d)) for c in l.src]
-
-  # assign to queues
-  peers = sorted({Device.canonicalize(d) for c in l.src if c.src[0].op is Ops.COPY
-                  for b in get_call_arg_uops(c) for d in to_tuple(b.device) if d.split(":")[0] == "AMD"})
-  num_queues = max(1, getenv("HCQ_NUM_SDMA", min(len(peers), 8) if ALL2ALL >= 1 else 1))
   queues = ["COMPUTE:0" if c.src[0].op is Ops.PROGRAM else "COPY:0" for c in l.src]
-  for i, c in enumerate(l.src):
-    if c.src[0].op is Ops.COPY and all(b.device in peers for b in get_call_arg_uops(c)):
-      queues[i] = f"COPY:{(peers.index(c.src[1].device) - peers.index(c.src[2].device) - 1) % len(peers) % num_queues}"
-
   srcs:list[UOp] = []
   for hcq, grp in itertools.groupby(zip(l.src, devs, queues), key=lambda e: bool(e[1])):
     srcs += [_finalize_batch(BatchCtx(list(grp), profile))] if hcq else [c for c, _, _ in grp]
@@ -401,8 +394,7 @@ def lower_call(call:UOp) -> UOp|None:
   # encode bodies
   ctx = EncodeCtx(call.arg.aux.device)
   devs = [Device[d] for d in dedup([d.split(":")[0] for d in ctx.devs])]
-  body = graph_rewrite(call.src[0], sum([d.pm_encode for d in devs if d.pm_encode is not None], PatternMatcher([])) + pm_hcq_encode,
-                       ctx=ctx, bpm=pm_patches, name="encode")
+  body = graph_rewrite(call.src[0], sum([d.pm_encode for d in devs], PatternMatcher([])) + pm_hcq_encode, ctx=ctx, bpm=pm_patches, name="encode")
   body = graph_rewrite(body, sum([d.pm_lower for d in devs if d.pm_lower is not None], PatternMatcher([])), ctx=ctx, bpm=pm_patches, name="lower")
 
   # resize table
@@ -428,21 +420,22 @@ def lower_call(call:UOp) -> UOp|None:
   return call.replace(src=(sink, *bufs), arg=replace(call.arg, aux=info)).after(*patches)
 pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call", allow_any_len=True), lower_call)])
 
-hcq_compile_cache:dict[tuple[UOp, bool, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
+hcq_compile_cache:dict[tuple[UOp, bool], UOp] = {} # eager templates: a buffer-free linear (uops are hash-consed) to its compiled form
 
-@rewrite_group(lambda linear,input_uops,profile,cache=False,ret=None: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool, cache=False) -> UOp:
+@rewrite_group(lambda linear,input_uops,profile,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
+def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
   if any(isinstance(getattr(c.without_after.arg, "aux", None), HCQInfo) for c in linear.src): return linear # compiled already
 
-  if cache and input_uops is not None:
+  if input_uops is not None:
     use_rt = len(linear.src) < HCQ_CACHE_THRESH # small schedules use runtime address patches so linked schedules can be cached without input buffers
     slots = {u:i for i,u in reversed(tuple(enumerate(input_uops)))}
     linear = graph_rewrite(linear, pm_replace_buffers, ctx=(use_rt, input_uops, slots), walk=True, name="replace buffers")
-  linear = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, ctx=tuple(input_uops or ()), name="prep calls")
-  if cache and input_uops is not None and (cached:=hcq_compile_cache.get(key:=(linear, profile, ALL2ALL >= 1))) is not None: return cached
-  lin = graph_rewrite(sched_batches(linear, profile), pm_encode, walk=True, name="encode")
+    if (cached:=hcq_compile_cache.get(key:=(linear, profile))) is not None: return cached
+  lin = graph_rewrite(linear, pm_unwrap_multi+pm_insert_copy_staging+pm_flatten_linear, name="prep calls")
+  lin = sched_batches(lin, profile)
+  lin = graph_rewrite(lin, pm_encode, walk=True, name="encode")
   with Context(EMULATED_DTYPES=""): final_linear = lower_and_compile(lin)
-  if cache and input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
+  if input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
   return final_linear
 
 # *****************
@@ -454,16 +447,14 @@ class LinkCtx: inputs:dict[UOp, UOp]; use_rt:bool; refs:list[UOp] = field(defaul
 def bufferize_buf(ctx:LinkCtx, b:UOp) -> UOp|None: # ctx: a kept link (the jit's) owns the linear's buffers, a one-shot borrows ring slots
   if b.tag is None: return None # a param, not a placeholder
 
-  dev = Device[to_tuple(b.device)[0]]
+  dev = cast(HCQ2Compiled, Device[to_tuple(b.device)[0]])
 
   # device owns the placeholders it names
   if (r:=cast(Buffer|None, dev.pm_bufferize.rewrite(b, ctx=dev))) is not None: pass
   elif not ctx.use_rt:
     spec = BufferSpec(host=b.arg.volatile, uncached=b.arg.volatile, cpu_access=True)
     r = Buffer(dev.device, b.max_numel(), b.dtype, options=spec, preallocate=True)
-  else:
-    off = dev.rt_allocator(True, b.arg.volatile).alloc(max(b.max_numel() * b.dtype.itemsize, 1), alignment=256)
-    r = dev.rt_buffer(True, b.arg.volatile).view(b.max_numel(), b.dtype, off).ensure_allocated()
+  else: r = dev.rt_view(b.max_numel() * b.dtype.itemsize, b.dtype, host=b.arg.volatile)
 
   return UOp.from_buffer(r, HCQ_RUNTIME_DEV.value)
 
@@ -471,7 +462,7 @@ def resolve_getaddr(ctx:LinkCtx, g:UOp) -> UOp|None:
   buf, off = unwrap_view(g.src[0])
   if buf.op not in {Ops.BUFFER, Ops.MSELECT}: return None
   ctx.refs.append(buf) # add to refs
-  return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]) + off, dtypes.uint64)
+  return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]).va_addr + off, dtypes.uint64)
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
   if getattr(b:=cast(Buffer, buf.buffer), '_hcq_written', None) is not blob.arg: # TODO: remove me
@@ -516,3 +507,129 @@ def hcq_link(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UO
   if ctx.refs: linked = linked.replace(src=(linked.src[0].after(*dedup(ctx.refs)), *linked.src[1:])) # attach refs to linear
   if cache and linked is not linear: link_linear_cache[linear] = linked
   return linked
+
+# *****************
+# Device classes
+
+class HCQ2Compiled(Compiled):
+  timestamp_divider: float = 1000.0
+  wait_timeout_ms: float = 30000.0
+  sleep_timeout_ms: int|None = None
+  rt_nbytes: int = 64 << 20 # the pool every per-linear buffer is carved out of
+  host_devs: frozenset[str] = frozenset({"CPU"})
+  pm_encode: PatternMatcher = PatternMatcher([]) # the backend's own encode rules, matched by its submit names
+  var_vals: dict[str, int] = {}
+
+  def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
+    self.can_recover = can_recover
+
+    self.pm_bufferize = PatternMatcher([
+      (UPat(Ops.PARAM, tag="timeline"), lambda ctx: ctx.timeline),
+      (UPat(Ops.PARAM, tag="program", name="b"),
+       lambda ctx, b: ctx.prog_bufs.setdefault(b, Buffer(ctx.device, b.max_numel(), b.dtype, options=BufferSpec(cpu_access=True, nolru=True)))),
+      (UPat(Ops.PARAM, name="b"), lambda b: cfunc_buf(*b.tag[1:]) if isinstance(b.tag, tuple) and b.tag[0] == "cfunc" else None),
+    ])
+    super().__init__(device, allocator, compilers, runtime, None, arch=arch)
+
+    self.prog_bufs:dict[UOp, Buffer] = {}
+    self.prof_ents:dict[tuple[Buffer, int], ProfileGraphEntry] = {} # (a batch's timestamps, start slot) -> entry, read at synchronize
+
+  @functools.cached_property
+  def timeline(self) -> Buffer: # [the signal, the value the last submitted batch signals]
+    buf = Buffer(self.device, 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    buf.host.view(fmt='B')[:16] = bytes(16)
+    return buf
+
+  def collect_prof(self):
+    if PROFILE:
+      es = list(self.prof_ents.items())
+      sigs = [buf.host.view(fmt='Q')[i]/decimal.Decimal(self.timestamp_divider) for (buf, _), e in es for i in (e.st_id, e.en_id)]
+      Compiled.profile_events.append(ProfileGraphEvent([replace(e, st_id=2*i, en_id=2*i+1) for i,(_, e) in enumerate(es)], [], sigs))
+    self.prof_ents.clear()
+
+  def _at_profile_finalize(self): # the device clock against the host's: the median offset over a few tiny kernels
+    from tinygrad.tensor import Tensor
+    tdiffs = []
+    for _ in range(5):
+      with Context(DEBUG=0, BEAM=0, TRACK_MATCH_STATS=0): Tensor.ones(1, device=self.device).contiguous().realize()
+      if not (ents:=list(self.prof_ents.items())): return
+      self.prof_ents.clear()
+      st = perf_counter_us()
+      self.synchronize()
+      gpu = max(buf.host.view(fmt='Q')[e.en_id] for (buf, _), e in ents)/decimal.Decimal(self.timestamp_divider)
+      tdiffs.append((st+perf_counter_us())/2 - gpu)
+    Compiled.profile_events.append(ProfileDeviceEvent(self.device, statistics.median(tdiffs), self.device_props()))
+
+  @functools.cache
+  def rt_allocator(self, uncached:bool=True, host:bool=False) -> BumpAllocator: return BumpAllocator(self.rt_nbytes)
+
+  @functools.cache
+  def rt_buffer(self, uncached:bool=True, host:bool=False) -> Buffer:
+    spec = BufferSpec(host=host, uncached=uncached, cpu_access=True)
+    return Buffer(self.device, self.rt_allocator(uncached, host).size, dtypes.uint8, options=spec, preallocate=True)
+
+  def rt_view(self, nbytes:int, dtype:DType=dtypes.uint8, uncached:bool=True, host:bool=False) -> Buffer: # a slot of the ring, wraps silently
+    off = self.rt_allocator(uncached, host).alloc(max(nbytes, 1), alignment=256)
+    return self.rt_buffer(uncached, host).view(nbytes // dtype.itemsize, dtype, off).ensure_allocated()
+
+  def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
+    timeout = timeout if timeout is not None and self.can_recover else None
+    st, done = time.perf_counter(), sig[0]
+    while done < value:
+      if done != (done:=sig[0]): st = time.perf_counter()
+      elif (elapsed:=time.perf_counter() - st) > (timeout or self.wait_timeout_ms) / 1000: raise RuntimeError(f"{self.device} signal wait timed out")
+      elif self.sleep_timeout_ms is not None and elapsed > self.sleep_timeout_ms / 1000: self.on_sleep()
+
+  def synchronize(self, timeout:int|None=None):
+    try: self._wait_signal(tl:=self.timeline.host.view(fmt='Q'), tl[1], timeout)
+    except RuntimeError:
+      self.on_device_hang()
+      raise
+    if self.prof_ents: self.collect_prof()
+
+  def on_device_hang(self): raise RuntimeError(f"{self.device} hang detected")
+
+  def on_sleep(self):
+    if (iface:=getattr(self, "iface", None)) is not None and hasattr(iface, "sleep"): iface.sleep(self.sleep_timeout_ms)
+
+  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
+
+  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
+
+  def finalize(self):
+    try: self.synchronize() # try to finalize the device in any case
+    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
+    super().finalize()
+
+@dataclass
+class HCQ2Buffer:
+  va_addr:sint
+  meta:Any=None
+  view:MMIOInterface|None=None
+
+  def offset(self, offset:int, size:int) -> HCQ2Buffer:
+    return HCQ2Buffer(self.va_addr+offset, meta=self.meta, view=(self.view.view(offset=offset, size=size) if self.view is not None else None))
+
+class HCQAllocator(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
+  def _as_buffer(self, buf:HCQBuffer) -> memoryview|None: return buf.view.mv if buf.view is not None else None
+  def _copyout(self, dest:memoryview, src:HCQBuffer): # TODO: remove with memcpy on cpu worker?
+    self.dev.synchronize()
+    with cpu_profile(f"{self.dev.device} -> TINY", f"{self.dev.device}:COPY"): ctypes.memmove(mv_address(dest), src.cpu_view().addr, dest.nbytes)
+
+  def _map(self, buf:HCQBuffer) -> tuple: # a mapping lives on the opaque, like hcq1: the lru hands the same one to many Buffers
+    if self.dev not in buf.mapped_devs:
+      if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
+      buf.mappings[self.dev] = self._do_map(buf)
+      buf.mapped_devs.append(self.dev)
+    return (mapped:=buf.mappings[self.dev]), mapped.meta
+
+  def _do_unmap(self, mb): getattr(self.dev, "iface").free(mb)
+
+  @suppress_finalizing
+  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
+    if options is not None and options.external_ptr is not None: return
+    for dev in buf.mapped_devs: dev.synchronize()
+    for d, mb in buf.mappings.items(): d.allocator._do_unmap(mb)
+    if hasattr(self, '_do_free'): self._do_free(buf, options)
+
+  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)

tinygrad/runtime/support/system.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
 import os, mmap, array, functools, ctypes, ctypes.util, select, contextlib, dataclasses, sys, struct, socket
-from tinygrad.device import BufferStorage, Buffer, Device
 from tinygrad.helpers import round_up, getenv, OSX, temp, ceildiv, DEBUG, pluralize
 from tinygrad.runtime.autogen import libc, pci, vfio
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, hcq_filter_visible_devices
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, HCQBuffer, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
 from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface
 
@@ -262,8 +261,7 @@ class PCIIfaceBase:
     self.dev_impl = dev_impl_t(self.pci_dev)
     self.dev, self.vram_bar, self.count = dev, vram_bar, len(hcq_filter_visible_devices(System.list_devices(vendor, devices, base_class), dn))
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False,
-            **kwargs) -> BufferStorage:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
     should_use_sysmem = host or ((cpu_access if self.is_bar_small() else (uncached and cpu_access)) and not force_devmem)
 
     # Align size to huge pages for large allocations, otherwise the unaligned tail falls back to 4KB pages, increasing TLB pressure.
@@ -273,35 +271,32 @@ class PCIIfaceBase:
       vaddr = self.dev_impl.mm.alloc_vaddr(size:=round_up(size, mmap.PAGESIZE), align=mmap.PAGESIZE)
       memview, paddrs = self.pci_dev.alloc_sysmem(size, vaddr=vaddr, contiguous=contiguous)
       mapping = self.dev_impl.mm.map_range(vaddr, size, [(paddr, 0x1000) for paddr in paddrs], aspace=AddrSpace.SYS, snooped=True, uncached=True)
-      return BufferStorage(vaddr, PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), memview)
+      return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
 
     mapping = self.dev_impl.mm.valloc(size:=round_up(size, 0x1000), uncached=uncached, contiguous=cpu_access, zero=zero)
     barview = self.pci_dev.map_bar(bar=self.vram_bar, off=mapping.paddrs[0][0], size=mapping.size) if cpu_access else None
-    return BufferStorage(mapping.va_addr, PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), barview)
+    return HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)
 
-  def free(self, storage:BufferStorage):
-    if storage.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(storage.meta.mapping)
-    if storage.meta.has_cpu_mapping: FileIOInterface.munmap(storage.buf, storage.meta.mapping.size)
-
-  def unmap(self, mapping:BufferStorage): self.dev_impl.mm.unmap_range(*mapping.meta)
+  def free(self, b:HCQBuffer):
+    if b.owner != self.dev: self.dev.iface.dev_impl.mm.unmap_range(b.va_addr, round_up(b.size, 0x1000))
+    if b.owner == self.dev and b.meta.mapping.aspace is AddrSpace.PHYS: self.dev_impl.mm.vfree(b.meta.mapping)
+    if b.owner == self.dev and b.meta.has_cpu_mapping: FileIOInterface.munmap(b.va_addr, b.size)
 
   def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
     return [(p + self.pci_dev.bar_info(self.vram_bar)[0], sz) for p, sz in paddrs], AddrSpace.SYS
 
-  def map(self, b:Buffer) -> BufferStorage:
-    if b.device.split(":")[0] in {"CPU", "PYTHON", "NPY"}:
-      if b._buf % 0x1000: raise RuntimeError("Host mapping requires a page-aligned address")
-      lo, size = b._buf, round_up(b.nbytes, 0x1000)
-      if not self.dev_impl.mm.va_base <= lo < lo + size <= self.dev_impl.mm.va_base + (1 << self.dev_impl.mm.va_bits):
-        raise RuntimeError(f"Host address {lo:#x} is outside the GPU virtual address range")
-      System.lock_memory(lo, size)
-      paddrs, aspace, snooped, uncached = [(x, 0x1000) for x in System.system_paddrs(lo, size)], AddrSpace.SYS, True, True
-    elif isinstance(ifa:=getattr(Device[b.device], "iface", None), PCIIfaceBase):
-      if ifa.is_bar_small(): raise RuntimeError(f"P2P mapping not supported for small bar devices: {b.device} -> {self.dev.device}")
-      lo, size, snooped, uncached = b._buf, b.meta.mapping.size, True, b.meta.mapping.uncached
+  def map(self, b:HCQBuffer):
+    if b.owner is not None and b.owner._is_cpu():
+      System.lock_memory(int(b.va_addr), b.size)
+      paddrs, aspace = [(x, 0x1000) for x in System.system_paddrs(int(b.va_addr), round_up(b.size, 0x1000))], AddrSpace.SYS
+      snooped, uncached = True, True
+    elif (ifa:=getattr(b.owner, "iface", None)) is not None and isinstance(ifa, PCIIfaceBase):
+      if ifa.is_bar_small(): raise RuntimeError(f"P2P mapping not supported for small bar devices: {b.owner} -> {self.dev}")
+
+      snooped, uncached = True, b.meta.mapping.uncached
       if b.meta.mapping.aspace is AddrSpace.SYS: paddrs, aspace = b.meta.mapping.paddrs, AddrSpace.SYS
       else: paddrs, aspace = ifa.p2p_paddrs(b.meta.mapping.paddrs)
-    else: raise RuntimeError(f"map failed: {b.device} -> {self.dev.device}")
+    else: raise RuntimeError(f"map failed: {b.owner} -> {self.dev}")
 
-    self.dev_impl.mm.map_range(lo, size, paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
-    return BufferStorage(b._buf, (lo, size))
+    self.dev_impl.mm.map_range(int(b.va_addr), round_up(b.size, 0x1000), paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
+    return HCQBuffer(b.va_addr, b.size, meta=b.meta, owner=b.owner)

tinygrad/runtime/support/usb.py

@@ -484,11 +484,14 @@ def _words(dev) -> Buffer: # zero the read signal and scratch
   b = Buffer(dev.device, 2, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True, nolru=True), preallocate=True)
   b.host.view(fmt='B')[:8] = bytes(8)
   return b
+@functools.cache
+def _asm24(dev) -> Buffer:
+  return Buffer(dev.device, 0x85000, dtypes.uint8, options=BufferSpec(external_ptr=dev.iface.ctrl.va_addr, nolru=True)).allocate(dev.iface.ctrl)
 pm_usb_bufferize = PatternMatcher([
   (UPat(Ops.PARAM, tag="usb_host"), lambda ctx: _host_block(ctx)),
   (UPat(Ops.PARAM, tag={"usb_xfer0", "usb_xfer1"}, name="b"), lambda ctx, b: _xfer(ctx, b.tag)),
   (UPat(Ops.PARAM, tag="usb_vram"), lambda ctx: _words(ctx)),
-  (UPat(Ops.PARAM, tag="usb_asm24"), lambda ctx: ctx.iface.ctrl),
+  (UPat(Ops.PARAM, tag="usb_asm24"), lambda ctx: _asm24(ctx)),
   (UPat(Ops.PARAM, name="b"), lambda b: Buffer("CPU", b.max_numel(), b.dtype, preallocate=True) if str(b.tag).startswith("cmdbuf_copy") else None),
 ])
 

tinygrad/schedule/indexing.py

@@ -34,10 +34,6 @@ def realize_srcs(ctx:IndexingContext, rb:UOp) -> None:
 def realize_store_after_src(ctx:IndexingContext, dest:UOp, src:UOp):
   # you don't usually have to do this for assign unless there's a WAR hazard like TestAssign.test_assign_double_diamond_reduce
   if dest.base in src.toposort(enter_calls=False): ctx.realize_map[src] = None
-  # the source of a cross device STORE is materialized on its own device first: the STORE itself is the copy
-  # NOTE: buffer identity views (shard views with max_shape != shape) must be materialized too, copies can't read them
-  if src.device is not None and dest.device != src.device:
-    ctx.realize_map[src] = ctx.non_removable[src] = None
 
 def realize_custom_kernel_srcs(ctx:IndexingContext, c:UOp) -> None:
   for s in c.src[1:]:
@@ -53,7 +49,7 @@ pm_generate_realize_map = PatternMatcher([
   (UPat({Ops.CONTIGUOUS, Ops.STORE}, name="tr"), realize),
   # realize srcs of these
   (UPat((Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
-  # sometimes we need to realize the src of STORE if there's a self-access, or if it's a cross device store
+  # sometimes we need to realize the src of STORE if there's a self-access
   (UPat(Ops.STORE, src=(UPat.var("dest"), UPat.var("src"))), realize_store_after_src),
 ])
 

tinygrad/schedule/prepare.py

@@ -1,3 +1,4 @@
+import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
 from tinygrad.uop.ops import graph_rewrite, rewrite_group, identity_element, resolve_returned_after
@@ -130,13 +131,6 @@ def expand_bitcast(bc:UOp) -> UOp|None:
   parts = [tmp>>8*i*ns for i in range(os//ns)]
   return parts[0].stack(*parts[1:], dim=-1).flatten(-2).cast(new_uint).bitcast(bc.dtype)
 
-def copy_to_anon_store(x:UOp, copy:UOp):
-  # the buffer created here is inside the call and is not persisted, like the buffers created for contiguous
-  # copies must read from a whole buffer, not a view: materialize anything lacking buffer identity (SDMA can't do offset copies)
-  if not x.has_buffer_identity(after_ok=True): x = x.contiguous()
-  buf = UOp.new_buffer(copy.device, prod(x.max_shape), copy.dtype).reshape(x.max_shape)
-  return buf.after(buf.store(x)).reshape(copy.shape)
-
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls with RETURNED inputs (inline the body)
   (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.has_unbound_outputs else None),
@@ -161,12 +155,8 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   # copy to same device is a no-op
   (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x if x.device == copy.device else None),
 
-  # a COPY in src[1] of a plain STORE can just be removed: a STORE to a buffer on a different device is a COPY
-  (UPat(Ops.STORE, src=(UPat.var("dst"), UPat(Ops.COPY, src=(UPat.var("x"),), name="cpy"))),
-   lambda dst,x,cpy: dst.store(x) if dst.device == cpy.device and dst.has_buffer_identity(after_ok=True) else None),
-
-  # a bare COPY is an anonymous store: realize it as a STORE into a fresh call-local buffer on the copy device
-  (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), copy_to_anon_store),
+  # copy on reshape is reshape on copy
+  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)),
 
   # reshaping on STORE can be a NOOP
   (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True),
@@ -203,10 +193,31 @@ earliest_rewrites = mop_cleanup+PatternMatcher([
   (UPat(Ops.AFTER, name="s"), lambda s: s.replace(src=(s.src[0],)+tuple(walk_mop(u) for u in s.src[1:] if u.op is not Ops.NOOP))),
 ])
 
+def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
+  input_src = copy.src[0]
+  # if it's a COPY, we need to give the input buffer identity
+  if not input_src.has_buffer_identity(after_ok=True) and copy.op is Ops.COPY: input_src = input_src.contiguous()
+  input_src = input_src.flatten()
+  if existing_buf is not None:
+    # if the existing buffer is not a full buffer, we can't use it
+    if not existing_buf.has_buffer_identity(after_ok=True): return None
+    # if there's already a buffer, we just use it
+    return existing_buf.flatten().store(input_src)
+  # create the output buffer
+  buf = UOp.new_buffer(copy.device, prod(input_src.max_shape), copy.dtype)
+  # reshape back to input
+  return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape)
+
+pm_copy_to_store = PatternMatcher([
+  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
+  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
+])
+
 @rewrite_group(new_ctx=False)
 def prepare_rangeify(sink:UOp) -> UOp:
   # prepare for rangeify
   tsink = graph_rewrite(sink, multi_pm, name="multi_pm")
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
+  tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
   return tsink

tinygrad/tensor.py

@@ -2,10 +2,10 @@
 from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
 from dataclasses import dataclass, field, replace
-from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
+from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
 from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, _from_np_dtype, _to_np_dtype, PyConst, AddrSpace
-from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey, is_numpy_ndarray
+from tinygrad.helpers import all_int, getenv, fetch, Metadata, TRACEMETA, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc, VIZ, pluralize, SPEC
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, Variable, ConstLike, UPat, PatternMatcher, GroupOp, graph_rewrite, rewrite_group
 from tinygrad.uop.ops import resolve_returned_after, remove_all_tags
@@ -43,10 +43,8 @@ def creation_copy_is_realized(u:UOp):
 # CONTIGUOUS and AFTER + parents are the only nodes that get updated
 add_tags = PatternMatcher([
   (UPat(Ops.COPY, name="u"), creation_copy_is_realized),
-  # no tag on copies that fill an AFTER's whole dest via STORE: merge COPY tag into AFTER (the copy reads that storage).
-  # a partial STORE keeps the tag: the copy mints its own storage like any bare creation copy
-  (UPat(Ops.AFTER, src=(UPat(name="dest"),
-    UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
+  # no tag on copies that are assigned via STORE+AFTER — merge COPY tag into AFTER
+  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE, src=(UPat(name="dest"), UPat(Ops.COPY, name="c")))), name="a"),
    lambda a,c,dest: a.replace(src=(a.src[0], a.src[1].replace(src=(dest, c.rtag(())))), tag=a.tag+c.tag) if a.tag and c.tag else None),
   (UPat(Ops.AFTER, name="x"), tag_uop),
   (UPat(GroupOp.All, name="x"), lambda ctx,x: tag_uop(x) if x in ctx.bases else None),
@@ -256,6 +254,8 @@ def _tensor_holds(u:UOp) -> bool: return any((t:=tref()) is not None and t.uop i
 
 # **** Tensor helper functions ****
 
+def is_numpy_ndarray(x) -> "TypeGuard[numpy.ndarray]": return str(type(x)) == "<class 'numpy.ndarray'>"
+
 def _fromnp(x: 'numpy.ndarray') -> UOp:
   ret = UOp.new_buffer("NPY", x.size, _from_np_dtype(x.dtype))
   # fake realize
@@ -312,7 +312,7 @@ class Tensor(RandMixin):
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
 
     # data might be on a different device
-    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device)
+    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device).clone()
     # cast on the target device, the source may not hold the dtype (numpy has no fp8/bfloat16) or be able to compute it (DISK)
     if _dtype is not None: self.uop = self.uop.cast(_dtype)
 
@@ -448,18 +448,15 @@ class Tensor(RandMixin):
       self.uop = (x.uop.src[0] if x.uop.op is Ops.CONTIGUOUS else x.uop).clone()
       return self
     # STORE+AFTER: STORE is the write effect (void), AFTER wraps the view for correct shape/ranging
-    assign = self.uop.after(store := self.uop.store(x.uop))
+    assign = self.uop.after(self.uop.store(x.uop))
     ib = self.uop
     while ib.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH} and not (ib.has_buffer_identity() and _tensor_holds(ib)): ib = ib.src[0]
     if ib is not self.uop:
       # a partial write needs storage to land in: a pending value gets explicit storage (a clone)
       target = ib if ib.has_buffer_identity(after_ok=True) else ib.clone()
-      if target is not ib:
-        assign = assign.substitute({ib: target}, walk=True)
-        store = assign.src[1]
-      # view assign: the base reads "after the store into the view" (one AFTER level). replace the node under the
-      # views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
-      _apply_map_to_tensors({ib: target.after(store)}, name="Embed View Assign")
+      if target is not ib: assign = assign.substitute({ib: target}, walk=True)
+      # view assign: replace the node under the views (e.g. RESHAPE(BUFFER)) so @function's substitution catches it
+      _apply_map_to_tensors({ib: target.after(assign)}, name="Embed View Assign")
     else:
       # simple assign
       self.uop = assign
@@ -548,9 +545,8 @@ class Tensor(RandMixin):
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
-    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
-    else: ret = Tensor(self.uop.copy_to_device(device))
+    # The transfer owns its destination from construction; COPY itself only describes the transfer.
+    ret = Tensor(self.uop.copy_to_device(device).clone())
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)
 

tinygrad/uop/ops.py

@@ -831,8 +831,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       assert bdtype.fmt is not None, f"{bdtype=} has None fmt"
       ret = UOp.empty(shape:=get_shape(x), dtype=bdtype, device="PYTHON")
       data = struct.pack(f"{prod(shape)}{bdtype.fmt}", *[truncate[bdtype](bdtype.const(xi)) for xi in fully_flatten(x)])
-    if not data: ret.buffer.allocate(memoryview(bytearray()))
-    else: ret.buffer.ensure_allocated().host[:] = data
+    ret.buffer.allocate(memoryview(bytearray(data))) # fake realize. buffer storage must be writable, and bytes isn't
     if ret.dtype != dtype: ret = ret.cast(dtype)
     return ret if ret.device == device else ret.copy_to_device(device)
   def clone(self, device=None) -> UOp:

tinygrad/viz/js/index.js

@@ -272,8 +272,6 @@ const canvasDims = () => {
   return [Math.round(document.querySelector("#profiler").clientWidth-sideRect.width), Math.round(sideRect.height)];
 }
 
-new ResizeObserver(([e]) => e.contentRect.width > 0 && e.target.dispatchEvent(new Event("resize"))).observe(document.getElementById("profiler"));
-
 function selectShape(key) {
   if (key == null) return {};
   const [t, idx] = key.split("-");
@@ -733,10 +731,9 @@ async function renderProfiler(path, opts) {
   let lastCanvasRect = null;
   function resize() {
     const [width, height] = canvasDims();
-    const pixelWidth = Math.trunc(width*dpr), pixelHeight = Math.trunc(height*dpr);
-    if (canvas.width === pixelWidth && canvas.height === pixelHeight) return;
-    canvas.width = pixelWidth;
-    canvas.height = pixelHeight;
+    if (canvas.width === width*dpr && canvas.height === height*dpr) return;
+    canvas.width = width*dpr;
+    canvas.height = height*dpr;
     canvas.style.height = `${height}px`;
     canvas.style.width = `${width}px`;
     ctx.scale(dpr, dpr);
@@ -751,9 +748,9 @@ async function renderProfiler(path, opts) {
   zoomLevel = getZoomIdentity();
   canvasZoom = d3.zoom().filter(vizZoomFilter).on("zoom", e => render(e.transform));
   d3.select(canvas).call(canvasZoom);
-  profiler.on("contextmenu", e => e.ctrlKey && e.preventDefault());
+  document.addEventListener("contextmenu", e => e.ctrlKey && e.preventDefault());
 
-  profiler.on("resize", (e) => resize()); resize();
+  new ResizeObserver(([e]) => e.contentRect.width > 0 && resize()).observe(profiler.node());
   profiler.on("scroll", () => render(zoomLevel));
 
   function findRectAtPosition(x, y) {

tinygrad/viz/serve.py

@@ -49,7 +49,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0",
                Ops.INDEX: "#CEF9B7", Ops.STACK: "#D8F9E4",
                Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff",
                **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80",
-               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#ff90c0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
                Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040",
                Ops.LINEAR: "#7DF4FF",
                Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D",