something

.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:
@@ -158,7 +154,7 @@ 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
@@ -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"
@@ -238,12 +230,9 @@ runs:
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
     - name: Add clang to PATH (Linux)
-      if: runner.os == 'Linux'
+      if: inputs.llvm == 'true' && runner.os == 'Linux'
       shell: bash
-      run: |
-        if [ -d /usr/lib/llvm-20/bin ]; then
-          echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
-        fi
+      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
 
     # **** AMD ****
     - name: Setup AMD (Linux)

.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

@@ -536,7 +536,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 14
+    timeout-minutes: 10
     defaults:
       run:
         shell: bash -e -o pipefail {0}

.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/platform.yml

@@ -70,10 +70,9 @@ jobs:
     - name: Run pytest (amd)
       env:
         DEV: MOCKKFD+AMD
-        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
     - name: Run pytest (ptx)
       env:
         DEV: "MOCK+NV:PTX"

.github/workflows/test.yml

@@ -253,7 +253,7 @@ jobs:
         deps: testing_unit
         llvm: 'true'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -478,7 +478,6 @@ jobs:
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
-      HCQ_RUNTIME_DEV: PYTHON
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -505,7 +504,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: SQTT_BUFFER_SIZE=16 VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0

docs/developer/developer.md

@@ -40,3 +40,7 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to
 Runtimes are responsible for device-specific interactions. They handle tasks such as initializing devices, allocating memory, loading/launching programs, and more. You can find more information about the runtimes API on the [runtime overview page](runtime.md).
 
 All runtime implementations can be found in the [runtime directory](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime).
+
+### HCQ Compatible Runtimes
+
+HCQ API is a lower-level API for defining runtimes. Interaction with HCQ-compatible devices occurs at a lower level, with commands issued directly to hardware queues. Some examples of such backends are [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) and [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py), which are userspace drivers for NVIDIA and AMD devices respectively. You can find more information about the API on [HCQ overview page](hcq.md)

docs/developer/hcq.md

@@ -0,0 +1,128 @@
+# HCQ Compatible Runtime
+
+## Overview
+
+The main aspect of HCQ-compatible runtimes is how they interact with devices. In HCQ, all interactions with devices occur in a hardware-friendly manner using [command queues](#command-queues). This approach allows commands to be issued directly to devices, bypassing runtime overhead such as HIP or CUDA. Additionally, by using the HCQ API, these runtimes can benefit from various optimizations and features, including [HCQGraph](#hcqgraph) and built-in profiling capabilities.
+
+### Command Queues
+
+To interact with devices you create a `HWQueue`. Some methods are required, like timestamp and synchronization methods like [signal](#tinygrad.runtime.support.hcq.HWQueue.signal) and [wait](#tinygrad.runtime.support.hcq.HWQueue.wait), while others are dependent on it being a compute or copy queue.
+
+For example, the following Python code enqueues a wait, execute, and signal command on the HCQ-compatible device:
+```python
+HWQueue().wait(signal_to_wait, value_to_wait) \
+         .exec(program, args_state, global_dims, local_dims) \
+         .signal(signal_to_fire, value_to_fire) \
+         .submit(your_device)
+```
+
+Each runtime should implement the required functions that are defined in the `HWQueue` classes.
+
+::: tinygrad.runtime.support.hcq.HWQueue
+    options:
+        members: [
+            "signal",
+            "wait",
+            "timestamp",
+            "bind",
+            "submit",
+            "memory_barrier",
+            "exec",
+            "copy",
+        ]
+        show_source: false
+
+### HCQ Compatible Device
+
+The `HCQCompiled` class defines the API for HCQ-compatible devices. This class serves as an abstract base class that device-specific implementations should inherit from and implement.
+
+::: tinygrad.runtime.support.hcq.HCQCompiled
+    options:
+        show_source: false
+
+#### Signals
+
+Signals are device-dependent structures used for synchronization and timing in HCQ-compatible devices. They should be designed to record both a `value` and a `timestamp` within the same signal. HCQ-compatible backend implementations should use `HCQSignal` as a base class.
+
+::: tinygrad.runtime.support.hcq.HCQSignal
+    options:
+        members: [value, timestamp, wait]
+        show_source: false
+
+The following Python code demonstrates the usage of signals:
+
+```python
+signal = your_device.new_signal(value=0)
+
+HWQueue().timestamp(signal) \
+         .signal(signal, value_to_fire) \
+         .submit(your_device)
+
+signal.wait(value_to_fire)
+signaled_value = signal.value # should be the same as `value_to_fire`
+timestamp = signal.timestamp
+```
+
+##### Synchronization signals
+
+Each HCQ-compatible device must allocate two signals for global synchronization purposes. These signals are passed to the `HCQCompiled` base class during initialization: an active timeline signal `self.timeline_signal` and a shadow timeline signal `self._shadow_timeline_signal` which helps to handle signal value overflow issues. You can find more about synchronization in the [synchronization section](#synchronization)
+
+### HCQ Compatible Allocator
+
+The `HCQAllocator` base class simplifies allocator logic by leveraging [command queues](#command-queues) abstractions. This class efficiently handles copy and transfer operations, leaving only the alloc and free functions to be implemented by individual backends.
+
+::: tinygrad.runtime.support.hcq.HCQAllocator
+    options:
+        members: [
+            "_alloc",
+            "_free",
+        ]
+        show_source: false
+
+#### HCQ Allocator Result Protocol
+
+Backends must adhere to the `HCQBuffer` protocol when returning allocation results.
+
+::: tinygrad.runtime.support.hcq.HCQBuffer
+    options:
+        members: true
+        show_source: false
+
+### HCQ Compatible Program
+
+`HCQProgram` is a base class for defining programs compatible with HCQ-enabled devices. It provides a flexible framework for handling different argument layouts (see `HCQArgsState`).
+
+::: tinygrad.runtime.support.hcq.HCQProgram
+    options:
+        members: true
+        show_source: false
+
+#### Arguments State
+
+`HCQArgsState` is a base class for managing the argument state for HCQ programs. Backend implementations should create a subclass of `HCQArgsState` to manage arguments for the given program.
+
+::: tinygrad.runtime.support.hcq.HCQArgsState
+    options:
+        members: true
+        show_source: false
+
+**Lifetime**: The `HCQArgsState` is passed to `HWQueue.exec` and is guaranteed not to be freed until `HWQueue.submit` for the same queue is called.
+
+### Synchronization
+
+HCQ-compatible devices use a global timeline signal for synchronizing all operations. This mechanism ensures proper ordering and completion of tasks across the device. By convention, `self.timeline_value` points to the next value to signal. So, to wait for all previous operations on the device to complete, wait for `self.timeline_value - 1` value. The following Python code demonstrates the typical usage of signals to synchronize execution to other operations on the device:
+
+```python
+HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
+         .exec(...)
+         .signal(your_device.timeline_signal, your_device.next_timeline()) \
+         .submit(your_device)
+
+# Optionally wait for execution
+your_device.timeline_signal.wait(your_device.timeline_value - 1)
+```
+
+## HCQGraph
+
+[HCQGraph](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/graph/hcq.py) is a core feature that implements `GraphRunner` for HCQ-compatible devices. `HCQGraph` builds static `HWQueue` for all operations per device. To optimize enqueue time, only the necessary parts of the queues are updated for each run using the symbolic variables, avoiding a complete rebuild.
+Optionally, queues can implement a `bind` API, which allows further optimization by eliminating the need to copy the queues into the device ring.

docs/developer/runtime.md

@@ -22,13 +22,18 @@ The `Compiled` class is responsible for initializing and managing a device.
 
 ### Allocator
 
-The `Allocator` class manages memory on the device and caches allocated buffers for reuse.
+The `Allocator` class is responsible for managing memory on the device. There is also a version called the `LRUAllocator`, which caches allocated buffers to optimize performance.
 
 ::: tinygrad.device.Allocator
     options:
         members: true
         show_source: false
 
+::: tinygrad.device.LRUAllocator
+    options:
+        members: true
+        show_source: false
+
 ### Program
 
 The `Program` class is created for each loaded program. It is responsible for executing the program on the device. As an example, here is a `CPUProgram` implementation which loads program and runs it.

examples/compile_tensorflow.py

@@ -97,3 +97,4 @@ if __name__ == "__main__":
   tf_output = keras_model(test_input).numpy()[0]
   print("keras:   ", tf_output, file=sys.stderr)
   np.testing.assert_allclose(tf_output, test_output, atol=1e-5, rtol=1e-5)
+

examples/gpt2.py

@@ -57,7 +57,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
     h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).clone()
+    return (h + self.mlp(self.ln_2(h))).contiguous()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

extra/dsp/invoke_bug.py

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

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)[0][0]
-b = hipallocator.alloc(N*N*2)[0][0]
-c = hipallocator.alloc(N*N*2)[0][0]
+a = hipallocator.alloc(N*N*4)
+b = hipallocator.alloc(N*N*2)
+c = hipallocator.alloc(N*N*2)
 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/hcq1/hcq.py

@@ -1,549 +0,0 @@
-from __future__ import annotations
-from typing import cast, Callable, Type, TypeVar, Generic, Any
-import contextlib, decimal, statistics, time, ctypes, array, collections, itertools
-from tinygrad.helpers import PROFILE, getenv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
-from tinygrad.helpers import suppress_finalizing, TracingKey
-from tinygrad.device import BufferSpec, Compiled, Allocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
-from tinygrad.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 HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
-  """
-  A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
-  """
-
-  def __init__(self):
-    self._q:Any = []
-    self.binded_device:HCQDeviceType|None = None
-    self.q_sints:list[tuple[int, int]] = []
-    self.mv_sints:list[tuple[MMIOInterface, int, int, int|None]] = []
-    self.syms:list[sint] = []
-    self._prev_resolved_syms:list[int|None] = []
-
-  def _new_sym(self, sym:sint) -> int:
-    if sym not in self.syms:
-      self.syms.append(sym)
-      self._prev_resolved_syms.append(None)
-    return self.syms.index(sym)
-
-  def q(self, *values):
-    """
-    Enqueues values in the queue.
-
-    Args:
-      values: The values to enqueue in the queue.
-    """
-
-    for v in values:
-      if isinstance(v, UOp):
-        self.q_sints.append((len(self._q), self._new_sym(v)))
-        self._q.append(0xbadc0ded)
-      else: self._q.append(v)
-
-  # *** common commands  ***
-
-  def timestamp(self, signal:SignalType):
-    """
-    Enqueues a timestamp command which records the current time in a signal after all previously enqueued commands are completed.
-
-    Args:
-      signal: The signal to store the timestamp
-    """
-
-  def signal(self, signal:SignalType, value:sint):
-    """
-    Enqueues a signal command which sets the signal to the given value, ensuring all previous operations are completed.
-
-    Args:
-      signal: The signal to set
-      value: The value to set the signal to
-    """
-
-  def wait(self, signal:SignalType, value:sint):
-    """
-    Enqueues a wait command which halts execution until the signal is greater than or equal to a specific value.
-
-    Args:
-      signal: The signal to wait on
-      value: The value to wait for
-    """
-
-  # *** commands for compute queues ***
-
-  def memory_barrier(self):
-    """
-    Enqueues a memory barrier command to ensure memory coherence between agents. Only on compute queues.
-    """
-
-  def exec(self, prg:ProgramType, args_state:ArgsStateType, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
-    """
-    Enqueues an execution command for a kernel program. Only on compute queues.
-
-    Args:
-      prg: The program to execute
-      args_state: The args state to execute program with
-      global_size: The global work size
-      local_size: The local work size
-    """
-
-  def write(self, b:HCQBuffer, val:sint, b64:bool=False):
-    """
-    Enqueues a command to write a value to a buffer address after all previously enqueued commands are completed.
-
-    Args:
-      b: The buffer to write to
-      val: The value to write
-      b64: If True, write a 64-bit value; otherwise write 32-bit
-    """
-    raise NotImplementedError("write not implemented")
-
-  def poll_bit(self, b:HCQBuffer, val:sint, mask:int):
-    """
-    Enqueues a poll command which halts execution until (mem[b] & mask) == val.
-    val must be 0 or mask (i.e. checks if masked bits are all clear or all set).
-
-    Args:
-      b: The buffer to poll
-      val: The expected value after masking (0 or mask)
-      mask: The bit mask to test
-    """
-    raise NotImplementedError("poll_bit not implemented")
-
-  # *** commands for copy queues ***
-
-  def copy(self, dest:HCQBuffer, src:HCQBuffer, copy_size:int):
-    """
-    Enqueues a copy command to transfer data. Only on copy queues.
-
-    Args:
-      dest: The destination buffer of the copy
-      src: The source buffer of the copy
-      copy_size: The size of data to copy
-    """
-
-  # *** submit and bind commands  ***
-
-  def bind(self, dev:HCQDeviceType):
-    """
-    Associates the queue with a specific device for optimized execution.
-
-    This optional method allows backend implementations to tailor the queue for efficient use on the given device. When implemented, it can eliminate
-    the need to copy queues into the device, thereby enhancing performance.
-
-    Args:
-      dev: The target device for queue optimization.
-
-    Note:
-      Implementing this method is optional but recommended for performance gains.
-    """
-
-  def bind_args_state(self, args_state:ArgsStateType):
-    for vals, mem, fmt in args_state.bind_data: self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt)
-
-  def bind_sints(self, *vals:sint, mem:MMIOInterface, struct_t:Type[ctypes.Structure], start_field:str, fmt, mask:int|None=None):
-    self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt, mask=mask, offset=getattr(struct_t, start_field).offset)
-
-  def bind_sints_to_mem(self, *vals:sint, mem:MMIOInterface, fmt, mask:int|None=None, offset:int=0):
-    mv = mem.view(offset=offset, size=len(vals)*8, fmt=fmt)
-    for i, val in enumerate(vals):
-      if isinstance(val, int): mv[i] = val if mask is None else ((mv[i] & ~mask) | val)
-      else: self.mv_sints.append((mv, i, self._new_sym(val), mask))
-
-  def _apply_var_vals(self, var_vals:dict[str, int]):
-    resolved_syms: list[int|None] = [sym_infer(sym, var_vals) for sym in self.syms]
-
-    for off, sym_idx in self.q_sints:
-      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
-      self._q[off] = resolved_syms[sym_idx]
-
-    for mv, off, sym_idx, mask in self.mv_sints:
-      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
-      mv[off] = resolved_syms[sym_idx] if mask is None else ((mv[off] & ~mask) | resolved_syms[sym_idx])
-
-    self._prev_resolved_syms = resolved_syms
-
-  def submit(self, dev:HCQDeviceType, var_vals:dict[str, int]|None=None):
-    """
-    Submits the command queue to a specific device for execution.
-
-    Args:
-      dev: The device to submit the queue to
-    """
-
-    if var_vals is not None: self._apply_var_vals(var_vals)
-    self._submit(dev)
-    return self
-  def _submit(self, dev:HCQDeviceType): raise NotImplementedError("need _submit")
-
-class HCQSignal(Generic[HCQDeviceType]):
-  def __init__(self, base_buf:HCQBuffer, value:int=0, owner:HCQDeviceType|None=None, is_timeline:bool=False, timestamp_divider=1000, virt=False):
-    self.base_buf, self.owner, self.is_timeline = base_buf, owner, is_timeline
-    self.should_return = isinstance(self.base_buf.va_addr, int) and self.owner is not None and not virt
-    self.timestamp_divider:decimal.Decimal = decimal.Decimal(timestamp_divider)
-    if isinstance(self.base_buf.va_addr, int) and not virt: self.value = value
-
-  def __del__(self):
-    if self.should_return: HCQCompiled.signal_pool[unwrap(self.owner).peer_group].append(self.base_buf)
-
-  @property
-  def value_addr(self) -> sint: return self.base_buf.va_addr
-
-  @property
-  def timestamp_addr(self) -> sint: return self.base_buf.va_addr + 8
-
-  @property
-  def value(self) -> int: return self.base_buf.cpu_view().view(0, 8, 'Q')[0]
-
-  @value.setter
-  def value(self, new_value:int): self.base_buf.cpu_view().view(0, 8, 'Q')[0] = new_value
-
-  @property
-  def timestamp(self) -> decimal.Decimal:
-    """
-    Get the timestamp field of the signal.
-
-    This property provides read-only access to the signal's timestamp.
-
-    Returns:
-      The timestamp in microseconds.
-    """
-    return self.base_buf.cpu_view().view(8, 8, 'Q')[0] / self.timestamp_divider
-
-  def _sleep(self, time_spent_since_last_sleep_ms:int):
-    """
-    Optional function which can implement sleep functionality for the signal.
-    Raises RuntimeError if a fault is detected.
-    """
-
-  def wait(self, value:int, timeout:int|None=None):
-    """
-    Waits the signal is greater than or equal to a specific value.
-
-    Args:
-      value: The value to wait for.
-      timeout: Maximum time to wait in milliseconds. Defaults to 30s.
-    """
-    timeout = timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)
-    start_time = int(time.perf_counter() * 1000)
-    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
-      self._sleep(cur_time - start_time)
-      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
-    if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
-
-@contextlib.contextmanager
-def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None,
-                profile_key:bytes|None=None):
-  st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
-  assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"
-
-  if enabled and queue is not None: queue.timestamp(st)
-  elif enabled and queue_type is not None:
-    queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
-
-  try: yield (st, en)
-  finally:
-    if enabled and queue is not None: queue.timestamp(en)
-    elif enabled and queue_type is not None:
-      queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
-
-    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device,
-                                                         profile_key))
-
-class HCQArgsState(Generic[ProgramType]):
-  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
-    self.buf, self.prg, self.bufs, self.vals = buf, prg, bufs, vals
-    self.bind_data:list[tuple[tuple[sint, ...], MMIOInterface, str]] = []
-
-  def bind_sints_to_buf(self, *vals:sint, buf:HCQBuffer, fmt, offset=0): self.bind_data.append((vals, buf.cpu_view().view(offset=offset), fmt))
-
-class CLikeArgsState(HCQArgsState[ProgramType]):
-  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=(), prefix:list[int]|None=None):
-    super().__init__(buf, prg, bufs, vals=vals)
-
-    if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix)
-
-    self.bind_sints_to_buf(*[b.va_addr for b in bufs], buf=self.buf, fmt='Q', offset=len(prefix or []) * 4)
-    for v,(val_offset,dt) in zip(vals, TinyELF.iter_sig(prg.signature[-len(vals):], len(bufs) * 8)):
-      assert v is not None
-      self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=len(prefix or []) * 4 + val_offset)
-
-class HCQProgram(Program[HCQDeviceType]):
-  def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None):
-    self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size
-    self.profile_key = obj.profile_key
-    self.prof_prg_counter = next(self.dev.prof_prg_counter)
-    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter, self.profile_key)]
-
-  @staticmethod
-  def _fini(dev, buf, spec): dev.allocator.free(((buf, buf.meta), buf.view), buf.size, spec)
-
-  def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int|None, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
-    """
-    Fills arguments for the kernel, optionally allocating space from the device if `kernargs_ptr` is not provided.
-    Args:
-      bufs: Buffers to be written to kernel arguments.
-      vals: Values to be written to kernel arguments.
-      kernargs_ptr: Optional pointer to pre-allocated kernel arguments memory.
-    Returns:
-      Arguments state with the given buffers and values set for the program.
-    """
-    argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size, 8),
-                                                       size=self.kernargs_alloc_size)
-    return self.args_state_t(argsbuf, self, bufs, vals=vals)
-
-  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:
-    """
-    Enqueues the program for execution with the given arguments and dimensions.
-
-    Args:
-      bufs: Buffer arguments to execute the kernel with.
-      global_size: Specifies the global work size for kernel execution (equivalent to CUDA's grid size).
-      local_size: Specifies the local work size for kernel execution (equivalent to CUDA's block size).
-      vals: Value arguments to execute the kernel with.
-      wait: If True, waits for the kernel to complete execution.
-
-    Returns:
-      Execution time of the kernel if 'wait' is True, otherwise None.
-    """
-
-    kernargs = self.fill_kernargs(bufs, vals)
-    q = unwrap(self.dev.hw_compute_queue_t)().wait(self.dev.timeline_signal, self.dev.timeline_value - 1).memory_barrier()
-
-    self.dev.prof_exec_counter += 1
-    with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE, profile_key=self.profile_key) as (sig_st, sig_en):
-      q.exec(self, kernargs, global_size, local_size)
-
-    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-
-    if wait: self.dev.synchronize(timeout=timeout)
-    return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None
-
-class HCQCompiled(Compiled, Generic[SignalType]):
-  """
-  A base class for devices compatible with the HCQ (Hardware Command Queue) API.
-  """
-  peer_groups: dict[str, list[HCQCompiled]] = collections.defaultdict(list)
-  signal_pages: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
-  signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
-  cpu_devices: list[HCQCompiled] = []
-
-  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
-               signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
-               kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
-    from extra.hcq1.graph import HCQGraph
-    super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
-
-    self.peer_group = getattr(getattr(self, 'iface', None), 'peer_group', device.split(":")[0])
-    HCQCompiled.peer_groups[self.peer_group].append(self)
-
-    self.signal_t, self.hw_compute_queue_t, self.hw_copy_queue_t = signal_t, comp_queue_t, copy_queue_t
-
-    self.timeline_value:int = 1
-    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str, bytes|None]] = []
-    self.prof_exec_counter:int = 0
-    self.prof_prg_counter = itertools.count(0)
-
-    if signal_t is not None:
-      # Map signals if any
-      for sig_page in HCQCompiled.signal_pages[self.peer_group]: cast(HCQAllocator, self.allocator)._map(sig_page)
-
-      self.sigalloc_size = sigalloc_size
-      self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
-
-    if comp_queue_t is not None:
-      self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))[0][0]
-      self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
-
-    self.can_recover = can_recover # Whether the device can recover from faults or timeouts
-    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
-
-    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
-
-  def synchronize(self, timeout:int|None=None):
-    if self.error_state is not None: raise self.error_state
-    if not hasattr(self, 'timeline_signal'): return
-
-    # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
-    if not self._is_cpu():
-      for dev in HCQCompiled.cpu_devices: dev.synchronize()
-
-    try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
-    except RuntimeError as e:
-      self.error_state = e
-      if hasattr(self, 'on_device_hang'): self.on_device_hang()
-      raise e
-
-    if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
-    if PROFILE:
-      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp, pk) for st,en,name,dev,pk in self.sig_prof_records]
-      self.sig_prof_records = []
-
-  def next_timeline(self):
-    self.timeline_value += 1
-    return self.timeline_value - 1
-
-  def new_signal(self, **kwargs) -> SignalType:
-    assert self.signal_t is not None, "Device does not support signals"
-    if not HCQCompiled.signal_pool[pg:=self.peer_group]:
-      HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True))[0][0])
-      HCQCompiled.signal_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)
-
-  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
-
-  def hw_compute_queues(self) -> list[tuple[str|None, Callable[[], HWQueue]]]:
-    return [(None, self.hw_compute_queue_t)] if self.hw_compute_queue_t is not None else []
-  def hw_copy_queues(self) -> list[tuple[str, Callable[[], HWQueue]]]:
-    return [("SDMA:0", self.hw_copy_queue_t)] if self.hw_copy_queue_t is not None else []
-
-  def _at_profile_finalize(self):
-    self.synchronize() # Expect device to be synchronizes
-
-    def _sync(d:HCQCompiled, q_t:Callable[[], HWQueue]):
-      q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d)
-      st = time.perf_counter_ns()
-      d.timeline_signal.wait(d.timeline_value - 1)  # average of the two
-      et = time.perf_counter_ns()
-      return (decimal.Decimal(et+st) / 2000) - d.timeline_signal.timestamp
-
-    for prefix, q_t in self.hw_compute_queues() + self.hw_copy_queues():
-      devname = f"{self.device}:{prefix}" if prefix else self.device
-      Compiled.profile_events += [ProfileDeviceEvent(devname, statistics.median([_sync(self, q_t) for _ in range(40)]), props=self.device_props())]
-
-  def _wrap_timeline_signal(self):
-    self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
-    self.timeline_signal.value = 0
-    cast(HCQAllocatorBase, self.allocator).b_timeline = [0] * len(cast(HCQAllocatorBase, self.allocator).b)
-
-  def _realloc(self, oldbuf:HCQBuffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQBuffer, bool]:
-    if oldbuf is not None: self.allocator.free(((oldbuf, oldbuf.meta), oldbuf.view), oldbuf.size, options=options)
-    try: buf, realloced = self.allocator.alloc(new_size, options=options)[0][0], True
-    except MemoryError:
-      if force: raise
-      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options)[0][0], False
-    return buf, realloced
-
-  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
-
-  def rdma_dev(self):
-    from extra.hcq1.ops_rdma import get_rdma_device
-    for i in itertools.count():
-      if (dev:=next((d for d in HCQCompiled.peer_groups[self.peer_group] if type(d).__name__ == 'RDMADevice'), None)): return dev
-      try: get_rdma_device(i)
-      except IndexError: raise RuntimeError(f"No RDMA found for peer group '{self.peer_group}'")
-
-  def finalize(self):
-    try: self.synchronize() # Try to finalize device in any case.
-    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
-    super().finalize()
-
-class HCQAllocatorBase(Allocator[HCQDeviceType], Generic[HCQDeviceType]):
-  """
-  A base allocator class compatible with the HCQ (Hardware Command Queue) API.
-
-  This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
-  """
-
-  def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, **kwargs):
-    super().__init__(dev, **kwargs)
-    self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True))[0][0] for _ in range(batch_cnt)]
-    self.b_timeline, self.b_next = [0] * len(self.b), 0
-
-  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")
-      if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
-      if (mb:=self._do_map(buf)) is not None: buf.mappings[self.dev] = mb
-      buf.mapped_devs.append(self.dev)
-    mapped = buf.mappings.get(self.dev, buf)
-    return mapped, mapped.meta
-
-  @suppress_finalizing
-  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
-    for dev in buf.mapped_devs: dev.synchronize()
-    for d, mb in buf.mappings.items(): d.allocator._do_unmap(mb)
-    if hasattr(self, '_do_free'): self._do_free(buf, options)
-
-  def _do_unmap(self, mb): self.dev.iface.free(mb)
-
-  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)
-
-class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
-  def _copyin(self, dest:HCQBuffer, src:memoryview):
-    if self.dev.hw_copy_queue_t is None:
-      self.dev.synchronize()
-      with cpu_profile(f'TINY -> {self.dev.device}', f"{self.dev.device}:COPY"): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
-      return
-
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {self.dev.device}", ret=src.nbytes), enabled=PROFILE,
-                     dev_suff="SDMA:0"):
-      for i in range(0, src.nbytes, self.b[0].size):
-        self.b_next = (self.b_next + 1) % len(self.b)
-        self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
-
-        lsize = min(self.b[self.b_next].size, src.nbytes - i)
-        self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src.cast('B')[i:i+lsize]
-        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
-                                  .copy(dest.offset(i), self.b[self.b_next], lsize) \
-                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-        self.b_timeline[self.b_next] = self.dev.timeline_value - 1
-
-  def copy_from_disk(self, dest:HCQBuffer, src, size):
-    def _get_temp_buf():
-      # Check if the next buffer is safe to be used (its signal has passed) and reserve it.
-      if self.b_timeline[(self.b_next + 1) % len(self.b)] <= self.dev.timeline_signal.value:
-        self.b_timeline[(self.b_next + 1) % len(self.b)], self.b_next = (1 << 64), (self.b_next + 1) % len(self.b)
-        return (self.b[self.b_next].cpu_view(), self.b_next)
-      return None
-
-    assert self.dev.hw_copy_queue_t is not None
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"DISK -> {self.dev.device}", ret=size), enabled=PROFILE,
-                     dev_suff="SDMA:0"):
-      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size,
-                                                                                             use_ioring=type(self.b[0].cpu_view()) is MMIOInterface):
-        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
-                                  .copy(dest.offset(dst_off), self.b[batch_info[1]].offset(src_off), copy_size) \
-                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-        self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1
-
-  def _copyout(self, dest:memoryview, src:HCQBuffer):
-    self.dev.synchronize()
-    if self.dev.hw_copy_queue_t is None:
-      with cpu_profile(f'{self.dev.device} -> TINY', f"{self.dev.device}:COPY"): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
-      return
-
-    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"{self.dev.device} -> TINY", ret=dest.nbytes), enabled=PROFILE,
-                     dev_suff="SDMA:0"):
-      for i in range(0, dest.nbytes, cp_size:=self.b[0].size):
-        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
-                                  .copy(self.b[0], src.offset(i), lsize:=min(cp_size, dest.nbytes-i)) \
-                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
-        self.dev.timeline_signal.wait(self.dev.timeline_value - 1)
-        dest.cast('B')[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
-
-  def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQDeviceType, dest_dev:HCQDeviceType):
-    if src_dev.peer_group != dest_dev.peer_group: return src_dev.rdma_dev().allocator._transfer(dest, src, sz, src_dev, dest_dev)
-
-    cast(HCQAllocator, src_dev.allocator)._map(dest)
-
-    assert src_dev.hw_copy_queue_t is not None
-    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=TracingKey(f"{src_dev.device} -> {dest_dev.device}", ret=sz), enabled=PROFILE,
-                     dev_suff="SDMA:0"):
-      src_dev.hw_copy_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
-                               .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
-                               .copy(dest, src, sz) \
-                               .signal(src_dev.timeline_signal, src_dev.next_timeline()).submit(src_dev)
-
-    if src_dev != dest_dev:
-      unwrap(dest_dev.hw_compute_queue_t)().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
-                                           .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
-                                           .signal(dest_dev.timeline_signal, dest_dev.next_timeline()).submit(dest_dev)

extra/introspection.py

@@ -9,7 +9,7 @@ def print_objects():
   tensors = [x for x in gc.get_objects() if isinstance(x, Tensor)]
   tensor_ram_used = sum([prod(x.shape)*4 for x in tensors])
   lazybuffers = [x for x in gc.get_objects() if isinstance(x, UOp)]
-  gpubuffers = [x for x in gc.get_objects() if isinstance(x, Buffer) and x.is_allocated()]
+  gpubuffers = [x for x in gc.get_objects() if isinstance(x, Buffer) and x.is_initialized()]
   realized_buffers = [x.realized for x in lazybuffers if x.base == x and x.realized]
   gpubuffers_orphaned = [x for x in gpubuffers if x not in realized_buffers]
 
@@ -31,7 +31,8 @@ def print_objects():
     cnt += 1
 
   for x in gpubuffers_orphaned:
-    if x.base.is_allocated(): x.base.deallocate()
+    if getattr(x, '_buf', None): del x._buf
+    if getattr(x, '_image', None): del x._image
 
   return len(gpubuffers_orphaned)
 

extra/mlx_driver/loopback.py

@@ -30,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 = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
-buf_dst = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))[0][0]
+buf_src = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))
+buf_dst = gpu.allocator.alloc(BUF_SIZE, BufferSpec(nolru=True))
 
 bar_base = gpu.iface.pci_dev.bar_info(gpu.iface.vram_bar)[0]
 src_paddr = buf_src.meta.mapping.paddrs[0][0] + bar_base

extra/models/llama.py

@@ -139,7 +139,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
     h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).clone().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):

extra/remote/bench.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os, sys, time
-from extra.hcq1.remote import RemotePCIDevice
+from extra.remote.hcq1_remote import RemotePCIDevice
 
 LAT_N_RUNS = 500
 THROUGHPUT_N_RUNS = 8
@@ -18,7 +18,7 @@ if __name__ == "__main__":
   print(f"connected to {os.environ['REMOTE']}, device: {name}\n")
 
   # ping (minimal server round-trip, no device I/O)
-  from extra.hcq1.remote import RemoteCmd
+  from extra.remote.hcq1_remote import RemoteCmd
   sock = pci.sock
   for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
   st = time.perf_counter()

extra/remote/serve.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 import socket, struct, sys
 from tinygrad.runtime.support.system import PCIDevice, System
-from extra.hcq1.remote import RemoteCmd
+from extra.remote.hcq1_remote import RemoteCmd
 from tinygrad.helpers import DEBUG, OSX
 
 def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)

mkdocs.yml

@@ -29,6 +29,7 @@ nav:
     - UOp: developer/uop.md
     - Runtime:
       - developer/runtime.md
+      - HCQ: developer/hcq.md
       - AM Driver: developer/am.md
   - tinybox: tinybox.md
 #- tinygrad: reference/

spec/tinyspec.tex

@@ -150,6 +150,18 @@ A value \op{Call} is void: its \op{Sink} body stores to output \op{Param}s bound
 
 \smallskip
 Assign is \op{Store} followed by \op{After}: write the value, then return the buffer with an ordering dependency.
+\op{After} orders consumers after its dependencies; it neither declares a write nor snapshots memory.
+In particular, \op{After}$(b, \op{Store}(d,v))$ returns $b$, not $v$, when $b$ and $d$ are disjoint.
+Views may share storage despite having different UOps. Differentiation follows the returned value:
+a matching unconditional full overwrite routes its gradient to the stored value; an unrelated write does not create a gradient path.
+Partial or uncertain aliased mutation gradients may be rejected.
+
+\smallskip
+\textbf{Tensor scheduling contract.} Within a lazy Tensor schedule, reads retain their assignment dependencies.
+A read must follow those dependencies and precede other writes that would destroy the required contents.
+Lowering must preserve these requirements until accesses are ordered, even when arguments share storage.
+Unsatisfiable requirements raise rather than read overwritten contents. This is a frontend requirement, not snapshot semantics for \op{After}.
+An executed \texttt{clone()} preserves data in fresh storage; \texttt{contiguous()} need not allocate.
 
 %% ============================================================
 \subsection*{{\color{elwyellow}Elementwise Ops} \normalfont\small--- all inputs same shape, output same shape, applied per-element}
@@ -236,7 +248,7 @@ Ternary & $(P, A, B)$
 \op{Custom} & (args\ldots) & fmt & Inject custom code string into generated source. \\
 \op{AtomicAdd} & (idx, val) & --- & Atomic read-modify-write: \texttt{buf[idx] += val}. \\[4pt]
 \op{CustomFunction} & (meta\ldots) & name & Opaque device function (e.g.\ HW decode). Via \op{Call}. \\
-\op{Program} & (linear, source, binary) & --- & Compiled kernel: instructions, source, and machine code. \\
+\op{Program} & (sink, \ldots) & metadata? & Kernel through compilation stages. \\
 \op{Source} & () & str & Human-readable rendered source code. \\
 \op{Binary} & () & bytes & Compiled machine code. \\
 \bottomrule
@@ -244,6 +256,11 @@ Ternary & $(P, A, B)$
 
 \smallskip
 These ops are not part of the core specification and are subject to change.
+\op{Program} contains a \op{Sink}, followed progressively by \op{Linear}, \op{Source}, and \op{Binary}.
+Access analysis derives reads and writes from the memory operands of \op{Load}/\op{Store}, resolving \op{Param}s through \op{Call} arguments.
+Compilation records these sets in \texttt{ProgramInfo.ins/outs} as zero-based argument slots; a read-modify-write belongs in both.
+Listing a parameter, returning an \op{After}, or declaring a write does not establish full initialization.
+Opaque code without computed access information is unsupported by assignment scheduling; its effects must not be guessed from its argument list.
 
 %% ============================================================
 \subsection*{Derived Properties}

test/amd/hw/helpers.py

@@ -225,7 +225,7 @@ amdhsa.kernels:
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
 
   buf_sz = _out_bytes(n_lanes)
-  out_gpu = dev.allocator.alloc(buf_sz)[0][0]
+  out_gpu = dev.allocator.alloc(buf_sz)
   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)
 

test/amd/hw/test_cdna_vop3.py

@@ -47,7 +47,7 @@ def _run_hw(instructions: list, out_reg: int = 2) -> int:
 
   dev = Device["AMD"]
   if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
-  out_gpu = dev.allocator.alloc(LANES * 4)[0][0]
+  out_gpu = dev.allocator.alloc(LANES * 4)
   code = _code(instructions, out_reg, out_gpu.va_addr)
   byte_str = ", ".join(f"0x{b:02x}" for b in code)
   asm_src = f""".text

test/amd/hw/test_dpp.py

@@ -84,7 +84,7 @@ amdhsa.kernels:
 """
   lib = compiler.compile(asm_src)
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  out_gpu = dev.allocator.alloc(WAVE64 * 4)[0][0]
+  out_gpu = dev.allocator.alloc(WAVE64 * 4)
   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)

test/amd/hw/test_rdna4_permlane_var.py

@@ -85,7 +85,7 @@ amdhsa.kernels:
 """
   lib = compiler.compile(asm_src)
   prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-  out_gpu = dev.allocator.alloc(LANES * 4)[0][0]
+  out_gpu = dev.allocator.alloc(LANES * 4)
   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)

test/amd/test_asm_kernel.py

@@ -3,7 +3,7 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import run_linear, estimate_uop, lower_and_compile
+from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
@@ -169,7 +169,7 @@ class TestAsmKernel(unittest.TestCase):
     if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    linear = lower_and_compile(a.schedule_linear())
+    linear = compile_linear(a.schedule_linear())
     est = estimate_uop(linear.src[-1])
     self.assertEqual(est.ops, a.numel())
     self.assertEqual(est.mem, a.nbytes()*2)

test/amd/test_mockgpu_invalid.py

@@ -1,29 +1,51 @@
 #!/usr/bin/env python3
 """Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, sys
+import unittest, subprocess, os, sys, time
 
 class TestMockGPUInvalidInstruction(unittest.TestCase):
   def test_unsupported_instruction_raises(self):
     """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
     test_code = '''
-import os, sys
-from tinygrad import Tensor
-from tinygrad.engine.realize import lower_and_compile, run_linear
+import struct
+from dataclasses import replace
+from tinygrad import Device, Tensor
+from tinygrad.engine.realize import compile_linear
 
-linear = lower_and_compile((Tensor.empty(1) + 1).schedule_linear())
-binary = linear.src[-1].src[0].src[3]
-lib = binary.arg.replace(bytes.fromhex("0000b0bf"), bytes.fromhex("00fe017e"), 1)
-try:
-  run_linear(linear.substitute({binary: binary.replace(arg=lib)}, enter_calls=True))
-except ValueError as error:
-  print(error, file=sys.stderr, flush=True)
-  os._exit(1)
+dev = Device["AMD"]
+a = Tensor([1.0]).realize()
+b = a + 1
+linear = compile_linear(b.schedule_linear())
+compiled_prg = linear.src[-1].src[0]
+lib = bytearray(compiled_prg.src[3].arg)
+
+# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
+# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
+found = False
+for i in range(0, len(lib) - 4, 4):
+  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
+    lib[i:i+4] = struct.pack("<I", 0x7E008400)
+    found = True
+    break
+assert found, "s_endpgm not found"
+
+patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
+b.uop.buffer.allocate()
+patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
+dev.synchronize()
 '''
 
-    env = {**os.environ, "DEV": "MOCKKFD+AMD", "HCQ_RUNTIME_DEV": "PYTHON"}
-    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=9)
-    self.assertEqual(result.returncode, 1)
-    self.assertIn("unknown rdna3 format word=0x7e01fe00", result.stderr)
+    env = os.environ.copy()
+    env["DEV"] = "MOCKKFD+AMD"
+    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
+
+    st = time.perf_counter()
+    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
+    elapsed = time.perf_counter() - st
+
+    self.assertNotEqual(result.returncode, 0, "should have raised")
+    self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
+    # Should exit immediately, not wait for the full timeout
+    self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_roundtrip.py

@@ -58,11 +58,11 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
   """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
   from tinygrad import Tensor
   from tinygrad.uop.ops import Ops
-  from tinygrad.engine.realize import lower_and_compile, resolve_params, unwrap_multi
+  from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
   from tinygrad.runtime.support.elf import elf_loader
 
   out = op_fn(Tensor)
-  linear = lower_and_compile(out.schedule_linear())
+  linear = compile_linear(out.schedule_linear())
   kernels = []
   buf_pool: dict[int, int] = {}  # buffer id -> size
   buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY

test/amd/test_sqtt_profiler.py

@@ -1,6 +1,5 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit, dtypes
-from tinygrad.dtype import AddrSpace
 from test.helpers import is_hcq2_device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Compiled, ProfileProgramEvent
@@ -9,7 +8,7 @@ from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.viz.serve import load_amd_counters, VizData
 from tinygrad.renderer.amd.sqtt import decode, print_packets
-from tinygrad.renderer.amd.dsl import s, v
+from tinygrad.renderer.amd.dsl import s
 
 @contextlib.contextmanager
 def save_sqtt():
@@ -28,46 +27,8 @@ def map_sqtt(profile:list) -> list[dict]:
 def custom_asm_cdna(A:UOp):
   import tinygrad.runtime.autogen.amd.cdna.ins as cdna
   WAVE_SIZE = 64
-  insts = [
-    cdna.s_barrier(),
-    cdna.s_getreg_b32(s[0], cdna.HWREG.HW_REG_HW_ID.value | (4 << 6) | (1 << 11)),
-
-    cdna.s_cmp_eq_u32(s[0], 0),
-    cdna.s_cbranch_scc1(16),
-
-    cdna.s_cmp_eq_u32(s[0], 1),
-    cdna.s_cbranch_scc1(9),
-
-    cdna.s_cmp_eq_u32(s[0], 2),
-    cdna.s_cbranch_scc1(3),
-
-    # SIMD 3
-    cdna.v_mov_b32_e32(v[0], 3),
-    cdna.s_nop(3),
-    cdna.s_endpgm(),
-
-    # SIMD 2
-    cdna.v_mov_b32_e32(v[0], 2),
-    cdna.s_nop(2),
-    cdna.s_nop(2),
-    cdna.s_endpgm(),
-
-    # SIMD 1
-    cdna.v_mov_b32_e32(v[0], 1),
-    cdna.s_nop(1),
-    cdna.s_nop(1),
-    cdna.s_nop(1),
-    cdna.s_endpgm(),
-
-    # SIMD 0
-    cdna.v_mov_b32_e32(v[0], 0),
-    cdna.s_nop(0),
-    cdna.s_nop(0),
-    cdna.s_nop(0),
-    cdna.s_nop(0),
-    cdna.s_endpgm(),
-  ]
-  return custom_asm(A, insts, WAVE_SIZE*4, 96*1024)
+  insts = [cdna.s_nop(0), cdna.s_mov_b32(s[0], 10)]
+  return custom_asm(A, insts+[cdna.s_endpgm()], WAVE_SIZE*2)
 
 def custom_asm_rdna(A:UOp):
   import tinygrad.runtime.autogen.amd.rdna3.ins as rdna3
@@ -75,9 +36,8 @@ def custom_asm_rdna(A:UOp):
   insts = [rdna3.s_nop(0), rdna3.s_mov_b32(s[0], 10)]
   return custom_asm(A, insts+[rdna3.s_endpgm()], WAVE_SIZE*2)
 
-def custom_asm(A, insts, num_threads, lds_size=0) -> UOp:
-  lds = UOp.placeholder((lds_size,), dtypes.uint8, addrspace=AddrSpace.LOCAL) if lds_size else None
-  return UOp(Ops.PROGRAM, src=(UOp.sink(A, lds, UOp.special(num_threads, "lidx0"), arg=KernelInfo("asm")), \
+def custom_asm(A, insts, num_threads) -> UOp:
+  return UOp(Ops.PROGRAM, src=(UOp.sink(A, UOp.special(num_threads, "lidx0"), arg=KernelInfo("asm")), \
       UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS,arg=(x,dtypes.void)) for x in insts]))))
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")

test/backend/test_assign.py

@@ -40,10 +40,10 @@ class TestAssign(unittest.TestCase):
   def test_assign_copy(self):
     a = Tensor([1.,2,3], device="PYTHON")
     c = Tensor.empty(3).assign(a.to(None))
-    # it should copy into the empty buffer
+    # The creation copy has its own storage, independent of the assignment destination.
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
+    assert_kernel_count(3 if is_hcq2_device() else 2)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -619,7 +619,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(5 if is_hcq2_device() else 3)  # TODO: first copy is dead, could be 2
+    assert_kernel_count(6 if is_hcq2_device() else 4)  # TODO: first copy is dead
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -629,7 +629,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
+    assert_kernel_count(3 if is_hcq2_device() else 2)
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_assign_temporary_copy_reshape(self):
@@ -637,7 +637,7 @@ class TestAssign(unittest.TestCase):
     c = Tensor.empty(2, 2).assign(a.to(None))
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(2 if is_hcq2_device() else 1)
+    assert_kernel_count(3 if is_hcq2_device() else 2)
     self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
 
 class TestAssignOrdering(unittest.TestCase):
@@ -828,6 +828,13 @@ class TestAssignOrdering(unittest.TestCase):
       b_np *= 0.9
     np.testing.assert_allclose(param.item(), p_np, atol=1e-5)
 
+  def test_after_store_to_other_buffer(self):
+    x, state = Tensor([2.]).realize(), Tensor([0.]).realize()
+    ordered = Tensor(x.uop.after(state.uop.store(x.uop * 3)))
+    self.assertEqual((ordered + x).tolist(), [4.])
+    self.assertEqual(state.tolist(), [6.])
+    self.assertEqual(x.tolist(), [2.])
+
   def test_war_reader_already_depends_on_write(self):
     x = Tensor([1.0]).contiguous().realize()
     y = Tensor([2.0]).contiguous().realize()
@@ -835,12 +842,8 @@ class TestAssignOrdering(unittest.TestCase):
     x.assign(x * 2)
     y.assign(y + x)
     z = y + x_expr
-    Tensor.realize(x, y, z)
-    try:
-      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 15.0])
-    except AssertionError:
-      # TODO: broken now, x_expr reads x after the assign
-      np.testing.assert_allclose([x.item(), y.item(), z.item()], [2.0, 4.0, 16.0])
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      Tensor.realize(x, y, z)
 
   def test_war_multi_read_then_assign(self):
     devices = ("CPU:0", "CPU:1")
@@ -875,8 +878,8 @@ class TestAssignOrdering(unittest.TestCase):
     a.assign(b + 1)                                    # a == 11
     v1 = a * 3                                         # reads 11 -> 33
     a.assign(b + 100)                                  # a == 110
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, ideally v1 is realized between the assigns
-      np.testing.assert_allclose((a + v1).numpy(), 143)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      (a + v1).numpy()
 
   def test_two_reads_between_three_assigns(self):
     a = Tensor.zeros(4).realize()
@@ -991,9 +994,8 @@ class TestAssignOrdering(unittest.TestCase):
       x.assign(x+1)
       return y+x
     a = Tensor([1.]).realize()
-    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, ideally y is realized between the assigns
-      out = outer(a).item()
-      self.assertEqual([out, a.item()], [7., 3.])
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      outer(a).item()
 
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
@@ -1014,24 +1016,6 @@ class TestAssignToUnrealizedView(unittest.TestCase):
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(c.tolist(), [[1,1],[2,1]])
 
-  def test_contiguous_partial_assign_realize(self):
-    x = Tensor([1., 2.]).realize()
-    y = (x + 1).contiguous()  # unrealized CONTIGUOUS
-    self.assertIs(y.uop.base.op, Ops.CONTIGUOUS)
-    # a partial write survives an explicit realize: the values are right, storage is an implementation detail
-    y[:1].assign(9.)
-    y.realize()
-    self.assertEqual(y.tolist(), [9., 3.])
-    # and it stays assigned across schedules
-    y[:1].assign(7.)
-    y.realize()
-    self.assertEqual(y.tolist(), [7., 3.])
-    # setitem syntax gives the same values, contiguous or not
-    for mk in (lambda xx: xx + 1, lambda xx: (xx + 1).contiguous()):
-      z = mk(Tensor([1., 2.]).realize())
-      z[:1] = 9.
-      self.assertEqual(z.tolist(), [9., 3.])
-
   def test_contiguous_backward(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
     cb = t.contiguous_backward()  # unrealized CONTIGUOUS_BACKWARD

test/backend/test_call.py

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

test/backend/test_const_folding.py

@@ -84,10 +84,10 @@ class TestReduceOpsConstFolding(unittest.TestCase):
       np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
 
   def test_zero_size_realize_folded(self):
-    # folded output doesn't realize on its own
+    # non contiguous folded output doesn't realize
     _check_ast_count(0, Tensor.empty(1, 0).sum())
-    # explicit storage of the folded const still schedules, and the value is usable
-    a = Tensor.empty(1, 0).sum().clone()
+    # contiguous folded const can still schedule
+    a = Tensor.empty(1, 0).sum().contiguous()
     _check_ast_count(2, a+2)
     self.assertIs(a.uop.base.op, Ops.BUFFER)
     np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)

test/backend/test_custom_kernel.py

@@ -105,6 +105,47 @@ def backward_gemm_custom(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]:
 # **** tests ****
 
 class TestCustomKernel(unittest.TestCase):
+  def test_readonly_after_args(self):
+    for chained in (False, True):
+      for corealize in (False, True):
+        with self.subTest(chained=chained, corealize=corealize):
+          x = Tensor([2.]).realize()
+          a, x1 = Tensor.empty(1).custom_kernel(x, fxn=custom_add_one_kernel)
+          b, x2 = Tensor.empty(1).custom_kernel(x1 if chained else x, fxn=custom_add_one_kernel)
+          if corealize:
+            y = x1 + b
+            Tensor.realize(y, x2)
+            self.assertEqual(y.tolist(), [5.])
+          else:
+            self.assertEqual((x1 + x2).tolist(), [4.])
+          self.assertEqual(a.tolist(), [3.])
+          self.assertEqual(b.tolist(), [3.])
+          self.assertEqual(x.tolist(), [2.])
+
+  def test_aliased_args_different_sizes(self):
+    def kernel(out:UOp, a:UOp, b:UOp):
+      i = UOp.range(4, 0)
+      return out[i].store(a[i] + b[0]).end(i).sink(arg=KernelInfo(name="aliased_sizes"))
+    x = Tensor([1., 2., 3., 4.]).realize()
+    out = Tensor.empty(4).custom_kernel(x, x[:1], fxn=kernel)[0]
+    self.assertEqual(out.tolist(), [2., 3., 4., 5.])
+
+  def test_unindexed_access_before_assign(self):
+    def kernel(out:UOp, x:UOp): return out.store(x + 1).sink(arg=KernelInfo(name="unindexed"))
+    x = Tensor([2.]).realize()
+    y = Tensor.empty(1).custom_kernel(x, fxn=kernel)[0]
+    x.assign(x * 2)
+    Tensor.realize(x, y)
+    self.assertEqual(x.tolist(), [4.])
+    self.assertEqual(y.tolist(), [3.])
+
+  def test_readonly_after_does_not_hide_write(self):
+    x = Tensor([2.]).realize()
+    _, before = Tensor.empty(1).custom_kernel(x, fxn=custom_add_one_kernel)
+    x.assign(x * 2)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):
+      (before + x).realize()
+
   def test_empty(self):
     a = Tensor.empty(1)
     a = Tensor.custom_kernel(a, fxn=lambda _: UOp.sink(arg=KernelInfo()))[0]

test/backend/test_jit.py

@@ -292,7 +292,7 @@ class TestJitGraphSplit(unittest.TestCase):
     if graph_t is None: return
 
     got = f.captured.linear.src
-    from extra.hcq1.graph import HCQGraph
+    from tinygrad.runtime.graph.hcq import HCQGraph
     from tinygrad.engine.jit import MultiGraphRunner
     if graph_t is HCQGraph:
       validate = hcqgraph

test/backend/test_ops.py

@@ -712,9 +712,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: 0**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: 0.7**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: (-2)**x, vals=[[-2.,-1,0,1,2,3]])
-    # 2**52+2 - 0.5 rounds back to itself
-    helper_test_op(None, lambda x: x**(2.0**52), vals=[[0.5, 1., 2.]], forward_only=True)
-    helper_test_op(None, lambda x: x**(2.0**52+2), vals=[[0.5, 1., 2.]], forward_only=True)
     # float to power of int
     helper_test_op(None, lambda x: 0.7**x, lambda x: (0.7**x).clone(), vals=[[-2,-1,0,1,2,3]], forward_only=True)
 
@@ -1089,8 +1086,6 @@ class TestOps(unittest.TestCase):
   def test_hardsigmoid_extreme(self):
     helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, low=300, high=400)
     helper_test_op([(45,65)], torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, low=-400, high=-300)
-    helper_test_op(None, torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, vals=[[1e7, 1e8, 2.68e8, 1e9]])
-    helper_test_op(None, torch.nn.functional.hardsigmoid, Tensor.hardsigmoid, vals=[[-3.1, -3., -2.9, 2.9, 3., 3.1]])
   def test_softplus(self):
     helper_test_op([(45,65)], torch.nn.functional.softplus, Tensor.softplus, grad_atol=1e-6)
     helper_test_op([(45,65)], lambda t: torch.nn.functional.softplus(t, beta=3), lambda t: Tensor.softplus(t, beta=3), grad_atol=1e-6)
@@ -1134,12 +1129,9 @@ class TestOps(unittest.TestCase):
   def test_relu6(self):
     helper_test_op([(45,65)], torch.nn.functional.relu6, Tensor.relu6)
     helper_test_op([()], torch.nn.functional.relu6, Tensor.relu6)
-    helper_test_op(None, torch.nn.functional.relu6, Tensor.relu6, vals=[[6.71089e7, 2.68435e8, 1e9]])
-    helper_test_op(None, torch.nn.functional.relu6, Tensor.relu6, vals=[[0., 6.]])
   def test_hardswish(self):
     helper_test_op([(45,65)], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
     helper_test_op([()], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
-    helper_test_op(None, torch.nn.functional.hardswish, Tensor.hardswish, vals=[[-3., 3.]], grad_atol=1e-6)
   def test_mish(self):
     helper_test_op([(45,65)], torch.nn.functional.mish, Tensor.mish)
     helper_test_op([()], torch.nn.functional.mish, Tensor.mish)

test/backend/test_profiler.py

@@ -2,7 +2,7 @@ import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
-from extra.hcq1.hcq import HCQCompiled
+from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
@@ -120,7 +120,7 @@ class TestProfiler(unittest.TestCase):
 
     for dev in [TestProfiler.d0.device, d1.device]:
       evs = [x for x in profile if isinstance(x, ProfileRangeEvent) and _dev_base(x.device) == dev]
-      assert len(evs) == (0 if buf1._host_mv() is not None else 1), "one kernel runs are expected"
+      assert len(evs) == (0 if hasattr(TestProfiler.d0.allocator, '_as_buffer') else 1), "one kernel runs are expected"
 
   def test_profile_multidev_transfer(self):
     try: d1 = Device[f"{Device.DEFAULT}:1"]

test/backend/test_setitem.py

@@ -1,5 +1,6 @@
 import unittest, operator
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
+from tinygrad.helpers import Context
 import numpy as np
 
 class TestSetitem(unittest.TestCase):
@@ -162,20 +163,21 @@ class TestSetitem(unittest.TestCase):
       np.testing.assert_allclose(t.numpy(), n)
 
   def test_jit_setitem_variable_offset(self):
-    @TinyJit
-    def f(t:Tensor, a:Tensor, v:Variable):
-      t.shrink(((v,v+1), None)).assign(a).realize()
+    with Context(CHECK_OOB=0):
+      @TinyJit
+      def f(t:Tensor, a:Tensor, v:Variable):
+        t.shrink(((v,v+1), None)).assign(a).realize()
 
-    t = Tensor.zeros(6, 6).contiguous().realize()
-    n = np.zeros((6, 6))
+      t = Tensor.zeros(6, 6).contiguous().realize()
+      n = np.zeros((6, 6))
 
-    for i in range(6):
-      v = Variable("v", 0, 6).bind(i)
-      a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
-      n[i, :] = i+1
-      f(t, a, v)
-      np.testing.assert_allclose(t.numpy(), n)
-    np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
+      for i in range(6):
+        v = Variable("v", 0, 6).bind(i)
+        a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
+        n[i, :] = i+1
+        f(t, a, v)
+        np.testing.assert_allclose(t.numpy(), n)
+      np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
 
   def test_setitem_overlapping_inplace1(self):
     t = Tensor([[3.0], [2.0], [1.0]]).contiguous()

test/backend/test_subbuffer.py

@@ -69,27 +69,41 @@ class TestSubBuffer(unittest.TestCase):
     buf = self.buf_unalloc
     sub_buf = buf.view(3, dtypes.uint8, offset=4)
     self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
     self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
 
     # base buffer alloc
     buf.allocate()
     self.assertTrue(buf.is_allocated())
-    self.assertFalse(sub_buf.is_allocated())
-    sub_buf.ensure_allocated()
+    self.assertTrue(buf.is_initialized())
     self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
+
+    # sub buffer alloc
+    sub_buf.allocate()
+    self.assertTrue(sub_buf.is_initialized())
+
+    # sub buffer dealloc
     sub_buf.deallocate()
     self.assertTrue(buf.is_allocated())
-    self.assertFalse(sub_buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
+    self.assertTrue(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
 
     # base buffer dealloc
     buf.deallocate()
     self.assertFalse(buf.is_allocated())
+    self.assertFalse(buf.is_initialized())
     self.assertFalse(sub_buf.is_allocated())
+    self.assertFalse(sub_buf.is_initialized())
 
-    # sub buffer alloc allocates the base
+    # sub buffer alloc
     sub_buf.ensure_allocated()
     self.assertTrue(buf.is_allocated())
+    self.assertTrue(buf.is_initialized())
     self.assertTrue(sub_buf.is_allocated())
+    self.assertTrue(sub_buf.is_initialized())
 
   def test_subbuffer_copy_in_out(self):
     sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]

test/backend/test_tensor.py

@@ -737,6 +737,25 @@ class TestZeroShapeTensor(unittest.TestCase):
     np.testing.assert_allclose(a.numpy(), b.numpy())
     self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
 
+  def test_clone_unrealized_copy_does_not_alias(self):
+    for realize_clone in (False, True):
+      with self.subTest(realize_clone=realize_clone):
+        a = Tensor([2.])
+        b = a.clone()
+        if realize_clone: b.realize()
+        b.assign(7.).realize()
+        self.assertEqual(a.tolist(), [2.])
+        self.assertEqual(b.tolist(), [7.])
+        self.assertIsNot(a.uop.base.buffer, b.uop.base.buffer)
+
+  def test_clone_preserves_creation_copy(self):
+    source = Tensor([2.], device="PYTHON")
+    copied = source.to("CPU")
+    cloned = copied.clone().realize()
+    source.assign(7.).realize()
+    self.assertEqual(copied.tolist(), [2.])
+    self.assertEqual(cloned.tolist(), [2.])
+
   def test_clone_deviceless_const(self):
     t = Tensor(UOp.const(2.0).cast(dtypes.float)).clone()
     np.testing.assert_equal(t.numpy(), 2.0)

test/device/test_hcq.py

@@ -3,8 +3,7 @@ from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import mv_address, DEBUG, DEV
 from test.helpers import slow, replace_opts
 from tinygrad.device import Buffer, BufferSpec
-from extra.hcq1.hcq import HCQCompiled
-from tinygrad.runtime.support.hcq import HCQBuffer
+from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runtime
@@ -236,7 +235,7 @@ class TestHCQ(unittest.TestCase):
       buf2 = Buffer(Device.DEFAULT, sz, dtypes.int8, options=BufferSpec(host=True, nolru=True)).ensure_allocated()
 
       ctypes.memset(buf2._buf.va_addr, 0x3e, sz)
-      buf2_q_view = buf2.host.view(fmt='Q')
+      buf2_q_view = buf2._buf.cpu_view().view(fmt='Q')
       for i in range(0, sz//8, 0x1000):
         for j in range(32): buf2_q_view[min(max(i + j - 16, 0), (sz // 8) - 1)] = random.randint(0, 0xffffffffffffffff)
 
@@ -568,7 +567,7 @@ class TestHCQ(unittest.TestCase):
 
     sz = 0x2000
     cpu_buffer = Buffer("CPU", sz, dtypes.uint8, options=BufferSpec(cpu_access=True)).ensure_allocated()
-    cpu_buffer.host.view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
+    cpu_buffer._buf.cpu_view().view(fmt='B')[:] = bytes([x & 0xff for x in range(sz)])
 
     for devid in range(6):
       if DEBUG >= 2: print(f"Testing map to device {Device.DEFAULT}:{devid}")

test/device/test_hcq2.py

@@ -1,12 +1,11 @@
 import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
 from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
-from tinygrad.device import Buffer
+from tinygrad.device import Buffer, BufferSpec
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import Context, dedup, partition, unwrap
+from tinygrad.helpers import Context, dedup, partition
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
 from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
-from tinygrad.codegen import do_to_program
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.c import init_c_struct_t
@@ -29,10 +28,8 @@ def chain(x:Tensor, n:int) -> Tensor:
 @contextlib.contextmanager
 def encoded_batches():
   batches, orig = [], hcq2.lower_and_compile
-  def track(l, *args, **kwargs):
-    batches.extend(c.without_after for c in l.src if call_is_hcq(c))
-    return orig(l, *args, **kwargs)
-  with patch.object(hcq2, "lower_and_compile", track): yield batches
+  with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
+    yield batches
 
 def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
   for _ in range(n): x = (x + 1).contiguous()
@@ -45,11 +42,6 @@ def patch_words(batch:UOp) -> list[UOp]:
 def rt_params(batch:UOp) -> list[str]:
   return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
 
-def cpu_buf(size:int=1, dtype=dtypes.uint8, **kwargs) -> UOp: return UOp.placeholder((size,), dtype, device="CPU", **kwargs)
-
-def lower_hcq(body:UOp) -> UOp:
-  return unwrap(hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test")).call(aux=hcq2.HCQInfo(("CPU",)))))
-
 class TestHCQ2Deps(unittest.TestCase):
   def test_disjoint_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="CPU")
@@ -69,8 +61,8 @@ class TestHCQ2Deps(unittest.TestCase):
       self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
       self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
 
-@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
-class TestHCQ2Schedule(unittest.TestCase):
+@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
+class TestHCQ2Core(unittest.TestCase):
   @staticmethod
   def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
 
@@ -101,33 +93,6 @@ class TestHCQ2Schedule(unittest.TestCase):
     linked = link_linear(compiled, input_uops=inputs)
     self.assertIs(link_linear(compiled, input_uops=inputs), linked)
 
-  def test_profile_slots_survive_indirect_access(self):
-    pm = PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat())),), allow_any_len=True),
-                          lambda buf: hcq2.rt_addr(buf, "CPU") if hcq2.unwrap_view(buf)[0].tag == "slots" else None)])
-    with patch.object(Device[Device.DEFAULT], "pm_lower", pm):
-      compiled = compile_linear(Tensor.ones(4).contiguous().schedule_linear(), profile=True)
-    self.assertFalse(any(param.op is Ops.PARAM and (param.arg.name or "").startswith("slots_")
-                         for param in compiled.src[0].without_after.src[0].toposort()))
-    call = link_linear(compiled).src[0].without_after
-    ((device, index),) = call.arg.aux.slots
-    self.assertEqual(device, Device.DEFAULT)
-    self.assertEqual(call.src[1 + index].buffer.dtype, dtypes.uint64)
-
-  def test_host_copies(self):
-    dev = Device[Device.DEFAULT]
-    if not dev.has_copy_queue: self.skipTest("copy queue required")
-    for host_device in ("CPU", "NPY", "DISK"):
-      for direct in (False, True):
-        for upload in (False, True):
-          with self.subTest(host_device=host_device, direct=direct, upload=upload):
-            host, gpu = UOp.new_buffer(host_device, 4, dtypes.uint8), UOp.new_buffer(dev.device, 4, dtypes.uint8)
-            src, dst = (host, gpu) if upload else (gpu, host)
-            linear = UOp(Ops.LINEAR, src=(src.copy_to_device(dst.device).call(dst, src),))
-            with patch.object(dev, "host_devs", frozenset({"CPU", host_device}) if direct else frozenset({"CPU"})):
-              compiled = compile_linear(linear, profile=False)
-            self.assertEqual(len(compiled.src), 1 if direct or host_device == "CPU" else 2)
-            self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
-
   def test_large_eager_not_cached(self):
     _, compiled, inputs = self.compiled(65)
     linked = link_linear(compiled, input_uops=inputs)
@@ -256,70 +221,62 @@ class TestHCQ2Schedule(unittest.TestCase):
     refs = [u.buffer for u in call.src[1:] if u.op is Ops.BUFFER]
     for n in names: self.assertTrue(any(r is getattr(dev, n) for r in refs), f"{n} is not a ref of the call")
 
-  def test_usb_renumbering(self):
-    programs = []
-    with Context(HCQ_RUNTIME_DEV="CPU"), patch("tinygrad.codegen.do_to_program", wraps=do_to_program) as build:
-      for ids in ((0, 1, 2, 3), (2, 0, 3, 1), (1, 0, 2, 3), (0, 1, 3, 2), (100, 101, 102, 103)):
-        with self.subTest(ids=ids):
-          regs = [UOp.placeholder((1,), dtypes.uint32, slot=i, addrspace=AddrSpace.REG) for i in ids[:2]]
-          a, b = [r.after(r.index(0).store(v)) for r, v in zip(regs, (3, 5))]
-          i, j = [UOp.range(UOp(Ops.NOOP), n, dtype=dtypes.void, src=(a, b)) for n in ids[2:]]
-          out = cpu_buf(dtype=dtypes.uint32, tag="out")
-          body = out.index(0).store(a.after(i, j).index(0).load()*10 + b.index(0).load()).end(j, UOp.const(False)).end(i, UOp.const(False))
-          compiled = lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),)))
-          programs.append(compiled.src[0].without_after.src[0])
-          self.assertIs(programs[-1], programs[0])
-          linear = hcq2.hcq_link(compiled, allow_cache=False)
-          run_linear(linear, jit=True)
-          self.assertEqual(linear.src[0].without_after.src[1].buffer.host.view(fmt='I')[0], 35)
-      self.assertLessEqual(build.call_count, 1)
-
-  def test_patched_view(self):
-    with Context(HCQ_RUNTIME_DEV="CPU"):
-      ctx = hcq2.EncodeCtx(("CPU",))
-      inner = hcq2.patch(cpu_buf(8, tag="inner"), [(4, UOp.const(42, dtypes.uint32))], bytes(8))
-      inner = unwrap(hcq2.hoist_links(ctx, inner))
-      outer = hcq2.patch(cpu_buf(8, tag="outer"), [(0, inner[4:8].getaddr("CPU"))])
-      with patch.object(hcq2, "EncodeCtx", return_value=ctx): call = lower_hcq(outer.bitcast(dtypes.uint64).index(0).load())
-      self.assertEqual(call.without_after.arg.aux.nargs, 1)
-      self.assertTrue(all(s.op is Ops.STORE for s in call.src[1:]))
-      linked = hcq2.hcq_link(UOp(Ops.LINEAR, src=(call,)), allow_cache=False).src[0]
-      inner_buf, outer_buf = linked.src[1].buffer, linked.without_after.src[1].buffer
-      self.assertEqual(inner_buf.host.view(fmt='I')[1], 42)
-      self.assertEqual(outer_buf.host.view(fmt='Q')[0], inner_buf._buf.va_addr + 4)
-
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):
   @staticmethod
   def _run(body:UOp) -> list[Buffer]:
-    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),))), allow_cache=False)
+    call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
+    assert call is not None
+    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), allow_cache=False)
     run_linear(linear, jit=True)
     return [u.buffer for u in linear.src[0].without_after.src[1:] if u.op is Ops.BUFFER]
 
   def test_ffi_ccall(self):
     with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = cpu_buf(dtype=dtypes.int32, slot=1, volatile=True, tag="ffi_result")
+      out = UOp.placeholder((1,), dtypes.int32, slot=1, device="CPU", volatile=True, tag="ffi_result")
       bufs = self._run(out.index(0).store(hcq2.ccall(libc.dll.ffs, 0x10)))
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int).host.view(fmt='i')[0], 5)
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 5)
 
   def test_ffi_cstruct(self):
     struct_t = init_c_struct_t(16, (("u8", ctypes.c_uint8, 0), ("u16", ctypes.c_uint16, 2),
                                   ("u32", ctypes.c_uint32, 4), ("u64", ctypes.c_uint64, 8)))
-    cpu_buf() # reserve slot zero for device-owned placeholders
+    UOp.placeholder((1,), dtypes.uint8, device="CPU") # reserve slot zero for device-owned placeholders
     with Context(HCQ_RUNTIME_DEV="CPU"):
       s = hcq2.cstruct(struct_t, u8=0x12, u16=UOp.const(0x3456, dtypes.uint16), u32=0x789ABCDE, u64=0xFEDCBA9876543210)
       bufs = self._run(s.index(0).load())
-    got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t)).host.view(fmt='B')))
+    got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
     self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
 
+  def test_device_lower_after_encode(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      out = UOp.placeholder((1,), dtypes.int32, device="CPU", tag="result")
+      encode = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_encode"), lambda: UOp.custom_function("test_lower"))])
+      lower = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="test_lower"), lambda out=out: out.index(0).store(42))])
+      with patch.object(Device["CPU"], "pm_encode", encode), patch.object(Device["CPU"], "pm_lower", lower):
+        bufs = self._run(UOp.custom_function("test_encode"))
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 42)
+
   def test_nested_cstruct_patches(self):
     with Context(HCQ_RUNTIME_DEV="CPU"):
       inner = hcq2.cstruct(init_c_struct_t(4, (("value", ctypes.c_uint32, 0),)), value=42)
       outer = hcq2.cstruct(init_c_struct_t(8, (("ptr", ctypes.c_uint64, 0),)), ptr=inner.getaddr("CPU"))
-      out = cpu_buf(dtype=dtypes.uint32, tag="result")
+      out = UOp.placeholder((1,), dtypes.uint32, device="CPU", tag="result")
       copied = hcq2.ccall(libc.memcpy, out.index(0), outer.bitcast(dtypes.uint64).index(0).load(), 4)
       bufs = self._run(out.after(copied).index(0).load())
-    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32).host.view(fmt='I')[0], 42)
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32)._buf.cpu_view().view(fmt='I')[0], 42)
+
+
+class TestHCQ2Timeline(unittest.TestCase):
+  def test_reused_timeline_is_zeroed(self):
+    buf = Buffer("CPU", 2, dtypes.uint64, options=BufferSpec(host=True, uncached=True, cpu_access=True), preallocate=True)
+    addr = buf._buf.va_addr
+    buf._buf.cpu_view().view(fmt='B')[:] = b'\xff' * 16
+    buf.deallocate()
+    dev = HCQ2Compiled.__new__(HCQ2Compiled)
+    dev.device = "CPU"
+    self.assertEqual(dev.timeline._buf.va_addr, addr)
+    self.assertEqual(bytes(dev.timeline._buf.cpu_view()), bytes(16))
+
 
 if __name__ == "__main__":
   unittest.main()

test/device/test_metal.py

@@ -58,5 +58,5 @@ kernel void r_5(device int* data0, const device int* data1, uint3 gid [[threadgr
 
     buf = device.allocator.alloc(size, BufferSpec(nolru=True))
     self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before+size, msg=f"{curr=} - {before=}")
-    device.allocator.free(buf, size, BufferSpec(nolru=True))
+    device.allocator.free(buf, buf.size, BufferSpec(nolru=True))
     self.assertEqual(curr:=device.sysdevice.currentAllocatedSize(), before, msg=f"{curr=} - {before=}")

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)[0][0], global_size=(1,1,1), local_size=(1,1,1), wait=True)')
+              'dev.allocator.alloc(64), 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)[0][0]
+buf = dev.allocator.alloc(64)
 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

@@ -13,7 +13,7 @@ class FakeProgram:
   def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
 
 class FakeAllocator(Allocator[Compiled]):
-  def _alloc(self, sz, options): return (None, None), None
+  def _alloc(self, sz, options): return None
   def _copyin(self, dest, src:memoryview): pass
 
 class TestLLaMASpeed(unittest.TestCase):

test/external/external_test_usb_asm24.py

@@ -1,7 +1,6 @@
 import unittest
 from tinygrad.helpers import Timing, getenv
-from tinygrad import Tensor, Device, TinyJit
-from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
+from tinygrad import Tensor, Device
 import numpy as np
 
 class USBTestCase(unittest.TestCase):
@@ -9,12 +8,7 @@ class USBTestCase(unittest.TestCase):
   def setUpClass(cls):
     cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
-    if not cls.dev.is_usb: raise unittest.SkipTest("only test this on USB devices")
-    cls.rng = np.random.default_rng(0)
-
-  def roundtrip(self, a:np.ndarray): # a copy in, a kernel, a copy out: the queue must order them
-    np.testing.assert_array_equal(a, Tensor(a, device="NPY").to(Device.DEFAULT).numpy())
-    np.testing.assert_array_equal(a + 1, (Tensor(a, device="NPY").to(Device.DEFAULT) + 1).numpy())
+    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
 
 class TestDevCopySpeeds(USBTestCase):
   def testCopyCPUtoDefault(self):
@@ -36,55 +30,73 @@ class TestUSBIntegrity(USBTestCase):
     t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
     Device[Device.DEFAULT].synchronize()
+
     y = x.to('CPU').realize()
+
     np.testing.assert_equal(t.numpy(), y.numpy())
+    del x, y, t
 
-  def testBoundaries(self): # around the slot, the chunk and the read window
-    for size in (1, 3, 508, 509, SLOT - 513, SLOT - 512, SLOT - 511, CHUNK - 1, CHUNK, CHUNK + 1, 2 * CHUNK - 1, 2 * CHUNK, 2 * CHUNK + 31, HALF,
-                 2 * HALF, 1 << 20):
-      with self.subTest(size=size): self.roundtrip(self.rng.integers(0, 256, size, dtype=np.uint8))
+  def testCopyinBoundaries(self):
+    rng, chunk = np.random.default_rng(0), 0x40000 - 4
+    for size in (1, 3, 508, 509, 0x3ffc, 0x3ffd, chunk, chunk+1, 2*chunk+31):
+      with self.subTest(size=size):
+        a = rng.integers(0, 256, size, dtype=np.uint8)
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
 
-  def testManyCopiesInABatch(self):
-    for n in (2, 7, 64, 300): # 300 chunks: the fence byte wraps
-      with self.subTest(n=n):
-        arrs = [self.rng.integers(0, 256, int(s), dtype=np.uint8) for s in self.rng.integers(1, 5000, n)]
-        ts = [Tensor(a, device="NPY").to(Device.DEFAULT) for a in arrs]
-        Tensor.realize(*ts)
-        for t, a in zip(ts, arrs): np.testing.assert_array_equal(a, t.numpy())
+  def testCopyinFenceWrap(self):
+    a = np.arange(2*(0x40000-4)+31, dtype=np.uint8)
+    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
+    self.dev.synchronize()
+    alloc, usb = self.dev.allocator, self.dev.iface.pci_dev.usb
+    clear = usb.read(0xA808, 1)
+    # Model a completed 256-chunk copy instead of the one-chunk warmup. The next clear tag must still change.
+    alloc._usb_seq += 255
+    usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
+    np.testing.assert_array_equal(a[:31], Tensor(a[:31], device="AMD").numpy())
+    self.assertNotEqual(clear, usb.read(0xA808, 1))
+    for bits in (8, 24):
+      with self.subTest(bits=bits):
+        alloc._usb_seq = ((alloc._usb_seq >> bits)+2)*(1 << bits)-2
+        usb.write(0xA800, bytes([alloc._usb_seq & 0xff]))
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
 
-  def testMixedBatch(self): # copies out and in, in one batch: runs of both directions
-    arrs = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
-    ts = [Tensor(a, device="NPY").to(Device.DEFAULT).realize() for a in arrs]
-    more = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
-    outs = [t.to("NPY") for t in ts] + [Tensor(a, device="NPY").to(Device.DEFAULT) for a in more]
-    Tensor.realize(*outs)
-    for o, a in zip(outs, arrs + more): np.testing.assert_array_equal(a, o.numpy())
-
-  def testRepeatedBatches(self): # a batch numbers its chunks from 0: the same batch again must not see what the last one left behind
-    a = self.rng.integers(0, 256, 2 * CHUNK + 31, dtype=np.uint8)
-    for _ in range(5): self.roundtrip(a)
-    @TinyJit
-    def step(x:Tensor) -> Tensor: return (x + 1).realize()
-    src = Tensor(a, device="NPY")
-    for i in range(5):
-      x = src.to(Device.DEFAULT)
-      np.testing.assert_array_equal(a + 1, step(x).numpy())
-
-  def testStaleSentinel(self): # payloads full of the tags the queue waits for, in both directions, before and around the real chunks
-    tags = np.array([0x51000000 | k for k in range(8)], dtype=np.uint32)
-    for tag in tags: # every dword of every chunk is the tag of some chunk of the copy
-      with self.subTest(payload=hex(tag)):
-        a = np.full((2 * CHUNK + 31) // 4, tag, dtype=np.uint32).view(np.uint8)
-        self.roundtrip(a)
-    with self.subTest(case="copyout residue"): # a read fills the sram with tags, then small chunks land in both halves
-      a = np.tile(tags, 2 * CHUNK // 32).view(np.uint8)
-      np.testing.assert_array_equal(a, (Tensor(a, device="NPY").to(Device.DEFAULT) * 1).numpy())
-      for size in (31, CHUNK + 31, 2 * CHUNK + 31): self.roundtrip(np.tile(tags, size // 32 + 1).view(np.uint8)[:size])
-
-  def testRingWrap(self): # 64MB of chunks: the sdma ring (1MB on usb) wraps within the copy
-    a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
-    t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
+  def testCopyinRingWrap(self):
+    rng = np.random.default_rng(0)
+    a = rng.integers(0, 256, 1 << 20, dtype=np.uint8)
+    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+    ring = self.dev.sdma_queue(0)
+    # A 16 MiB copyin needs more than 4 KiB of SDMA packets, forcing the submission to wrap.
+    target = ring.ring.nbytes - 0x1000
+    padding = target - ring.put_value % ring.ring.nbytes - 16  # four-dword timeline fence
+    self.assertGreaterEqual(padding, 0)
+    q = self.dev.hw_copy_queue_t()
+    q.q(*([0] * (padding // 4)))
+    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+    self.dev.synchronize()
+    before = ring.put_value // ring.ring.nbytes
+    a = rng.integers(0, 256, 16 << 20, dtype=np.uint8)
+    t = Tensor(a, device="AMD").realize()
+    self.assertGreater(ring.put_value // ring.ring.nbytes, before)
     np.testing.assert_array_equal(a, t.numpy())
 
+  def testCopyinStaleSentinel(self):
+    a = np.arange(16, dtype=np.uint8)
+    np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+    chunk = 0x40000 - 4
+    for case in ("copyout", "reuse"):
+      with self.subTest(case=case):
+        if case == "copyout":
+          # A 512 KiB copyin takes three chunks. Copyout then fills both SRAM windows with the next expected tag.
+          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 3) & 0xFFFFFF)
+          a = np.full(0x80000 // 4, tag, dtype=np.uint32)
+          np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+          a = np.arange(31, dtype=np.uint8)
+        else:
+          # The first full chunk contains the tag expected by the short third chunk in the same window.
+          tag = 0x51000000 | ((self.dev.allocator._usb_seq + 2) & 0xFFFFFF)
+          a = np.arange(2 * chunk + 31, dtype=np.uint8)
+          a[:chunk].view(np.uint32)[:] = tag
+        np.testing.assert_array_equal(a, Tensor(a, device="AMD").numpy())
+
 if __name__ == "__main__":
   unittest.main()

test/external/fuzz_fast_idiv.py

@@ -12,8 +12,7 @@ if __name__ == "__main__":
     if i % 1000 == 0:
       print(f"Progress: {i}")
     dt = random.choice(dtypes.ints)
-    vmax = random.randint(1, 2**random.randint(1, dt.max.bit_length()))
-    u = UOp.variable('x', random.randint(0, vmax-1) if vmax > 1 else 0, vmax, dtype=dt)
+    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
     d = random.randint(1, max(1, u.vmax)*2)
     if d in powers_of_two: continue
     expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)

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))[0][0]) for i in range(100)]
+  buffers = [(rd:=random.choice(dev), rd.allocator.alloc(random.randint(1, 10000))) for i in range(100)]
 
   for _ in trange(100000):
     d1, b1 = random.choice(buffers)

test/helpers.py

@@ -67,7 +67,7 @@ def assert_kernel_count(expected:int):
 
 def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
   from tinygrad.runtime.support.hcq2 import HCQ_DEVS
-  return Device.DEFAULT.split(":")[0] in HCQ_DEVS
+  return Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}
 
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
@@ -125,12 +125,12 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   allocator = dev.allocator
   bufs = []
   for buf_dt, data in inputs or []:
-    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)[0][0])
+    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
     allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
   g = UOp.param(0, uop.dtype, 1)
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
-  prog(out_buf:=allocator.alloc(uop.dtype.itemsize)[0][0], *bufs, vals=vals)
+  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *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/null/test_helpers.py

@@ -78,13 +78,6 @@ class TestContextVars(unittest.TestCase):
     test()
     self.assertEqual(VARIABLE.value, 0)
 
-  def test_decorator_recursive(self):
-    @Context(VARIABLE=1)
-    def test(n):
-      if n: test(n-1)
-    test(2)
-    self.assertEqual(VARIABLE.value, 0)
-
   def test_context_exit_reverts_updated_values(self):
     D = ContextVar("D", 1)
     D.value = 2

test/null/test_schedule.py

@@ -582,8 +582,8 @@ class TestSchedule(unittest.TestCase):
     p = P[0]
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
-    # assign on a pending contiguous overwrites the whole value, no store hazard
-    check_schedule(p, 3)
+    # TODO: this should be 3 if fix store hazard worked correctly
+    check_schedule(p, 4)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
     self.enterContext(Context(DEFAULT_FLOAT=dtype))

test/null/test_tensor_uop_mixin.py

@@ -246,10 +246,8 @@ class TestTensorUOpRand(unittest.TestCase):
     self.assertIs(Tensor._threefry_random_bits(Tensor(key), Tensor(c0), Tensor(c1)).uop, UOp._threefry_random_bits(key, c0, c1))
   def test_rand(self):
     k, c = UOp.empty((2,), dtype=dtypes.uint32), UOp.zeros(2, dtype=dtypes.uint32)
-    self.assertIs(_strip_unique(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop),
-                  _strip_unique(UOp._rand(k, c, (2, 2), dtypes.float32)))
-    self.assertIs(_strip_unique(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop),
-                  _strip_unique(UOp._rand(k, c, (0, 3), dtypes.float32)))
+    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop, UOp._rand(k, c, (2, 2), dtypes.float32))
+    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop, UOp._rand(k, c, (0, 3), dtypes.float32))
 
 class TestTensorUOpGather(unittest.TestCase):
   def _check(self, t, dim, idx):

test/null/test_uop_symbolic.py

@@ -919,11 +919,6 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a % -8) // 2, -4, 0, "(a%-8//2)")
     self.helper_test_variable((a % -8) % 2, 0, 1, "(a%2)")
 
-  def test_nested_div_mod_symbolic_inner_divisor(self):
-    a = Variable("a", 0, 100)
-    self.helper_test_variable((a % (Variable("n", 1, 10)*4)) // 2, 0, 19, "(a//2%(n*2))")
-    check_uop_against_string(self, (a % (Variable("n", 0, 10)*4) // 2).simplify(), "(a%(n*4)//2)")
-
   def test_floordiv_lt_negative_c(self):
     # x//d<c with negative c also reduces to x<c*d for d>0
     idx = Variable("idx", -20, 20)

test/null/test_uops.py

@@ -269,6 +269,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
     for x in gated_uops: self.assertIs(x.op, Ops.STORE)
     for x in gated_uops: self.assertEqual(len(x.src), 2)
 
+@unittest.skipIf(Device.DEFAULT == "METAL", "compiler bug")
 @unittest.skipUnless(Ops.SHR in Device[Device.DEFAULT].renderer.code_for_op, "fast_idiv requires SHR")
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
@@ -309,7 +310,7 @@ class TestFastIdiv(unittest.TestCase):
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @Context(DISABLE_FAST_IDIV=0)
-  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fast_idiv widens uint32 to uint64")
+  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
@@ -328,25 +329,6 @@ class TestFastIdiv(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.CMOD, ops)
 
-  @Context(DISABLE_FAST_IDIV=0)
-  def test_fast_idiv_nonpositive_divisor(self):
-    ridx = UOp.range(20, 0)
-    for d in (-3, 0):
-      for op in (Ops.CDIV, Ops.CMOD):
-        ops = [x.op for x in to_uops_list([ridx.alu(op, UOp.const(d))], ren=Device[Device.DEFAULT].renderer)]
-        self.assertNotIn(Ops.SHR, ops, f"fast_idiv fired on {op} by {d}")
-
-  @Context(DISABLE_FAST_IDIV=0)
-  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "needs a uint64 buffer")
-  def test_fast_idiv_cmod_kept_when_idiv_declines(self):
-    ren = Device[Device.DEFAULT].renderer
-    d = UOp.param(0, dtypes.int32, 4).index(UOp.const(0))
-    ops = [x.op for x in to_uops_list([UOp.range(30, 0).alu(Ops.CMOD, d)], ren=ren)]
-    self.assertIn(Ops.CMOD, ops, "CMOD by a non-const divisor should be left alone")
-    big = UOp.param(1, dtypes.uint64, 4).index(UOp.const(0))
-    ops = [x.op for x in to_uops_list([big.alu(Ops.CMOD, UOp.const(3, dtypes.uint64))], ren=ren)]
-    self.assertIn(Ops.CMOD, ops, "CMOD should be left alone when fast_idiv declines")
-
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
@@ -360,7 +342,6 @@ class TestFastIdiv(unittest.TestCase):
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
     self.assertNotIn(dtypes.long, [x.dtype for x in uops])
-    self.assertNotIn(Ops.CDIV, [x.op for x in uops])
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/null/test_validate_oob.py

@@ -186,20 +186,40 @@ class TestValidateOOB(unittest.TestCase):
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8
 
-  # local memory
-  def test_gated_local(self):
-    with Context(CHECK_OOB=1, SPEC=2):
+  # skipped tests (moved from test_uop_graph.py)
+  @unittest.skip("if not allowed in graph")
+  def test_in_bounds_access_gated_local(self):
+    with Context(CHECK_OOB=1):
+      # Define buffers
       gbuf = UOp.param(0, dtypes.uint, 400)
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
+
+      # Define indices, valids and barrier
       gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
       lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
-      store = sbuf.index(lidx.valid(lidx < 8)).store(UOp.const(1))
-      load = sbuf.after(store).index(lidx.valid(lidx < 8)).load()
-      to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(load)])  # valid: local store and load gated to 8, global store gated to 400
-      with self.assertRaises(RuntimeError):
-        to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(sbuf.after(store).index(lidx).load())])  # lidx 0..9 into 8
-      with self.assertRaises(RuntimeError):
-        to_uops_list([gbuf.index(gidx).store(load)])  # gidx 0..415 into 400
+
+      gate = (gidx<400) & (lidx<8)
+
+      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
+
+      barrier = UOp(Ops.BARRIER, src=(local_store,))
+      if_barrier = UOp(Ops.IF, src=(gate, barrier))
+
+      # Load from local memory (after the IF/barrier)
+      local_load = UOp(Ops.LOAD, src=(sbuf.index(lidx), if_barrier))
+
+      # Store to global memory
+      global_store = UOp(Ops.STORE, src=(gbuf.index(gidx), local_load))
+      to_uops_list([global_store])
+
+  @unittest.skip("Bool load is not supported yet")
+  def test_load_mask(self):
+    with Context(CHECK_OOB=1):
+      glbl0 = UOp.param(0, dtypes.int, 16)
+      mask = UOp.param(0, dtypes.bool, 16)
+      ridx = UOp.range(20, 0)
+      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
+      to_uops_list([ld0])
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_viz.py

@@ -454,7 +454,7 @@ class TestVizIntegration(unittest.TestCase):
   def test_jit(self):
     with save_viz():
       @TinyJit
-      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(a.to(c.device)).contiguous(), b.assign(c.to(b.device))
+      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
       a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
       for _ in range(3): Tensor.realize(*f(a, b, c))
     out = load_profile(cpu_events)

test/testextra/test_hevc.py

@@ -98,7 +98,7 @@ class TestHevc(unittest.TestCase):
       Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
 
     compiled = compile_linear(decoded.linear_with_vars()[0])
-    self.assertTrue(any(call.without_after.src[0].op is Ops.PROGRAM for call in compiled.src))
+    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
     encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
     self.assertEqual(len(encdec_calls), 1)
 

test/unit/test_after.py

@@ -1,91 +0,0 @@
-import unittest
-from tinygrad import Tensor
-
-class TestAfterCounterexamples(unittest.TestCase):
-  def test_ordered_writes_allowed(self):
-    x = Tensor([0.]).realize().uop
-    a = x.after(x.store(1))
-    b = a.after(a.store(2))
-    self.assertEqual(Tensor(b).tolist(), [2.])
-
-  def test_disjoint_writes_allowed(self):
-    x = Tensor([0., 0.]).realize().uop
-    y = Tensor(x.after(x[:1].store(1), x[1:].store(2)))
-    self.assertEqual(y.tolist(), [1., 2.])
-
-  def test_read_modify_write_chain(self):
-    x = Tensor([2.]).clone()
-    x.assign(x + 1)
-    x.assign(x * 2)
-    self.assertEqual(x.tolist(), [6.])
-
-  def test_overwrite_cuts_gradient(self):
-    x = Tensor([2.])
-    y = x.clone()
-    y.assign(3)  # overwriting with a constant makes y independent of x
-    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0.])
-
-  def test_shared_state_readers(self):
-    x = Tensor([2.]).clone()
-    x.assign(x + 1)
-    a, b = x + 1, x * 2
-    Tensor.realize(a, b)
-    self.assertEqual(a.tolist(), [4.])
-    self.assertEqual(b.tolist(), [6.])
-
-  @unittest.expectedFailure
-  def test_chained_square_assign_gradient(self):
-    x = Tensor([2.0])
-    y = x.clone()
-    y.assign(y*y)
-    y.assign(y*y)
-    # y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
-    self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
-
-  @unittest.expectedFailure
-  def test_partial_store_gradient(self):
-    x = Tensor([2., 3.]).realize()
-    y = Tensor(x.uop.after(x[:1].uop.store(4)))
-    # y = [4, x[1]]. Currently returns [0., 0.].
-    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
-
-  @unittest.expectedFailure
-  def test_partial_store_source_gradient(self):
-    x = Tensor([4.])
-    y = Tensor([2., 3.]).realize()
-    z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
-    # x contributes once, not twice. Currently returns [2.].
-    self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
-
-  def test_unrelated_store_gradient(self):
-    x = Tensor([2.]).realize()
-    y = x.clone()
-    z = Tensor(x.uop.after(y.uop.store(0)))
-    # 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.
-    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
-
-  @unittest.expectedFailure
-  def test_unordered_overlapping_stores_rejected(self):
-    x = Tensor([0.]).realize().uop
-    # No ordering between the writes. Currently succeeds with [2.].
-    with self.assertRaises(RuntimeError):
-      Tensor(x.after(x.store(1), x.store(2))).realize()
-
-  @unittest.expectedFailure
-  def test_gradient_after_callify(self):
-    x = Tensor([2.]).realize()
-    y = x * 2
-    y.callify()
-    # Currently raises: "expected a CALL with unbound BUFFER outputs or a grad_fxn".
-    self.assertEqual(y.sum().gradient(x)[0].tolist(), [2.])
-
-if __name__ == "__main__":
-  unittest.main()

test/unit/test_buffer.py

@@ -1,46 +0,0 @@
-import unittest
-from tinygrad.device import Buffer
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import Context
-
-class TestBuffer(unittest.TestCase):
-  def test_host_view(self):
-    b = Buffer("CPU", 4, dtypes.uint32)
-    v = b.view(2, dtypes.uint16, 4)
-    host = v.host
-    host.view(fmt='H')[0] = 0x1234
-    self.assertEqual(b.host.view(fmt='H')[2], 0x1234)
-    self.assertEqual(v._buf.va_addr, b._buf.va_addr + 4)
-    self.assertIs(v.host, host)
-    self.assertIs(v.meta, b.meta)
-
-  def test_mapping(self):
-    b = Buffer("CPU", 8, dtypes.uint8, initial_value=b"abcdefgh")
-    self.assertIs(b.get_storage("PYTHON")[0][1], b.get_buf("PYTHON"))
-    v = b.view(4, dtypes.uint8, 2)
-    mapped = v.get_storage("PYTHON")
-    self.assertEqual(bytes(mapped[0][0]), b"cdef")
-    self.assertIs(mapped[1], v.host)
-    self.assertIsNone(mapped[0][1])
-    self.assertIs(v.get_storage("PYTHON")[0], mapped[0])
-
-  def test_view_reallocation(self):
-    b = Buffer("CPU", 8, dtypes.uint8)
-    v = b.view(4, dtypes.uint8, 2)
-    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")[0], old)
-    self.assertEqual(bytes(v.get_buf("PYTHON")), b"test")
-
-  def test_cache_owned_storage_only(self):
-    for opaque in (None, memoryview(bytearray(8))):
-      with self.subTest(imported=opaque is not None), Context(LRU=1):
-        b = Buffer("PYTHON", 8, dtypes.uint8, opaque=opaque)
-        buf = b._buf
-        b.deallocate()
-        self.assertEqual(b._buf is buf, opaque is None)
-
-if __name__ == "__main__": unittest.main()

test/unit/test_call_access.py

@@ -0,0 +1,66 @@
+import unittest
+from tinygrad import Tensor, Context, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
+
+
+class TestCallAccess(unittest.TestCase):
+  def test_computed_reads_writes_and_unused_arguments(self):
+    out, x, unused = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(3))
+    body = out.store(x + 1).sink(arg=KernelInfo())
+    self.assertEqual(body.call(out, x, unused).call_access(), ((x,), (out,)))
+
+  def test_computed_read_modify_write(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    body = out.store(out + x).sink(arg=KernelInfo())
+    self.assertEqual(body.call(out, x).call_access(), ((out, x), (out,)))
+
+  def test_computed_empty_effects(self):
+    x = UOp.param(0, dtypes.float, (1,), "CPU")
+    self.assertEqual(UOp.sink(x, arg=KernelInfo()).call(x).call_access(), ((), ()))
+
+  def test_computed_program_accesses(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    sink = out.store(x.load()).sink(arg=KernelInfo())
+    program = UOp(Ops.PROGRAM, src=(sink,), arg=ProgramInfo.from_sink(sink))
+    self.assertEqual(program.call(out, x).call_access(), ((x,), (out,)))
+
+  def test_nested_linear_parameter_scopes(self):
+    a, b, c = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(3))
+    inner = a.store(b + 1).sink(arg=KernelInfo()).call(b, a)
+    body = UOp(Ops.LINEAR, src=(inner,))
+    self.assertEqual(body.call(a, b, c).call_access(), ((a,), (b,)))
+
+  def test_copy_accesses(self):
+    out, x = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    self.assertEqual(UOp(Ops.COPY, src=(x,), arg=out.device).call(out, x).call_access(), ((x,), (out,)))
+
+  def test_unknown_opaque_accesses_reject(self):
+    x = UOp.param(0, dtypes.float, (1,), "CPU")
+    bodies = (UOp(Ops.PROGRAM, src=(UOp.sink(x),)),
+              UOp(Ops.CUSTOM, src=(x,), arg=("", dtypes.void)).sink(arg=KernelInfo()))
+    for body in bodies:
+      with self.assertRaisesRegex(RuntimeError, "cannot compute accesses"): body.call(x).call_access()
+
+  @Context(DEV="CPU")
+  def test_unknown_effects_do_not_replace_tensors_on_failure(self):
+    def kernel(x): return UOp(Ops.PROGRAM, src=(UOp.sink(x, arg=KernelInfo()),))
+    x = Tensor([2.]).realize().custom_kernel(fxn=kernel)[0]
+    before = x.uop
+    for _ in range(2):
+      with self.assertRaisesRegex(RuntimeError, "cannot compute accesses"): x.realize()
+      self.assertIs(x.uop, before)
+
+  def test_bad_access_slots(self):
+    arg = UOp.param(0, dtypes.float, (1,), "CPU")
+    for slot in (-1, 1):
+      p = UOp(Ops.PROGRAM, src=(UOp.sink(arg),), arg=ProgramInfo(globals=(0,), ins=(slot,), outs=()))
+      with self.assertRaisesRegex(RuntimeError, "invalid CALL access slot"): p.call(arg, arg).call_access()
+
+  def test_compiled_writable_alias_rejects(self):
+    a, b = (UOp.param(i, dtypes.float, (1,), "CPU") for i in range(2))
+    sink = a.store(b.load()).sink(arg=KernelInfo())
+    program = UOp(Ops.PROGRAM, src=(sink,), arg=ProgramInfo.from_sink(sink))
+    with self.assertRaisesRegex(RuntimeError, "aliased opaque"): program.call(a, a).call_access()
+
+
+if __name__ == "__main__": unittest.main()

test/unit/test_gradient.py

@@ -88,6 +88,72 @@ class TestTensorGradient(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])     # gradient flows through clone
     np.testing.assert_allclose(base.grad.numpy(), [0.0, 0.0, 0.0, 0.0])  # ...but detach blocks it from base
 
+  def test_gradient_through_single_assign(self):
+    x = Tensor([2., 3.]).realize()
+    y = x.clone()
+    y.assign(y.square())
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [4., 6.])
+
+  def test_gradient_through_assign_requires_old_versions(self):
+    for count in (2, 3):
+      with self.subTest(count=count):
+        x = Tensor([2., 3.]).realize()
+        y = x.clone()
+        for _ in range(count): y.assign(y.square())
+        g = y.sum().gradient(x)[0]
+        before = (x.uop, y.uop, g.uop)
+        # Reject incompatible versions, including on retry: failed scheduling must not replace them with buffers.
+        for _ in range(2):
+          with self.assertRaisesRegex(RuntimeError, "cycle"): g.realize()
+          self.assertEqual((x.uop, y.uop, g.uop), before)
+
+  def test_gradient_through_assign_with_snapshots(self):
+    x = Tensor([2., 3.]).realize()
+    y = x.clone()
+    for _ in range(2): y.assign(y.clone().square())
+    g = y.sum().gradient(x)[0]
+    gg = g.sum().gradient(x)[0]
+    Tensor.realize(g, gg)
+    self.assertEqual(g.tolist(), [32., 108.])
+    self.assertEqual(gg.tolist(), [48., 108.])
+
+  def test_gradient_after_unrelated_store(self):
+    x, v, dst = Tensor([2.]).realize(), Tensor([3.]).realize(), Tensor.empty(1)
+    y = Tensor(x.uop.after(dst.uop.store(v.uop)))
+    self.assertEqual([g.tolist() for g in y.sum().gradient(x, v)], [[1.], [0.]])
+    self.assertEqual(y.tolist(), [2.])
+    self.assertEqual(dst.tolist(), [3.])
+
+  def test_gradient_after_multiple_unrelated_stores(self):
+    x, a, b = Tensor([2.]).realize(), Tensor.empty(1), Tensor.empty(1)
+    y = Tensor(x.uop.after(a.uop.store(x.uop * 3), b.uop.store(x.uop * 4)))
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [1.])
+
+  def test_gradient_after_readonly_call(self):
+    x = Tensor([2.]).realize()
+    def kernel(dst, src): return dst.store(src * 3).sink(arg=KernelInfo())
+    for grad_fxn in (None, lambda g, k: (None, g * 3)):
+      _, unchanged = Tensor.empty(1).custom_kernel(x, fxn=kernel, grad_fxn=grad_fxn)
+      self.assertEqual(unchanged.sum().gradient(x)[0].tolist(), [1.])
+
+  def test_gradient_after_unrelated_call(self):
+    x, v, dst = Tensor([2.]).realize(), Tensor([3.]).realize(), Tensor.empty(1)
+    p, q = dst.uop.param_like(0), v.uop.param_like(1)
+    call = p.store(q * 3).sink(arg=KernelInfo()).call(dst.uop, v.uop, grad_fxn=lambda g, k: (None, g * 3))
+    y = Tensor(x.uop.after(call))
+    self.assertEqual([g.tolist() for g in y.sum().gradient(x, v)], [[1.], [0.]])
+
+  def test_gradient_after_aliased_store_view_rejects(self):
+    x = Tensor([2., 3.]).realize()
+    y = Tensor(x.uop.after(x.uop.shrink(((0, 1),)).store(4.)))
+    with self.assertRaisesRegex(RuntimeError, "aliased write"): y.sum().gradient(x)
+
+  def test_gradient_after_duplicate_call_output_rejects(self):
+    x = Tensor([2.]).realize()
+    def kernel(a, b): return a.store(b * 2).sink(arg=KernelInfo())
+    y = x.custom_kernel(x, fxn=kernel, grad_fxn=lambda g, k: (g, g))[0]
+    with self.assertRaisesRegex(RuntimeError, "ambiguous CALL"): y.sum().gradient(x)
+
   def test_setitem_on_grad_used_tensor_raises(self):
     x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
     _ = (x * 2.0).sum()
@@ -135,11 +201,22 @@ class TestTensorGradient(unittest.TestCase):
     (Tensor.rand(()) + w).backward()
     self.assertIsNone(w.grad)
 
-  def test_max_backward_many_ties(self):
-    t = Tensor.ones(70000, dtype=dtypes.half).contiguous()
-    np.testing.assert_allclose(t.max().gradient(t)[0].sum().numpy(), 1.0, atol=1e-3)
-
 class TestMultiOutputGradient(unittest.TestCase):
+  def test_custom_kernel_inplace_gradient(self):
+    def double(x:UOp): return x[0].store(x[0]*2).sink(arg=KernelInfo(name="double_inplace"))
+    def backward(g:UOp, call:UOp): return (g*2,)
+    x = Tensor([2.]).realize()
+    y = x.custom_kernel(fxn=double, grad_fxn=backward)[0]
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [2.])
+    self.assertEqual(y.tolist(), [4.])
+
+  def test_custom_kernel_unchanged_output_gradient(self):
+    def noop(x:UOp): return x[0].store(x[0]).sink(arg=KernelInfo(name="identity"))
+    def backward(g:UOp, call:UOp): return (g,)
+    x = Tensor([2.]).realize()
+    y = x.custom_kernel(fxn=noop, grad_fxn=backward)[0]
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [1.])
+
   @staticmethod
   def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
     C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()

test/unit/test_hcq_graph.py

@@ -3,8 +3,8 @@ from tinygrad import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
-from extra.hcq1.graph import HCQGraph
-from extra.hcq1.hcq import HCQCompiled
+from tinygrad.runtime.graph.hcq import HCQGraph
+from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 

tinygrad/codegen/__init__.py

@@ -7,7 +7,7 @@ from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.renderer.isa import ISARenderer, IselContext
+from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
 from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
@@ -439,13 +439,12 @@ def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
   lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
   # isa renderers need to allocate registers
   if isinstance(ctx, ISARenderer):
-    lin_ctx = ctx.linear_ctx_type(ctx)
-    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, lin_ctx)
+    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
     # register definitions (INS without srcs) move to the top so regalloc sees their live ranges span the whole program (callee saved regs)
     lst = sorted(lst, key=lambda u: u.op is not Ops.INS or bool(u.src))
-    regalloc_ctx = LinearScanRegallocContext(lin_ctx, lst, ctx)
+    regalloc_ctx = LinearScanRegallocContext(lst, ctx)
     lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx)
-    lst = line_rewrite(lst, ctx.post_regalloc_matcher, lin_ctx)
+    lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx)
     if DEBUG >= 4: print(ctx.asm_str(lst, sink.arg.function_name))
   return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
 

tinygrad/codegen/decomp/op.py

@@ -18,19 +18,29 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
   assert False
 
 def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
-  if d <= 0 or x.vmin < 0: return None
-  if (vmax:=min(x.vmax, x.dtype.max)) < d: return x.const_like(0)
-  m,s = magicgu(vmax, d)
-  if m*vmax <= x.dtype.max: return (x*m) >> s
+  from tinygrad.renderer.cstyle import MetalRenderer
+  # NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization
+  if isinstance(ren, MetalRenderer): return None
+  # If d is a power of two this is not valid for signed ints!
+  is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
+  assert d>0, "Sign should have been taken out of divisor"
+  vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
+  if vmin > -d and vmax < d: return x.const_like(0)
+  m,s = magicgu(max(vmax, abs(vmin)), d)
+  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
+    return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
   # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
-  if (k := (d & -d).bit_length()-1) > 0:
-    if (ret:=fast_idiv(ren, x >> k, d >> k, dont_cast=True)) is not None: return ret
+  # use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout
+  if (largest_factor_of_two_in_d := (d & -d)) > 1:
+    if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)),
+                       d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
   if dont_cast: return None
   # the next integer width that holds x*m
   widen = {dtypes.int8:dtypes.int16, dtypes.int16:dtypes.int32, dtypes.int32:dtypes.int64, dtypes.int64:dtypes.uint64,
            dtypes.uint8:dtypes.uint16, dtypes.uint16:dtypes.uint32, dtypes.uint32:dtypes.uint64}
   if (next_dtype := widen.get(x.dtype)) is not None and next_dtype in ren.supported_dtypes():
-    if m*vmax <= next_dtype.max: return ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
+    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
+      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
   return None
 
 # ***** threefry *****
@@ -95,12 +105,13 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
       lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
         if (v:=powers_of_two.get(c.val, 0)) else None)]
     if not disable_fast_idiv:
-      # fast_idiv handles non-pow2 divisors on non-negative inputs
-      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.val))]
-      # rewrite raw CMOD -> x - d*fast_idiv(x,d), only when fast_idiv can actually divide;
+      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
+      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
+        lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+      # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
       # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
-      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
-        lambda ctx, x, d: x - d * q if (q:=fast_idiv(ctx, x, d.val)) is not None else None)]
+      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
+        lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
   if Ops.NEG in ops:
     pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
     if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]

tinygrad/codegen/late/regalloc.py

@@ -1,18 +1,20 @@
 import itertools
 from tinygrad.helpers import dedup
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
-from tinygrad.renderer.isa import ISARenderer, Register, rdef, LinearContext
-from typing import Any
+from tinygrad.renderer.isa import ISARenderer, Register, greg
+from tinygrad.dtype import dtypes
 
 PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
   def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]]
-  def __init__(self, ctx:LinearContext, uops:list[UOp], ren:ISARenderer):
+  def __init__(self, uops:list[UOp], ren:ISARenderer):
     self.uops = uops
     self.ren = ren
     self.idx = itertools.count()
+    # the label associated with each loop NOTE: this is only used post regalloc and should be removed
+    self.loop_label: dict[UOp, str] = {}
 
     # compute live ranges
     self.live_range: dict[Register, list[int]] = {}
@@ -21,15 +23,16 @@ class LinearScanRegallocContext:
     for idx,u in reversed(list(enumerate(uops))):
       if u.op in PSEUDO_OPS: continue
       defs = u.tag if isinstance(u.tag, tuple) else ()
-      for v in defs + tuple(rdef(s) for s in dedup(u.src)):
+      for v in defs + tuple(greg(s) for s in dedup(u.src)):
         if isinstance(v, Register): lr.setdefault(v, []).insert(0, idx)
       for v in defs:
         if v in lr and (n:=max((e for s,e in loops.items() if s <= lr[v][-1] < e), default=None)): lr[v].append(n)
       if u.op is Ops.RANGE: loops[idx] = max(j for j,x in enumerate(uops) if u in x.src)
 
     # allocate registers
+    self.stack_size: int = 0
     self.locals: dict[UOp, UOp] = {}
-    self.spills: dict[Register, Any] = {} # mapping from virtual to arbitrary spill slot
+    self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
     self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point
     self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point
     live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it
@@ -46,7 +49,11 @@ class LinearScanRegallocContext:
     # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
     def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
       if v not in self.spills:
-        self.spills[v] = ctx.assign_spill_slot(v, self.vdef(v))
+        # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
+        sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
+        offset = self.stack_size + (sz - self.stack_size % sz) % sz
+        self.spills[v] = UOp.cconst(offset, dtypes.int32)
+        self.stack_size = offset + sz
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
       return r
@@ -57,7 +64,7 @@ class LinearScanRegallocContext:
       for s in u.src:
         # HACK: cause of later hacks to lower range
         if u.op is Ops.END: continue
-        if not isinstance(v:=rdef(s), Register): continue
+        if not isinstance(v:=greg(s), Register): continue
         if v not in live: live[v] = fill(v, i)
         self.reals.setdefault(i, {})[v] = live[v]
 
@@ -69,12 +76,17 @@ class LinearScanRegallocContext:
           cons = v.cons
           # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
           if ren.is_two_address(u) and j == 0:
-            uses = tuple(live.get(rdef(s)) for s in u.src)
+            uses = tuple(live.get(greg(s)) for s in u.src)
             cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses)
           # HACK: cause the range is missing the comparison
           live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
           self.reals.setdefault(i, {})[v] = live[v]
 
+      # allocate stack array
+      if u.op is Ops.BUFFER:
+        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
+        self.stack_size += u.max_numel() * u.dtype.itemsize
+
       # loop prologue, avoid loading inside the loop
       if u.op is Ops.RANGE:
         # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
@@ -101,14 +113,22 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   nsrc = []
   for j,s in enumerate(x.src):
     # v here is the virtual defined by the original s as s is the rewritten version
-    if i in ctx.reals and (v:=rdef(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
+    if i in ctx.reals and (v:=greg(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
     else: nsrc.append(s)
   ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
-  nx = x.replace(src=tuple(nsrc), tag=ndefs)
+  if x.op is Ops.BUFFER: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], tag=ndefs))
+  else: nx = x.replace(src=tuple(nsrc), tag=ndefs)
 
   before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])]
   after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
 
+  # alloc/dealloc stack
+  if ctx.stack_size > 0:
+    sp = ctx.ren.stack_pointer()
+    offset = UOp.cconst(ctx.stack_size, sp.dtype)
+    if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
+    elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
+
   return nx, before + [nx] + after
 
 pm_regalloc_rewrite = PatternMatcher([

tinygrad/device.py

@@ -7,7 +7,6 @@ from tinygrad.helpers import LRU, getenv, diskcache_get, diskcache_put, DEBUG, G
 from tinygrad.helpers import Context, CCACHE, ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE, cpu_events, ProfileEvent, ProfilePointEvent, suppress_finalizing
 from tinygrad.helpers import select_by_name, select_first_inited, DEV, TracingKey, size_to_str, pluralize, Target, unwrap, round_up
 from tinygrad.dtype import DType, _to_np_dtype
-from tinygrad.runtime.support.memory import MMIOInterface
 if TYPE_CHECKING: from tinygrad.renderer import Renderer
 
 # **************** Device ****************
@@ -102,12 +101,11 @@ class Buffer:
   def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None,
                initial_value:bytes|pickle.PickleBuffer|None=None, base:Buffer|None=None, offset:int=0, preallocate=False):
     assert isinstance(dtype, DType)
-    self.device, self.size, self.dtype, self.offset, self.allocated_views, self._base = Device.canonicalize(device), size, dtype, offset, 0, base
-    self.options = options if options is not None else BufferSpec()
-    self._storage:tuple|None = None
-    self._maps:dict[str, tuple] = {}
+    self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = Device.canonicalize(device), size, dtype, options, offset, 0
+    self._bufs: dict[str, Any] = {}
     if base is None:
       assert offset == 0, "base buffers can't have offset"
+      self._base = None
       if opaque is not None: self.allocate(opaque)
       if initial_value is not None:
         self.allocate()
@@ -116,76 +114,56 @@ class Buffer:
     else:
       assert base._base is None, "base can't have a base"
       assert self.device == base.device, "base must have the same device"
+      self._base = base
     if preallocate: self.allocate()
-
-  @suppress_finalizing
-  def __del__(self): self._storage is None or self.deallocate()
-
-  def __repr__(self):
-    return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
-           (f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options != BufferSpec() else "") + ">"
-
   @property
   def base(self) -> Buffer: return self._base if self._base is not None else self
-  @functools.cached_property
-  def allocator(self) -> Allocator: return self.base.allocator if self._base is not None else Device[self.device].allocator
   @property
-  def _buf(self) -> Any: return self.get_storage()[0][0]
-  @property
-  def host(self) -> MMIOInterface: return unwrap(self.get_storage()[1])
-  @property
-  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) -> 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 self._maps:
+  def _buf(self) -> Any: return self._bufs[self.device]
+  # check if the underlying buffer is allocated and the current buffer/view is initialized
+  def is_initialized(self) -> bool: return self.is_allocated() and self.device in self._bufs
+  # check if the underlying buffer is allocated, possibly from the base object
+  def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs
+  def get_buf(self, device: str) -> Any:
+    if device not in self._bufs and (device:=Device.canonicalize(device)) not in self._bufs:
       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)[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
+      if device == self.device: self.ensure_allocated()
+      elif self._base is not None: self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset)
+      else: self._bufs[device] = allocator.map(self.ensure_allocated())
+    return self._bufs[device]
+  def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self
   def allocate(self, opaque=None, external_ptr=None) -> Buffer:
-    assert not self.is_allocated(), "can't allocate already allocated buffer"
+    assert not self.is_initialized(), "can't allocate already allocated buffer"
     if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
-    if not self.device.startswith("NULL") and self.size > MAX_BUFFER_SIZE > 0 and self.options.external_ptr is None:
+    if not self.device.startswith("NULL") and self.size > MAX_BUFFER_SIZE > 0 and (self.options is None or self.options.external_ptr is None):
       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)
+    self.allocator:Allocator = Device[self.device].allocator
+    if external_ptr is not None:
+      self.options = replace(self.options, external_ptr=external_ptr) if self.options else BufferSpec(external_ptr=external_ptr)
     if self._base is not None:
-      (buf, meta), host = self.base.get_storage()
-      mapping = self.allocator._offset(buf, self.nbytes, self.offset), meta
+      self._base.ensure_allocated()
+      self._base.allocated_views += 1
+      self._bufs[self.device] = self.allocator._offset(self.base._buf, self.nbytes, self.offset)
     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:
+      self._bufs[self.device] = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options)
+      if not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None):
         GlobalCounters.mem_used += self.nbytes
         GlobalCounters.mem_used_per_device[self.device] += self.nbytes
       if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size}))
-    elif self._storage is None: self.base.allocated_views += 1
-    self._storage, self._maps, self._base_storage = storage, {}, self.base._storage if self._base else None
     return self
-
   def deallocate(self):
-    assert self._storage is not None, "buffer must be allocated to deallocate"
+    assert self.device in self._bufs, "buffer must be allocated to deallocate"
     if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
     if self._base is None:
-      if GlobalCounters is not None and not self.device.startswith("DISK") and self.options.external_ptr is None:
+      if GlobalCounters is not None and not self.device.startswith("DISK") and (self.options is None or 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, "free", self.trace_num))
-      for dev, mb in self._maps.items(): Device[dev].allocator._unmap(mb[0])
-      self.allocator.free(self._storage, self.nbytes, self.options)
-    else: self.base.allocated_views -= 1
-    self._storage, self._maps, self._base_storage = None, {}, None
-
+      for dev, mb in self._bufs.items():
+        if dev != self.device: Device[dev].allocator._unmap(mb)
+      self.allocator.free(self._buf, self.nbytes, self.options)
+    elif self._base is not None: self._base.allocated_views -= 1
+    self._bufs.clear()
   def __reduce_ex__(self, protocol):
     buf:bytearray|pickle.PickleBuffer|None = None
     if self._base is not None:
@@ -194,40 +172,37 @@ class Buffer:
     if self.is_allocated():
       buf = pickle.PickleBuffer(self.as_memoryview()) if protocol >= 5 else bytearray(self.as_memoryview())
     return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf)
-
   @property
   def trace_num(self) -> int:
     if not hasattr(self, '_trace_num'): self._trace_num = len(Buffer.profile_events)
     return self._trace_num
-
-  def _host_mv(self) -> memoryview|None:
-    if self.is_allocated() and hasattr(host:=self.get_storage()[1], 'mv'): return unwrap(host).view(fmt='B').mv
-    if self.is_allocated() and hasattr(self.allocator, '_as_buffer'): return self.allocator._as_buffer(self._buf)
-    return None
-
+  @property
+  def nbytes(self): return self.size*self.dtype.itemsize
+  @suppress_finalizing
+  def __del__(self): (self.device not in self._bufs) or self.deallocate()
+  def __repr__(self):
+    return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
+           (f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
   def as_memoryview(self, allow_zero_copy=False, force_zero_copy=False, no_sync=False) -> memoryview:
     # zero copy with as_memoryview (disabled by default due to use after free)
-    if (force_zero_copy or allow_zero_copy) and (mv:=self._host_mv()) is not None:
+    if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'):
       if not no_sync: self.allocator.dev.synchronize()
-      return mv
+      if (mv:=self.allocator._as_buffer(self._buf)) is not None: return mv
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv
-
   def numpy(self) -> 'np.ndarray': # type: ignore [name-defined] # noqa: F821
     import numpy as np
     assert _to_np_dtype(self.dtype) is not None, f"no np dtype for {self.dtype}"
     return np.frombuffer(self.as_memoryview(), dtype=_to_np_dtype(self.dtype))
-
   def copy_from(self, src:Buffer) -> Buffer:
     assert self.nbytes == src.nbytes, f"copy size mismatch, {self.nbytes} != {src.nbytes}"
-    assert self.is_allocated() and src.is_allocated(), "copy requires allocated buffers"
+    assert self.is_initialized() and src.is_initialized(), "copy requires allocated buffers"
     from tinygrad.engine.realize import run_linear
     from tinygrad.uop.ops import UOp, Ops
     du, su = UOp.from_buffer(self), UOp.from_buffer(src)
     run_linear(UOp(Ops.LINEAR, src=(su.param_like(1).copy_to_device(self.device).call(du, su),)), update_stats=False)
     return self
-
   def view(self, size:int, dtype:DType, offset:int) -> Buffer:
     assert offset < self.nbytes, "offset must be less than nbytes"
     return Buffer(self.device, size, dtype, base=self.base, offset=self.offset+offset)
@@ -236,47 +211,55 @@ DeviceType = TypeVar('DeviceType', bound='Compiled')
 
 # TODO: size, dest, src are the same type. can we enforce this?
 class Allocator(Generic[DeviceType]):
-  lru = True
-
   def __init__(self, dev:DeviceType, supports_copy_from_disk:bool=True, supports_transfer:bool=True):
     self.dev: DeviceType = dev
     self.default_buffer_spec: BufferSpec = BufferSpec()
-    self.cache:dict[tuple[int, BufferSpec|None], list[tuple]] = defaultdict(list)
     self.supports_copy_from_disk, self.supports_transfer = supports_copy_from_disk, supports_transfer
-
-  def alloc(self, size:int, options:BufferSpec|None=None) -> tuple:
+  # overridden in LRUAllocator
+  def alloc(self, size:int, options:BufferSpec|None=None):
     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
-    try: return self._alloc(size, spec)
-    except (RuntimeError, MemoryError): self.free_cache()
-    try: return self._alloc(size, spec)
+    try: return self._alloc(size, options if options is not None else self.default_buffer_spec)
     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
+                                                               f"Used: {size_to_str(GlobalCounters.mem_used_per_device[self.dev.device])}") from e
+  def free(self, opaque, size:int, options:BufferSpec|None=None):
+    self._free(opaque, options if options is not None else self.default_buffer_spec)
 
-  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._free(storage[0][0], spec)
-
-  def free_cache(self):
-    for (_, options), storages in self.cache.items():
-      for storage in storages: self._free(storage[0][0], options if options is not None else self.default_buffer_spec)
-      storages.clear()
-
-  def map(self, buf:Buffer) -> tuple: return self._map(buf.ensure_allocated()._buf)
+  def map(self, buf:Buffer): return self._map(buf.ensure_allocated()._buf)
 
   # implemented by the runtime
-  def _alloc(self, size:int, options:BufferSpec) -> tuple: raise NotImplementedError("need alloc")
+  def _alloc(self, size:int, options:BufferSpec): 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) -> tuple: raise NotImplementedError("need map")
+  def _map(self, buf): raise NotImplementedError("need map")
   def _unmap(self, mb): pass  # default no-op; override if _map allocates iface-side state
+  # def _as_buffer(self, src) -> memoryview:
   def _offset(self, buf, size:int, offset:int): raise NotImplementedError("need offset")
   # def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
   def _encode_decode(self, bufout, bufin, desc, hist:list, shape:tuple[int,...], frame_pos:int): raise NotImplementedError("need encdec") # optional
 
+class LRUAllocator(Allocator, Generic[DeviceType]):
+  """
+  The LRU Allocator is responsible for caching buffers.
+  It ensures that buffers are not freed until it is absolutely necessary, optimizing performance.
+  """
+  def __init__(self, dev:DeviceType, **kwargs):
+    self.cache: dict[tuple[int, BufferSpec|None], Any] = defaultdict(list)
+    super().__init__(dev, **kwargs)
+  def alloc(self, size:int, options:BufferSpec|None=None):
+    if len(c := self.cache[(size, options)]): return c.pop()
+    try: return super().alloc(size, options)
+    except (RuntimeError, MemoryError):
+      self.free_cache()
+      return super().alloc(size, options)
+  def free_cache(self):
+    for (sz,options),opaques in self.cache.items():
+      for opaque in opaques: super().free(opaque, sz, options)
+      opaques.clear()
+  def free(self, opaque:Any, size:int, options:BufferSpec|None=None):
+    if LRU and (options is None or (not (options.nolru or options.zero) and options.external_ptr is None)): self.cache[(size, options)].append(opaque)
+    else: super().free(opaque, size, options)
+
 class DepsTracker:
   def __init__(self):
     # tracks (offset, end, dep) ranges per base buffer id to handle suballocated buffers correctly.

tinygrad/engine/jit.py

@@ -187,7 +187,7 @@ class CapturedJit(Generic[ReturnType]):
     for u in self._written_uops:
       if u.op is not Ops.BUFFER or (buf:=u.arg.buffer) is None: continue
       for b in (buf.bufs if isinstance(buf, MultiBuffer) else (buf,)):
-        if b.is_allocated(): b.deallocate()
+        if b.is_initialized(): b.deallocate()
         if (base:=b._base) is not None and base.allocated_views == 0 and base.is_allocated(): base.deallocate()
 
 def _prepare_jit_inputs(args, kwargs):

tinygrad/engine/realize.py

@@ -52,8 +52,6 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
-  call = call.without_after
-  if isinstance(call.arg.aux, HCQInfo): return call.arg.aux.estimates
   if (ast:=call.src[0]).op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
   if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
     return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
@@ -156,7 +154,7 @@ def exec_copy(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
     elif src.device.startswith("DISK") and getattr(src.allocator.dev, 'fd', None) is not None \
          and hasattr(dest.allocator, 'copy_from_disk') and src.nbytes >= 4096 and dest.allocator.supports_copy_from_disk:
       dest.allocator.copy_from_disk(dest._buf, src._buf, src.nbytes)
-    elif dest._host_mv() is not None: src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
+    elif hasattr(dest.allocator, '_as_buffer'): src.allocator._copyout(dest.as_memoryview(force_zero_copy=True), src._buf)
     else: dest.allocator._copyin(dest._buf, src.as_memoryview(allow_zero_copy=True))
   return []
 
@@ -194,21 +192,20 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   if (info:=call.arg.aux).inputs:
-    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr + off for u, dev, off in info.inputs]
-    cast(Buffer, call.src[1 + info.table].buffer).host.view(fmt='Q')[:] = array.array('Q', addrs)
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr for u, dev in info.inputs]
+    cast(Buffer, call.src[1 + info.table].buffer)._buf.cpu_view().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,))
   if not (ctx.wait or PROFILE): return ets
 
   slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots}
-  for devs, name, _, prof, pkey in info.kernels:
-    for d in (devs if prof else ()): cast(Any, Device[d]).prof_ents[(slots[d], prof[0])] = ProfileGraphEntry(d, name, prof[0], prof[1], pkey)
-  if ctx.wait:
-    for device in info.device: cast(Any, Device[device]).synchronize(timeout=ctx.timeout)
-  def _prof_tm(device:str, prof:tuple[int, ...]) -> float:
-    st, en = (slots[device].host.view(fmt='Q')[x] for x in prof)
-    return float(en-st) / cast(Any, Device[device]).timestamp_divider / 1e6
-  return ets + [_prof_tm(device, prof) if ctx.wait else None for devices, _, _, prof, _ in info.kernels if prof for device in devices]
+  def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
+    (d:=cast(Any, Device[device])).prof_ents[(slots[device], prof[0])] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
+    if not ctx.wait: return None
+    d.synchronize(timeout=ctx.timeout)
+    st, en = (slots[device]._buf.cpu_view().view(fmt='Q')[x] for x in prof)
+    return float(en-st) / d.timestamp_divider / 1e6
+  return ets + [_prof_tm(device, name, prof, profile_key) for devices,name,_,prof,profile_key in info.kernels if prof for device in devices]
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([

tinygrad/helpers.py

@@ -166,8 +166,6 @@ def stderr_log(msg:str): print(msg, end='', file=sys.stderr, flush=True)
 
 class Context(contextlib.ContextDecorator):
   def __init__(self, **kwargs): self.kwargs = kwargs
-  # ContextDecorator otherwise reuses self, so recursive calls overwrite old_context.
-  def _recreate_cm(self): return Context(**self.kwargs)
   def __enter__(self):
     self.old_context:dict[str, Any] = {k: ContextVar._cache[k].value for k in self.kwargs}
     for k,v in self.kwargs.items(): ContextVar._cache[k].value = v
@@ -241,7 +239,7 @@ TRANSCENDENTAL = ContextVar("TRANSCENDENTAL", 1)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 1)
+VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)

tinygrad/mixin/elementwise.py

@@ -56,14 +56,14 @@ class ElementwiseMixin(CreationMixin):
     """
     return self.cast(dtypes.bool).ne(True)
 
-  def contiguous(self) -> Self:
+  def contiguous(self, **kwargs) -> Self:
     """
     Returns a contiguous tensor.
     """
     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)
-    return self._wrap_uop(uop.alu(Ops.CONTIGUOUS))
+    return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
 
   def contiguous_backward(self) -> Self:
     """
@@ -705,7 +705,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([-9., -6., -3., 0., 3., 6., 9.]).relu6().numpy())
     ```
     """
-    return ((r:=self.relu()) < 6).where(r, 6)
+    return self.relu() - (self-6).relu()
 
   def hardswish(self) -> Self:
     """
@@ -730,7 +730,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).hardsigmoid().numpy())
     ```
     """
-    return ((y:=(alpha * self + beta).relu()) < 1).where(y, 1)
+    return (alpha * self + beta).relu() - (alpha * self + beta - 1).relu()
 
   def hardtanh(self, min_val=-1, max_val=1) -> Self:
     """

tinygrad/mixin/gradient.py

@@ -3,14 +3,12 @@ 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.device import Buffer
 from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
-  if op == Ops.MAX:
-    # count the ties in the acc dtype, the count can overflow the gradient dtype
-    mask = ret.src[0].eq(ret).cast(sum_acc_dtype(ctx.dtype))
-    return ((mask/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))).cast(ctx.dtype) * ctx,)
+  if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
   if op == Ops.MUL:
     # d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others
     safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1]))
@@ -67,6 +65,25 @@ 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 after_gradient(ctx:UOp, ret:UOp):
+  value, *deps = ret.src
+  if len(deps) == 1:
+    dep = deps[0]
+    if dep.op is Ops.STORE and len(dep.src) == 2 and value is dep.src[0]: return (None, ctx)
+    if dep.op is Ops.CALL and (value.unsharded_base.is_unbound or value in dep.call_access()[1]):
+      if dep.src[1:].count(value) != 1: raise RuntimeError("ambiguous CALL output gradient")
+      return (None, UOp.sink(*(ctx if a is value else UOp(Ops.NOOP) for a in dep.src[1:])))
+  for dep in deps:
+    if dep.op is Ops.STORE: writes = dep.src[:1]
+    elif dep.op is Ops.CALL: _, writes = dep.call_access()
+    else: raise RuntimeError(f"gradient through {dep.op} ordering is unsupported")
+    for w in writes:
+      a, b = (u.storage_base.arg.buffer if u.storage_base.op is Ops.BUFFER else None for u in (value, w))
+      if not isinstance(a, Buffer) or not isinstance(b, Buffer) or a.base is b.base or \
+         any(buf.base.options is not None and buf.base.options.external_ptr is not None for buf in (a, b)):
+        raise RuntimeError("gradient through an aliased write is unsupported")
+  return (ctx,) + (None,)*len(deps)
+
 # ctx is grad_output
 pm_gradient = PatternMatcher([
   (UPat(Ops.CAST, name="ret"), lambda ctx, ret: (ctx.cast(ret.src[0].dtype),)),
@@ -97,13 +114,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.COPY, name="ret"), lambda ctx, ret: (ctx.copy_to_device(ret.src[0].device),)),
   (UPat(Ops.UNSHARD, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   (UPat(Ops.SINK), lambda ctx: ctx.src),
-  (UPat(Ops.AFTER, src=(UPat.var("d"), UPat(Ops.CALL, name="k"))), lambda ctx, d, k:
-    (ctx, UOp.sink(*([ctx if i == k.src.index(d)-1 else UOp(Ops.NOOP) for i in range(len(k.src)-1)])))),
-  # ordering-only AFTER: store target is a different buffer, gradient flows straight through to dest
-  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="t"), UPat())))),
-   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, name="ret"), after_gradient),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/mixin/rand.py

@@ -40,7 +40,7 @@ class RandMixin(OpMixin):
   def _rand(cls, key:Self, counter:Self, shape:tuple[int, ...], dtype:DType, contiguous:bool=True) -> Self:
     bits = cls.random_bits(key, counter, ceildiv(prod(shape) * dtype.itemsize, 4))
     out = cls._bits_to_rand(bits, shape, dtype)
-    return out.clone() if contiguous else out
+    return out.contiguous() if contiguous else out
 
   @staticmethod
   def _next_counter(device:str, num:int):
@@ -293,7 +293,7 @@ class RandMixin(OpMixin):
     if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]")
     if not TRAINING or p == 0: return self
     if p == 1: return self.const_like(0)
-    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).clone().where(self, 0) / (1.0 - p)
+    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)
 
   def scaled_dot_product_attention(self, key:Self, value:Self, attn_mask:Self|None=None, dropout_p:float=0.0,
                                    is_causal:bool=False, enable_gqa:bool=False) -> Self:

tinygrad/renderer/amd/sqtt.py

@@ -670,17 +670,8 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
       for wave in range(10):
         if (p.inst >> (wave * 2)) & 3 == 3:
           inst = pc_map[pc:=wave_pc[(p.simd, wave)]]
-          if getattr(inst, 'op_name', '') not in {'S_NOP', 'S_WAITCNT'}: continue
           wave_pc[(p.simd, wave)] += inst.size()
           yield (p, InstructionInfo(pc, wave, inst))
-    elif isinstance(p, CDNA_INST):
-      inst = pc_map[pc:=wave_pc[(p.simd, p.wave)]]
-      if p.op == InstOpCDNA.JUMP:
-        x = getattr(inst, 'simm16') & 0xffff
-        wave_pc[(p.simd, p.wave)] += inst.size() + (x - 0x10000 if x & 0x8000 else x)*4
-      else:
-        wave_pc[(p.simd, p.wave)] += inst.size()
-      yield (p, InstructionInfo(pc, p.wave, inst))
     # map INST events on this SIMD to the program counter, we know the waves
     elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):
       inst = pc_map[pc:=wave_pc[(simd, p.wave)]]
@@ -744,6 +735,5 @@ if __name__ == "__main__":
   sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
   evt_num = getenv("SQTT_EVENT", -1)
   for i, event in enumerate(sqtt_events):
-    if evt_num == -1 or i == evt_num:
-      print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
-      print_packets(decode(event.blob))
+    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
+    print_packets(decode(event.blob))

tinygrad/renderer/cstyle.py

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

tinygrad/renderer/isa/__init__.py

@@ -3,7 +3,6 @@ import itertools
 from dataclasses import dataclass, field
 from tinygrad.renderer import Renderer
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops
-from typing import Any
 
 @dataclass(frozen=True)
 class Register:
@@ -24,24 +23,23 @@ class IselContext:
   def vreg(self, cons:tuple[Register, ...]|Register):
     return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
 
-def rdef(u:UOp):
-  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return rdef(u.src[0])
-  return u.tag[0] if isinstance(u.tag, tuple) else u.tag
+def greg(u:UOp):
+  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0])
+  if isinstance(u.tag, tuple): return u.tag[0]
+  return u.tag
 
-class LinearContext:
-  def __init__(self, ren:ISARenderer):
-    self.ren, self.stack_size = ren, 0
-    self.loop_label: dict[UOp, str] = {}
-  def assign_spill_slot(self, r:Register, u:UOp) -> Any: raise NotImplementedError("arch specific")
+@dataclass
+class PreRegAllocContext:
+  lock: UOp|None = None
 
 class ISARenderer(Renderer):
   pre_isel_matcher: PatternMatcher
   isel_matcher: PatternMatcher
   pre_regalloc_matcher: PatternMatcher
   post_regalloc_matcher: PatternMatcher
-  linear_ctx_type: type = LinearContext
 
   def is_two_address(self, x:UOp) -> bool: return False
-  def spill(self, spill_slot:Any, x:UOp) -> UOp: raise NotImplementedError("arch specific")
-  def fill(self, spill_slot:Any, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
+  def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific")
+  def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
+  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
   def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")

tinygrad/renderer/isa/x86.py

@@ -1,4 +1,3 @@
-from __future__ import annotations
 # flake8: noqa: E702
 # allow semicolons to put multiple ops on one line
 import sys, struct, functools
@@ -7,7 +6,7 @@ from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
-from tinygrad.renderer.isa import ISARenderer, IselContext, Register, LinearContext, rdef
+from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
 from tinygrad.helpers import unwrap, Target
 
 # ***** X86 Ops *****
@@ -159,9 +158,6 @@ pre_isel_matcher = PatternMatcher([
 ])
 
 # ***** X86 registers *****
-def def_reg(dt:DType, reg:Register) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=(reg,))
-# undefined operand, used for VEX instructions
-def undef(): return UOp(Ops.NOOP)
 
 RAX = Register("rax", 0)
 RCX = Register("rcx", 1)
@@ -182,12 +178,11 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
         "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
         **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}}
 
-stack_pointer = def_reg(dtypes.uint64, RSP)
-
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
 def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
+def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=None if reg is None else (reg,))
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
   if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
@@ -211,13 +206,13 @@ def vinsertps(x:UOp) -> UOp:
   def _insert(ret:UOp, i:int) -> UOp:
     s, v = base(x, i), lane(x, i)
     return x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
-  return functools.reduce(_insert, range(len(x.src)), undef())
+  return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
 
 # vpinsrd xmm2, xmm0, eax, imm
 # inserts the element in eax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
   op = {2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD}[x.dtype.itemsize]
-  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), undef())
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
 def idiv(ctx:IselContext, x:UOp) -> UOp:
@@ -270,7 +265,7 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None:
   # the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
   def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
   def _stack_arg(disp:int):
-    return (stack_pointer, UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), src=(imm(dtypes.int32, disp),)), imm(dtypes.uint8, 8))
+    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
   if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
   else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8)
   # this move "cleanses" the abi register constraint
@@ -325,9 +320,9 @@ isel_matcher = PatternMatcher([
     lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])),
   # **** Op -> X86Op ****
   # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc
-  # so regalloc builds the prologue/epilogue naturally. they all share the stack pointer define's dtype so the the stack pointer define is first
+  # so regalloc builds the prologue/epilogue naturally
   (UPat(Ops.SINK, name="x"), lambda x:
-   x.replace(src=(x.ins(X86Ops.RET, src=x.src + (stack_pointer,) + tuple(def_reg(dtypes.uint64, r) for r in CALLEE_SAVED)),))
+   x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64, r) for r in CALLEE_SAVED)),)) \
     if not x.src or x.src[0].op is not Ops.INS or x.src[0].arg[0] is not X86Ops.RET else None),
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
@@ -422,8 +417,8 @@ isel_matcher = PatternMatcher([
   (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)),
   (UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))),
   (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(undef(), y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(undef(), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
   (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
    x.ins(X86Ops.MOVZX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
   (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD)),
@@ -441,7 +436,7 @@ isel_matcher = PatternMatcher([
   # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
   # load, store
   (UPat(Ops.LOAD, dtypes.floats, src=(UPat(name="a"),), name="x"), lambda x,a:
-   x.ins(X86Ops.VPINSRW, src=(undef(),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
+   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
    x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a:
    x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else x.ins(_xmm_sz(x), src=fold_address(a))),
@@ -459,21 +454,14 @@ isel_matcher = PatternMatcher([
 # the flags belong to the last instruction that wrote them. x86 has no good way to store/restore them (then regalloc would
 # handle it), so a consumer that no longer owns its compare re-emits it. Unlike a regalloc rematerialization this is not
 # optional, there is no fallback load from stack
-def flag_rematerialize(ctx:X86LinearContext, x:UOp):
+def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
   if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags: ctx.lock = x
   elif x.arg[0] in X86GroupOp.ReadFlags and ctx.lock is not (flag_def:=x.src[-1]):
     ctx.lock = flag_def
     return (x, [flag_def, x])
   return None
 
-# TODO: dont use rewrite
-def alloc_buffer(ctx:X86LinearContext, x:UOp):
-  nx = isel_matcher.rewrite(stack_pointer.index(UOp.cconst(ctx.stack_size, dtypes.uint32), tag=x.tag))
-  ctx.stack_size += x.max_numel() * x.dtype.itemsize
-  return nx, [nx]
-
 pre_regalloc_matcher = PatternMatcher([
-  (UPat(Ops.BUFFER, name="x"), alloc_buffer),
   (UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize),
 ])
 
@@ -504,14 +492,8 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
-  # the frame is allocated after the stack pointer define at the top of the program and freed before RET
-  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [x, x.ins(X86Ops.SUBi, src=(imm(dtypes.int32, ctx.stack_size),))])
-    if ctx.stack_size and x.arg[0] is X86Ops.DEFINE and rdef(x) == RSP else None),
-  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [stack_pointer.ins(X86Ops.ADDi, src=(imm(dtypes.int32, ctx.stack_size),)), x])
-    if ctx.stack_size and x.arg[0] is X86Ops.RET else None),
   # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, src=(UPat.cvar("disp").cast(),), name="x"), lambda ctx,disp,x:
-    (nx:=UOp.cconst(ctx.stack_size + disp.val, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
   # expand the cmp here so we can preserve rng src edge to get label from ctx
   (UPat(Ops.INS, name="x"), lambda ctx,x: lower_loop(ctx, x) if x.arg[0] is X86Ops.LOOP_CMP else None),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
@@ -520,7 +502,7 @@ post_regalloc_matcher = PatternMatcher([
   (UPat(Ops.END, name="x"), lower_end),
   # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
   (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
-   [ctx.ren.copy(x.src[0], rdef(x)), nx] if rdef(x) != rdef(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
+   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
 ])
 
 # ***** X86 instruction encoding *****
@@ -530,9 +512,9 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
               vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes:
     nonlocal reg, opc
     # get the encoding values of the different fields
-    reg = cast(int, cast(Register, rdef(reg_uop)).index if reg_uop is not None else reg)
-    rm = cast(Register, rdef(rm_uop)).index
-    idx = cast(Register, rdef(idx_uop)).index if idx_uop is not None and rdef(idx_uop) is not None else 4
+    reg = cast(int, cast(Register, greg(reg_uop)).index if reg_uop is not None else reg)
+    rm = cast(Register, greg(rm_uop)).index
+    idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4
     # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register
     rm_sz = sz_uop.src[0].val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
@@ -544,7 +526,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # r extends reg field, x extends index field, b extends rm or base field
     r, _x, b = reg >> 3, idx >> 3, rm >> 3
     if sel: # VEX bytes
-      vvvv = (vd.index if isinstance(vd := rdef(vvvv_uop), Register) else reg) if vvvv_uop is not None else 0
+      vvvv = cast(Register, greg(vvvv_uop)).index if vvvv_uop is not None else 0
       if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | pp])
       else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | pp])
     else: # optional PREFIX and REX bytes
@@ -589,7 +571,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # IMM byte
     if imm_uop is not None:
       if imm_uop.op is Ops.CAST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.src[0].val)
-      elif isinstance(rdef(imm_uop), Register): inst += bytes([(rdef(imm_uop).index & 0b1111) << 4 | 0b0000])
+      elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000])
     return inst
 
   # get the encoding structure of the uop
@@ -622,7 +604,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
 encodings = {
   # moves
   X86Ops.MOVABS: lambda x:
-   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | rdef(x).index >> 3, 0xB8 + (rdef(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
   X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
   X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
   X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
@@ -684,16 +666,6 @@ encodings = {
   X86Ops.RET: lambda x: bytes([0xC3]),
 }
 
-class X86LinearContext(LinearContext):
-  def __init__(self, ren:X86Renderer):
-    super().__init__(ren)
-    self.lock: UOp|None = None
-  def assign_spill_slot(self, r:Register, u:UOp) -> int:
-    sz = r.cons[0].size
-    offset = self.stack_size + (sz - self.stack_size % sz) %sz
-    self.stack_size = offset + sz
-    return offset
-
 class X86Renderer(ISARenderer):
   device = "CPU"
   has_local = False
@@ -704,36 +676,34 @@ class X86Renderer(ISARenderer):
   pre_regalloc_matcher = pre_regalloc_matcher
   post_regalloc_matcher = post_regalloc_matcher
   code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
-  linear_ctx_type = X86LinearContext
   def __init__(self, target:Target):
     if target.arch.split(",")[0] != "x86_64": raise RuntimeError(f"X86Renderer only supports x86_64, got {target.arch}")
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()
   def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
+  def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
   def copy(self, x:UOp, reg:Register) -> UOp: return x.ins(X86Ops.MOV, src=(x,), tag=reg)
 
-  def spill(self, spill_slot:int, x:UOp) -> UOp:
+  def spill(self, disp:UOp, x:UOp) -> UOp:
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    disp = UOp.cconst(spill_slot, dtypes.int32)
-    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
 
   # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
-  def fill(self, spill_slot:int, x:UOp, reg:Register) -> UOp:
+  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16
     dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    disp = UOp.cconst(spill_slot, dtypes.int32)
-    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
+    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
 
   def asm_str(self, uops:list[UOp], function_name:str) -> str:
     def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg[0]))[-1] in ('i', 'm') else o[7:]).lower():7s}"
     def _format_operands(x:UOp) -> str:
       def _format(src:tuple[UOp, ...]) -> list[str]:
         return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
-                (o:=str(rdef(s))) in reg_strs else o for s in src if rdef(s) is not None]
+                (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None]
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
-        return [f"[{rdef(base)}" + (f" + {rdef(idx)}*{sz.src[0].val}" if rdef(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
+        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
 
       if len(x.src) > 4 and x.arg[0] in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg[0] in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

tinygrad/runtime/graph/hcq.py

@@ -1,8 +1,7 @@
 import collections, time
 from typing import Any, cast
 from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, suppress_finalizing, TracingKey, unwrap
-from extra.hcq1.hcq import HCQCompiled, HCQAllocator, HCQSignal, HWQueue, HCQArgsState
-from tinygrad.runtime.support.hcq import HCQBuffer, BumpAllocator, MMIOInterface
+from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, 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
@@ -30,7 +29,7 @@ class HCQGraph(MultiGraphRunner):
     for runtime in self.runtimes:
       if runtime is None: continue
       kernargs_size[runtime.dev] += round_up(runtime.kernargs_alloc_size, 16)
-    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True))[0][0] for d,sz in kernargs_size.items()}
+    self.kernargs_bufs: dict[Compiled, HCQBuffer] = {d:d.allocator._alloc(max(sz, 1), BufferSpec(cpu_access=True)) for d,sz in kernargs_size.items()}
 
     # Fill initial arguments.
     self.ji_args: dict[int, HCQArgsState] = {}
@@ -103,7 +102,7 @@ class HCQGraph(MultiGraphRunner):
       elif is_rdma:
         enqueue_queue = self.comp_queues[enqueue_dev]
         rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
-        from extra.hcq1.ops_rdma import RDMACopyQueue
+        from tinygrad.runtime.ops_rdma import RDMACopyQueue
         self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
       else:
         assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"

tinygrad/runtime/graph/metal.py

@@ -4,7 +4,7 @@ from tinygrad.helpers import dedup, getenv, unwrap, PROFILE
 from tinygrad.device import Buffer, Device, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.runtime.ops_metal import MetalDevice, wait_check, to_ns_str
+from tinygrad.runtime.ops_metal import MetalDevice, MetalAllocator, wait_check, to_ns_str
 from tinygrad.runtime.autogen import metal
 
 class MetalGraph(GraphRunner):
@@ -26,8 +26,8 @@ class MetalGraph(GraphRunner):
 
     self.var_bind_data = []
     if len(self.vars):
-      (self.var_buf, _), host = self.dev.allocator.alloc(sum(dt.itemsize for r in self.runtimes for (_,_,dt,s) in unwrap(r).signature if s == ()))
-      self.var_buf_view, var_buf_offset = unwrap(host).mv, 0
+      self.var_buf = 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 = cast(MetalAllocator, self.dev.allocator)._as_buffer(self.var_buf), 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_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 BufferSpec, Allocator, Compiled, Compiler, CompileError, TinyELF, Program
+from tinygrad.device import BufferSpec, LRUAllocator, 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]]]
@@ -75,10 +75,9 @@ class CLProgram(Program['CLDevice']):
       return float(end.value-start.value) * OSX_TIMING_RATIO * 1e-9
     return None
 
-class CLAllocator(Allocator['CLDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
-    return (checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status), None), None
-
+class CLAllocator(LRUAllocator['CLDevice']):
+  def _alloc(self, size:int, options:BufferSpec) -> cl.cl_mem:
+    return checked(cl.clCreateBuffer(self.dev.context, cl.CL_MEM_READ_WRITE, size, None, status := ctypes.c_int32()), status)
   @suppress_finalizing
   def _free(self, opaque:cl.cl_mem, options:BufferSpec): check(cl.clReleaseMemObject(opaque))
   def _copyin(self, dest:cl.cl_mem, src:memoryview):

tinygrad/runtime/ops_cpu.py

@@ -76,19 +76,18 @@ class CPUProgram(Program['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) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
     if options.external_ptr is not None: addr, buf = options.external_ptr, None
     elif WIN: addr = mv_address(buf:=mmap.mmap(-1, size, access=mmap.ACCESS_WRITE))
     else: addr = mv_address(buf:=mmap.mmap(-1, size, mmap.MAP_ANON | mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE))
-    return (opaque:=HCQBuffer(addr, size, meta=buf, view=MMIOInterface(addr, size, fmt='B'), owner=self.dev), opaque.meta), opaque.view
-
+    return HCQBuffer(va:=addr, sz:=size, meta=buf, view=MMIOInterface(va, sz, fmt='B'), owner=self.dev)
   def _as_buffer(self, src) -> memoryview: return to_mv(src.va_addr, src.size)
   def _copyin(self, dest:HCQBuffer, src:memoryview):
     self.dev.synchronize()
     ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
-    dest[:] = to_mv(int(src.va_addr), dest.nbytes)[:]
+    dest[:] = self._as_buffer(src)[:len(dest)]
   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)

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 MMIOInterface, Compiled, BufferSpec, Allocator, Program, TinyELF
+from tinygrad.device import Compiled, BufferSpec, LRUAllocator, Program, TinyELF
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.runtime.autogen import cuda
@@ -64,14 +64,12 @@ class CUDAProgram(Program['CUDADevice']):
       for i in range(len(vals)): self.c_args.__setattr__(f'v{i}', vals[i])
     return cu_time_execution(lambda: check(cuda.cuLaunchKernel(self.prg, *global_size, *local_size, self.smem, None, None, self.vargs)), enable=wait)
 
-class CUDAAllocator(Allocator['CUDADevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+class CUDAAllocator(LRUAllocator['CUDADevice']):
+  def _alloc(self, size, options:BufferSpec):
     check(cuda.cuCtxSetCurrent(self.dev.context))
-    if options.external_ptr: opaque = cuda.CUdeviceptr_v2(options.external_ptr)
-    elif options.host: opaque = init_c_var(ctypes.c_void_p, lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
-    else: opaque = init_c_var(cuda.CUdeviceptr, lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
-    return (opaque, None), MMIOInterface(opaque.value, size) if options.host else None
-
+    if options.external_ptr: return cuda.CUdeviceptr_v2(options.external_ptr)
+    if options.host: return init_c_var(ctypes.c_void_p, lambda x: check(cuda.cuMemHostAlloc(ctypes.byref(x), size, 0x01)))
+    return init_c_var(cuda.CUdeviceptr, lambda x: check(cuda.cuMemAlloc_v2(ctypes.byref(x), size)))
   @suppress_finalizing
   def _free(self, opaque, options:BufferSpec):
     if options.external_ptr: return
@@ -81,8 +79,8 @@ class CUDAAllocator(Allocator['CUDADevice']):
     check(cuda.cuCtxSetCurrent(self.dev.context))
     host_mem = self.alloc(len(src), BufferSpec(host=True))
     self.dev.pending_copyin.append((host_mem, len(src), BufferSpec(host=True)))
-    ctypes.memmove(host_mem[0][0], mv_address(src), len(src))
-    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem[0][0], len(src), None))
+    ctypes.memmove(host_mem, mv_address(src), len(src))
+    check(cuda.cuMemcpyHtoDAsync_v2(dest, host_mem, len(src), None))
   def _copyout(self, dest:memoryview, src):
     CUDADevice.synchronize_system()
     check(cuda.cuCtxSetCurrent(self.dev.context))
@@ -116,7 +114,7 @@ class CUDADevice(Compiled):
       check(cuda.cuCtxEnablePeerAccess(dev.context, 0))
       CUDADevice.peer_access = True
 
-    self.pending_copyin: list[tuple[tuple, int, BufferSpec|None]] = []
+    self.pending_copyin: list[tuple[int, 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 MMIOInterface, Compiled, Allocator
+from tinygrad.helpers import OSX, round_up
+from tinygrad.device import Compiled, Allocator
 with contextlib.suppress(ImportError):
   import _posixshmem
   from tinygrad.runtime.autogen import io_uring, libc
@@ -78,11 +78,10 @@ class DiskBuffer:
 
 MAP_LOCKED, MAP_POPULATE = 0 if OSX else 0x2000, getattr(mmap, "MAP_POPULATE", 0 if OSX else 0x008000)
 class DiskAllocator(Allocator):
-  lru = False
-  def _alloc(self, size:int, options) -> tuple:
+  def __init__(self, dev:DiskDevice): super().__init__(dev)
+  def _alloc(self, size:int, options):
     self.dev._might_open(size)
-    return (opaque:=DiskBuffer(self.dev, size), None), MMIOInterface(mv_address(opaque._buf()), size)
-
+    return DiskBuffer(self.dev, size)
   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

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 MMIOInterface, BufferSpec, Compiled, Allocator, Compiler, Program, TinyELF
+from tinygrad.device import 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,13 +75,12 @@ class DSPBuffer:
     self.va_addr, self.size, self.share_info, self.offset = va_addr, size, share_info, offset
 
 class DSPAllocator(Allocator['DSPDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec):
     if getenv("MOCKDSP"): fd, share_info, flags = -1, None, mmap.MAP_SHARED|mmap.MAP_ANONYMOUS
     else:
       b = qcom_dsp.ION_IOC_ALLOC(self.dev.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
       fd, flags = (share_info:=qcom_dsp.ION_IOC_SHARE(self.dev.ion_fd, handle=b.handle)).fd, mmap.MAP_SHARED
-    opaque = DSPBuffer(libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, flags, fd, 0), size, share_info, offset=0)
-    return (opaque, opaque.share_info), MMIOInterface(opaque.va_addr, size)
+    return DSPBuffer(libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, flags, fd, 0), size, share_info, offset=0)
 
   @suppress_finalizing
   def _free(self, opaque:DSPBuffer, options:BufferSpec):
@@ -132,7 +131,7 @@ class DSPDevice(Compiled):
       self.ion_fd = os.open('/dev/ion', os.O_RDONLY)
       super().__init__(device, DSPAllocator(self), [DSPRenderer], DSPProgram)
       fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
-      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))[0][0]
+      self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
       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 Compiled, Allocator, BufferSpec, Program, TinyELF
+from tinygrad.device import Compiled, LRUAllocator, 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
@@ -56,11 +56,10 @@ class HIPProgram(Program[HIPDevice]):
       check(hip.hipEventElapsedTime(ctypes.byref(ret := ctypes.c_float()), self.dev.time_event_st, self.dev.time_event_en))
       return ret.value * 1e-3
 
-class HIPAllocator(Allocator[HIPDevice]):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+class HIPAllocator(LRUAllocator[HIPDevice]):
+  def _alloc(self, size:int, options:BufferSpec):
     check(hip.hipSetDevice(self.dev.device_id))
-    return (init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size))), None), None
-
+    return init_c_var(hip.hipDeviceptr_t, lambda x: check(hip.hipMalloc(ctypes.byref(x), size)))
   def _free(self, opaque, options:BufferSpec): check(hip.hipFree(opaque))
   def _copyin(self, dest, src: memoryview):
     check(hip.hipSetDevice(self.dev.device_id))

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 MMIOInterface, Compiled, Compiler, CompileError, Program, TinyELF, Allocator, ProfileDeviceEvent
+from tinygrad.device import Compiled, Compiler, CompileError, Program, TinyELF, LRUAllocator, ProfileDeviceEvent
 from tinygrad.renderer.cstyle import MetalRenderer
 from tinygrad.runtime.autogen import metal
 from tinygrad.runtime.support.c import DLL
@@ -154,14 +154,15 @@ class MetalProgram(Program[MetalDevice]):
 class MetalBuffer:
   def __init__(self, buf:metal.MTLBuffer, size:int, offset=0): self.buf, self.size, self.offset = buf, size, offset
 
-class MetalAllocator(Allocator[MetalDevice]):
-  def _alloc(self, size:int, options) -> tuple:
-    ret = metal.MTLBuffer(options.external_ptr) if options.external_ptr else \
-      self.dev.sysdevice.newBufferWithLength_options(size, metal.MTLResourceStorageModeShared)
-    setattr(ret, "retain", False) # Buffer is explicitly released in _free()
-    if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
-    return (MetalBuffer(ret, size), None), MMIOInterface(addr, size) if (addr:=ret.contents()) is not None else None
+class MetalAllocator(LRUAllocator[MetalDevice]):
+  def _alloc(self, size:int, options) -> MetalBuffer:
+    if options.external_ptr: return MetalBuffer(metal.MTLBuffer(options.external_ptr), size)
 
+    # Buffer is explicitly released in _free() rather than garbage collected via reference count
+    ret = self.dev.sysdevice.newBufferWithLength_options(size, metal.MTLResourceStorageModeShared)
+    ret.retain = False
+    if ret.value is None: raise MemoryError(f"Metal OOM while allocating {size=}")
+    return MetalBuffer(ret, size)
   @suppress_finalizing
   def _free(self, opaque:MetalBuffer, options):
     if not options.external_ptr: opaque.buf.release()
@@ -188,6 +189,6 @@ class MetalAllocator(Allocator[MetalDevice]):
     self.dev.synchronize()
     with cpu_profile(prof_desc, f"{self.dev.device}:COPY"): dst[:] = src
   def _as_buffer(self, src:MetalBuffer) -> memoryview: return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]
-  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(to_mv(dest.buf.contents()+dest.offset, dest.size), src, "TINY -> METAL")
-  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, to_mv(src.buf.contents()+src.offset, src.size), "METAL -> TINY")
+  def _copyin(self, dest:MetalBuffer, src:memoryview): self._cp_mv(self._as_buffer(dest), src, "TINY -> METAL")
+  def _copyout(self, dest:memoryview, src:MetalBuffer): self._cp_mv(dest, self._as_buffer(src), "METAL -> TINY")
   def _offset(self, buf:MetalBuffer, size:int, offset:int): return MetalBuffer(buf.buf, size, offset)

tinygrad/runtime/ops_npy.py

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

tinygrad/runtime/ops_null.py

@@ -22,8 +22,7 @@ class NullProgram(Program['NullDevice']):
     with cpu_profile(self.name, self.device, profile_key=self.profile_key): return 1e-3
 
 class NullAllocator(Allocator['NullDevice']):
-  def _alloc(self, size:int, options) -> tuple: return (None, None), None
-
+  def _alloc(self, size, options): pass
   def _copyin(self, dest, src:memoryview): pass
   def _copyout(self, dest:memoryview, src):
     if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")

tinygrad/runtime/ops_nv.py

@@ -95,6 +95,11 @@ class QMD:
 
 class NVQueue(HWQueue):
   dev:NVDevice
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
 
   def nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
@@ -122,6 +127,11 @@ class NVQueue(HWQueue):
     return doorbell.after(queued).index(0).store(UOp.const(fifo.token, dtypes.uint32))
 
 class NVComputeQueue(NVQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+  ]) + NVQueue.q_rewrite
+
   def __init__(self, ctx, submit):
     super().__init__(ctx, submit)
 
@@ -180,6 +190,11 @@ class NVComputeQueue(NVQueue):
     self.prev_qmd = qmd
 
 class NVCopyQueue(NVQueue):
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ()),
+  ]) + NVQueue.q_rewrite
+
   def copy(self, call:UOp):
     dest, src = (a.getaddr(self.devs) for a in call.src[1:3])
     for off in range(0, sz:=call.src[2].max_numel() * call.src[2].dtype.itemsize, step:=(1 << 31)):
@@ -291,8 +306,8 @@ def nv_build_program(dev:NVDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[NVPro
   return cached
 
 class NVAllocator(HCQAllocator['NVDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
-    return (opaque:=self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero), opaque.meta), opaque.view
+  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
+    return self.dev.iface.alloc(size, cpu_access=options.cpu_access, host=options.host, zero=options.zero)
 
   def _do_free(self, opaque:HCQBuffer, options:BufferSpec): self.dev.iface.free(opaque)
 
@@ -661,9 +676,9 @@ class NVDevice(HCQ2Compiled):
     return [x.data for x in infos]
 
   def _push(self, fifo:GPFifo, cmds:list[int]): # a pushbuffer built in python: channel setup and video decode
-    (buf:=self.rt_view(len(cmds) * 4)).host.view(fmt='I')[:] = array.array('I', cmds)
+    (buf:=self.rt_view(len(cmds) * 4))._buf.cpu_view().view(fmt='I')[:] = array.array('I', cmds)
 
-    put = fifo.put_value.host.view(fmt='Q')
+    put = fifo.put_value._buf.view.view(fmt='Q')
     fifo.ring._buf.cpu_view().view(fmt='Q')[put[0] % fifo.entries] = buf._buf.va_addr | (len(cmds) << 42) | (1 << 41)
     fifo.gpput._buf.cpu_view().view(fmt='I')[0] = (put[0] + 1) % fifo.entries
 
@@ -672,7 +687,7 @@ class NVDevice(HCQ2Compiled):
     put[0] += 1
 
   def _submit_cmds(self, fifo:GPFifo, *cmds:int): # runs cmds once everything already submitted is done, then bumps the timeline
-    tl, addr = self.timeline.host.view(fmt='Q'), self.timeline._buf.va_addr
+    tl, addr = self.timeline._buf.cpu_view().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),

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 MMIOInterface, Buffer, Compiled, Compiler, Allocator, Program, TinyELF
+from tinygrad.device import Buffer, Compiled, Compiler, Allocator, Program, TinyELF
 from tinygrad.renderer import tc
 from tinygrad.uop.ops import exec_alu, python_alu, Ops, UOp, GroupOp
 from tinygrad.renderer import Renderer
@@ -237,14 +237,13 @@ class PythonRenderer(Renderer):
   def supported_dtypes(self): return {d for d in super().supported_dtypes() if d != dtypes.half or sys.version_info >= (3, 12)}
 
 class PythonAllocator(Allocator['PythonDevice']):
-  def _alloc(self, size:int, options) -> tuple: return (buf:=memoryview(bytearray(size)), buf), MMIOInterface(mv_address(buf), size)
-
+  def _alloc(self, size, options): return memoryview(bytearray(size))
   def _as_buffer(self, src) -> memoryview: return src
   def _copyin(self, dest, src:memoryview):
     with cpu_profile("TINY -> PYTHON", f"{self.dev.device}:COPY"): dest[:] = src
   def _copyout(self, dest:memoryview, src):
     with cpu_profile("PYTHON -> TINY", f"{self.dev.device}:COPY"): dest[:] = src
-  def map(self, buf:Buffer) -> tuple: return (mv:=buf.as_memoryview(force_zero_copy=True)), mv
+  def map(self, buf:Buffer): return buf.as_memoryview(force_zero_copy=True)
   def _offset(self, buf:memoryview, size:int, offset:int): return buf[offset:offset+size]
 
 class PythonDevice(Compiled):

tinygrad/runtime/ops_qcom.py

@@ -53,6 +53,14 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 
 class QCOMComputeQueue(HWQueue):
   dev:QCOMDevice
+  q_rewrite = PatternMatcher([
+    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
+    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
+    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
+    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
+    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
+  ])
+
   def cmd(self, opcode:int, *vals): self.q(pkt7_hdr(opcode, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 
   def reg(self, reg:int, *vals): self.q(pkt4_hdr(reg, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
@@ -302,8 +310,8 @@ def qcom_build_program(dev:QCOMDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[Q
   return cached
 
 class QCOMAllocator(HCQAllocator['QCOMDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
-    return (opaque:=self.dev._gpu_map(options.external_ptr, size) if options.external_ptr else self.dev._gpu_alloc(size), opaque.meta), opaque.view
+  def _alloc(self, size:int, opts:BufferSpec) -> HCQBuffer:
+    return self.dev._gpu_map(opts.external_ptr, size) if opts.external_ptr else self.dev._gpu_alloc(size)
 
   def _do_free(self, opaque, options:BufferSpec): self.dev._gpu_free(opaque)
 

tinygrad/runtime/ops_rdma.py

@@ -1,9 +1,8 @@
 from __future__ import annotations
-import mmap, struct, functools, atexit
+import mmap, struct, functools
 from typing import cast
 from tinygrad.uop.ops import sint
-from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue
-from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
+from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, 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
@@ -82,7 +81,7 @@ class RDMAAllocator(HCQAllocatorBase):
                      meta=self.dev.iface.mlx_dev.register_mem(pages, len(pages) * page_sz, page_sz.bit_length() - 1))
 
   def _do_free(self, buf:HCQBuffer, options): self.dev.iface.mlx_dev.unregister_mem(buf.meta)
-  def _do_unmap(self, mb): self.dev.iface.mlx_dev.unregister_mem(mb.meta)
+  def _unmap(self, mb): self.dev.iface.mlx_dev.unregister_mem(mb.meta)
 
   def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQCompiled, dest_dev:HCQCompiled):
     # sync device
@@ -104,9 +103,3 @@ class RDMADevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.iface = MLXIface(self, int(device.split(":")[1]) if ":" in device else 0)
     super().__init__(device, RDMAAllocator(self), [], None, signal_t=None)
-
-@functools.cache
-def get_rdma_device(index:int) -> RDMADevice:
-  dev = RDMADevice(f"RDMA:{index}")
-  atexit.register(dev.finalize)
-  return dev

tinygrad/runtime/ops_webgpu.py

@@ -147,10 +147,10 @@ class WebGPUProgram(Program['WebGpuDevice']):
     return None
 
 class WebGpuAllocator(Allocator['WebGpuDevice']):
-  def _alloc(self, size:int, options:BufferSpec) -> tuple:
+  def _alloc(self, size:int, options:BufferSpec) -> webgpu.WGPUBuffer:
     # WebGPU buffers have to be 4-byte aligned
-    return (webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
-      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc)), None), None
+    return webgpu.wgpuDeviceCreateBuffer(self.dev.device_res, webgpu.WGPUBufferDescriptor(size=round_up(size, 4),
+      usage=webgpu.WGPUBufferUsage_Storage | webgpu.WGPUBufferUsage_CopyDst | webgpu.WGPUBufferUsage_CopySrc))
   def _copyin(self, dest:webgpu.WGPUBuffer, src:memoryview):
     if src.nbytes % 4:
       padded_src = bytearray(round_up(src.nbytes, 4))

tinygrad/runtime/support/hcq.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
-from typing import Any
-import ctypes, os
+from typing import cast, Callable, Type, TypeVar, Generic, Any
+import contextlib, decimal, statistics, time, ctypes, array, os, collections, itertools
 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.helpers import DEV, PROFILE, getenv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
+from tinygrad.helpers import suppress_finalizing, pluralize, TracingKey
+from tinygrad.device import Device, BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
+from tinygrad.uop.ops import sym_infer, sint, UOp
 from tinygrad.runtime.autogen import libc
-from tinygrad.runtime.support.memory import MMIOInterface as MMIOInterface, BumpAllocator as BumpAllocator
+from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
+from tinygrad.renderer import Renderer
 
 class FileIOInterface:
   """
@@ -56,11 +58,447 @@ def hcq_filter_visible_devices(devs, device):
   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
 
+SignalType = TypeVar('SignalType', bound='HCQSignal')
+HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled')
+ProgramType = TypeVar('ProgramType', bound='HCQProgram')
+ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState')
+
+class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
+  """
+  A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
+  """
+
+  def __init__(self):
+    self._q:Any = []
+    self.binded_device:HCQDeviceType|None = None
+    self.q_sints:list[tuple[int, int]] = []
+    self.mv_sints:list[tuple[MMIOInterface, int, int, int|None]] = []
+    self.syms:list[sint] = []
+    self._prev_resolved_syms:list[int|None] = []
+
+  def _new_sym(self, sym:sint) -> int:
+    if sym not in self.syms:
+      self.syms.append(sym)
+      self._prev_resolved_syms.append(None)
+    return self.syms.index(sym)
+
+  def q(self, *values):
+    """
+    Enqueues values in the queue.
+
+    Args:
+      values: The values to enqueue in the queue.
+    """
+
+    for v in values:
+      if isinstance(v, UOp):
+        self.q_sints.append((len(self._q), self._new_sym(v)))
+        self._q.append(0xbadc0ded)
+      else: self._q.append(v)
+
+  # *** common commands  ***
+
+  def timestamp(self, signal:SignalType):
+    """
+    Enqueues a timestamp command which records the current time in a signal after all previously enqueued commands are completed.
+
+    Args:
+      signal: The signal to store the timestamp
+    """
+
+  def signal(self, signal:SignalType, value:sint):
+    """
+    Enqueues a signal command which sets the signal to the given value, ensuring all previous operations are completed.
+
+    Args:
+      signal: The signal to set
+      value: The value to set the signal to
+    """
+
+  def wait(self, signal:SignalType, value:sint):
+    """
+    Enqueues a wait command which halts execution until the signal is greater than or equal to a specific value.
+
+    Args:
+      signal: The signal to wait on
+      value: The value to wait for
+    """
+
+  # *** commands for compute queues ***
+
+  def memory_barrier(self):
+    """
+    Enqueues a memory barrier command to ensure memory coherence between agents. Only on compute queues.
+    """
+
+  def exec(self, prg:ProgramType, args_state:ArgsStateType, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
+    """
+    Enqueues an execution command for a kernel program. Only on compute queues.
+
+    Args:
+      prg: The program to execute
+      args_state: The args state to execute program with
+      global_size: The global work size
+      local_size: The local work size
+    """
+
+  def write(self, b:HCQBuffer, val:sint, b64:bool=False):
+    """
+    Enqueues a command to write a value to a buffer address after all previously enqueued commands are completed.
+
+    Args:
+      b: The buffer to write to
+      val: The value to write
+      b64: If True, write a 64-bit value; otherwise write 32-bit
+    """
+    raise NotImplementedError("write not implemented")
+
+  def poll_bit(self, b:HCQBuffer, val:sint, mask:int):
+    """
+    Enqueues a poll command which halts execution until (mem[b] & mask) == val.
+    val must be 0 or mask (i.e. checks if masked bits are all clear or all set).
+
+    Args:
+      b: The buffer to poll
+      val: The expected value after masking (0 or mask)
+      mask: The bit mask to test
+    """
+    raise NotImplementedError("poll_bit not implemented")
+
+  # *** commands for copy queues ***
+
+  def copy(self, dest:HCQBuffer, src:HCQBuffer, copy_size:int):
+    """
+    Enqueues a copy command to transfer data. Only on copy queues.
+
+    Args:
+      dest: The destination buffer of the copy
+      src: The source buffer of the copy
+      copy_size: The size of data to copy
+    """
+
+  # *** submit and bind commands  ***
+
+  def bind(self, dev:HCQDeviceType):
+    """
+    Associates the queue with a specific device for optimized execution.
+
+    This optional method allows backend implementations to tailor the queue for efficient use on the given device. When implemented, it can eliminate
+    the need to copy queues into the device, thereby enhancing performance.
+
+    Args:
+      dev: The target device for queue optimization.
+
+    Note:
+      Implementing this method is optional but recommended for performance gains.
+    """
+
+  def bind_args_state(self, args_state:ArgsStateType):
+    for vals, mem, fmt in args_state.bind_data: self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt)
+
+  def bind_sints(self, *vals:sint, mem:MMIOInterface, struct_t:Type[ctypes.Structure], start_field:str, fmt, mask:int|None=None):
+    self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt, mask=mask, offset=getattr(struct_t, start_field).offset)
+
+  def bind_sints_to_mem(self, *vals:sint, mem:MMIOInterface, fmt, mask:int|None=None, offset:int=0):
+    mv = mem.view(offset=offset, size=len(vals)*8, fmt=fmt)
+    for i, val in enumerate(vals):
+      if isinstance(val, int): mv[i] = val if mask is None else ((mv[i] & ~mask) | val)
+      else: self.mv_sints.append((mv, i, self._new_sym(val), mask))
+
+  def _apply_var_vals(self, var_vals:dict[str, int]):
+    resolved_syms: list[int|None] = [sym_infer(sym, var_vals) for sym in self.syms]
+
+    for off, sym_idx in self.q_sints:
+      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
+      self._q[off] = resolved_syms[sym_idx]
+
+    for mv, off, sym_idx, mask in self.mv_sints:
+      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
+      mv[off] = resolved_syms[sym_idx] if mask is None else ((mv[off] & ~mask) | resolved_syms[sym_idx])
+
+    self._prev_resolved_syms = resolved_syms
+
+  def submit(self, dev:HCQDeviceType, var_vals:dict[str, int]|None=None):
+    """
+    Submits the command queue to a specific device for execution.
+
+    Args:
+      dev: The device to submit the queue to
+    """
+
+    if var_vals is not None: self._apply_var_vals(var_vals)
+    self._submit(dev)
+    return self
+  def _submit(self, dev:HCQDeviceType): raise NotImplementedError("need _submit")
+
+class HCQSignal(Generic[HCQDeviceType]):
+  def __init__(self, base_buf:HCQBuffer, value:int=0, owner:HCQDeviceType|None=None, is_timeline:bool=False, timestamp_divider=1000, virt=False):
+    self.base_buf, self.owner, self.is_timeline = base_buf, owner, is_timeline
+    self.should_return = isinstance(self.base_buf.va_addr, int) and self.owner is not None and not virt
+    self.timestamp_divider:decimal.Decimal = decimal.Decimal(timestamp_divider)
+    if isinstance(self.base_buf.va_addr, int) and not virt: self.value = value
+
+  def __del__(self):
+    if self.should_return: HCQCompiled.signal_pool[unwrap(self.owner).peer_group].append(self.base_buf)
+
+  @property
+  def value_addr(self) -> sint: return self.base_buf.va_addr
+
+  @property
+  def timestamp_addr(self) -> sint: return self.base_buf.va_addr + 8
+
+  @property
+  def value(self) -> int: return self.base_buf.cpu_view().view(0, 8, 'Q')[0]
+
+  @value.setter
+  def value(self, new_value:int): self.base_buf.cpu_view().view(0, 8, 'Q')[0] = new_value
+
+  @property
+  def timestamp(self) -> decimal.Decimal:
+    """
+    Get the timestamp field of the signal.
+
+    This property provides read-only access to the signal's timestamp.
+
+    Returns:
+      The timestamp in microseconds.
+    """
+    return self.base_buf.cpu_view().view(8, 8, 'Q')[0] / self.timestamp_divider
+
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
+    """
+    Optional function which can implement sleep functionality for the signal.
+    Raises RuntimeError if a fault is detected.
+    """
+
+  def wait(self, value:int, timeout:int|None=None):
+    """
+    Waits the signal is greater than or equal to a specific value.
+
+    Args:
+      value: The value to wait for.
+      timeout: Maximum time to wait in milliseconds. Defaults to 30s.
+    """
+    timeout = timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)
+    start_time = int(time.perf_counter() * 1000)
+    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
+      self._sleep(cur_time - start_time)
+      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
+    if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
+
+@contextlib.contextmanager
+def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None,
+                profile_key:bytes|None=None):
+  st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
+  assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"
+
+  if enabled and queue is not None: queue.timestamp(st)
+  elif enabled and queue_type is not None:
+    queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
+
+  try: yield (st, en)
+  finally:
+    if enabled and queue is not None: queue.timestamp(en)
+    elif enabled and queue_type is not None:
+      queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
+
+    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device,
+                                                         profile_key))
+
+class HCQArgsState(Generic[ProgramType]):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
+    self.buf, self.prg, self.bufs, self.vals = buf, prg, bufs, vals
+    self.bind_data:list[tuple[tuple[sint, ...], MMIOInterface, str]] = []
+
+  def bind_sints_to_buf(self, *vals:sint, buf:HCQBuffer, fmt, offset=0): self.bind_data.append((vals, buf.cpu_view().view(offset=offset), fmt))
+
+class CLikeArgsState(HCQArgsState[ProgramType]):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=(), prefix:list[int]|None=None):
+    super().__init__(buf, prg, bufs, vals=vals)
+
+    if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix)
+
+    self.bind_sints_to_buf(*[b.va_addr for b in bufs], buf=self.buf, fmt='Q', offset=len(prefix or []) * 4)
+    for v,(val_offset,dt) in zip(vals, TinyELF.iter_sig(prg.signature[-len(vals):], len(bufs) * 8)):
+      assert v is not None
+      self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=len(prefix or []) * 4 + val_offset)
+
+class HCQProgram(Program[HCQDeviceType]):
+  def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None):
+    self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size
+    self.profile_key = obj.profile_key
+    self.prof_prg_counter = next(self.dev.prof_prg_counter)
+    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter, self.profile_key)]
+
+  @staticmethod
+  def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)
+
+  def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int|None, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
+    """
+    Fills arguments for the kernel, optionally allocating space from the device if `kernargs_ptr` is not provided.
+    Args:
+      bufs: Buffers to be written to kernel arguments.
+      vals: Values to be written to kernel arguments.
+      kernargs_ptr: Optional pointer to pre-allocated kernel arguments memory.
+    Returns:
+      Arguments state with the given buffers and values set for the program.
+    """
+    argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size, 8),
+                                                       size=self.kernargs_alloc_size)
+    return self.args_state_t(argsbuf, self, bufs, vals=vals)
+
+  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:
+    """
+    Enqueues the program for execution with the given arguments and dimensions.
+
+    Args:
+      bufs: Buffer arguments to execute the kernel with.
+      global_size: Specifies the global work size for kernel execution (equivalent to CUDA's grid size).
+      local_size: Specifies the local work size for kernel execution (equivalent to CUDA's block size).
+      vals: Value arguments to execute the kernel with.
+      wait: If True, waits for the kernel to complete execution.
+
+    Returns:
+      Execution time of the kernel if 'wait' is True, otherwise None.
+    """
+
+    kernargs = self.fill_kernargs(bufs, vals)
+    q = unwrap(self.dev.hw_compute_queue_t)().wait(self.dev.timeline_signal, self.dev.timeline_value - 1).memory_barrier()
+
+    self.dev.prof_exec_counter += 1
+    with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE, profile_key=self.profile_key) as (sig_st, sig_en):
+      q.exec(self, kernargs, global_size, local_size)
+
+    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+
+    if wait: self.dev.synchronize(timeout=timeout)
+    return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None
+
+class HCQCompiled(Compiled, Generic[SignalType]):
+  """
+  A base class for devices compatible with the HCQ (Hardware Command Queue) API.
+  """
+  peer_groups: dict[str, list[HCQCompiled]] = collections.defaultdict(list)
+  signal_pages: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
+  signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
+  cpu_devices: list[HCQCompiled] = []
+
+  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
+               signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
+               kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
+    from tinygrad.runtime.graph.hcq import HCQGraph
+    super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
+
+    self.peer_group = getattr(getattr(self, 'iface', None), 'peer_group', device.split(":")[0])
+    HCQCompiled.peer_groups[self.peer_group].append(self)
+
+    self.signal_t, self.hw_compute_queue_t, self.hw_copy_queue_t = signal_t, comp_queue_t, copy_queue_t
+
+    self.timeline_value:int = 1
+    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str, bytes|None]] = []
+    self.prof_exec_counter:int = 0
+    self.prof_prg_counter = itertools.count(0)
+
+    if signal_t is not None:
+      # Map signals if any
+      for sig_page in HCQCompiled.signal_pages[self.peer_group]: cast(HCQAllocator, self.allocator)._map(sig_page)
+
+      self.sigalloc_size = sigalloc_size
+      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))
+      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
+    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
+
+    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
+
+  def synchronize(self, timeout:int|None=None):
+    if self.error_state is not None: raise self.error_state
+    if not hasattr(self, 'timeline_signal'): return
+
+    # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
+    if not self._is_cpu():
+      for dev in HCQCompiled.cpu_devices: dev.synchronize()
+
+    try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
+    except RuntimeError as e:
+      self.error_state = e
+      if hasattr(self, 'on_device_hang'): self.on_device_hang()
+      raise e
+
+    if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
+    if PROFILE:
+      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp, pk) for st,en,name,dev,pk in self.sig_prof_records]
+      self.sig_prof_records = []
+
+  def next_timeline(self):
+    self.timeline_value += 1
+    return self.timeline_value - 1
+
+  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)))
+      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)
+
+  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
+
+  def hw_compute_queues(self) -> list[tuple[str|None, Callable[[], HWQueue]]]:
+    return [(None, self.hw_compute_queue_t)] if self.hw_compute_queue_t is not None else []
+  def hw_copy_queues(self) -> list[tuple[str, Callable[[], HWQueue]]]:
+    return [("SDMA:0", self.hw_copy_queue_t)] if self.hw_copy_queue_t is not None else []
+
+  def _at_profile_finalize(self):
+    self.synchronize() # Expect device to be synchronizes
+
+    def _sync(d:HCQCompiled, q_t:Callable[[], HWQueue]):
+      q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d)
+      st = time.perf_counter_ns()
+      d.timeline_signal.wait(d.timeline_value - 1)  # average of the two
+      et = time.perf_counter_ns()
+      return (decimal.Decimal(et+st) / 2000) - d.timeline_signal.timestamp
+
+    for prefix, q_t in self.hw_compute_queues() + self.hw_copy_queues():
+      devname = f"{self.device}:{prefix}" if prefix else self.device
+      Compiled.profile_events += [ProfileDeviceEvent(devname, statistics.median([_sync(self, q_t) for _ in range(40)]), props=self.device_props())]
+
+  def _wrap_timeline_signal(self):
+    self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
+    self.timeline_signal.value = 0
+    cast(HCQAllocatorBase, self.allocator).b_timeline = [0] * len(cast(HCQAllocatorBase, self.allocator).b)
+
+  def _realloc(self, oldbuf:HCQBuffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQBuffer, bool]:
+    if oldbuf is not None: self.allocator.free(oldbuf, oldbuf.size, options=options)
+    try: buf, realloced = self.allocator.alloc(new_size, options=options), True
+    except MemoryError:
+      if force: raise
+      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
+    return buf, realloced
+
+  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
+
+  def rdma_dev(self):
+    for i in itertools.count():
+      if (dev:=next((d for d in HCQCompiled.peer_groups[self.peer_group] if type(d).__name__ == 'RDMADevice'), None)): return dev
+      try: Device[f'RDMA:{i}']
+      except IndexError: raise RuntimeError(f"No RDMA found for peer group '{self.peer_group}'")
+
+  def finalize(self):
+    try: self.synchronize() # Try to finalize device in any case.
+    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
+    super().finalize()
+
 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
+    self._mappings:dict[HCQCompiled, 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,
@@ -78,3 +516,107 @@ class HCQBuffer:
 
   @property
   def mapped_devs(self): return self._devs if self._base is None else self._base._devs
+
+class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
+  """
+  A base allocator class compatible with the HCQ (Hardware Command Queue) API.
+
+  This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
+  """
+
+  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)) for _ in range(batch_cnt)]
+    self.b_timeline, self.b_next = [0] * len(self.b), 0
+
+  def _map(self, buf:HCQBuffer) -> HCQBuffer:
+    if self.dev in buf.mapped_devs: return buf
+    if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
+    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
+
+    # Since it's unified memory space, any buffer mapping is valid for all devices after successful map.
+    # Devices can save mappings and internal metadata as a new buffer.
+    if (mb:=self._do_map(buf)) is not None: buf.mappings[self.dev] = mb
+    buf.mapped_devs.append(self.dev)
+    return buf
+
+  @suppress_finalizing
+  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
+    for dev in buf.mapped_devs: dev.synchronize()
+    for d, mb in buf.mappings.items(): d.allocator._unmap(mb)
+    if hasattr(self, '_do_free'): self._do_free(buf, options)
+
+  def _unmap(self, mb): self.dev.iface.free(mb)
+
+  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)
+
+class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    if self.dev.hw_copy_queue_t is None:
+      self.dev.synchronize()
+      with cpu_profile(f'TINY -> {self.dev.device}', f"{self.dev.device}:COPY"): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      return
+
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {self.dev.device}", ret=src.nbytes), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for i in range(0, src.nbytes, self.b[0].size):
+        self.b_next = (self.b_next + 1) % len(self.b)
+        self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
+
+        lsize = min(self.b[self.b_next].size, src.nbytes - i)
+        self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src.cast('B')[i:i+lsize]
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(dest.offset(i), self.b[self.b_next], lsize) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.b_timeline[self.b_next] = self.dev.timeline_value - 1
+
+  def copy_from_disk(self, dest:HCQBuffer, src, size):
+    def _get_temp_buf():
+      # Check if the next buffer is safe to be used (its signal has passed) and reserve it.
+      if self.b_timeline[(self.b_next + 1) % len(self.b)] <= self.dev.timeline_signal.value:
+        self.b_timeline[(self.b_next + 1) % len(self.b)], self.b_next = (1 << 64), (self.b_next + 1) % len(self.b)
+        return (self.b[self.b_next].cpu_view(), self.b_next)
+      return None
+
+    assert self.dev.hw_copy_queue_t is not None
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"DISK -> {self.dev.device}", ret=size), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size,
+                                                                                             use_ioring=type(self.b[0].cpu_view()) is MMIOInterface):
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(dest.offset(dst_off), self.b[batch_info[1]].offset(src_off), copy_size) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1
+
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    self.dev.synchronize()
+    if self.dev.hw_copy_queue_t is None:
+      with cpu_profile(f'{self.dev.device} -> TINY', f"{self.dev.device}:COPY"): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      return
+
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"{self.dev.device} -> TINY", ret=dest.nbytes), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for i in range(0, dest.nbytes, cp_size:=self.b[0].size):
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(self.b[0], src.offset(i), lsize:=min(cp_size, dest.nbytes-i)) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.dev.timeline_signal.wait(self.dev.timeline_value - 1)
+        dest.cast('B')[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
+
+  def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQDeviceType, dest_dev:HCQDeviceType):
+    if src_dev.peer_group != dest_dev.peer_group: return src_dev.rdma_dev().allocator._transfer(dest, src, sz, src_dev, dest_dev)
+
+    cast(HCQAllocator, src_dev.allocator)._map(dest)
+
+    assert src_dev.hw_copy_queue_t is not None
+    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=TracingKey(f"{src_dev.device} -> {dest_dev.device}", ret=sz), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      src_dev.hw_copy_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
+                               .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
+                               .copy(dest, src, sz) \
+                               .signal(src_dev.timeline_signal, src_dev.next_timeline()).submit(src_dev)
+
+    if src_dev != dest_dev:
+      unwrap(dest_dev.hw_compute_queue_t)().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
+                                           .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
+                                           .signal(dest_dev.timeline_signal, dest_dev.next_timeline()).submit(dest_dev)