delete output_shape [PR] (#18090 )

start to clean up reduce related axes [pr] (#18089 )
* start to clean up reduce related axes [pr] * revert
2026-09-09 15:25:29 -04:00 · 2026-09-09 13:42:32 -04:00 · 2026-09-09 20:10:44 +03:00 · 2026-09-09 08:21:48 -07:00 · 2026-09-09 08:11:40 -07:00 · 2026-09-09 12:10:21 +03:00
64 changed files with 740 additions and 479 deletions
@@ -137,7 +137,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')
+      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')
      shell: bash
      run: |
        sudo mkdir -p /var/cache/apt/archives
@@ -165,7 +165,7 @@ runs:
        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')
+      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')
      id: apt-pkgs
      shell: bash
      run: |
@@ -211,20 +211,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') && 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' || inputs.autogen == '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') && 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' || inputs.autogen == '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') && 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' || inputs.autogen == 'true') && steps.apt-pkgs.outputs.installed == 'false'
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -431,6 +431,8 @@ 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
@@ -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: 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
+        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
      - 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 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
+        run: VIZ=-1 PMA=1 PMA_BUFFER_SIZE=8 DEBUG=5 python3 test/backend/test_ops.py TestOps.test_add
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -10,7 +10,6 @@ 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):
@@ -47,8 +46,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._buf.size});")
-    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl.nbytes});")
+    weights += cl.as_memoryview()
  cprog.append("}")

  # write the weights to disk
@@ -15,6 +15,7 @@ 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
@@ -151,7 +152,10 @@ class GPTOSS:

    # output
    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    if getenv("GPTOSS_EMBEDDING", 0):
+      self.tok_embeddings = GPTOSSEmbedding(vocab_size, dim)
+    else:
+      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)
@@ -315,7 +315,7 @@ return 0; }
 if __name__ == "__main__":
  dev = DSPDevice()

-  bufs = [dev.allocator.alloc(0x60000)[0][0] for _ in range(4)]
+  bufs = [dev.allocator.alloc(0x60000).buf for _ in range(4)]

  only_entry = dev.compiler.compile(entry)
  app1 = dev.runtime("test", only_entry)
@@ -268,7 +268,7 @@ return HAP_perf_get_time_us() == 1 ? 4 : 0;
 if __name__ == "__main__":
  dev = DSPDevice()

-  bufs = [dev.allocator.alloc(0x60000)[0][0] for _ in range(4)]
+  bufs = [dev.allocator.alloc(0x60000).buf for _ in range(4)]

  only_entry = dev.compiler.compile(entry)
  app1 = dev.runtime("test", only_entry)
@@ -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, to_mv, prod
+from tinygrad.helpers import Context, 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 bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
+      weight = ''.join(["\\x%02X"%x for x in cl.as_memoryview()])
      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] + ["}"]
@@ -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)[0][0]
-b = hipallocator.alloc(N*N*2)[0][0]
-c = hipallocator.alloc(N*N*2)[0][0]
+a = hipallocator.alloc(N*N*4).buf
+b = hipallocator.alloc(N*N*2).buf
+c = hipallocator.alloc(N*N*2).buf
 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)
@@ -0,0 +1,90 @@
+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)
@@ -0,0 +1,46 @@
+#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
@@ -0,0 +1,37 @@
+#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;
+  }
+}
@@ -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 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 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 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))[0][0] for d,sz in kernargs_size.items()}
+    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()}

    # 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(buf, BufferSpec(cpu_access=True))
+    for fdev, buf in self.kernargs_bufs.items(): fdev.allocator._free(BufferStorage(buf, buf.meta, buf.view), BufferSpec(cpu_access=True))

  @staticmethod
  def supports_uop(batch_devs:list[Compiled], new_call:UOp) -> bool:
@@ -3,17 +3,39 @@ 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 BufferSpec, Compiled, Allocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
+from tinygrad.device import BufferStorage, Buffer, 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.
@@ -282,7 +304,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(((buf, buf.meta), buf.view), buf.size, spec)
+  def _fini(dev, buf, spec): dev.allocator.free(BufferStorage(buf, buf.meta, buf.view), buf.size, spec)

  def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int|None, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
    """
@@ -359,7 +381,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))[0][0]
+      self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True)).buf
      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
@@ -393,7 +415,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))[0][0])
+      HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True)).buf)
      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)
@@ -425,11 +447,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(((oldbuf, oldbuf.meta), oldbuf.view), oldbuf.size, options=options)
-    try: buf, realloced = self.allocator.alloc(new_size, options=options)[0][0], True
+    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
    except MemoryError:
      if force: raise
-      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options)[0][0], False
+      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options).buf, False
    return buf, realloced

  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
@@ -455,9 +477,11 @@ 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))[0][0] for _ in range(batch_cnt)]
+    self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)).buf 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")
@@ -468,10 +492,10 @@ class HCQAllocatorBase(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
    return mapped, mapped.meta

  @suppress_finalizing
-  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 _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 _do_unmap(self, mb): self.dev.iface.free(mb)

@@ -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 HCQCompiled, HCQAllocator, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, hcq_profile
-from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
+from extra.hcq1.hcq import HCQBuffer, HCQCompiled, HCQAllocator, HWQueue, CLikeArgsState, HCQSignal, HCQProgram, hcq_profile
+from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.hcq import MMIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec, TinyELF
+from tinygrad.device import BufferStorage, 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(((self.hw_page, self.hw_page.meta), self.hw_page.view), self.hw_page.size,
+      self.binded_device.allocator.free(BufferStorage(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))[0][0]
+    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
    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))[0][0]
+    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
    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))[0][0]
+    self.hw_page = dev.allocator.alloc((qsz:=round_up(len(self._q), 8)) * 4, BufferSpec(cpu_access=True, nolru=True, uncached=True)).buf
    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))[0][0]
+    self.lib_gpu = self.dev.allocator.alloc(round_up(image.nbytes, 0x1000), buf_spec:=BufferSpec(nolru=True)).buf
    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) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
    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 (opaque, opaque.meta), opaque.view
+    return BufferStorage(opaque, opaque.meta, opaque.view)

  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)

@@ -981,6 +981,19 @@ 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")):
@@ -1086,7 +1099,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))[0][0]
+      self.pmc_buffer = self.allocator.alloc(self.pmc_sched[-1].off + self.pmc_sched[-1].size, BufferSpec(nolru=True, uncached=True)).buf
      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)
@@ -1095,8 +1108,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))[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_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_next_cmd_id = itertools.count(0)

    if self.is_am():
@@ -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
-from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
+from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, HCQBuffer
+from tinygrad.runtime.support.hcq import 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
@@ -3,8 +3,7 @@ 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
-from tinygrad.runtime.support.hcq import HCQBuffer
+from extra.hcq1.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runtime
@@ -2,7 +2,8 @@
 # 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
+from tinygrad.device import Device, BufferSpec, Buffer
+from tinygrad.dtype import dtypes
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.memory import AddrSpace
 from tinygrad.runtime.ops_amd import AMDComputeQueue
@@ -30,8 +31,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 = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
-buf_dst = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
+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)

 bar_base = gpu.iface.pci_dev.bar_info(gpu.iface.vram_bar)[0]
 src_paddr = buf_src.meta.mapping.paddrs[0][0] + bar_base
@@ -40,8 +41,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!"
-gpu.allocator._copyin(buf_src, memoryview(bytearray(test_msg.ljust(BUF_SIZE, b'\x00'))))
-gpu.allocator._copyin(buf_dst, memoryview(bytearray(BUF_SIZE)))
+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.synchronize()

 # post recv WQE on RQ from CPU (scatter entry: byte_count, lkey, addr)
@@ -90,7 +91,7 @@ qp.poll_cq()

 # read back
 result = bytearray(BUF_SIZE)
-gpu.allocator._copyout(memoryview(result), buf_dst)
+result[:] = buf_dst.as_memoryview()
 gpu.synchronize()

 got = bytes(result[:len(test_msg)])
@@ -57,18 +57,34 @@ class USBGPUDebug:
    self.provisioned = True
    print("Provisioning complete.")

-  def reset(self, bootloader=False):
+  def reset(self, bootloader=False, wait=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:
-      time.sleep(1)
-      self.ftdi.set_cbus_gpio(0)
+    if bootloader or wait: self._wait_for(bootloader)
+    if bootloader: 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')

@@ -79,7 +95,9 @@ 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()

@@ -88,15 +106,17 @@ if __name__ == "__main__":
      dbg.provision()

    if args.reset:
-      dbg.reset(bootloader=False)
+      dbg.reset(bootloader=False, wait=args.wait)

    if args.bootloader:
-      dbg.reset(bootloader=True)
+      dbg.reset(bootloader=True, wait=args.wait)

    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)
-
@@ -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}
+\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace, anonymous store}

 \begin{tabular}{@{}l l l l@{}}
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Load} & (buf, alt?, gate?) & device, addrspace & Read (pull) from buffer into a new anonymous buffer. \\
-  & & & Note: this replaces \op{Copy} and \op{Contiguous}. \\
+\op{Load} & (buf, alt?, gate?) & --- & Read from buffer into AddrSpace.ALU. \\
+\op{Copy} & (buf,) & device & Copy buf to device arg. \\
 \bottomrule
 \end{tabular}

@@ -170,9 +170,10 @@ 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
+  from tinygrad.device import Device, TinyELF, Buffer
  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import Target, flat_mv
+  from tinygrad.helpers import Target
+  from tinygrad.dtype import dtypes

  dev = Device["AMD"]
  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
@@ -225,12 +226,11 @@ amdhsa.kernels:
  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))

  buf_sz = _out_bytes(n_lanes)
-  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_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_buf = bytearray(buf_sz)
-  dev.allocator._copyout(flat_mv(memoryview(out_buf)), out_gpu)
+  out_buf = out_gpu.as_memoryview()

  return parse_output(bytes(out_buf), n_lanes)

@@ -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, flat_mv
+from tinygrad.helpers import Target
 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,13 +42,14 @@ 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
+  from tinygrad.device import Device, TinyELF, Buffer
+  from tinygrad.dtype import dtypes
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device["AMD"]
  if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
-  out_gpu = dev.allocator.alloc(LANES * 4)[0][0]
-  code = _code(instructions, out_reg, out_gpu.va_addr)
+  out_gpu = Buffer(dev.device, LANES * 4, dtypes.uint8, preallocate=True)
+  code = _code(instructions, out_reg, out_gpu._buf)
  byte_str = ", ".join(f"0x{b:02x}" for b in code)
  asm_src = f""".text
 .globl test
@@ -87,8 +88,7 @@ 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 = bytearray(LANES * 4)
-  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
+  out = out_gpu.as_memoryview()
  return struct.unpack("<I", out)[0]

 def run_cdna(instructions: list, out_reg: int = 2) -> int:
@@ -6,7 +6,7 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import Target, flat_mv
+from tinygrad.helpers import Target
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm

@@ -37,7 +37,8 @@ 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
+  from tinygrad.device import Device, TinyELF, Buffer
+  from tinygrad.dtype import dtypes
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device["AMD"]
@@ -84,10 +85,9 @@ amdhsa.kernels:
 """
  lib = compiler.compile(asm_src)
  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  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)
+  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()
  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]:
@@ -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, flat_mv
+from tinygrad.helpers import Target
 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,7 +36,8 @@ 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
+  from tinygrad.device import Device, TinyELF, Buffer
+  from tinygrad.dtype import dtypes
  from tinygrad.runtime.support.compiler_amd import HIPCompiler

  dev = Device['AMD']
@@ -85,10 +86,9 @@ amdhsa.kernels:
 """
  lib = compiler.compile(asm_src)
  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  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)
+  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()
  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]:
@@ -359,8 +359,7 @@ class TestUOpPrograms(unittest.TestCase):
    prog = C.end(i, j, k)

    # run program
-    # TODO: make this work with opts_to_apply
-    self._run(prog.sink(arg=KernelInfo(opts_to_apply=())), a, b, c)
+    self._run(prog.sink(arg=KernelInfo()), a, b, c)

    with Context(DEBUG=0): self.assertLessEqual((c-ref).square().mean().item(), 1e-6)

@@ -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), inputs
+    return out, compile_linear(out.schedule_linear(), input_uops=inputs, cache=True), 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.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
+    ranges = [(b._buf, b._buf + 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.va_addr
+            addr = buf._buf
            self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))

  def test_small_eager_cached(self):
@@ -143,7 +143,7 @@ class TestHCQ2Schedule(unittest.TestCase):
          before = tuple(inputs)
          with rt_views() as borrowed:
            for linear in (compiled, linked):
-              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
+              self.assertIs(compile_linear(linear, input_uops=inputs, cache=not jit), linear)
          self.assertEqual(tuple(inputs), before)
          self.assertFalse(borrowed)
          run_linear(linked, input_uops=inputs, jit=True, wait=True)
@@ -286,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.va_addr + 4)
+      self.assertEqual(outer_buf.host.view(fmt='Q')[0], inner_buf._buf + 4)

@unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):
@@ -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)[0][0], global_size=(1,1,1), local_size=(1,1,1), wait=True)')
+              'dev.allocator.alloc(64).buf, 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]')
@@ -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)[0][0]
+buf = dev.allocator.alloc(64).buf
 try:
  broken_prg(buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)
  print("  ERROR: Kernel did not fault!")
@@ -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 Allocator, Compiled
+from tinygrad.device import BufferStorage, 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 (None, None), None
+  def _alloc(self, sz, options): return BufferStorage(None)
  def _copyin(self, dest, src:memoryview): pass

 class TestLLaMASpeed(unittest.TestCase):
@@ -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))[0][0]) for i in range(100)]
+  buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000)).buf) for i in range(100)]

  for _ in trange(100000):
    d1, b1 = random.choice(buffers)
@@ -125,12 +125,12 @@ 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)[0][0])
+    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize).buf)
    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())
-  prog(out_buf:=allocator.alloc(uop.dtype.itemsize)[0][0], *bufs, vals=vals)
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize).buf, *bufs, vals=vals)
  return out_buf.cast(uop.dtype.fmt or "").tolist()[0]

 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:
@@ -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.va_addr, 0, vgpr_size * 4)
+      ctypes.memset(self.accvgpr_buf._buf, 0, vgpr_size * 4)
    else:
      self.accvgpr_buf = self.vgpr_buf
    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.va_addr, 0, vgpr_size * 4)
-    ctypes.memset(self.sgpr_buf._buf.va_addr, 0, SGPR_COUNT * 4)
+    ctypes.memset(self.vgpr_buf._buf, 0, vgpr_size * 4)
+    ctypes.memset(self.sgpr_buf._buf, 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,11 +1992,10 @@ 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.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)]))
+      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)]))
    done = [False] * len(waves)
    for _ in range(10_000_000):
      if all(done): return
@@ -2027,7 +2026,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.va_addr, 0, max(lds_size, 4))  # reset LDS for next workgroup
+      if lds_size > 0: ctypes.memset(lds_buf._buf, 0, max(lds_size, 4))  # reset LDS for next workgroup

  if PROFILE: sqtt_traces.append(sqtt_finalize())
  return 0
@@ -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

@@ -249,6 +249,32 @@ 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)),
@@ -208,6 +208,24 @@ 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
@@ -42,19 +42,17 @@ 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]]. Currently returns [0., 0.].
+    # y = [4, x[1]]; only the untouched element depends on x.
    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. Currently returns [2.].
+    # x contributes once, not twice.
    self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])

  def test_unrelated_store_gradient(self):
@@ -64,14 +62,26 @@ 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 creates a nested AFTER; currently raises in backward.
+    # View assign is an AFTER on a partial STORE; only the untouched element depends on x.
    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
@@ -10,29 +10,29 @@ 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.va_addr, b._buf.va_addr + 4)
+    self.assertEqual(v._buf, b._buf + 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.assertIs(b.get_storage("PYTHON")[0][1], b.get_buf("PYTHON"))
+    self.assertIs(b.get_storage("PYTHON").meta, b.get_buf("PYTHON"))
    v = b.view(4, dtypes.uint8, 2)
    mapped = v.get_storage("PYTHON")
-    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])
+    self.assertEqual(bytes(mapped.buf), b"cdef")
+    self.assertIs(mapped.host, v.host)
+    self.assertIsNone(mapped.meta)
+    self.assertIs(v.get_storage("PYTHON"), mapped)

  def test_view_reallocation(self):
    b = Buffer("CPU", 8, dtypes.uint8)
    v = b.view(4, dtypes.uint8, 2)
-    old = v.get_storage("PYTHON")[0]
+    old = v.get_storage("PYTHON")
    b.deallocate()
    b.allocate()
    self.assertFalse(v.is_allocated())
    v.host[:] = b"test"
-    self.assertIsNot(v.get_storage("PYTHON")[0], old)
+    self.assertIsNot(v.get_storage("PYTHON"), old)
    self.assertEqual(bytes(v.get_buf("PYTHON")), b"test")

  def test_cache_owned_storage_only(self):
@@ -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.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE)) == 1 or (TC_OPT.value >= 1)):
+  if USE_TC > 0 and (len(k.reduce_axes) == 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.output_shape[i] for i in k.upcastable_dims) <= (240 if k.ren.target.device == "QCOM" else 2048), False):
+  if resolve(prod(k.full_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.output_shape[i] for i in k.upcastable_dims) >= 1024) and (k.upcast_size() < 32):
+  while resolve(prod(k.full_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]):
@@ -99,14 +99,19 @@ 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.axes_of(AxisType.GROUP_REDUCE, AxisType.REDUCE) \
-                                                  if isinstance(s:=self.full_shape[i], int) and s > 1]
+  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 real_axis(self, op:OptOps, axis:int|None) -> int:
    if axis is None or op is OptOps.TC: return -1
@@ -291,9 +296,6 @@ 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))
@@ -1,5 +1,5 @@
 from __future__ import annotations
-from dataclasses import dataclass, replace
+from dataclasses import dataclass, replace, field
 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
@@ -97,6 +97,9 @@ 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,
@@ -104,14 +107,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:tuple|None = None
-    self._maps:dict[str, tuple] = {}
+    self._storage:BufferStorage|None = None
    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()
-        self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
+        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))))
        if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
    else:
      assert base._base is None, "base can't have a base"
@@ -130,24 +133,25 @@ 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()[0][0]
+  def _buf(self) -> Any: return self.get_storage().buf
  @property
-  def host(self) -> MMIOInterface: return unwrap(self.get_storage()[1])
+  def host(self) -> MMIOInterface: return unwrap(self.get_storage().host)
  @property
-  def meta(self) -> Any: return self.get_storage()[0][1]
+  def meta(self) -> Any: return self.get_storage().meta
  @property
  def nbytes(self): return self.size * self.dtype.itemsize

-  def get_storage(self, device:str|None=None) -> tuple:
+  def get_storage(self, device:str|None=None) -> BufferStorage:
    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 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]
+    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]

-  def get_buf(self, device:str) -> Any: return self.get_storage(device)[0][0]
+  def get_buf(self, device:str) -> Any: return self.get_storage(device).buf

  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
@@ -158,19 +162,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:
-      (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
+      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)
+      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)
    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._maps, self._base_storage = storage, {}, self.base._storage if self._base else None
+    self._storage, self._base_storage = storage, self.base._storage if self._base else None
    return self

  def deallocate(self):
@@ -181,10 +185,9 @@ 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._maps, self._base_storage = None, {}, None
+    self._storage, self._base_storage = None, None

  def __reduce_ex__(self, protocol):
    buf:bytearray|pickle.PickleBuffer|None = None
@@ -201,7 +204,7 @@ class Buffer:
    return self._trace_num

  def _host_mv(self) -> memoryview|None:
-    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'): return unwrap(host).view(fmt='B').mv
+    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(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
    return None

@@ -241,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[tuple]] = defaultdict(list)
+    self.cache:dict[tuple[int, BufferSpec|None], list[BufferStorage]] = 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) -> tuple:
+  def alloc(self, size:int, options:BufferSpec|None=None) -> BufferStorage:
    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
@@ -254,24 +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:tuple, size:int, options:BufferSpec|None=None):
+  def free(self, storage:BufferStorage, 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._free(storage[0][0], spec)
+    else: self.do_free(storage, spec)

  def free_cache(self):
    for (_, options), storages in self.cache.items():
-      for storage in storages: self._free(storage[0][0], options if options is not None else self.default_buffer_spec)
+      for storage in storages: self.do_free(storage, options if options is not None else self.default_buffer_spec)
      storages.clear()

-  def map(self, buf:Buffer) -> tuple: return self._map(buf.ensure_allocated()._buf)
+  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())

  # implemented by the runtime
-  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 _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 _copyin(self, dest, src:memoryview): raise NotImplementedError("need copyin")
  def _copyout(self, dest:memoryview, src): raise NotImplementedError("need copyout")
-  def _map(self, buf) -> tuple: raise NotImplementedError("need map")
+  def _map(self, buf) -> BufferStorage: 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):
@@ -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))
+  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value), input_uops=input_uops, cache=False)
  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
@@ -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
+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.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,6 +156,9 @@ 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 src.device.split(":")[0] in HCQ_DEVS and dest._host_mv() is not None:
+      dst_mv, src_mv = dest.as_memoryview(force_zero_copy=True), src.as_memoryview(force_zero_copy=True)
+      with cpu_profile(f"{src.device} -> TINY", f"{src.device}:COPY"): dst_mv[:] = src_mv[:]
    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 []
@@ -194,7 +197,7 @@ 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).va_addr + off for u, dev, off in info.inputs]
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev) + 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,))
@@ -277,13 +280,13 @@ pm_exec = PatternMatcher([
  (UPat(Ops.CALL, src=(UPat(Ops.CUSTOM_FUNCTION, arg="validate", name="ast"),), name="call", allow_any_len=True), exec_validate),
 ])

-from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV, HCQInfo # noqa: E402 # down here, hcq2 imports realize
+from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV, HCQInfo, HCQ_DEVS # 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) -> UOp:
+def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None, cache=False) -> 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)
+  linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile, cache=cache)
  return linear

 def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
@@ -291,13 +294,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), input_uops=inputs)
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs, cache=True), 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), allow_cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True, cache=False), allow_cache=ctx.cache)
  while True:
    if clear_l2:
      if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()
@@ -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 sum_acc_dtype
+from tinygrad.dtype import dtypes, sum_acc_dtype
 from tinygrad.function import renumber_invalid_outputs

 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
@@ -67,6 +67,19 @@ 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),)),
@@ -104,6 +117,7 @@ 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,)),
@@ -26,8 +26,8 @@ class MetalGraph(GraphRunner):

    self.var_bind_data = []
    if len(self.vars):
-      (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
+      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

    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)):
@@ -3,9 +3,9 @@ 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 HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
 from tinygrad.uop.ops import sint, UOp, ProgramInfo
-from tinygrad.device import BufferSpec, Buffer, Device, Compiled, ProfileProgramEvent
+from tinygrad.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, 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, mv_address, ContextVar, VIZ
@@ -14,7 +14,7 @@ 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, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
+from tinygrad.runtime.support.hcq import FileIOInterface, 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
@@ -555,19 +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(HCQAllocator['AMDDevice']):
+class AMDAllocator(Allocator['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 _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 _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)

-  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)
+  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

@dataclass
 class AMDQueueDesc:
@@ -611,7 +611,7 @@ 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) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> BufferStorage:
    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
@@ -637,53 +637,52 @@ 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

-    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
+    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)

-  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 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 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 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):
+  def map(self, buf:Buffer) -> BufferStorage:
+    if buf.device.split(":")[0] == "CPU":
+      return replace(mem:=self.alloc(buf.nbytes, host=True, cpu_addr=buf._buf), meta=(mem.meta.handle, True))
+    self._map_handle(buf.meta.handle)
+    return BufferStorage(buf._buf, (buf.meta.handle, False))
+
+  def _map_handle(self, handle):
    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
    assert stm.n_success == 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)
+  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)
    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._buf.meta.handle)
+        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.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
+    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf, ring_size=ring.nbytes, 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.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)
+      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)

    if not hasattr(self, 'doorbells'):
      self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
@@ -742,7 +741,6 @@ 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
@@ -767,11 +765,10 @@ 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.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
-        gart._buf.va_addr+wptr, idx)))
+      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr, gart._buf+wptr, idx)))
    else:
-      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)))
+      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)))

    (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)
@@ -803,9 +800,10 @@ 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) -> tuple:
+  def map(self, buf:Buffer) -> BufferStorage:
    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
+    return BufferStorage(mv_address(mv), mv)
+  def _unmap(self, mapping:BufferStorage): pass

 class USBIface(PCIIface):
  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
@@ -823,10 +821,10 @@ class USBIface(PCIIface):
    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, nolru=True),
-                  opaque=HCQBuffer(vaddr, 0x85000, view=view, owner=self.dev))
+    return Buffer(self.dev.device, 0x85000, dtypes.uint8, options=BufferSpec(external_ptr=vaddr), opaque=BufferStorage(vaddr, host=view))

-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False,
+            **kwargs) -> BufferStorage:
    # 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)

@@ -995,7 +993,7 @@ class AMDDevice(HCQ2Compiled):
    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.va_addr
+    base = self.scratch._buf
    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'),
@@ -1024,13 +1022,13 @@ class AMDDevice(HCQ2Compiled):
      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.va_addr, b.arg.slot, key))
+        Compiled.profile_events.append(ProfileProgramEvent(self.device, name, lib, buf._buf, 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.va_addr + off) // 32) & 0x1FFFFFFF) * 32
+    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf + 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"))
@@ -5,7 +5,7 @@ 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, mv_address, suppress_finalizing, unwrap, round_up, is_image_shape
 from tinygrad.renderer.cstyle import OpenCLRenderer
-from tinygrad.device import BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
+from tinygrad.device import BufferStorage, 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,11 +76,11 @@ class CLProgram(Program['CLDevice']):
    return None

 class CLAllocator(Allocator['CLDevice']):
-  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
+  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))

  @suppress_finalizing
-  def _free(self, opaque:cl.cl_mem, options:BufferSpec): check(cl.clReleaseMemObject(opaque))
+  def _free(self, storage:BufferStorage, options:BufferSpec): check(cl.clReleaseMemObject(storage.buf))
  def _copyin(self, dest:cl.cl_mem, src:memoryview):
    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))
@@ -2,9 +2,9 @@ from __future__ import annotations
 import platform, sys, ctypes, mmap, struct, time
 from typing import cast
 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.device import BufferStorage, BufferSpec, TinyELF, Program, Device, Buffer, Allocator
+from tinygrad.runtime.support.memory import MMIOInterface
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
@@ -56,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:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+  def __call__(self, *bufs:int, 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, *[b.va_addr for b in bufs])
+      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *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 = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
+      args = [*bufs, *cast(tuple[int, ...], vals)]
      self.fxn(*[ctypes.c_uint64(x) for x in args])
    return time.perf_counter() - st if wait else None

@@ -74,25 +74,24 @@ 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 CPUAllocator(HCQAllocator['CPUDevice']):
+class CPUAllocator(Allocator['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:
+  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 (opaque:=HCQBuffer(addr, size, meta=buf, view=MMIOInterface(addr, size, fmt='B'), owner=self.dev), opaque.meta), opaque.view
+    return BufferStorage(addr, buf, MMIOInterface(addr, size, fmt='B'))

-  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
-  def _copyin(self, dest:HCQBuffer, src:memoryview):
+  def _copyin(self, dest:int, src:memoryview):
    self.dev.synchronize()
-    ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
-  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    ctypes.memmove(dest, from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:int):
    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
+    dest[:] = to_mv(src, dest.nbytes)[:]
+  def _map(self, buf:Buffer) -> BufferStorage:
+    if not isinstance(host:=buf.get_storage().host, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
+    return BufferStorage(host.addr)
+  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset

 class CPUDevice(HCQ2Compiled):
  wait_timeout_ms, has_copy_queue = 30000, False
@@ -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 MMIOInterface, Compiled, BufferSpec, Allocator, Program, TinyELF
+from tinygrad.device import BufferStorage, 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,24 +65,23 @@ 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) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
    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 (opaque, None), MMIOInterface(opaque.value, size) if options.host else None
+    return BufferStorage(opaque, None, MMIOInterface(opaque.value, size) if options.host else None)

  @suppress_finalizing
-  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 _free(self, storage:BufferStorage, options:BufferSpec):
+    if options.host: check(cuda.cuMemFreeHost(storage.buf))
+    else: check(cuda.cuMemFree_v2(storage.buf))
  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[0][0], mv_address(src), len(src))
-    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem[0][0], len(src), None))
+    ctypes.memmove(host_mem.buf, mv_address(src), len(src))
+    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem.buf, len(src), None))
  def _copyout(self, dest:memoryview, src):
    CUDADevice.synchronize_system()
    check(cuda.cuCtxSetCurrent(self.dev.context))
@@ -116,7 +115,7 @@ class CUDADevice(Compiled):
      check(cuda.cuCtxEnablePeerAccess(dev.context, 0))
      CUDADevice.peer_access = True

-    self.pending_copyin: list[tuple[tuple, int, BufferSpec|None]] = []
+    self.pending_copyin: list[tuple[BufferStorage, int, BufferSpec|None]] = []
    CUDADevice.devices.append(self)

    from tinygrad.runtime.graph.cuda import CUDAGraph
@@ -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 MMIOInterface, Compiled, Allocator
+from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator
 with contextlib.suppress(ImportError):
  import _posixshmem
  from tinygrad.runtime.autogen import io_uring, libc
@@ -79,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) -> tuple:
+  def _alloc(self, size:int, options) -> BufferStorage:
    self.dev._might_open(size)
-    return (opaque:=DiskBuffer(self.dev, size), None), MMIOInterface(mv_address(opaque._buf()), size)
+    return BufferStorage(opaque:=DiskBuffer(self.dev, size), None, MMIOInterface(mv_address(opaque._buf()), size))

-  def _free(self, opaque, options): self.dev._might_close()
+  def _free(self, storage:BufferStorage, 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):
@@ -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 MMIOInterface, BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
+from tinygrad.device import BufferStorage, 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) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
    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 (opaque, opaque.share_info), MMIOInterface(opaque.va_addr, size)
+    return BufferStorage(opaque, opaque.share_info, MMIOInterface(opaque.va_addr, size))

  @suppress_finalizing
-  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 _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 _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))[0][0]
+      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True)).buf
      ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)

      self.init_dsp()
@@ -1,6 +1,6 @@
 import ctypes
 from tinygrad.helpers import mv_address, getenv, suppress_finalizing
-from tinygrad.device import Compiled, Allocator, BufferSpec, Program, TinyELF
+from tinygrad.device import BufferStorage, 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
@@ -57,11 +57,11 @@ class HIPProgram(Program[HIPDevice]):
      return ret.value * 1e-3

 class HIPAllocator(Allocator[HIPDevice]):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
    check(hip.hipSetDevice(self.dev.device_id))
-    return (init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size))), None), None
+    return BufferStorage(init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size))))

-  def _free(self, opaque, options:BufferSpec): check(hip.hipFree(opaque))
+  def _free(self, storage:BufferStorage, options:BufferSpec): check(hip.hipFree(storage.buf))
  def _copyin(self, dest, src: memoryview):
    check(hip.hipSetDevice(self.dev.device_id))
    check(hip.hipMemcpy(dest, mv_address(src), len(src), hip.hipMemcpyHostToDevice))
@@ -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 MMIOInterface, Compiled, Compiler, CompileError, Program, TinyELF, Allocator, ProfileDeviceEvent
+from tinygrad.device import BufferStorage, 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
@@ -155,16 +155,15 @@ 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) -> tuple:
+  def _alloc(self, size:int, options) -> BufferStorage:
    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 (MetalBuffer(ret, size), None), MMIOInterface(addr, size) if (addr:=ret.contents()) is not None else None
+    return BufferStorage(MetalBuffer(ret, size), None, MMIOInterface(addr, size) if (addr:=ret.contents()) is not None else None)

  @suppress_finalizing
-  def _free(self, opaque:MetalBuffer, options):
-    if not options.external_ptr: opaque.buf.release()
+  def _free(self, storage:BufferStorage, options): storage.buf.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()
@@ -1,9 +1,10 @@
 import numpy as np
 from tinygrad.helpers import flat_mv
-from tinygrad.device import MMIOInterface, Compiled, Allocator
+from tinygrad.device import BufferStorage, 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 _alloc(self, size:int, options) -> BufferStorage:
+    return BufferStorage(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)
@@ -1,5 +1,5 @@
 import inspect, math
-from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
+from tinygrad.device import BufferStorage, 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) -> tuple: return (None, None), None
+  def _alloc(self, size:int, options) -> BufferStorage: return BufferStorage(None)

  def _copyin(self, dest, src:memoryview): pass
  def _copyout(self, dest:memoryview, src):
@@ -2,12 +2,12 @@ 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
-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 dataclasses import dataclass, replace
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HWQueue, encode_submit, patch, to_name, unwrap_view
+from tinygrad.runtime.support.hcq import 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.device import BufferStorage, Buffer, BufferSpec, Allocator, 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
@@ -290,16 +290,19 @@ 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(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
+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)

-  def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
+  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_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"
+  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"

    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)
@@ -308,15 +311,15 @@ class NVAllocator(HCQAllocator['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.va_addr >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin.va_addr >> 8)
+    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)
    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.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_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_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)
@@ -326,7 +329,7 @@ class NVAllocator(HCQAllocator['NVDevice']):
 # device

@dataclass
-class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; entries: int; token: int # noqa: E702
+class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; notifier: Buffer; entries: int; token: int # noqa: E702

 class NVKIface:
  root = None
@@ -436,7 +439,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) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> BufferStorage:
    # 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))
@@ -476,17 +479,20 @@ class NVKIface:

    return self._gpu_uvm_map(va_addr, size, mem_handle, has_cpu_mapping=cpu_access or host)

-  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)
+  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)
      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)

-    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 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))

-  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False, owner=None) -> HCQBuffer:
+  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False) -> BufferStorage:
    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,
@@ -500,14 +506,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 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)
+    return BufferStorage(va_base, uvm_map, MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None)

-  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 map(self, buf:Buffer) -> BufferStorage:
+    mem = buf.meta
+    if buf.device.split(":")[0] == "CPU":
+      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))
+    return replace(mapping:=self._gpu_uvm_map(buf._buf, mem.length, mem.hMemory, create_range=False), meta=(mapping.meta, False))

  def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):
    return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment)
@@ -576,9 +582,6 @@ class NVDevice(HCQ2Compiled):
    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))

@@ -600,8 +603,9 @@ class NVDevice(HCQ2Compiled):

  @functools.cached_property
  def fifos(self) -> dict[str, GPFifo]:
-    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)
+    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)
+
    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))
@@ -615,11 +619,11 @@ class NVDevice(HCQ2Compiled):
    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 = 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,
+    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,
      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)

@@ -642,7 +646,7 @@ class NVDevice(HCQ2Compiled):
    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), entries=entries, token=ws_token_params.workSubmitToken)
+      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), notifier=notifier, 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
@@ -664,15 +668,15 @@ class NVDevice(HCQ2Compiled):
    (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
+    fifo.ring.host.view(fmt='Q')[put[0] % fifo.entries] = buf._buf | (len(cmds) << 42) | (1 << 41)
+    fifo.gpput.host.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
+    tl, addr = self.timeline.host.view(fmt='Q'), self.timeline._buf
    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),
@@ -687,7 +691,7 @@ class NVDevice(HCQ2Compiled):
    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(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf.va_addr)),
+    self._submit_cmds(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf)),
                       *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):
@@ -750,12 +754,14 @@ class NVDevice(HCQ2Compiled):
      (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 = 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_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_rptr = 0

    pma_stream = nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS(hMemPmaBuffer=self.pma_buf.meta.hMemory,
-      pmaBufferSize=self.pma_buf.size, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf.va_addr)
+      pmaBufferSize=self.pma_buf.nbytes, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf._buf)
    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))
@@ -822,9 +828,9 @@ class NVDevice(HCQ2Compiled):
    if params.bOverflowStatus: raise RuntimeError("PMA profiler: buffer overflow detected")
    if params.bytesAvailable == 0: return None

-    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
+    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

    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))
@@ -7,7 +7,7 @@ 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, cpu_profile, mv_address
-from tinygrad.device import MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
+from tinygrad.device import BufferStorage, 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
@@ -237,14 +237,15 @@ 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 _alloc(self, size:int, options) -> BufferStorage:
+    return BufferStorage(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 map(self, buf:Buffer) -> BufferStorage: return BufferStorage(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):
@@ -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 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.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, TinyELF
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HWQueue, HCQ_RUNTIME_DEV, encode_submit, ccall, cstruct, patch, unwrap_view
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
 from tinygrad.runtime.autogen import kgsl, mesa, libc
 from tinygrad.renderer.cstyle import QCOMCLRenderer
 from tinygrad.renderer.nir import IR3Renderer
@@ -301,11 +301,14 @@ 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(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
+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)

-  def _do_free(self, opaque, options:BufferSpec): self.dev._gpu_free(opaque)
+  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 flag(nm, val): return (val << getattr(kgsl, f"{nm}_SHIFT")) & getattr(kgsl, f"{nm}_MASK")

@@ -354,11 +357,9 @@ class QCOMDevice(HCQ2Compiled):

  @functools.cached_property
  def border_color(self) -> Buffer: # zeros: the samplers clamp to a black border
-    (b:=Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)) \
-      .as_memoryview(force_zero_copy=True)[:] = bytes(0x1000)
-    return b
+    return Buffer(self.device, 0x1000, dtypes.uint8, options=BufferSpec(nolru=True), initial_value=bytes(0x1000))

-  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> HCQBuffer:
+  def _gpu_alloc(self, size:int, flags:int=0, uncached=False, fill_zeroes=False) -> BufferStorage:
    flags |= flag("KGSL_MEMALIGN", alignment_hint:=12) | kgsl.KGSL_MEMFLAGS_USE_CPU_MAP
    if uncached: flags |= flag("KGSL_CACHEMODE", kgsl.KGSL_CACHEMODE_UNCACHED)

@@ -366,24 +367,24 @@ class QCOMDevice(HCQ2Compiled):
    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 HCQBuffer(va_addr=va_addr, size=size, meta=(alloc, True), view=MMIOInterface(va_addr, size, fmt='B'), owner=self)
+    return BufferStorage(va_addr, (alloc, True), MMIOInterface(va_addr, size, fmt='B'))

-  def _gpu_map(self, ptr:int, size:int) -> HCQBuffer:
+  def _gpu_map(self, ptr:int, size:int) -> BufferStorage:
    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 HCQBuffer(mi.gpuaddr + (ptr - ptr_aligned), size=size, meta=(mi, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
+      return BufferStorage(mi.gpuaddr + (ptr - ptr_aligned), (mi, False), MMIOInterface(ptr, size, fmt='B'))
    except OSError as e:
-      if e.errno == 14: return HCQBuffer(va_addr=ptr, size=size, meta=(None, False), view=MMIOInterface(ptr, size, fmt='B'), owner=self)
+      if e.errno == 14: return BufferStorage(ptr, (None, False), MMIOInterface(ptr, size, fmt='B'))
      raise RuntimeError("Failed to map external pointer to GPU memory") from e

-  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
+  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
    else:
-      kgsl.IOCTL_KGSL_GPUOBJ_FREE(self.fd, id=mem.meta[0].id)
-      FileIOInterface.munmap(mem.va_addr, mem.meta[0].mmapsize)
+      kgsl.IOCTL_KGSL_GPUOBJ_FREE(self.fd, id=storage.meta[0].id)
+      FileIOInterface.munmap(storage.buf, storage.meta[0].mmapsize)

  def _wait_signal(self, sig:MMIOInterface|memoryview, value:int, timeout:int|None=None):
    if sig[0] < value:
@@ -1,5 +1,5 @@
 import functools, struct
-from tinygrad.device import Compiled, Allocator, BufferSpec, Program, TinyELF
+from tinygrad.device import BufferStorage, 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) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
    # WebGPU buffers have to be 4-byte aligned
-    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
+    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)))
  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, opaque:webgpu.WGPUBuffer, options:BufferSpec): self.dev.free(opaque)
+  def _free(self, storage:BufferStorage, options:BufferSpec): self.dev.free(storage.buf)

 class WebGpuDevice(Compiled):
  def __init__(self, device:str):
@@ -1,11 +1,8 @@
 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

@@ -55,26 +52,3 @@ 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
@@ -1,24 +1,21 @@
 from __future__ import annotations
-from typing import cast, TypeVar, Generic, Any, TYPE_CHECKING
+from typing import cast, Any
 import functools, time, itertools, decimal, weakref, statistics, ctypes, importlib
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2, cpu_profile, mv_address
+from tinygrad.helpers import dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2
 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.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
 from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
 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())
@@ -422,11 +419,11 @@ pm_encode = PatternMatcher([(UPat(Ops.CALL, src=(UPat(Ops.SINK),), name="call",

 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,ret: f"HCQ Compile {pluralize('Kernel', len(ret.src))}")
-def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
+@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:
  if any(isinstance(getattr(c.without_after.arg, "aux", None), HCQInfo) for c in linear.src): return linear # compiled already

-  if input_uops is not None:
+  if cache and 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")
@@ -435,7 +432,7 @@ def hcq_compile(linear:UOp, input_uops:list[UOp]|None, profile:bool) -> UOp:
  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 input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
+  if cache and input_uops is not None and final_linear is not linear: hcq_compile_cache[key] = final_linear
  return final_linear

 # *****************
@@ -462,7 +459,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]).va_addr + off, dtypes.uint64)
+  return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]) + 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
@@ -520,7 +517,7 @@ class HCQ2Compiled(Compiled):
  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):
+  def __init__(self, device:str, allocator:Allocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
    self.can_recover = can_recover

    self.pm_bufferize = PatternMatcher([
@@ -594,42 +591,7 @@ class HCQ2Compiled(Compiled):

  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)
@@ -1,8 +1,9 @@
 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, HCQBuffer, hcq_filter_visible_devices
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, hcq_filter_visible_devices
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace, BumpAllocator
 from tinygrad.runtime.support.usb import USB3, CustomASM24Controller, USBMMIOInterface

@@ -261,7 +262,8 @@ 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) -> HCQBuffer:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False,
+            **kwargs) -> BufferStorage:
    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.
@@ -271,32 +273,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 HCQBuffer(vaddr, size, meta=PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), view=memview, owner=self.dev)
+      return BufferStorage(vaddr, PCIAllocationMeta(mapping, has_cpu_mapping=True, hMemory=paddrs[0]), memview)

    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 HCQBuffer(mapping.va_addr, size, view=barview, meta=PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), owner=self.dev)
+    return BufferStorage(mapping.va_addr, PCIAllocationMeta(mapping, cpu_access, hMemory=mapping.paddrs[0][0]), barview)

-  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 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 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: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
+  def map(self, b:Buffer) -> BufferStorage:
+    if b.device.split(":")[0] == "CPU":
+      lo, size = b._buf & ~0xfff, round_up(b._buf + b.nbytes, 0x1000) - (b._buf & ~0xfff)
+      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
      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.owner} -> {self.dev}")
+    else: raise RuntimeError(f"map failed: {b.device} -> {self.dev.device}")

-    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)
+    self.dev_impl.mm.map_range(lo, size, paddrs, aspace=aspace, snooped=snooped, uncached=uncached)
+    return BufferStorage(b._buf, (lo, size))
@@ -43,8 +43,10 @@ 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 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"),
+  # 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"),
   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),
@@ -448,15 +450,18 @@ 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(self.uop.store(x.uop))
+    assign = self.uop.after(store := 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)
-      # 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")
+      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")
    else:
      # simple assign
      self.uop = assign
@@ -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: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6",
+               Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#ff90c0", 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",
Author	SHA1	Message	Date
chenyuandGitHub	b6ac3838b5	delete output_shape [PR] (#18090 )	2026-09-09 15:25:29 -04:00
chenyuandGitHub	803890b6f5	start to clean up reduce related axes [pr] (#18089 ) * start to clean up reduce related axes [pr] * revert	2026-09-09 13:42:32 -04:00
nimlgenandGitHub	d22b12b4d5	hcq2: explicit cache flag (#18088 ) * hcq: make compilation caching explicit * hcq: require input buffers for compilation caching * hcq: match rewrite label arguments to compile signature	2026-09-09 20:10:44 +03:00
geohot	85d64af69f	hotfix: add copy to spec, make the VIZ color of copy closer to load	2026-09-09 08:21:48 -07:00
George HotzandGitHub	f07260ea15	partial view assign grad (#18067 ) * gradient: partial grad for a write through a non-overlapping view (AFTER of AFTER) * assign: view writes are a single AFTER on a partial STORE, not a nested AFTER * one line * callify: merged COPY tags map through the STORE's target view, not the AFTER dest * assign: stored creation values own their storage at store construction * callify: merge COPY tags only when the STORE fills the AFTER's whole dest * precreate all copy targets at construction, never tag creation copies * cut the copy-storage churn: partial store merges only when the STORE fills the AFTER's whole dest	2026-09-09 08:11:40 -07:00
nimlgenandGitHub	fec703cbec	remove hcqbuf (#18073 ) * runtime: remove HCQBuffer from active runtimes * minify * x * f	2026-09-09 12:10:21 +03:00
nimlgenandGitHub	2b84a5e5b1	buffer: BufferStorage dataclass (#18083 ) * runtime: keep mappings in storage and pass storage to free * buffer: return complete storage from get_storage * use dc * frozen * frozen	2026-09-09 10:50:22 +03:00
nimlgenandGitHub	aed30173c9	remove hcqbufs from runtimes (#18082 ) * x * x * smaller pma	2026-09-09 10:17:41 +03:00
qazalandGitHub	d3a44990b3	remove HCQ2=0 in llama 8b CI (#18080 )	2026-09-09 13:16:56 +09:00
sirhcmandGitHub	fb9a356014	benchmarks: reset mac chestnut (#18079 )	2026-09-08 23:16:13 -04:00
wozeparrotandGitHub	461b15ab89	gptoss: padded embedding (#18078 )	2026-09-08 19:03:33 -07:00
chenyuandGitHub	7595039187	failing test for PADTO+TC (#18077 )	2026-09-08 21:08:52 -04:00
sirhcmandGitHub	822d8c4682	ci: more setup-tinygrad (#18076 )	2026-09-08 18:33:15 -04:00