simpler

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

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

.github/workflows/autogen.yml

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

.github/workflows/benchmark.yml

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

.github/workflows/docs.yml

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

.github/workflows/test.yml

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

examples/compile_tensorflow.py

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

examples/mlperf/models/gpt_oss.py

@@ -15,7 +15,6 @@ 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
@@ -152,10 +151,7 @@ class GPTOSS:
 
     # output
     self.norm = nn.RMSNorm(dim, norm_eps)
-    if getenv("GPTOSS_EMBEDDING", 0):
-      self.tok_embeddings = GPTOSSEmbedding(vocab_size, dim)
-    else:
-      self.tok_embeddings = nn.Embedding(vocab_size, dim)
+    self.tok_embeddings = nn.Embedding(vocab_size, dim)
     self.tok_embeddings.weight = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.output = Tensor.normal(vocab_size, dim, mean=0.0, std=INIT_STD, dtype=dtypes.bfloat16)
     self.freqs_cis = precompute_freqs_cis(head_dim, max_context * 2, rope_theta).contiguous().is_param_(False)

extra/dsp/invoke_bug.py

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

extra/dsp/invoke_bug_2.py

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

extra/export_model.py

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

extra/gemm/hip_matmul.py

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

extra/gptoss_kernels/embedding/__init__.py

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

extra/gptoss_kernels/embedding/embedding_bwd.cpp

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

extra/gptoss_kernels/embedding/embedding_fwd.cpp

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

extra/hcq1/graph.py

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

extra/hcq1/hcq.py

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

extra/hcq1/ops_amd_old.py

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

extra/hcq1/ops_rdma.py

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

extra/hcq1/test_hcq.py

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

extra/mlx_driver/loopback.py

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

extra/usbgpu/debug.py

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

test/amd/hw/helpers.py

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

test/amd/hw/test_cdna_vop3.py

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

test/amd/hw/test_dpp.py

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

test/amd/hw/test_rdna4_permlane_var.py

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

test/backend/test_assign.py

@@ -998,14 +998,10 @@ class TestAssignOrdering(unittest.TestCase):
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
-    c = t.to("CPU:1")  # unrealized COPY
-    self.assertIs(c.uop.base.op, Ops.COPY)
+    c = t.to("CPU:1")  # the pending transfer already owns its destination
+    self.assertIs(c.uop.base.op, Ops.AFTER)
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
-    try:
-      self.assertEqual(c.tolist(), [[0,1],[0,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(c.tolist(), [[0,0],[0,0]])
+    self.assertEqual(c.tolist(), [[0,1],[0,1]])
 
   def test_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
@@ -1045,14 +1041,10 @@ class TestAssignToUnrealizedView(unittest.TestCase):
 
   def test_detach_copy(self):
     t = Tensor.zeros(2,2, dtype=dtypes.int).to("CPU:0").contiguous().realize()
-    d = t.to("CPU:1").detach()  # DETACH(unrealized COPY)
-    self.assertIs(d.uop.base.op, Ops.COPY)
+    d = t.to("CPU:1").detach()
+    self.assertIs(d.uop.base.op, Ops.AFTER)
     d[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).to("CPU:1").contiguous().realize())
-    try:
-      self.assertEqual(d.tolist(), [[0,1],[0,1]])
-    except AssertionError:
-      # TODO: broken now
-      self.assertEqual(d.tolist(), [[0,0],[0,0]])
+    self.assertEqual(d.tolist(), [[0,1],[0,1]])
 
   def test_detach_contiguous(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()

test/backend/test_call.py

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

test/device/test_hcq2.py

@@ -86,14 +86,14 @@ class TestHCQ2Schedule(unittest.TestCase):
   def test_jit_has_no_rt_buffers(self):
     dev = Device[Device.DEFAULT]
     rings = [dev.rt_buffer(True, host) for host in (False, True)]
-    ranges = [(b._buf, b._buf + b.nbytes) for b in rings]
+    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
     for n in (1, 65):
       with self.subTest(kernels=n):
         x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
         for _ in range(2): f(x)
         for u in f.captured.linear.toposort():
           if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
-            addr = buf._buf
+            addr = buf._buf.va_addr
             self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
 
   def test_small_eager_cached(self):
@@ -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 + 4)
+      self.assertEqual(outer_buf.host.view(fmt='Q')[0], inner_buf._buf.va_addr + 4)
 
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):

test/device/test_metal.py

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

test/external/external_test_am_fault_recovery.py

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

test/external/external_test_hive_reset.py

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

test/external/external_test_speed_llama.py

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

test/external/fuzz_kfd.py

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

test/helpers.py

@@ -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).buf)
+    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)[0][0])
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
-  prog(out_buf:=allocator.alloc(uop.dtype.itemsize).buf, *bufs, vals=vals)
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize)[0][0], *bufs, vals=vals)
   return out_buf.cast(uop.dtype.fmt or "").tolist()[0]
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:

test/mockgpu/amd/emu.py

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

test/mockgpu/mockgpu.py

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

test/mockgpu/nv/nvdriver.py

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

test/opt/test_kernel_opts.py

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

test/opt/test_tensor_cores.py

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

test/unit/test_after.py

@@ -64,14 +64,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 creates a nested AFTER; 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

test/unit/test_buffer.py

@@ -1,4 +1,4 @@
-import unittest
+import gc, unittest, weakref
 from tinygrad.device import Buffer
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import Context
@@ -10,29 +10,47 @@ class TestBuffer(unittest.TestCase):
     host = v.host
     host.view(fmt='H')[0] = 0x1234
     self.assertEqual(b.host.view(fmt='H')[2], 0x1234)
-    self.assertEqual(v._buf, b._buf + 4)
+    self.assertEqual(v._buf.va_addr, b._buf.va_addr + 4)
     self.assertIs(v.host, host)
     self.assertIs(v.meta, b.meta)
 
+  def test_memoryview_keeps_allocation_alive(self):
+    for device in ("CPU", "PYTHON", "NPY"):
+      with self.subTest(device=device), Context(LRU=0):
+        b = Buffer(device, 8, dtypes.uint8).ensure_allocated()
+        b.host[:] = b"abcdefgh"
+        v = b.view(4, dtypes.uint8, 2).ensure_allocated()
+        mv = v.as_memoryview(force_zero_copy=True)[1:]
+        b_ref, v_ref = weakref.ref(b), weakref.ref(v)
+        del b, v
+        gc.collect()
+        self.assertIsNotNone(b_ref())
+        self.assertIsNotNone(v_ref())
+        self.assertEqual(bytes(mv), b"def")
+        del mv
+        gc.collect()
+        self.assertIsNone(v_ref())
+        self.assertIsNone(b_ref())
+
   def test_mapping(self):
     b = Buffer("CPU", 8, dtypes.uint8, initial_value=b"abcdefgh")
-    self.assertIs(b.get_storage("PYTHON").meta, b.get_buf("PYTHON"))
+    self.assertIs(b.get_storage("PYTHON")[0][1], b.get_buf("PYTHON"))
     v = b.view(4, dtypes.uint8, 2)
     mapped = v.get_storage("PYTHON")
-    self.assertEqual(bytes(mapped.buf), b"cdef")
-    self.assertIs(mapped.host, v.host)
-    self.assertIsNone(mapped.meta)
-    self.assertIs(v.get_storage("PYTHON"), mapped)
+    self.assertEqual(bytes(mapped[0][0]), b"cdef")
+    self.assertIs(mapped[1], v.host)
+    self.assertIsNone(mapped[0][1])
+    self.assertIs(v.get_storage("PYTHON")[0], mapped[0])
 
   def test_view_reallocation(self):
     b = Buffer("CPU", 8, dtypes.uint8)
     v = b.view(4, dtypes.uint8, 2)
-    old = v.get_storage("PYTHON")
+    old = v.get_storage("PYTHON")[0]
     b.deallocate()
     b.allocate()
     self.assertFalse(v.is_allocated())
     v.host[:] = b"test"
-    self.assertIsNot(v.get_storage("PYTHON"), old)
+    self.assertIsNot(v.get_storage("PYTHON")[0], old)
     self.assertEqual(bytes(v.get_buf("PYTHON")), b"test")
 
   def test_cache_owned_storage_only(self):

test/unit/test_call.py

@@ -5,7 +5,7 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.tensor import transform_to_call
 
-def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop))[0].src[0].key
+def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop.materialize()))[0].src[0].key
 
 class TestCall(unittest.TestCase):
   def test_call_plus(self):

test/unit/test_callify.py

@@ -1,7 +1,186 @@
 import unittest
-from tinygrad import Tensor, dtypes
+from unittest.mock import patch
+from tinygrad import Tensor, dtypes, function
+from tinygrad.tensor import transform_to_call
+from tinygrad.uop.ops import UOp, Ops, ParamArg
+from tinygrad.uop.render import pyrender
+from tinygrad.uop.spec import eval_pyrender
 
 class TestCallify(unittest.TestCase):
+  def test_no_buffer_creation_in_callify(self):
+    x = Tensor([1., 2.]).realize()
+    for precompile in (False, True):
+      @function(precompile=precompile)
+      def f(x): return x + 1
+      called = f(x)
+      roots = ((x + 2).uop.materialize(), called.uop, x.clone().uop)
+      with patch.object(UOp, "new_buffer", side_effect=AssertionError("callify created storage")), \
+           patch.object(UOp, "empty_like", side_effect=AssertionError("callify replaced storage")), \
+           patch.object(UOp, "bind_buffer", side_effect=AssertionError("callify bound storage")):
+        call, mapped = transform_to_call(UOp.sink(*roots))
+      self.assertIs(mapped[called.uop].storage_base, called.uop.storage_base)
+      self.assertIn(called.uop.storage_base, call.src[1:])
+
+  def test_unbound_store_binds_original_declaration(self):
+    buf = UOp(Ops.BUFFER, arg=ParamArg(next(UOp.unique_num), dtypes.float32, size=2, device="CPU"))
+    alias = Tensor(buf)
+    t = Tensor(buf.after(buf.store(buf.const_like(7.))))
+    t.callify().callify().realize()
+    self.assertEqual(t.uop.storage_base.arg.slot, buf.arg.slot)
+    self.assertFalse(t.uop.storage_base.is_unbound)
+    self.assertIs(alias.uop.buffer, t.uop.buffer)
+    self.assertEqual(t.tolist(), [7., 7.])
+    self.assertEqual(t.tolist(), [7., 7.])
+
+  def test_symbolic_view_keeps_bindings(self):
+    start, size = UOp.variable("start", 0, 4).bind(2), UOp.variable("size", 1, 4).bind(3)
+    t = Tensor.arange(8).float().realize()[start:start+size].clone()
+    shape = t.shape
+    t.callify().realize()
+    self.assertEqual(t.shape, shape)
+    self.assertEqual(t[:3].tolist(), [2., 3., 4.])
+    self.assertEqual(t[:3].tolist(), [2., 3., 4.])
+
+  def test_effect_only_call_body(self):
+    # An opaque tensor-level body needs no returned AFTERs to make its root stores execute.
+    for shape in ((6,), (2, 3)):
+      with self.subTest(shape=shape):
+        x = Tensor.arange(6).float().reshape(shape).realize()
+        a, b = Tensor.zeros(shape).contiguous().realize(), Tensor.zeros(shape).contiguous().realize()
+        a_buf, b_buf = a.uop.buffer, b.uop.buffer
+        a, b = Tensor.custom_kernel(a, b, x, fxn=lambda a,b,x: UOp.sink(a.store(x+1), b.store(x*2)))[:2]
+        a.realize(b)
+        self.assertIs(a.uop.buffer, a_buf)
+        self.assertIs(b.uop.buffer, b_buf)
+        self.assertEqual(a.flatten().tolist(), [1., 2., 3., 4., 5., 6.])
+        self.assertEqual(b.flatten().tolist(), [0., 2., 4., 6., 8., 10.])
+
+  def test_effect_only_slice_store(self):
+    x = Tensor.zeros(4, 4).contiguous().realize()
+    y = Tensor.ones(2, 2).contiguous().realize()
+    out = Tensor.custom_kernel(x, y, fxn=lambda x,y: x.shrink(((1, 3), (1, 3))).store(y).sink())[0]
+    self.assertEqual(out.tolist(), [[0., 0., 0., 0.], [0., 1., 1., 0.], [0., 1., 1., 0.], [0., 0., 0., 0.]])
+
+  def test_empty_declaration_binds(self):
+    buf = UOp(Ops.BUFFER, arg=ParamArg(next(UOp.unique_num), dtypes.float32, size=2, device="CPU"))
+    t = Tensor(buf).realize()
+    self.assertEqual(t.uop.arg.slot, buf.arg.slot)
+    self.assertFalse(t.uop.is_unbound)
+
+  def test_declaration_pyrender(self):
+    for size in (None, 2):
+      buf = UOp(Ops.BUFFER, arg=ParamArg(next(UOp.unique_num), dtypes.float32, size=size, device="CPU"))
+      self.assertIs(eval_pyrender(pyrender(buf)), buf)
+
+  def test_scalar_declaration_binds(self):
+    buf = UOp(Ops.BUFFER, arg=ParamArg(next(UOp.unique_num), dtypes.float32, device="CPU"))
+    t = Tensor(buf.after(buf.store(buf.const_like(7.)))).realize()
+    self.assertEqual(t.shape, ())
+    self.assertEqual(t.uop.storage_base.arg.slot, buf.arg.slot)
+    self.assertEqual(t.uop.buffer.size, 1)
+    self.assertEqual(t.item(), 7.)
+
+  def test_call_output_identity_and_cache(self):
+    for precompile in (False, True):
+      @function(precompile=precompile)
+      def f(x): return x + 1, x * 2
+      x = Tensor([1., 2.]).realize()
+      a, b = f(x)
+      decls = (a.uop.storage_base, b.uop.storage_base)
+      a.callify(b).realize(b)
+      self.assertEqual((a.uop.storage_base.arg.slot, b.uop.storage_base.arg.slot), tuple(d.arg.slot for d in decls))
+      self.assertEqual(a.tolist(), [2., 3.])
+      self.assertEqual(b.tolist(), [2., 4.])
+      c, d = f(x)
+      c.realize(d)
+      self.assertIsNot(a.uop.buffer, c.uop.buffer)
+      self.assertIsNot(b.uop.buffer, d.uop.buffer)
+      self.assertEqual(c.tolist(), [2., 3.])
+      self.assertEqual(d.tolist(), [2., 4.])
+
+  def test_call_read_materializes_declared_output(self):
+    for precompile in (False, True):
+      @function(precompile=precompile)
+      def f(x): return x + 1
+      x = Tensor([1., 2.]).realize()
+      y = f(x)
+      slot = y.uop.storage_base.arg.slot
+      self.assertEqual(y.tolist(), [2., 3.])
+      self.assertEqual(y.uop.storage_base.arg.slot, slot)
+      x.assign(0).realize()
+      self.assertEqual(y.tolist(), [2., 3.])
+
+  def test_output_aliases_share_materialization(self):
+    x = Tensor([1., 2.]).realize() + 1
+    y, z = x.contiguous_backward(), x.contiguous()
+    x.realize(y, z, x)
+    self.assertIs(x.uop.buffer, y.uop.buffer)
+    self.assertIs(x.uop.buffer, z.uop.buffer)
+    self.assertEqual(x.tolist(), [2., 3.])
+
+  def test_output_slots_survive_binding(self):
+    x = Tensor([1., 2.]).realize()
+    p = x.uop.param_like(1)
+    (out,) = UOp.call_with_outputs((p + 1,), x.uop, output_pos=(0,))
+    c = out.src[1]
+    bound = c.substitute({out.storage_base: out.storage_base.bind_buffer()})
+    self.assertTrue(bound.is_value_call)
+    self.assertFalse(bound.has_unbound_outputs)
+    self.assertEqual(bound.arg.output_pos, (0,))
+    self.assertEqual(Tensor(bound.call_outputs[0]).tolist(), [2., 3.])
+
+  def test_output_scoping_preserves_storage_targets(self):
+    x = Tensor([1., 2.]).realize()
+    y = x.clone()
+    x.assign(0)
+    y.realize(x)
+    self.assertEqual(y.tolist(), [1., 2.])
+    self.assertEqual(x.tolist(), [0., 0.])
+    self.assertIsNot(y.uop.buffer, x.uop.buffer)
+
+  def test_shared_output_order(self):
+    for reverse in (False, True):
+      x = Tensor([1., 2.]).realize()
+      a = (x + 1).sum()
+      b = a * 2
+      roots = (b, a) if reverse else (a, b)
+      Tensor.realize(*roots)
+      x.assign(0).realize()
+      self.assertEqual(a.item(), 5.)
+      self.assertEqual(b.item(), 10.)
+
+  def test_transfers_own_storage(self):
+    a = Tensor([1., 2.], device="CPU:0")
+    self.assertIs(a.uop.op, Ops.AFTER)
+    b = a.to("CPU:1")
+    self.assertIs(b.uop.op, Ops.AFTER)
+    self.assertIsNot(a.uop.storage_base, b.uop.storage_base)
+    c = Tensor.empty(2, device="CPU:1").assign(b).realize()
+    a.assign(0).realize()
+    self.assertEqual(b.tolist(), [1., 2.])
+    self.assertEqual(c.tolist(), [1., 2.])
+    b.assign(3).realize()
+    self.assertEqual(c.tolist(), [1., 2.])
+
+  def test_virtual_output_does_not_allocate(self):
+    t = Tensor(2.)
+    with patch.object(UOp, "new_buffer", side_effect=AssertionError("virtual storage")):
+      t.callify().realize()
+    self.assertEqual(t.item(), 2.)
+
+  def test_contiguous_through_wrapper_keeps_copy(self):
+    for wrapper in ("detach", "contiguous_backward"):
+      with self.subTest(wrapper=wrapper):
+        x = Tensor([1., 2.]).realize()
+        y = getattr(x.flip(0), wrapper)().contiguous().realize()
+        x.assign(0).realize()
+        self.assertEqual(y.tolist(), [2., 1.])
+
+  def test_intermediate_contiguous_through_wrapper_is_view(self):
+    x = Tensor([1., 2., 3., 4.], device="CPU").realize()
+    y = x[:2].contiguous_backward().contiguous() + 1
+    self.assertEqual(len(y.schedule_linear().src), 1)
+
   def test_basic(self):
     a = Tensor([1.,2,3])
     b = Tensor([4.,5,6])

tinygrad/device.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from dataclasses import dataclass, replace, field
+from dataclasses import dataclass, replace
 from collections import defaultdict
 from typing import Any, Callable, Generic, TypeVar, Iterator, Generator, Self, TYPE_CHECKING
 import importlib, inspect, functools, pathlib, os, contextlib, re, atexit, pickle, decimal, subprocess, struct
@@ -97,9 +97,6 @@ class MultiBuffer:
   def is_allocated(self): return all(x.is_allocated() for x in self.bufs)
   def __repr__(self): return f"<multibuf real:{self.is_allocated()} device:{tuple(x.device for x in self.bufs)} size:{self.size} dtype:{self.dtype}>"
 
-@dataclass(frozen=True)
-class BufferStorage: buf:Any; meta:Any=None; host:MMIOInterface|None=None; maps:dict[str, BufferStorage]=field(default_factory=dict) # noqa: E702
-
 class Buffer:
   profile_events:list[ProfileEvent] = []
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
@@ -107,14 +104,14 @@ class Buffer:
     assert isinstance(dtype, DType)
     self.device, self.size, self.dtype, self.offset, self.allocated_views, self._base = Device.canonicalize(device), size, dtype, offset, 0, base
     self.options = options if options is not None else BufferSpec()
-    self._storage:BufferStorage|None = None
+    self._storage:tuple|None = None
+    self._maps:dict[str, tuple] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
       if opaque is not None: self.allocate(opaque)
       if initial_value is not None:
         self.allocate()
-        if (host:=self.get_storage().host) is not None: host[:] = memoryview(initial_value).cast('B')
-        else: self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
+        self.copy_from(Buffer("PYTHON", self.size, self.dtype, opaque=memoryview(bytearray(initial_value))))
         if isinstance(initial_value, pickle.PickleBuffer): initial_value.release()
     else:
       assert base._base is None, "base can't have a base"
@@ -133,25 +130,24 @@ class Buffer:
   @functools.cached_property
   def allocator(self) -> Allocator: return self.base.allocator if self._base is not None else Device[self.device].allocator
   @property
-  def _buf(self) -> Any: return self.get_storage().buf
+  def _buf(self) -> Any: return self.get_storage()[0][0]
   @property
-  def host(self) -> MMIOInterface: return unwrap(self.get_storage().host)
+  def host(self) -> MMIOInterface: return unwrap(self.get_storage()[1])
   @property
-  def meta(self) -> Any: return self.get_storage().meta
+  def meta(self) -> Any: return self.get_storage()[0][1]
   @property
   def nbytes(self): return self.size * self.dtype.itemsize
 
-  def get_storage(self, device:str|None=None) -> BufferStorage:
+  def get_storage(self, device:str|None=None) -> tuple:
     storage = unwrap(self.ensure_allocated()._storage)
     device = Device.canonicalize(device) if device is not None else self.device
     if device == self.device: return storage
-    if device not in storage.maps:
-      alloc = Device[device].allocator
-      storage.maps[device] = BufferStorage(alloc._offset(self.base.get_buf(device), self.nbytes, self.offset)) if self._base else alloc.map(self)
-    if storage.maps[device].host is not storage.host: storage.maps[device] = replace(storage.maps[device], host=storage.host)
-    return storage.maps[device]
+    if device not in self._maps:
+      allocator = Device[device].allocator
+      self._maps[device] = (allocator._offset(self.base.get_buf(device), self.nbytes, self.offset), None) if self._base else allocator.map(self)
+    return self._maps[device], storage[1]
 
-  def get_buf(self, device:str) -> Any: return self.get_storage(device).buf
+  def get_buf(self, device:str) -> Any: return self.get_storage(device)[0][0]
 
   def is_allocated(self) -> bool: return self._storage is not None and (self._base is None or self._base_storage is self.base._storage)
   def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self
@@ -162,19 +158,19 @@ class Buffer:
       raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
     if external_ptr is not None: self.options = replace(self.options, external_ptr=external_ptr)
     if self._base is not None:
-      storage = replace(self.base.get_storage(), buf=self.allocator._offset(self.base._buf, self.nbytes, self.offset), maps={})
-    elif opaque is not None:
-      self.options = replace(self.options, nolru=True)
-      storage = opaque if isinstance(opaque, BufferStorage) else BufferStorage(opaque)
-    else: storage = self.allocator.alloc(self.nbytes, self.options)
-    storage = replace(storage, host=storage.host.view(self.offset, self.nbytes, fmt='B') if storage.host is not None else None)
+      (buf, meta), host = self.base.get_storage()
+      mapping = self.allocator._offset(buf, self.nbytes, self.offset), meta
+    else:
+      if opaque is not None: self.options = replace(self.options, nolru=True)
+      mapping, host = ((opaque, None), None) if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
+    storage = mapping, host.view(self.offset, self.nbytes, fmt='B') if host is not None else None
     if self._base is None:
       if not self.device.startswith("DISK") and self.options.external_ptr is None:
         GlobalCounters.mem_used += self.nbytes
         GlobalCounters.mem_used_per_device[self.device] += self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
     elif self._storage is None: self.base.allocated_views += 1
-    self._storage, self._base_storage = storage, self.base._storage if self._base else None
+    self._storage, self._maps, self._base_storage = storage, {}, self.base._storage if self._base else None
     return self
 
   def deallocate(self):
@@ -185,9 +181,10 @@ class Buffer:
         GlobalCounters.mem_used -= self.nbytes
         GlobalCounters.mem_used_per_device[self.device] -= self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "free", self.trace_num))
+      for dev, mb in self._maps.items(): Device[dev].allocator._unmap(mb[0])
       self.allocator.free(self._storage, self.nbytes, self.options)
     else: self.base.allocated_views -= 1
-    self._storage, self._base_storage = None, None
+    self._storage, self._maps, self._base_storage = None, {}, None
 
   def __reduce_ex__(self, protocol):
     buf:bytearray|pickle.PickleBuffer|None = None
@@ -204,7 +201,10 @@ class Buffer:
     return self._trace_num
 
   def _host_mv(self) -> memoryview|None:
-    if self.is_allocated() and hasattr(host:=self.get_storage().host, 'mv'): return unwrap(host).view(fmt='B').mv
+    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'):
+      mv = unwrap(host).view(fmt='B').mv
+      mv.obj._buffer = self  # raw ctypes views do not own their memory; keep the allocation alive for asynchronous copies
+      return mv
     if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
     return None
 
@@ -244,10 +244,10 @@ class Allocator(Generic[DeviceType]):
   def __init__(self, dev:DeviceType, supports_copy_from_disk:bool=True, supports_transfer:bool=True):
     self.dev: DeviceType = dev
     self.default_buffer_spec: BufferSpec = BufferSpec()
-    self.cache:dict[tuple[int, BufferSpec|None], list[BufferStorage]] = defaultdict(list)
+    self.cache:dict[tuple[int, BufferSpec|None], list[tuple]] = defaultdict(list)
     self.supports_copy_from_disk, self.supports_transfer = supports_copy_from_disk, supports_transfer
 
-  def alloc(self, size:int, options:BufferSpec|None=None) -> BufferStorage:
+  def alloc(self, size:int, options:BufferSpec|None=None) -> tuple:
     assert size > 0, f"alloc size must be positive, getting {size}"
     if len(c:=self.cache[(size, options)]): return c.pop()
     spec = options if options is not None else self.default_buffer_spec
@@ -257,29 +257,24 @@ class Allocator(Generic[DeviceType]):
     except (RuntimeError, MemoryError) as e: raise MemoryError(f"Allocation of {size_to_str(size)} failed on {self.dev.device}. "
                                                             f"Used: {size_to_str(GlobalCounters.mem_used_per_device[self.dev.device])}") from e
 
-  def free(self, storage:BufferStorage, size:int, options:BufferSpec|None=None):
+  def free(self, storage:tuple, size:int, options:BufferSpec|None=None):
     spec = options if options is not None else self.default_buffer_spec
     if LRU and self.lru and not (spec.nolru or spec.zero) and spec.external_ptr is None: self.cache[(size, options)].append(storage)
-    else: self.do_free(storage, spec)
+    else: self._free(storage[0][0], spec)
 
   def free_cache(self):
     for (_, options), storages in self.cache.items():
-      for storage in storages: self.do_free(storage, options if options is not None else self.default_buffer_spec)
+      for storage in storages: self._free(storage[0][0], options if options is not None else self.default_buffer_spec)
       storages.clear()
 
-  def do_free(self, storage:BufferStorage, options:BufferSpec):
-    for dev in storage.maps: Device[dev].synchronize()
-    for dev, mb in storage.maps.items(): Device[dev].allocator._unmap(mb)
-    if options.external_ptr is None: self._free(storage, options)
-
-  def map(self, buf:Buffer) -> BufferStorage: return self._map(buf.ensure_allocated())
+  def map(self, buf:Buffer) -> tuple: return self._map(buf.ensure_allocated()._buf)
 
   # implemented by the runtime
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage: raise NotImplementedError("need alloc")
-  def _free(self, storage:BufferStorage, options:BufferSpec): pass  # if opaque is a Python object, you don't need a free
+  def _alloc(self, size:int, options:BufferSpec) -> tuple: raise NotImplementedError("need alloc")
+  def _free(self, opaque, options:BufferSpec): pass  # if opaque is a Python object, you don't need a free
   def _copyin(self, dest, src:memoryview): raise NotImplementedError("need copyin")
   def _copyout(self, dest:memoryview, src): raise NotImplementedError("need copyout")
-  def _map(self, buf) -> BufferStorage: raise NotImplementedError("need map")
+  def _map(self, buf) -> tuple: raise NotImplementedError("need map")
   def _unmap(self, mb): pass  # default no-op; override if _map allocates iface-side state
   def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
   # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):

tinygrad/engine/realize.py

@@ -3,7 +3,7 @@ from typing import cast, Iterator, Any, Sequence
 import weakref, decimal, array
 from dataclasses import dataclass, replace, field
 from tinygrad.helpers import colored, DEBUG, GlobalCounters, ansipad, prod, flatten, Context, to_tuple, tqdm, dedup
-from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us, cpu_profile
+from tinygrad.helpers import BEAM, size_to_str, time_to_str, VALIDATE_WITH_CPU, PROFILE, ProfilePointEvent, cpu_events, perf_counter_us
 from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, AxisType, sym_infer, graph_rewrite, ProgramInfo
 from tinygrad.device import Device, Buffer, MultiBuffer, ProfileGraphEntry
 from tinygrad.renderer import Estimates, Renderer
@@ -156,9 +156,6 @@ 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 []
@@ -197,7 +194,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) + off for u, dev, off in info.inputs]
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr + off for u, dev, off in info.inputs]
     cast(Buffer, call.src[1 + info.table].buffer).host.view(fmt='Q')[:] = array.array('Q', addrs)
   ctx = replace(ctx, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}})
   ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,))
@@ -280,7 +277,7 @@ 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, HCQ_DEVS # noqa: E402 # down here, hcq2 imports realize
+from tinygrad.runtime.support.hcq2 import hcq_compile, hcq_link, HCQ_RUNTIME_DEV, HCQInfo # noqa: E402 # down here, hcq2 imports realize
 
 def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:list[UOp]|None=None, profile:bool|None=None) -> UOp:
   if validate: linear = graph_rewrite(linear, pm_validate, name="validate", walk=True)

tinygrad/mixin/elementwise.py

@@ -1,6 +1,6 @@
 import math, functools, operator
 from typing import TYPE_CHECKING, Literal, Self
-from tinygrad.uop import Ops
+from tinygrad.uop import Ops, GroupOp
 from tinygrad.dtype import dtypes, ConstType, DType, PyConst, least_upper_dtype, least_upper_float, weak_dtype
 from tinygrad.helpers import argfix, polyN
 from tinygrad.mixin.creation import CreationMixin
@@ -63,6 +63,7 @@ class ElementwiseMixin(CreationMixin):
     if self.dtype in dtypes.weaks: return self
     uop = self._uop
     if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
+    if uop.op in GroupOp.Movement|{Ops.BITCAST} and (view:=uop.buffer_view()) is not None: return self._wrap_uop(view)
     return self._wrap_uop(uop.alu(Ops.CONTIGUOUS))
 
   def contiguous_backward(self) -> Self:

tinygrad/mixin/gradient.py

@@ -2,7 +2,7 @@ from typing import cast
 import math, dataclasses
 from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, all_metadata, broadcast_axes
 from tinygrad.helpers import argsort
-from tinygrad.dtype import 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):
@@ -32,14 +32,14 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
     # grads align with the call's src positions (None for the body and for RETURNED outputs, wherever they are)
     def arg_grads(g):
       git = iter(g)
-      return (None,) + tuple(next(git) if not a.unsharded_base.is_unbound else None for a in k.src[1:])
+      return (None,) + tuple(None if i in (k.arg.output_pos or ()) else next(git) for i in range(len(args)))
     if ctx.op is Ops.SINK:
       real = [on_dev(g, i) for i,g in enumerate(ctx.src) if g.op is not Ops.NOOP]
       return arg_grads(k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k))
     return arg_grads(k.arg.grad_fxn(on_dev(ctx, 0), k))
   # the RETURNED inputs are the call outputs: their positions in the args get the output gradients from the AFTER rule
-  assert fxn.op is Ops.SINK and k.has_unbound_outputs, f"expected a CALL with unbound BUFFER outputs or a grad_fxn, got {fxn.op}"
-  ret_pos = [i for i, a in enumerate(args) if a.unsharded_base.is_unbound]
+  assert fxn.op is Ops.SINK and k.is_value_call, f"expected a value CALL or a grad_fxn, got {fxn.op}"
+  ret_pos = k.arg.output_pos
   # the body stores the outputs into output PARAMs: the values are the stored values in slot order
   values = UOp.sink(*[st.src[1] for st in fxn.src if st.op is Ops.STORE])
   params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
@@ -53,7 +53,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
   grads = compute_gradient(values, root_grad, set(params.values()))
   # for precompiled calls, substitute forward outputs with params so intermediates aren't recomputed
   fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(values.src)} if k.arg.precompile else {}
-  fwd_outs = k.unbound_outputs if k.arg.precompile else ()
+  fwd_outs = k.call_outputs if k.arg.precompile else ()
   # collect needed gradient bodies, compact unused params, create a single backward CALL
   grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
   bwd_body = UOp.sink(*[gb for _, gb in grad_bodies]).substitute(fwd_subs, walk=True)
@@ -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_after_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,8 @@ 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.AFTER, src=(UPat(name="view"),
+    UPat(Ops.STORE, src=(UPat(name="view"), UPat())))))), partial_after_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/renderer/cstyle.py

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

tinygrad/runtime/graph/metal.py

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

tinygrad/runtime/ops_amd.py

@@ -3,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, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name, patch, unwrap_view, rt_addr
 from tinygrad.uop.ops import sint, UOp, ProgramInfo
-from tinygrad.device import BufferStorage, BufferSpec, Buffer, Device, Allocator, Compiled, ProfileProgramEvent
+from tinygrad.device import BufferSpec, Buffer, Device, Compiled, ProfileProgramEvent
 from tinygrad.dtype import dtypes
 from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, prod, colored
 from tinygrad.helpers import ceildiv, unwrap, pluralize, HCQ2, mv_address, ContextVar, VIZ
@@ -14,10 +14,10 @@ from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
 from tinygrad.runtime.autogen.am import am
 from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface, hcq_filter_visible_devices
+from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
 from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
-from tinygrad.runtime.support.system import PCIIfaceBase, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
+from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
 from tinygrad.runtime.support.usb import USB3, pm_usb_batch, pm_usb_lower, pm_usb_bufferize
 from tinygrad.runtime.support.memory import AddrSpace
 if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
@@ -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(Allocator['AMDDevice']):
+class AMDAllocator(HCQAllocator['AMDDevice']):
   def __init__(self, dev:AMDDevice):
     super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
 
-  def _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 _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 _free(self, storage:BufferStorage, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(storage)
-  def _map(self, buf:Buffer) -> BufferStorage: return self.dev.iface.map(buf)
-  def _unmap(self, mapping:BufferStorage): self.dev.iface.unmap(mapping)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
+
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
 
 @dataclass
 class AMDQueueDesc:
@@ -576,7 +576,7 @@ class AMDQueueDesc:
 
 class KFDIface:
   kfd:FileIOInterface|None = None
-  event_page:Buffer
+  event_page:HCQBuffer|None = None
   gpus:list[FileIOInterface] = []
   count:int = 0
 
@@ -611,7 +611,25 @@ class KFDIface:
     kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
     if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> BufferStorage:
+    # Set these for our device.
+    if KFDIface.event_page is None:
+      KFDIface.event_page = self.alloc(0x8000, uncached=True)
+      kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
+    else: self.map(KFDIface.event_page)
+
+    # Event to wait for queues completion
+    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
+    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
+
+    # OS events to collect memory and hardware faults
+    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
+    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)
+
+    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
+      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
+    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
+
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
     flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
 
     if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
@@ -637,52 +655,43 @@ class KFDIface:
       buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset)
       assert addr == buf == mem.va_addr
 
-    self._map_handle(mem.handle)
-    return BufferStorage(mem.va_addr, mem, MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None)
+    view = MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None
+    self.map(hcqbuf:=HCQBuffer(mem.va_addr, mem.size, meta=mem, view=view, owner=self.dev))
+    return hcqbuf
 
-  def free(self, storage:BufferStorage):
-    self._unmap_handle(storage.meta.handle)
-    if storage.buf: FileIOInterface.munmap(storage.buf, storage.meta.size)
-    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=storage.meta.handle)
+  def free(self, mem):
+    self._unmap(mem)
+    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
+    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
-  def unmap(self, mapping:BufferStorage):
-    handle, owned = mapping.meta
-    self._unmap_handle(handle)
-    if owned: kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=handle)
+  def unmap(self, mem):
+    self._unmap(mem)
+    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
 
-  def map(self, buf:Buffer) -> BufferStorage:
-    if buf.device.split(":")[0] == "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):
+  def _unmap(self, mem):
     gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
     assert stm.n_success == 1
 
-  def _unmap_handle(self, handle):
-    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
+  def map(self, mem):
+    if mem.owner is not None and mem.owner._is_cpu():
+      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      cast(Any, mapped)._owns_kfd_handle = True
+      return mapped
+
+    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
+    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
     assert stm.n_success == 1
+    return HCQBuffer(mem.va_addr, mem.size, meta=mem.meta, owner=mem.owner)
 
   def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
                    xcc_id=0, idx=0):
-    if not hasattr(self, 'queue_event_arr'):
-      if not hasattr(KFDIface, 'event_page'):
-        KFDIface.event_page = Buffer(self.dev.device, 0x8000, dtypes.uint8, options=BufferSpec(uncached=True), preallocate=True)
-        kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
-
-      KFDIface.event_page.get_buf(self.dev.device)
-      self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(*[kfd.struct_kfd_event_data(event_id=kfd.AMDKFD_IOC_CREATE_EVENT(
-        KFDIface.kfd, event_type=t, auto_reset=int(t == kfd.KFD_IOC_EVENT_SIGNAL)).event_id)
-        for t in (kfd.KFD_IOC_EVENT_SIGNAL, kfd.KFD_IOC_EVENT_MEMORY, kfd.KFD_IOC_EVENT_HW_EXCEPTION)])
-
-    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf, ring_size=ring.nbytes, gpu_id=self.gpu_id,
+    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
       queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
-      eop_buffer_address=eop_buffer._buf if eop_buffer else 0, eop_buffer_size=eop_buffer.nbytes if eop_buffer else 0,
-      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
-      write_pointer_address=gart._buf+wptr, read_pointer_address=gart._buf+rptr+8*xcc_id)
+      eop_buffer_address=eop_buffer._buf.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer._buf.size if eop_buffer else 0,
+      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf.va_addr if cwsr_buffer else 0,
+      ctx_save_restore_size=ctx_save_restore_size,
+      write_pointer_address=gart._buf.va_addr+wptr, read_pointer_address=gart._buf.va_addr+rptr+8*xcc_id)
 
     if not hasattr(self, 'doorbells'):
       self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
@@ -695,7 +704,7 @@ class KFDIface:
       write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer)
 
   def sleep(self, tm:int):
-    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=ctypes.addressof(self.queue_event_arr), num_events=3, wait_for_all=0, timeout=tm)
+    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
     if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang()
 
   def on_device_hang(self):
@@ -741,6 +750,7 @@ class PCIIface(PCIIfaceBase):
 
   def require_profile_mode(self): return True
   def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
+  def unmap(self, mem): self.free(mem)
 
   def _compute_props(self):
     self.ip_versions = self.dev_impl.ip_ver
@@ -765,10 +775,11 @@ class PCIIface(PCIIfaceBase):
 
     rcvr_params: tuple
     if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr, gart._buf+wptr, idx)))
+      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, idx)))
     else:
-      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf, ring.nbytes, gart._buf+rptr,
-        gart._buf+wptr, eop_buffer._buf, eop_buffer.nbytes, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
+      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
+        gart._buf.va_addr+wptr, eop_buffer._buf.va_addr, eop_buffer._buf.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
 
     (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True)).host.view(fmt='Q')[0] = 0
     doorbell = Buffer("CPU", 1, dtypes.uint64, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True)
@@ -800,10 +811,9 @@ class PCIIface(PCIIfaceBase):
   def device_fini(self): self.dev_impl.fini()
 
 class USBAllocator(AMDAllocator): # the host program reads another device's memory in place: its bytes are the mapping
-  def map(self, buf:Buffer) -> BufferStorage:
+  def map(self, buf:Buffer) -> tuple:
     mv = buf.ensure_allocated().as_memoryview(force_zero_copy=True, no_sync=True)
-    return BufferStorage(mv_address(mv), mv)
-  def _unmap(self, mapping:BufferStorage): pass
+    return HCQBuffer(addr:=mv_address(mv), mv.nbytes, meta=mv, view=MMIOInterface(addr, mv.nbytes, fmt='B'), owner=self.dev), mv
 
 class USBIface(PCIIface):
   def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
@@ -812,19 +822,15 @@ class USBIface(PCIIface):
     self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
     self.dev_impl = AMDev(self.pci_dev)
     self._compute_props()
-
-  @functools.cached_property
-  def ctrl(self) -> Buffer:
     # the controller's memory the queue and the host share, one range (usb.py slices it): the sys page at 0, the cq page at 0x1000, the sram at
     # 0x5000. the host's view starts at the sys page's controller address 0xa000, which puts the sram on its scsi window 0xf000
     vaddr, pieces = self.dev_impl.mm.alloc_vaddr(size=0x85000), [(0x0, 0x820000, 0x1000), (0x1000, 0x822000, 0x1000), (0x5000, 0x200000, 0x80000)]
-    for off, paddr, n in pieces: self.dev_impl.mm.map_range(vaddr + off, n, [(paddr, n)], aspace=AddrSpace.SYS, uncached=True)
-    view = self.pci_dev.dma_view(0xa000, 0x85000)
-    for off, n in ((0x800, 4), (0x5000, 0x80000)): view.view(off, n)[:] = bytes(n) # no stale fence or sentinel
-    return Buffer(self.dev.device, 0x85000, dtypes.uint8, options=BufferSpec(external_ptr=vaddr), opaque=BufferStorage(vaddr, host=view))
+    maps = [self.dev_impl.mm.map_range(vaddr + off, n, [(sys, n)], aspace=AddrSpace.SYS, uncached=True) for off, sys, n in pieces]
+    self.ctrl = HCQBuffer(vaddr, 0x85000, meta=PCIAllocationMeta(maps[0], has_cpu_mapping=False), view=self.pci_dev.dma_view(0xa000, 0x85000),
+                          owner=self.dev)
+    for off, n in ((0x800, 4), (0x5000, 0x80000)): unwrap(self.ctrl.view).view(off, n)[:] = bytes(n) # no stale fence or sentinel
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False,
-            **kwargs) -> BufferStorage:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, zero=False, **kwargs) -> HCQBuffer:
     # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
     return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
 
@@ -993,7 +999,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
+    base = self.scratch._buf.va_addr
     self.aql_desc.scratch_backing_memory_location = base
     self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size
     self.aql_desc.scratch_resource_descriptor[:] = [lo32(base), int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(base), SWIZZLE_ENABLE=1), 'little'),
@@ -1022,13 +1028,13 @@ class AMDDevice(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, b.arg.slot, key))
+        Compiled.profile_events.append(ProfileProgramEvent(self.device, name, lib, buf._buf.va_addr, b.arg.slot, key))
     return self.prog_bufs[b]
 
   def sqtt_trace(self, slot:int, se:int) -> bytes:
     off = (se * self.prof_slots + slot) * self.sqtt_win
     wptr = (self.sqtt_wptrs.host.view(fmt='I')[slot * self.sqtt_ses + se] & 0x1FFFFFFF) * 32
-    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf + off) // 32) & 0x1FFFFFFF) * 32
+    if self.target[:2] == (11, 0): wptr -= (((self.sqtt_buf._buf.va_addr + off) // 32) & 0x1FFFFFFF) * 32
     assert 0 <= wptr <= self.sqtt_win, f"{wptr} > {self.sqtt_win}, should never happen"
     if wptr >= self.sqtt_win - 32: # the wptr stops at the last dword when the window overflows
       print(colored(f"{self.device}: Warning: SQTT buffer is full (SE {se})! Increase SQTT buffer with SQTT_BUFFER_SIZE=X (in MB)", "yellow"))

tinygrad/runtime/ops_cl.py

@@ -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 BufferStorage, BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
+from tinygrad.device import BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
 
 CC_CB = c.CFUNCTYPE[None, [c.POINTER[ctypes.c_char], c.POINTER[None], cl.size_t, c.POINTER[None]]]
 BP_CB = c.CFUNCTYPE[None, [cl.cl_program, c.POINTER[None]]]
@@ -76,11 +76,11 @@ class CLProgram(Program['CLDevice']):
     return None
 
 class CLAllocator(Allocator['CLDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return BufferStorage(checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status))
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), None), None
 
   @suppress_finalizing
-  def _free(self, storage:BufferStorage, options:BufferSpec): check(cl.clReleaseMemObject(storage.buf))
+  def _free(self, opaque:cl.cl_mem, options:BufferSpec): check(cl.clReleaseMemObject(opaque))
   def _copyin(self, dest:cl.cl_mem, src:memoryview):
     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))

tinygrad/runtime/ops_cpu.py

@@ -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 BufferStorage, BufferSpec, TinyELF, Program, Device, Buffer, Allocator
-from tinygrad.runtime.support.memory import MMIOInterface
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled
+from tinygrad.device import BufferSpec, TinyELF, Program, Device
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator
 from tinygrad.runtime.support.c import DLL
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import CPULLVMRenderer
@@ -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:int, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
                vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
     st = time.perf_counter()
     if self.lvp:
       lvp_args = bytearray(12 + (len(bufs) + len(vals)) * 8)
       addr = mv_address(lvp_args)
-      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *bufs)
+      struct.pack_into(f'<3I{len(bufs)}Q', lvp_args, 0, *data64_le(addr+12), (len(bufs)+len(vals))*2, *[b.va_addr for b in bufs])
       for v,(off,dt) in zip(vals, TinyELF.iter_sig(self.signature[-len(vals):], len(bufs)*8)): struct.pack_into(f'<{dt.fmt}', lvp_args, 12+off, v)
       self.fxn(addr)
     else:
-      args = [*bufs, *cast(tuple[int, ...], vals)]
+      args = [*[cast(int, b.va_addr) for b in bufs], *cast(tuple[int, ...], vals)]
       self.fxn(*[ctypes.c_uint64(x) for x in args])
     return time.perf_counter() - st if wait else None
 
@@ -74,24 +74,25 @@ 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(Allocator['CPUDevice']):
+class CPUAllocator(HCQAllocator['CPUDevice']):
   def __init__(self, dev:CPUDevice): super().__init__(dev, supports_copy_from_disk=False, supports_transfer=False)
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
     if options.external_ptr is not None: addr, buf = options.external_ptr, None
     elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
     else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
-    return BufferStorage(addr, buf, MMIOInterface(addr, size, fmt='B'))
+    return (opaque:=HCQBuffer(addr, size, meta=buf, view=MMIOInterface(addr, size, fmt='B'), owner=self.dev), opaque.meta), opaque.view
 
-  def _copyin(self, dest:int, src:memoryview):
+  def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
     self.dev.synchronize()
-    ctypes.memmove(dest, from_mv(src), len(src))
-  def _copyout(self, dest:memoryview, src:int):
+    ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
-    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
+    dest[:] = to_mv(int(src.va_addr), dest.nbytes)[:]
+  def _do_map(self, buf:HCQBuffer):
+    if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
+    return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)
+  def _do_unmap(self, mb): pass  # CPU _do_map returns a view wrapper, nothing to release
 
 class CPUDevice(HCQ2Compiled):
   wait_timeout_ms, has_copy_queue = 30000, False

tinygrad/runtime/ops_cuda.py

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

tinygrad/runtime/ops_disk.py

@@ -1,7 +1,7 @@
 import os, sys, mmap, io, ctypes, contextlib, pathlib
 from typing import Generator, Callable
 from tinygrad.helpers import OSX, round_up, mv_address
-from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator
+from tinygrad.device import MMIOInterface, Compiled, Allocator
 with contextlib.suppress(ImportError):
   import _posixshmem
   from tinygrad.runtime.autogen import io_uring, libc
@@ -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) -> BufferStorage:
+  def _alloc(self, size:int, options) -> tuple:
     self.dev._might_open(size)
-    return BufferStorage(opaque:=DiskBuffer(self.dev, size), None, MMIOInterface(mv_address(opaque._buf()), size))
+    return (opaque:=DiskBuffer(self.dev, size), None), MMIOInterface(mv_address(opaque._buf()), size)
 
-  def _free(self, storage:BufferStorage, options): self.dev._might_close()
+  def _free(self, opaque, options): self.dev._might_close()
   def _as_buffer(self, src:DiskBuffer): return src._buf()
   def _copyin(self, dest:DiskBuffer, src:memoryview): dest._buf()[:] = src
   def _copyout(self, dest:memoryview, src:DiskBuffer):

tinygrad/runtime/ops_dsp.py

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

tinygrad/runtime/ops_hip.py

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

tinygrad/runtime/ops_metal.py

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

tinygrad/runtime/ops_npy.py

@@ -1,10 +1,9 @@
 import numpy as np
 from tinygrad.helpers import flat_mv
-from tinygrad.device import BufferStorage, MMIOInterface, Compiled, Allocator
+from tinygrad.device import MMIOInterface, Compiled, Allocator
 
 class NpyAllocator(Allocator['NpyDevice']):
-  def _alloc(self, size:int, options) -> BufferStorage:
-    return BufferStorage(arr:=np.empty(size, dtype=np.uint8), arr, MMIOInterface(arr.ctypes.data, size))
+  def _alloc(self, size:int, options) -> tuple: return (arr:=np.empty(size, dtype=np.uint8), arr), MMIOInterface(arr.ctypes.data, size)
 
   def _as_buffer(self, src:np.ndarray) -> memoryview: return flat_mv(np.require(src, requirements='C').data)
   def _copyout(self, dest:memoryview, src:np.ndarray): dest[:] = self._as_buffer(src)

tinygrad/runtime/ops_null.py

@@ -1,5 +1,5 @@
 import inspect, math
-from tinygrad.device import BufferStorage, Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
+from tinygrad.device import Compiled, Allocator, ProfileGraphEntry, ProfileGraphEvent, Program, TinyELF
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.renderer import Renderer, cstyle, nir, ptx, llvmir, wgsl
 from tinygrad.renderer.cstyle import CStyleLanguage
@@ -22,7 +22,7 @@ class NullProgram(Program['NullDevice']):
     with cpu_profile(self.name, self.device, profile_key=self.profile_key): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):
-  def _alloc(self, size:int, options) -> BufferStorage: return BufferStorage(None)
+  def _alloc(self, size:int, options) -> tuple: return (None, None), None
 
   def _copyin(self, dest, src:memoryview): pass
   def _copyout(self, dest:memoryview, src):

tinygrad/runtime/ops_nv.py

@@ -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, 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 dataclasses import dataclass
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, patch, to_name, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface, FileIOInterface, BumpAllocator, hcq_filter_visible_devices
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher
 from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
-from tinygrad.device import BufferStorage, Buffer, BufferSpec, Allocator, Compiled, Device, TinyELF
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, TinyELF
 from tinygrad.dtype import dtypes, DType
 from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, PROFILE, ContextVar, VIZ
 from tinygrad.helpers import ProfileEvent
@@ -290,19 +290,16 @@ def nv_build_program(dev:NVDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[NVPro
     cached = _nv_program_cache[key] = (data, patch(buf, rows, data.image))
   return cached
 
-class NVAllocator(Allocator['NVDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> BufferStorage:
-    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero)
+class NVAllocator(HCQAllocator['NVDevice']):
+  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+    return (opaque:=self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero), opaque.meta), opaque.view
 
-  def _free(self, storage:BufferStorage, options:BufferSpec):
-    self.dev.synchronize()
-    self.dev.iface.free(storage)
-  def _map(self, buf:Buffer) -> BufferStorage: return self.dev.iface.map(buf)
-  def _unmap(self, mapping:BufferStorage): self.dev.iface.unmap(mapping)
-  def _offset(self, buf:int, size:int, offset:int) -> int: return buf + offset
+  def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
 
-  def _encode_decode(self, bufout:int, bufin:int, desc_buf:int, hist:list[int], shape:tuple[int,...], frame_pos:int):
-    assert all(h % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
+  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
+
+  def _encode_decode(self, bufout:HCQBuffer, bufin:HCQBuffer, desc_buf:HCQBuffer, hist:list[HCQBuffer], shape:tuple[int,...], frame_pos:int):
+    assert all(h.va_addr % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
 
     h, w = ((2 * shape[0]) // 3 if shape[0] % 3 == 0 else (2 * shape[0] - 1) // 3), shape[1]
     dev, chroma_off = self.dev, round_up(w, 64) * round_up(h, 64)
@@ -311,15 +308,15 @@ class NVAllocator(Allocator['NVDevice']):
     cmds = nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
     cmds += nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
                 err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin.va_addr >> 8)
     for pos, buf in zip([(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)] + [frame_pos], hist + [bufout]):
-      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf >> 8)
-      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, (buf + chroma_off) >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, (desc_buf + 0x200) >> 8)
-    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf) >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, desc_buf.offset(0x200).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf.va_addr) >> 8)
     cmds += nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_addr + dev.intra_top_off) >> 8)
     if dev.intra_unk_off is not None: cmds += nvm(4, 0x4dc, (filter_addr + dev.intra_unk_off) >> 8)
     cmds += nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
@@ -329,7 +326,7 @@ class NVAllocator(Allocator['NVDevice']):
 # device
 
 @dataclass
-class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; notifier: Buffer; entries: int; token: int # noqa: E702
+class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; entries: int; token: int # noqa: E702
 
 class NVKIface:
   root = None
@@ -439,7 +436,7 @@ class NVKIface:
     if made.params.status != 0: raise RuntimeError(f"_gpu_map_to_cpu returned {get_error_str(made.params.status)}")
     return fd_dev.mmap(target, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED | (MAP_FIXED if target is not None else 0), 0)
 
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> BufferStorage:
+  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, map_flags=0, cpu_addr=None, **kwargs) -> HCQBuffer:
     # Uncached memory is "system". Use huge pages only for gpu memory.
     page_size = mmap.PAGESIZE if uncached or host else ((2 << 20) if size >= (8 << 20) else (mmap.PAGESIZE if isinstance(self, MOCKIface) else
                                                                                              4 << 10))
@@ -479,20 +476,17 @@ class NVKIface:
 
     return self._gpu_uvm_map(va_addr, size, mem_handle, has_cpu_mapping=cpu_access or host)
 
-  def free(self, storage:BufferStorage):
-    mem = storage.meta
-    if mem.hMemory > NVKIface.host_object_enumerator: # not a host object, clear phys mem.
-      made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, hObjectOld=mem.hMemory)
+  def free(self, mem:HCQBuffer):
+    if mem.owner != self.dev: return
+    if mem.meta.hMemory > NVKIface.host_object_enumerator: # not a host object, clear phys mem.
+      made = nv_gpu.NVOS00_PARAMETERS(hRoot=self.root, hObjectParent=self.dev.nvdevice, hObjectOld=mem.meta.hMemory)
       nv_iowr(self.fd_ctl, nv_gpu.NV_ESC_RM_FREE, made)
       if made.status != 0: raise RuntimeError(f"_gpu_free returned {get_error_str(made.status)}")
-    self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=storage.buf, length=mem.length))
-    if storage.host is not None: FileIOInterface.munmap(storage.buf, mem.length)
 
-  def unmap(self, mapping:BufferStorage):
-    mem, owns_range = mapping.meta
-    if owns_range: self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=mapping.buf, length=mem.length))
+    self.uvm(nv_gpu.UVM_FREE, nv_gpu.UVM_FREE_PARAMS(base=int(mem.va_addr), length=mem.size))
+    if mem.view is not None: FileIOInterface.munmap(int(mem.va_addr), mem.size)
 
-  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False) -> BufferStorage:
+  def _gpu_uvm_map(self, va_base, size, mem_handle, create_range=True, has_cpu_mapping=False, owner=None) -> HCQBuffer:
     if create_range:
       self.uvm(nv_gpu.UVM_CREATE_EXTERNAL_RANGE, nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS(base=va_base, length=size))
       made = nv_gpu.NVOS46_PARAMETERS(hClient=self.root, hDevice=self.dev.nvdevice, hDma=self.dev.virtmem, hMemory=mem_handle, length=size,
@@ -506,14 +500,14 @@ class NVKIface:
 
     self.uvm(nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION, uvm_map:=nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS(base=va_base, length=size,
       rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hMemory=mem_handle, gpuAttributesCount=1, perGpuAttributes=attrs, mapped_gpu_ids=[self.gpu_uuid]))
-    return BufferStorage(va_base, uvm_map, MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None)
+    return HCQBuffer(va_base, size, meta=uvm_map, view=MMIOInterface(va_base, size, fmt='B') if has_cpu_mapping else None,
+                     owner=self.dev if owner is None else owner)
 
-  def map(self, buf:Buffer) -> BufferStorage:
-    mem = buf.meta
-    if buf.device.split(":")[0] == "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 map(self, mem:HCQBuffer):
+    if mem.owner is not None and mem.owner._is_cpu():
+      if not any(x.device.startswith("NV") for x in mem.mapped_devs): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
+      mem = mem.mappings[next(x for x in mem.mapped_devs if x.device.startswith("NV"))]
+    return self._gpu_uvm_map(mem.va_addr, mem.size, mem.meta.hMemory, create_range=False, owner=mem.owner)
 
   def _alloc_gpu_vaddr(self, size, alignment=(4 << 10), force_low=False):
     return NVKIface.low_uvm_vaddr_allocator.alloc(size, alignment) if force_low else NVKIface.uvm_vaddr_allocator.alloc(size, alignment)
@@ -582,6 +576,9 @@ 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))
 
@@ -603,9 +600,8 @@ class NVDevice(HCQ2Compiled):
 
   @functools.cached_property
   def fifos(self) -> dict[str, GPFifo]:
-    mem = self.iface.alloc(3<<20, contiguous=True, cpu_access=True, force_devmem=True, map_flags=nv_gpu.NVOS33_FLAGS_CACHING_TYPE_WRITECOMBINED<<23)
-    self.gpfifo_buf = Buffer(self.device, 3<<20, dtypes.uint8, opaque=mem)
-
+    self.gpfifo_buf = Buffer(self.device, self.gpfifo_mem.size, dtypes.uint8, options=BufferSpec(external_ptr=self.gpfifo_mem.va_addr, nolru=True)) \
+                        .allocate(opaque=self.gpfifo_mem)
     compute = self._new_gpu_fifo("COMPUTE:0", self.ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
     copy = self._new_gpu_fifo("COPY:0", self.ctxshare, self.channel_group, offset=0x100000, entries=0x10000)
     self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
@@ -619,11 +615,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 = Buffer(self.device, size:=48 << 20, dtypes.uint8, opaque=self.iface.alloc(size, uncached=True))
-    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_buf._buf+offset, gpFifoEntries=entries,
-      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_buf.meta.hMemory,
-      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_buf.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
-      engineType=19 if video else 0, hContextShare=ctxshare,
+    notifier = self.iface.alloc(48 << 20, uncached=True)
+    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_mem.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
+      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_mem.meta.hMemory,
+      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_mem.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
+      engineType=19 if video else 0,
       hVASpace=self.vaspace if video and self.is_nvd() else 0) # gsp has no default vaspace, rm maps the decoder ctx into its own
     gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params)
 
@@ -646,7 +642,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), notifier=notifier, entries=entries, token=ws_token_params.workSubmitToken)
+      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), entries=entries, token=ws_token_params.workSubmitToken)
     self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag=to_name(n, name)), lambda ctx, b=getattr(fifo, n): b)
                                         for n in ("ring", "gpput", "doorbell", "put_value")]) + self.pm_bufferize
     return fifo
@@ -668,15 +664,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.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
+    fifo.ring._buf.cpu_view().view(fmt='Q')[put[0] % fifo.entries] = buf._buf.va_addr | (len(cmds) << 42) | (1 << 41)
+    fifo.gpput._buf.cpu_view().view(fmt='I')[0] = (put[0] + 1) % fifo.entries
 
     System.memory_barrier()
     self.gpu_mmio[0x90 // 4] = fifo.token
     put[0] += 1
 
   def _submit_cmds(self, fifo:GPFifo, *cmds:int): # runs cmds once everything already submitted is done, then bumps the timeline
-    tl, addr = self.timeline.host.view(fmt='Q'), self.timeline._buf
+    tl, addr = self.timeline.host.view(fmt='Q'), self.timeline._buf.va_addr
     self._push(fifo, nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1]),
                          nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit")) + list(cmds) +
                      nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1] + 1),
@@ -691,7 +687,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)),
+    self._submit_cmds(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf.va_addr)),
                        *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(bytes_per_tpc), 0xff))
 
   def _ensure_has_vid_hw(self, w, h):
@@ -754,14 +750,12 @@ 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 = Buffer(self.device, size:=getenv("PMA_BUFFER_SIZE", 512) << 20, dtypes.uint8,
-                          opaque=self.iface.alloc(size, uncached=True, cpu_cached=True, cpu_access=True))
-    self.pma_bytes = Buffer(self.device, size:=0x1000, dtypes.uint8,
-                            opaque=self.iface.alloc(size, uncached=True, cpu_cached=True, cpu_access=self.is_nvd(), read_only=True))
+    self.pma_buf = self.iface.alloc(getenv("PMA_BUFFER_SIZE", 512) << 20, uncached=True, cpu_cached=True, cpu_access=True)
+    self.pma_bytes = self.iface.alloc(0x1000, uncached=True, cpu_cached=True, cpu_access=self.is_nvd(), read_only=True)
     self.pma_rptr = 0
 
     pma_stream = nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS(hMemPmaBuffer=self.pma_buf.meta.hMemory,
-      pmaBufferSize=self.pma_buf.nbytes, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf._buf)
+      pmaBufferSize=self.pma_buf.size, hMemPmaBytesAvailable=self.pma_bytes.meta.hMemory, pmaBufferVA=self.pma_buf.va_addr)
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_ALLOC_PMA_STREAM, pma_stream, extra=(self.pma_buf, self.pma_bytes))
 
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_RESERVE_HWPM_LEGACY, nv_gpu.struct_NVB0CC_CTRL_RESERVE_HWPM_LEGACY_PARAMS(ctxsw=0))
@@ -828,9 +822,9 @@ class NVDevice(HCQ2Compiled):
     if params.bOverflowStatus: raise RuntimeError("PMA profiler: buffer overflow detected")
     if params.bytesAvailable == 0: return None
 
-    start, end, view = self.pma_rptr, self.pma_rptr + params.bytesAvailable, self.pma_buf.host
-    pma_data = bytes(view[start:min(end, self.pma_buf.nbytes)]) + bytes(view[:max(0, end - self.pma_buf.nbytes)])
-    self.pma_rptr = end % self.pma_buf.nbytes
+    start, end = self.pma_rptr, self.pma_rptr + params.bytesAvailable
+    pma_data = bytes(self.pma_buf.cpu_view()[start:min(end, self.pma_buf.size)]) + bytes(self.pma_buf.cpu_view()[:max(0, end - self.pma_buf.size)])
+    self.pma_rptr = end % self.pma_buf.size
 
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_PMA_STREAM_UPDATE_GET_PUT,
       nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS(bytesConsumed=params.bytesAvailable))

tinygrad/runtime/ops_python.py

@@ -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 BufferStorage, MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
+from tinygrad.device import MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.renderer import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -237,15 +237,14 @@ 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) -> BufferStorage:
-    return BufferStorage(buf:=memoryview(bytearray(size)), buf, MMIOInterface(mv_address(buf), size))
+  def _alloc(self, size:int, options) -> tuple: return (buf:=memoryview(bytearray(size)), buf), MMIOInterface(mv_address(buf), size)
 
   def _as_buffer(self, src) -> memoryview: return src
   def _copyin(self, dest, src:memoryview):
     with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
   def _copyout(self, dest:memoryview, src):
     with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
-  def map(self, buf:Buffer) -> BufferStorage: return BufferStorage(mv:=buf.as_memoryview(force_zero_copy=True), mv)
+  def map(self, buf:Buffer) -> tuple: return (mv:=buf.as_memoryview(force_zero_copy=True)), mv
   def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
 
 class PythonDevice(Compiled):

tinygrad/runtime/ops_qcom.py

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

tinygrad/runtime/ops_webgpu.py

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

tinygrad/runtime/support/hcq.py

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

tinygrad/runtime/support/hcq2.py

@@ -1,21 +1,24 @@
 from __future__ import annotations
-from typing import cast, Any
+from typing import cast, TypeVar, Generic, Any, TYPE_CHECKING
 import functools, time, itertools, decimal, weakref, statistics, ctypes, importlib
 from dataclasses import replace, dataclass, field
-from tinygrad.helpers import dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2
+from tinygrad.helpers import suppress_finalizing, dedup, pluralize, unwrap, PROFILE, VIZ, HCQ2, cpu_profile, mv_address
 from tinygrad.helpers import to_tuple, ContextVar, Context, panic, partition, perf_counter_us, DEV
 from tinygrad.device import Device, Buffer, BufferSpec, Compiled, Allocator, DepsTracker
 from tinygrad.device import ProfileGraphEntry, ProfileGraphEvent, ProfileDeviceEvent
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
+from tinygrad.uop.ops import Ops, sint, UOp, UPat, PatternMatcher, KernelInfo, GroupOp, graph_rewrite, rewrite_group, exec_alu
 from tinygrad.dtype import dtypes, DType, DTYPES_DICT, AddrSpace
 from tinygrad.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())
@@ -459,7 +462,7 @@ def resolve_getaddr(ctx:LinkCtx, g:UOp) -> UOp|None:
   buf, off = unwrap_view(g.src[0])
   if buf.op not in {Ops.BUFFER, Ops.MSELECT}: return None
   ctx.refs.append(buf) # add to refs
-  return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]) + off, dtypes.uint64)
+  return UOp.const(cast(Buffer, buf.buffer).get_buf(to_tuple(g.arg)[0]).va_addr + off, dtypes.uint64)
 
 def fold_binary(buf:UOp, blob:UOp) -> UOp:
   if getattr(b:=cast(Buffer, buf.buffer), '_hcq_written', None) is not blob.arg: # TODO: remove me
@@ -517,7 +520,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:Allocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
+  def __init__(self, device:str, allocator:HCQAllocator, compilers:list[type[Renderer]], runtime, can_recover:bool=False, arch=None):
     self.can_recover = can_recover
 
     self.pm_bufferize = PatternMatcher([
@@ -591,7 +594,42 @@ 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)

tinygrad/runtime/support/system.py

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

tinygrad/runtime/support/usb.py

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

tinygrad/schedule/__init__.py

@@ -121,7 +121,7 @@ def lower_sink_to_linear(call:UOp) -> UOp|None:
   function = call.src[0]
   if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None
   # value calls (with unbound outputs) are inlined positionally during prepare: their bodies are not programs to schedule
-  if call.has_unbound_outputs: return None
+  if call.is_value_call: return None
   st = time.perf_counter()
   cache_key = function.key
   if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None:
@@ -181,8 +181,14 @@ pm_copy_from_store = PatternMatcher([
   (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices),
 ])
 
-@rewrite_group(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
-def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]:
+@rewrite_group(lambda _,ret,**kwargs: f"Schedule {pluralize('Kernel', len(ret[0].src))}")
+def create_linear_with_vars(big_sink:UOp, buffer_bindings:dict[UOp, UOp]|None=None) -> tuple[UOp, dict[str, int]]:
+  # Only bind external declarations here. BUFFERs inside a body remain lexical schedule temporaries.
+  bindings = buffer_bindings if buffer_bindings is not None else {}
+  for arg in big_sink.src[1:]:
+    for b in arg.toposort(enter_calls=False):
+      if b.is_unbound and b not in bindings: bindings[b] = b.bind_buffer()
+  big_sink = big_sink.replace(src=(big_sink.src[0],)+tuple(a.substitute(bindings) for a in big_sink.src[1:]))
   # big_sink srcs are all the Tensors
   linear_call = graph_rewrite(big_sink, pm_schedule, name="schedule to linear", enter_calls=True)
 

tinygrad/schedule/multi.py

@@ -294,11 +294,11 @@ multi_pm = PatternMatcher([
     lambda multi,red: multi.src[0].allreduce(*red.arg).unshard(multi.arg, multi.src[1:])),
 
   # rewrite value-producing calls explicitly for UNSHARD
-  (UPat(Ops.CALL, name="call"), lambda call: rewrite_into_function(call) if call.has_unbound_outputs else None),
+  (UPat(Ops.CALL, name="call"), lambda call: rewrite_into_function(call) if call.is_value_call else None),
   (UPat((Ops.CALL, Ops.AFTER), src=(UPat(Ops.UNSHARD, name="multi"), ), name="root", allow_any_len=True), passthrough_multi),
   # just strip the UNSHARD from non-value-producing CALLs (custom kernels, etc.) — value-producing CALLs are handled by rewrite_into_function
   (UPat(Ops.CALL, dtype=dtypes.void, name="root", custom_early_reject=set([Ops.UNSHARD])), lambda root:
-    UOp(root.op, src=tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), arg=root.arg) if not root.has_unbound_outputs else None),
+    UOp(root.op, src=tuple(x.src[0] if x.op is Ops.UNSHARD else x for x in root.src), arg=root.arg) if not root.is_value_call else None),
   (UPat((Ops.CAST, Ops.BITCAST, Ops.CONTIGUOUS, Ops.DETACH, Ops.CONTIGUOUS_BACKWARD),
         src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), passthrough_multi),
   # STORE of a sharded value into an unsharded dest (e.g. a fragment into a full output tile)

tinygrad/schedule/prepare.py

@@ -113,7 +113,8 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
     if p.arg.size is not None:
       n, flat = flat_storage(a)
       if p.arg.size != n: raise TypeError(f"arg {i} shape mismatch: expected size {p.arg.size}, got {a.shape}")
-      dict_map[p] = flat
+      # Output PARAMs address storage, not padded values: padding a symbolic output view would put WHERE on a STORE destination.
+      dict_map[p] = a.storage_base.reshape((n,)) if p.arg.slot in (c.arg.output_pos or ()) else flat
     elif a.shape != ():
       raise TypeError(f"arg {i} shape mismatch: expected scalar, got {a.shape}")
     if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
@@ -133,10 +134,10 @@ def expand_bitcast(bc:UOp) -> UOp|None:
 
 earliest_rewrites = mop_cleanup+PatternMatcher([
   # resolve calls with RETURNED inputs (inline the body)
-  (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.has_unbound_outputs else None),
+  (UPat(Ops.CALL, name="c"), lambda c: resolve_function(c) if c.is_value_call else None),
 
   # resolve AFTER on RETURNED (call outputs)
-  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
+  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True, name="a"), resolve_returned_after),
 
   # resolve allreduce (must be bottom up)
   (UPat(Ops.ALLREDUCE, src=(UPat.var("buf"),), name="red"), create_allreduce_function),
@@ -220,4 +221,6 @@ def prepare_rangeify(sink:UOp) -> UOp:
   if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters")
   tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites")
   tsink = graph_rewrite(tsink, pm_copy_to_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store")
-  return tsink
+  # An effect-only body still produces buffer states. Root stores must participate in RAW/WAR scheduling
+  # just like stores already carried by AFTER; their destination and value are unchanged.
+  return tsink.replace(src=tuple(walk_mop(s.src[0]).after(s) if s.op is Ops.STORE else s for s in tsink.src))

tinygrad/tensor.py

@@ -12,7 +12,6 @@ from tinygrad.uop.ops import resolve_returned_after, remove_all_tags
 from tinygrad.uop.spec import type_verify, spec_tensor
 from tinygrad.mixin.rand import RandMixin
 from tinygrad.schedule import create_linear_with_vars
-from tinygrad.schedule.multi import multi_pm
 from tinygrad.device import Buffer, canonicalize_device
 from tinygrad.engine.realize import run_linear
 
@@ -21,14 +20,11 @@ from tinygrad.engine.realize import run_linear
 @dataclass
 class AllocCtx:
   buffer_map: dict[UOp, UOp] = field(default_factory=dict)
-  bases: set[UOp] = field(default_factory=set)
   stores: list[UOp] = field(default_factory=list)
   replacements: list[UOp] = field(default_factory=list)
   unbound: dict[UOp, UOp] = field(default_factory=dict)
   views: set[UOp] = field(default_factory=set)
-
-# a tag is the tuple of original pre-rewrite UOps a node provides storage for
-def tag_uop(x:UOp): return None if x.tag is not None else x.replace(tag=(x,))
+  outputs: set[UOp] = field(default_factory=set)
 
 # a base needs storage of its own if it can back a buffer and doesn't already have one
 def needs_storage(u:UOp) -> bool: return not u.is_virtual and not u.has_buffer_identity()
@@ -36,108 +32,29 @@ def needs_storage(u:UOp) -> bool: return not u.is_virtual and not u.has_buffer_i
 def on_disk(u:UOp): return isinstance(u.device, str) and u.device.startswith("DISK")
 def is_creation_device(u:UOp): return isinstance(u.device, str) and u.device.startswith(("DISK", "NPY", "PYTHON"))
 
-def creation_copy_is_realized(u:UOp):
-  # all copies from disk/numpy are realized into a real buffer
-  if is_creation_device(u.src[0]): return tag_uop(u)
-
-# 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"),
-   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),
-])
-
-def mint_tagged_storage(x:UOp):
-  if x.tag is None: return None          # untouched
-  # empty tag from rtag(()): a COPY already handled via buffer_map or merged into a parent AFTER.
-  # () is falsy but not None, so it isn't re-tagged like a bare (tag=None) node would be; just strip it here
-  if not x.tag: return x.rtag(None)
-  # a tagged CONTIGUOUS is consumed by the mint: the buffer stores its source directly
-  src = x.src[0] if x.op is Ops.CONTIGUOUS else x.rtag(None)
-  # virtual values and DISK tensors don't get real buffers: keep the (single) annotation, drop the tag
-  if x.is_virtual or on_disk(x): return src.alu(Ops.CONTIGUOUS)
-  # if size is 0, remove the contig
-  if 0 in x.shape: return src
-  buf = x.empty_like()
-  return buf.after(buf.store(src)).replace(tag=x.tag)
-
 def contiguous_mops_to_view(ctx:AllocCtx, c:UOp, src:UOp):
-  """MOPS(BUFFER) → SHRINK when movement ops collapse to a contiguous range."""
-  buf = src.base
-  while buf.op is Ops.BITCAST: buf = buf.src[0].base
-  # no symbolic shape
-  if buf.op not in {Ops.BUFFER, Ops.UNSHARD} or not all_int(c.shape): return None
-
-  # for UNSHARD tensors, use multi_pm to resolve per-shard movement ops, then view the resolved shard
-  unshard = None
-  if buf.op is Ops.UNSHARD:
-    if isinstance(c.device, str): return None
-    if (unshard := graph_rewrite(src, multi_pm, name="multi_buffer_view")).op is not Ops.UNSHARD: return None
-    src = unshard.src[0]
-
-  # offset the base buffer by the collapsed movement ops and view it
-  if (cv := src.contiguous_view()) is None or (buf := cv[0]).op is not Ops.BUFFER: return None
-  # NB: make offset a UOp.variable here to do the offset computation in the kernels
-  view = buf[cv[1]:cv[1] + src.max_numel() * src.element_size() // buf.element_size()].bitcast(src.dtype)
-  ctx.views.add(view)
-  if unshard is not None: return view.reshape(src.shape).unshard(unshard.arg, unshard.src[1:])
-  view = view.reshape(c.shape)
+  if (view:=src.buffer_view()) is None: return None
+  buf = view
+  while buf.op in {Ops.RESHAPE, Ops.UNSHARD}: buf = buf.src[0]
+  ctx.views.add(buf)
   return c.replace(src=(view,)+c.src[1:]) if c.op in {Ops.COPY, Ops.STORE} else view
 
 def transform_precompiled_call(c:UOp) -> UOp|None:
-  if c.arg is None or not c.arg.precompile or not c.has_unbound_outputs: return None
+  if not c.is_value_call or not c.arg.precompile: return None
   assert c.src[0].op is Ops.SINK, "precompiled call bodies are SINKs of stores into the output PARAMs"
-  # the RETURNED srcs are the call outputs (slots are src positions)
-  ret_pos = [p for p,a in enumerate(c.src[1:]) if a.unsharded_base.is_unbound]
-  srcs = tuple(st.src[1] for st in c.src[0].src if st.op is Ops.STORE)
-
-  # add the outputs to the call
-  outs = tuple(c.src[1+p].empty_like() for p in ret_pos)
-  targets = [o.param_like(p).shrink_to(s.shape) for p,o,s in zip(ret_pos, outs, srcs)]
-
-  # how each stored value lands in its output PARAM target: a CONTIGUOUS materializes straight into the target and
-  # a real buffer/UNSHARD rebinds its storage to the target (once per unique value); everything else is copied into it
-  placed:dict[UOp, UOp] = {}
-  items:list[UOp] = []
-  for s, t in zip(srcs, targets):
-    deps:list[UOp] = []
-    while s.op is Ops.AFTER:
-      deps.extend(s.src[1:])
-      s = s.src[0]
-    if s not in placed:
-      if s.op is Ops.CONTIGUOUS: placed[s] = t.after(t.store(s.src[0]))
-      elif s.op in {Ops.BUFFER, Ops.UNSHARD} and s.has_buffer_identity(): placed[s] = t
-      if s in placed:
-        items.append(s.after(*deps))
-        continue
-    items.append(t.after(t.store(s.after(*deps))))
-  # swap every placed value for its target storage, also inside other stores' AFTER deps
-  fxn = UOp.sink(*(x.substitute(placed) for x in items))
-
-  # all bodies are SINKs now, the node just becomes an opaque CALL: outs take the RETURNEDs' places; afters on real
-  # buffers are the input storage, afters on RETURNED placeholders have no storage yet, materialize them
-  rmap = dict(zip(ret_pos, outs))
-  new_call = c.replace(src=(fxn, *[rmap.get(i, a if a.has_buffer_identity(after_ok=True) else a.contiguous())
-                                   for i, a in enumerate(c.src[1:])]))
-  rets = tuple(o.after(new_call) for o in outs)
-
-  # if the CALL has symbolic shapes, shrink the max-sized output to the actual symbolic shape
-  # NOTE: must use the resolved shapes of the RETURNED placeholders (which substitute PARAMs with external args), not raw body shapes
-  rets = tuple(r.shrink_to(rs.shape) for r,rs in zip(rets, (c.src[1+p] for p in ret_pos)))
-
-  # the AFTER outputs resolve against this: stores of each real output into its RETURNED placeholder
-  return UOp.sink(*[c.src[1+p].store(v) for p, v in zip(ret_pos, rets)])
+  # The body already stores into the output PARAMs. Preserve it and its declared destinations.
+  ret_pos = c.arg.output_pos
+  new_call = c.replace(src=(c.src[0], *[a if i in ret_pos or a.has_buffer_identity(after_ok=True) else a.contiguous()
+                                      for i, a in enumerate(c.src[1:])]), arg=replace(c.arg, output_pos=None))
+  return UOp.sink(*(c.src[1+p].store(c.src[1+p].after(new_call)) for p in ret_pos))
 
 # NOTE: adding rules to here is bad. these all need to run before the schedule cache
 pm_early_transform_tensor_graph = PatternMatcher([
-  # transform precompiled value-producing calls into opaque CALLs (outputs become real buffers)
+  # lower precompiled value-producing calls into opaque CALLs using their declared output storage
   (UPat(Ops.CALL, name="c"), transform_precompiled_call),
 
   # resolve AFTER on RETURNED placeholders (for precompiled calls)
-  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True), resolve_returned_after),
+  (UPat(Ops.AFTER, src=(UPat(name="r"), UPat(Ops.SINK, name="t")), allow_any_len=True, name="a"), resolve_returned_after),
 
   # fold MOPS+BITCAST over BUFFER into SHRINK when movement ops collapse to contiguous range
   (UPat((Ops.COPY, Ops.CONTIGUOUS), src=(UPat(GroupOp.Movement|{Ops.BITCAST}, name="src"),), name="c"), contiguous_mops_to_view),
@@ -150,19 +67,12 @@ pm_early_transform_tensor_graph = PatternMatcher([
   (UPat(GroupOp.Movement-{Ops.SHRINK, Ops.RESHAPE}, name="x").f(Ops.COPY, name="copy"), lambda x,copy:
    x.replace(src=(copy.replace(src=(x.src[0],), tag=None),)+x.src[1:]) if on_disk(x) else None),
 
-  # strip DETACH/CONTIGUOUS_BACKWARD before minting (tags carry over)
-  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"),
-   lambda x: x.src[0].replace(tag=(x.src[0].tag or ())+(x.tag or ())) if x.tag else x.src[0]),
-  # contiguous of an already-materialized value is a no-op (tags carry over for held values)
-  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),), name="c"),
-   lambda a,c: a.replace(tag=(a.tag or ())+(c.tag or ())) if a.src[0].has_buffer_identity() else None),
-  # mint buffers for tagged values; an untagged CONTIGUOUS flows through to the scheduler, which bufferizes it
-  (UPat(GroupOp.All-{Ops.AFTER, Ops.STORE}, name="x"), mint_tagged_storage),
+  # strip graph-only wrappers
+  (UPat((Ops.DETACH, Ops.CONTIGUOUS_BACKWARD), name="x"), lambda x: x.src[0]),
+  # contiguous of an already-materialized value is a no-op
+  (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.AFTER, name="a"),)), lambda a: a if a.src[0].has_buffer_identity() else None),
 ])
 
-# a store's storage keeps the views and drops AFTERs (they only sequence stores)
-pm_drop_after = PatternMatcher([(UPat(Ops.AFTER, name="a"), lambda a: a.src[0])])
-
 def replace_input_buffer(ctx:AllocCtx, b:UOp):
   ctx.replacements.append(b)
   return b.param_like(len(ctx.replacements)-1)
@@ -170,6 +80,7 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp):
 # unbound BUFFERs get canonical scope-local id slots here so structurally identical calls hash identically for the
 # schedule cache (fresh slots are all positive from the global counter; negative slots are already canonical)
 def canonicalize_unbound_buffer(ctx:AllocCtx, b:UOp):
+  if b in ctx.outputs: return None
   if b.arg.slot >= 0 and b not in ctx.unbound: ctx.unbound[b] = b.replace(arg=replace(b.arg, slot=-1-len(ctx.unbound)))
   return ctx.unbound.get(b)
 
@@ -185,52 +96,53 @@ pm_canonicalize_unbound = PatternMatcher([
 pm_replace_buf = pm_canonicalize_unbound+PatternMatcher([
   # replace BUFFER with PARAM for cache key normalization (ALU addrspace buffers are Variables, they stay, and unbound BUFFERs too)
   (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
-   replace_input_buffer(ctx, b) if b.addrspace is AddrSpace.GLOBAL and not b.is_unbound else None),
+   replace_input_buffer(ctx, b) if b.addrspace is AddrSpace.GLOBAL and (not b.is_unbound or b in ctx.outputs) else None),
   # replace buffer views (SHRINK/BITCAST) with PARAM (only the views created by contiguous_mops_to_view)
   (UPat((Ops.SHRINK, Ops.BITCAST), name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b in ctx.views else None),
   # strip the stored value from bound Variables for cache key normalization, so different values hit same cache
   (UPat(Ops.AFTER, name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.is_bound_var else None),
 ])
 
+def is_persistent_effect(u:UOp, outputs:set[UOp]) -> bool:
+  if u.op is Ops.COPY: return on_disk(u)
+  return u.op is Ops.AFTER and not u.is_bound_var and (
+    not u.src[0].unsharded_base.is_unbound or u.src[0].storage_base in outputs or u.src[1].op is Ops.STORE or
+    (u.src[1].op is Ops.CALL and (not u.src[1].is_value_call or u.src[1].arg.precompile)))
+
 @rewrite_group(lambda _,ret: f"Callify {pluralize('Buffer', len(ret[1]))}")
 def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]:
   if VIZ: graph_rewrite(big_sink, PatternMatcher([]), name="View Tensor Graph")
   if SPEC: type_verify(big_sink, spec_tensor)
-  # bases to realize. an AFTER already names the storage its store writes into
-  ctx = AllocCtx(bases={base for x in big_sink.src if needs_storage(base:=x.base) and base.op is not Ops.AFTER})
-
-  # this rewrite is "read-only", it adds simple things to buffer_map and may sink things on big_sink, bottom_up
-  # this is the only one where we have to be careful to not break the tensor graph
-  big_sink = graph_rewrite(big_sink, add_tags, ctx=ctx, bottom_up=True, name="add tags")
-
-  # final outputs of value calls materialize with fresh storage
-  srcs:list[UOp] = []
-  for u in big_sink.src:
-    if u.op is Ops.AFTER and u.src[0].unsharded_base.is_unbound and u.src[1].op is Ops.CALL:
-      # precompiled calls don't need this: transform_precompiled_call gives their outputs real buffers
-      call = u.src[1]
-      if not (call.arg is not None and call.arg.precompile):
-        buf = u.empty_like()
-        u = buf.after(buf.store(u.rtag(None))).replace(tag=u.tag)
-    srcs.append(u)
-  big_sink = big_sink.replace(src=tuple(srcs))
-
-  # here we can break the tensor graph. tags propagate through replaces so we can still find the original UOps
+  # Escaping declarations become parameters of the schedule, not scope-local temporaries.
+  ctx = AllocCtx(outputs={x.storage_base for x in big_sink.src if not x.is_virtual})
+  # Tensor replacements name the original destinations, independently of how their effects are lowered.
+  for u in big_sink.toposort(enter_calls=False):
+    if u.op is Ops.AFTER and is_persistent_effect(u, ctx.outputs):
+      ctx.buffer_map[u] = u.src[0].storage_view
   big_sink = graph_rewrite(big_sink, pm_early_transform_tensor_graph, ctx=ctx, name="early transform tensor graph")
 
-  # collect the stores (never entering call bodies) and map tagged AFTERs to their storage; tags are stripped at the end
-  # copies to disk are stores to the disk buffer; bound Variables are call inputs and RETURNEDs are call outputs
-  # AFTERs on unbound STORAGE (clones) are collected too: the clone's own buffer is the storage, no fresh copy
+  # Collect effects without entering call bodies. Escaping declarations become schedule parameters.
   for u in big_sink.toposort(enter_calls=False):
-    if (u.op is Ops.COPY and on_disk(u)) or (u.op is Ops.AFTER and not u.is_bound_var and
-        (not u.src[0].unsharded_base.is_unbound or u.src[1].op is Ops.STORE)):
+    if is_persistent_effect(u, ctx.outputs):
       ctx.stores.append(u)
-      if u.tag: ctx.buffer_map.update({t:graph_rewrite(u.src[0], pm_drop_after).shrink_to(t.shape) for t in u.tag})
-  ret = graph_rewrite(UOp.sink(*ctx.stores), pm_replace_buf+remove_all_tags, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements)
+      if u.op is Ops.AFTER: ctx.outputs.add(u.src[0].storage_base)
+  body = graph_rewrite(UOp.sink(*ctx.stores), pm_replace_buf+remove_all_tags, ctx=ctx, bottom_up=True, name="replace bufs")
+  # An empty declaration may have no effects but still needs a binding when it escapes.
+  for b in ctx.outputs:
+    if b.is_unbound and b not in ctx.replacements: ctx.replacements.append(b)
+  ret = body.call(*ctx.replacements)
   assert not any(x in ctx.buffer_map for x in ctx.buffer_map.values())
   if VIZ: graph_rewrite(ret, PatternMatcher([]), name="View Call")
   return ret, ctx.buffer_map
 
+def outputs_to_call(*xs:UOp) -> tuple[UOp, dict[UOp, UOp]]:
+  # Build output requests, not a preparation pass over the graph. Intermediate storage is already declared.
+  memo:dict[UOp, UOp] = {}
+  outputs = {x.base:x.base.materialize(memo) for x in xs}
+  big_sink, becomes_map = transform_to_call(UOp.sink(*outputs.values()))
+  becomes_map.update({x:y.substitute(becomes_map) for x,y in outputs.items() if x is not y})
+  return big_sink, becomes_map
+
 # *** all in scope Tensors are here. this gets relevant UOps ***
 
 all_tensors: dict[weakref.ref[Tensor], None] = {}
@@ -312,7 +224,7 @@ class Tensor(RandMixin):
     if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
 
     # data might be on a different device
-    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device)
+    self.uop:UOp = data if data.device is None or data.device == _device else data.copy_to_device(_device).clone()
     # cast on the target device, the source may not hold the dtype (numpy has no fp8/bfloat16) or be able to compute it (DISK)
     if _dtype is not None: self.uop = self.uop.cast(_dtype)
 
@@ -386,8 +298,7 @@ class Tensor(RandMixin):
     return [Tensor(u) for u in UOp.custom_kernel(*[t.uop for t in (self,)+lst], fxn=fxn, grad_fxn=grad_fxn)]
 
   def callify(self, *lst:Tensor) -> Tensor:
-    big_sink = UOp.sink(*[x.uop for x in (self,)+lst])
-    big_sink, buffer_map = transform_to_call(big_sink)
+    big_sink, buffer_map = outputs_to_call(*[x.uop for x in (self,)+lst])
     _apply_map_to_tensors({x:y.after(big_sink) for x,y in buffer_map.items()}, name="callify")
     return self
 
@@ -396,9 +307,11 @@ class Tensor(RandMixin):
     # weakness ends where storage begins
     if any(t.dtype in dtypes.weaks and t.uop.device is not None for t in (self,)+lst):
       raise RuntimeError("cannot realize a weak dtype; cast to a concrete dtype first")
-    big_sink, becomes_map = transform_to_call(UOp.sink(*[x.uop for x in (self,)+lst]))
-    _apply_map_to_tensors(becomes_map, name="buffers")
-    return create_linear_with_vars(big_sink)
+    big_sink, becomes_map = outputs_to_call(*[x.uop for x in (self,)+lst])
+    bindings:dict[UOp, UOp] = {}
+    ret = create_linear_with_vars(big_sink, buffer_bindings=bindings)
+    _apply_map_to_tensors({**bindings, **{x:y.substitute(bindings) for x,y in becomes_map.items()}}, name="buffers")
+    return ret
 
   def schedule_linear(self, *lst:Tensor) -> UOp:
     """Creates the schedule needed to realize these Tensor(s)."""
@@ -455,7 +368,6 @@ class Tensor(RandMixin):
       # 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")
     else:
       # simple assign
@@ -545,9 +457,8 @@ class Tensor(RandMixin):
     """
     if self.uop.device is None: return self
     if (device:=canonicalize_device(device)) == self.device: return self
-    # a copy to disk wants to persist, so it inserts a clone: the disk buffer is the storage of the copied value
-    if isinstance(device, str) and device.startswith("DISK"): ret = Tensor(self.uop.clone(device))
-    else: ret = Tensor(self.uop.copy_to_device(device))
+    # The transfer owns its destination from construction; COPY itself only describes the transfer.
+    ret = Tensor(self.uop.copy_to_device(device).clone())
     if self.grad is not None: ret.grad = self.grad.to(device)
     return ret.is_param_(self.is_param)
 

tinygrad/uop/ops.py

@@ -537,14 +537,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     if len(srcs) == 1 and isinstance(srcs[0], UOp): return srcs[0]
     return UOp(Ops.GROUP, src=tuple([x for x in srcs if x is not None]), **kwargs)
   @property
+  def is_value_call(self) -> bool:
+    return self.op is Ops.CALL and self.arg is not None and self.arg.output_pos is not None
+  @property
+  def call_outputs(self) -> tuple[UOp, ...]:
+    assert self.is_value_call
+    return tuple(self.src[1+i].after(self) for i in self.arg.output_pos)
+  @property
   def has_unbound_outputs(self) -> bool:
-    """does this call still have unresolved outputs: unbound BUFFERs among its inputs (minted by call_with_outputs,
-    resolved when the call is inlined or the outputs are materialized). a lifecycle query, not a call type"""
-    return self.op is Ops.CALL and any(x.unsharded_base.is_unbound for x in self.src[1:])
+    """Whether any declared value-call outputs still lack backing storage (not a call-kind query)."""
+    return self.is_value_call and any(self.src[1+i].unsharded_base.is_unbound for i in self.arg.output_pos)
   @property
   def unbound_outputs(self) -> tuple[UOp, ...]:
     """the unresolved outputs of this call: an AFTER on each unbound BUFFER input, usable like a normal buffer"""
-    return tuple(x.after(self) for x in self.src[1:] if x.unsharded_base.is_unbound)
+    return tuple(x for x in self.call_outputs if x.src[0].unsharded_base.is_unbound) if self.is_value_call else ()
   def index(self, *srcs:UOp|int|None, **kwargs):
     new_srcs: list[UOp] = [UOp.const(x) if isinstance(x, int) else x for x in srcs if x is not None]
     if len(new_srcs) == 1 and new_srcs[0].op is Ops.CONST and self.op is Ops.STACK: return self.src[new_srcs[0].val]
@@ -764,6 +770,15 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     while b.op in {Ops.BITCAST, Ops.AFTER, Ops.UNSHARD}: b = b.src[0].unsharded_base
     return b
 
+  @property
+  def storage_view(self) -> UOp:
+    """The addressed view without storage-state dependencies. Shape expressions retain their bindings."""
+    if self.op is Ops.AFTER: return self.src[0].storage_view
+    if self.op in GroupOp.Movement|{Ops.BITCAST, Ops.DETACH, Ops.UNSHARD, Ops.MSELECT}:
+      return self.replace(src=(self.src[0].storage_view,)+self.src[1:])
+    if self.op is Ops.MSTACK: return self.replace(src=tuple(s.storage_view for s in self.src))
+    return self
+
   # cached property here makes external_uop_gc fail, why?
   @property
   def as_shape(self) -> tuple[sint, ...]:
@@ -812,6 +827,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     slot = next(UOp.unique_num) if num is None else num
     buf = MultiBuffer(device, size, dtype) if isinstance(device, tuple) else Buffer(device, size, dtype)
     return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=size, device=device, buffer=buf))
+  def bind_buffer(self) -> UOp:
+    """Attach backing storage to an existing declaration without minting a new storage slot."""
+    assert self.is_unbound and not self.is_virtual
+    buf = MultiBuffer(self.device, self.max_numel(), self.dtype) if isinstance(self.device, tuple) else \
+          Buffer(self.device, self.max_numel(), self.dtype)
+    return self.replace(arg=replace(self.arg, buffer=buf))
   @staticmethod
   def from_buffer(opaque:Buffer, device:str|tuple[str, ...]|None=None):
     # the opaque Buffer goes straight in the arg: the ucache dedups because the arg (and thus the Buffer) is part of the key
@@ -833,7 +854,27 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
       data = struct.pack(f"{prod(shape)}{bdtype.fmt}", *[truncate[bdtype](bdtype.const(xi)) for xi in fully_flatten(x)])
     ret.buffer.allocate(memoryview(bytearray(data))) # fake realize. buffer storage must be writable, and bytes isn't
     if ret.dtype != dtype: ret = ret.cast(dtype)
-    return ret if ret.device == device else ret.copy_to_device(device)
+    return ret if ret.device == device else ret.clone(device)
+  def materialize(self, memo:dict[UOp, UOp]|None=None) -> UOp:
+    """Build an explicit output request. Share destinations within a multi-output request."""
+    if memo is None: memo = {}
+    if self not in memo: memo[self] = self._materialize(memo)
+    return memo[self]
+
+  def _materialize(self, memo:dict[UOp, UOp]) -> UOp:
+    if self.is_virtual or (isinstance(self.device, str) and self.device.startswith("DISK")): return self
+    if self.op in {Ops.DETACH, Ops.CONTIGUOUS_BACKWARD}: return self.src[0].materialize(memo)
+    if self.op is Ops.AFTER or self.storage_base.op in {Ops.BUFFER, Ops.PARAM}: return self
+    if self.op in GroupOp.Movement:
+      return self.replace(src=(self.src[0].materialize(memo),)+self.src[1:])
+    if self.op is Ops.CONTIGUOUS:
+      src = self.src[0]
+      while src.op in {Ops.DETACH, Ops.CONTIGUOUS_BACKWARD}: src = src.src[0]
+      if src.op is Ops.CONTIGUOUS: return src.materialize(memo)
+      if src.has_buffer_identity(after_ok=True, unbound_ok=True): return src
+      if (view:=src.buffer_view()) is not None: return view
+      return src.clone() if src.op in GroupOp.Movement|{Ops.AFTER, Ops.BITCAST} else src.materialize(memo)
+    return self.clone()
   def clone(self, device=None) -> UOp:
     device = device or self.device
     ret = self.empty_like(device=device)
@@ -883,6 +924,22 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     while len(s.src) and s.op not in {Ops.BUFFER, Ops.PARAM, Ops.STAGE, Ops.MSTACK}: s = s.src[0]
     return s
 
+  def buffer_view(self) -> UOp|None:
+    """Construct a zero-copy view when movement/bitcast operations describe a contiguous buffer range."""
+    if not all_int(self.shape): return None
+    src, buf = self, self.base
+    while buf.op is Ops.BITCAST: buf = buf.src[0].base
+    if buf.op not in {Ops.BUFFER, Ops.UNSHARD}: return None
+    unshard = None
+    if buf.op is Ops.UNSHARD:
+      from tinygrad.schedule.multi import multi_pm
+      if isinstance(self.device, str): return None
+      if (unshard := graph_rewrite(src, multi_pm, name="multi buffer view")).op is not Ops.UNSHARD: return None
+      src = unshard.src[0]
+    if (cv := src.contiguous_view()) is None or (buf := cv[0]).op is not Ops.BUFFER: return None
+    view = buf[cv[1]:cv[1] + src.max_numel() * src.element_size() // buf.element_size()].bitcast(src.dtype).reshape(src.shape)
+    return view.unshard(unshard.arg, unshard.src[1:]) if unshard is not None else view.reshape(self.shape)
+
   def contiguous_view(self) -> tuple[UOp, int]|None:
     from tinygrad.schedule.prepare import pm_mops
     from tinygrad.uop.symbolic import symbolic
@@ -901,12 +958,12 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
   def contiguous_view_offset(self) -> int|None: return None if (view := self.contiguous_view()) is None else view[1]
 
-  def has_buffer_identity(self, after_ok=False):
-    """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/UNSHARD -> BUFFER chain)."""
+  def has_buffer_identity(self, after_ok=False, unbound_ok=False):
+    """Check for storage through shape wrappers; unbound_ok also accepts declarations without backing buffers."""
     # TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer
-    if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok)
-    if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok)
-    return self.op in {Ops.BUFFER, Ops.PARAM} and not self.is_unbound
+    if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok, unbound_ok)
+    if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok, unbound_ok)
+    return self.op in {Ops.BUFFER, Ops.PARAM} and (unbound_ok or not self.is_unbound)
   @property
   def is_unbound(self) -> bool:
     # an unbound GLOBAL BUFFER has no storage bound yet: it's a declaration of storage (call output, scheduler temp)
@@ -1209,8 +1266,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     output_pos gives the position of each output in the arg list (default: a block after the inputs), the inputs take
     the remaining positions in order; when it's given, input params must already be slotted at their final positions.
     output_pos must be strictly ascending: the body's stores and the call args pair positionally by values order"""
+    # Precompiled outputs are storage, including otherwise virtual constant results.
+    if precompile: values = tuple(v.cast(v.commit_dtype()) for v in values)
     # the device defaults to the first device in the values or args, like srcs-based device resolution
     default_dev = next((x.device for x in itertools.chain(values, srcs) if x.device is not None), None)
+    if precompile and default_dev is None: default_dev = canonicalize_device(None)
     pos = tuple(range(len(srcs), len(srcs)+len(values))) if output_pos is None else output_pos
     assert len(pos) == len(values) and len(set(pos)) == len(pos), "output_pos must be one distinct position per output"
     assert all(a < b for a, b in zip(pos, pos[1:])), f"output_pos {output_pos} must be strictly ascending"
@@ -1241,6 +1301,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     it = iter(srcs)
     call = body.call(*[r if r is not None else next(it) for r in args], grad_fxn=grad_fxn, name=name, precompile=precompile,
                      precompile_backward=precompile_backward, aux=aux)
+    call = call.replace(arg=replace(call.arg, output_pos=pos))
     return tuple(r.after(call) for r in rets)
 
   # one-line convenience for the single-output case: self is the value
@@ -1323,12 +1384,14 @@ class CallInfo:
   precompile_backward: bool = False
   aux: Any = None
   dtype: DType = dtypes.void
+  # None for opaque calls; value-call outputs are positional, independent of their backing-buffer bindings.
+  output_pos: tuple[int, ...]|None = None
   # grad_fxn can't be pickled
-  def __reduce__(self): return (CallInfo, (None, self.name, self.precompile, self.precompile_backward, self.aux, self.dtype))
+  def __reduce__(self): return (CallInfo, (None, self.name, self.precompile, self.precompile_backward, self.aux, self.dtype, self.output_pos))
   def __repr__(self):
     gf = id(self.grad_fxn) if self.grad_fxn else None
     return f"CallInfo({gf}, {repr(self.name)}, {self.precompile}, {self.precompile_backward})" + \
-      (f", {self.dtype}" if self.dtype is not dtypes.void else "")
+      (f", {self.dtype}" if self.dtype is not dtypes.void else "") + (f", output_pos={self.output_pos}" if self.output_pos is not None else "")
 
 # ******** ops in python ********
 
@@ -1807,12 +1870,16 @@ def to_max_shape(shape:tuple[sint, ...]) -> tuple[int, ...]: return tuple(int(x.
 _substitute = PatternMatcher([(UPat(tuple(Ops), name="x"), lambda ctx,x: ctx.get(x,None))])
 _pm_resolve_params = PatternMatcher([(UPat(Ops.PARAM, name="p"), lambda ctx,p: ctx[p.arg.slot])])
 
-def resolve_returned_after(r:UOp, t:UOp) -> UOp|None:
+def resolve_returned_after(a:UOp, r:UOp, t:UOp) -> UOp|None:
   """AFTER on a RETURNED placeholder extracts the call output value: the value of its matching store in a SINK body
   (called from patterns that bind t to a SINK)"""
-  vals = [st.src[1] for st in t.src if st.op is Ops.STORE and st.src[0].unsharded_base is r.unsharded_base] \
-    if r.unsharded_base.is_unbound else []
-  return vals[0] if len(vals) == 1 else None
+  stores = [st for st in t.src if st.op is Ops.STORE and st.src[0].unsharded_base is r.unsharded_base]
+  if len(stores) != 1: return None
+  # Unbound, scope-local outputs are values and can fuse. Escaping outputs have been scoped as PARAMs:
+  # keep their STORE instead of extracting its value and losing the declared destination.
+  val = stores[0].src[1]
+  ret = val if r.unsharded_base.is_unbound or (val.op is Ops.AFTER and val.src[0] is r) else r.after(stores[0])
+  return ret.replace(tag=(ret.tag or ()) + a.tag) if a.tag else ret
 remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 def gate_kernel_sink(x:UOp) -> bool:

tinygrad/uop/render.py

@@ -1,4 +1,4 @@
-from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.dtype import dtypes
 from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.ops import ParamArg, UOp, PatternMatcher, UPat, multirange_str, range_str, consumer_map_from_toposort, sint
 from tinygrad.helpers import strip_parens
@@ -84,9 +84,6 @@ pm_pyrender_extra = PatternMatcher([
   (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val})"),
   (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
   (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)})"),
-  (UPat(Ops.BUFFER, src=(), name="x"), lambda x:
-    f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
-    if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
   (UPat(Ops.COPY, src=(UPat(name="x"),), name="copy"), lambda ctx,x,copy: f"{ctx[x]}.copy_to_device({repr(copy.arg)})"),
   (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda ctx,x: f"UOp(Ops.CUSTOM_FUNCTION, src={srcs(ctx, x.src)}, arg={x.arg!r})"),
   (UPat(Ops.REDUCE, name="r"), lambda ctx,r: f"{ctx[r.src[0]]}._rop({r.arg[0]}, {tuple(range(r.arg[1]))})" if r.arg[1] else None),

tinygrad/uop/spec.py

@@ -135,7 +135,7 @@ spec_tensor = PatternMatcher([
 
   # BUFFER
   (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf:
-   True if buf.is_unbound else (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device))
+   True if buf.is_unbound else (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, (int, type(None))) and is_device(buf.arg.device))
    if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),
 
   # a Variable is a 0-d ALU BUFFER with a value range and no device