tinygrad 0.9.0__py3-none-any.whl → 0.9.1__py3-none-any.whl

{tinygrad-0.9.0.dist-info → tinygrad-0.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinygrad
-Version: 0.9.0
+Version: 0.9.1
 Summary: You like pytorch? You like micrograd? You love tinygrad! <3
 Author: George Hotz
 License: MIT
@@ -10,7 +10,6 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
-Requires-Dist: tqdm
 Requires-Dist: pyobjc-framework-Metal ; platform_system == "Darwin"
 Requires-Dist: pyobjc-framework-libdispatch ; platform_system == "Darwin"
 Provides-Extra: arm
@@ -39,6 +38,7 @@ Requires-Dist: onnx ==1.16.0 ; extra == 'testing'
 Requires-Dist: onnx2torch ; extra == 'testing'
 Requires-Dist: opencv-python ; extra == 'testing'
 Requires-Dist: tabulate ; extra == 'testing'
+Requires-Dist: tqdm ; extra == 'testing'
 Requires-Dist: safetensors ; extra == 'testing'
 Requires-Dist: transformers ; extra == 'testing'
 Requires-Dist: sentencepiece ; extra == 'testing'
@@ -47,6 +47,7 @@ Requires-Dist: librosa ; extra == 'testing'
 Requires-Dist: networkx ; extra == 'testing'
 Requires-Dist: hypothesis ; extra == 'testing'
 Requires-Dist: nibabel ; extra == 'testing'
+Requires-Dist: bottle ; extra == 'testing'
 Provides-Extra: testing_tf
 Requires-Dist: tensorflow ==2.15.1 ; extra == 'testing_tf'
 Requires-Dist: tensorflow-addons ; extra == 'testing_tf'
@@ -64,7 +65,7 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an
 
 <h3>
 
-[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](/docs) | [Examples](/examples) | [Showcase](/docs/showcase.md) | [Discord](https://discord.gg/ZjZadyC7PK)
+[Homepage](https://github.com/tinygrad/tinygrad) | [Documentation](https://docs.tinygrad.org/) | [Discord](https://discord.gg/ZjZadyC7PK)
 
 </h3>
 
@@ -139,7 +140,8 @@ tinygrad already supports numerous accelerators, including:
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
-- [x] [HSA](tinygrad/runtime/ops_hsa.py)
+- [x] [AMD](tinygrad/runtime/ops_amd.py)
+- [x] [NV](tinygrad/runtime/ops_nv.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 
@@ -163,7 +165,7 @@ python3 -m pip install git+https://github.com/tinygrad/tinygrad.git
 
 ## Documentation
 
-Documentation along with a quick start guide can be found in the [docs/](/docs) directory.
+Documentation along with a quick start guide can be found on the [docs website](https://docs.tinygrad.org/) built from the [docs/](/docs) directory.
 
 ### Quick example comparing to PyTorch
 
@@ -209,7 +211,7 @@ Now, what we want:
 - Bug fixes (with a regression test) are great! This library isn't 1.0 yet, so if you stumble upon a bug, fix it, write a test, and submit a PR, this is valuable work.
 - Solving bounties! tinygrad [offers cash bounties](https://docs.google.com/spreadsheets/d/1WKHbT-7KOgjEawq5h5Ic1qUWzpfAzuD_J06N1JwOCGs/edit?usp=sharing) for certain improvements to the library. All new code should be high quality and well tested.
 - Features. However, if you are adding a feature, consider the line tradeoff. If it's 3 lines, there's less of a bar of usefulness it has to meet over something that's 30 or 300 lines. All features must have regression tests. In general with no other constraints, your feature's API should match torch or numpy.
-- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win.
+- Refactors that are clear wins. In general, if your refactor isn't a clear win it will be closed. But some refactors are amazing! Think about readability in a deep core sense. A whitespace change or moving a few functions around is useless, but if you realize that two 100 line functions can actually use the same 110 line function with arguments while also improving readability, this is a big win. Refactors should pass [process replay](#process-replay-tests).
 - Tests/fuzzers. If you can add tests that are non brittle, they are welcome. We have some fuzzers in here too, and there's a plethora of bugs that can be found with them and by improving them. Finding bugs, even writing broken tests (that should pass) with `@unittest.expectedFailure` is great. This is how we make progress.
 - Dead code removal from core `tinygrad/` folder. We don't care about the code in extra, but removing dead code from the core library is great. Less for new people to read and be confused by.
 
@@ -225,3 +227,9 @@ python3 -m pip install -e '.[testing]'  # install extra deps for testing
 python3 test/test_ops.py                # just the ops tests
 python3 -m pytest test/                 # whole test suite
 ```
+
+#### Process replay tests
+
+[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/process_replay.py) detects changes in the generated kernels of CI tests by comparing them against tinygrad master. If your PR is a refactor or speedup without any expected behavior change, it should include a green process replay pass to get merged.
+
+You can enable process replay by adding [run_process_replay] to your PR title. [example](https://github.com/tinygrad/tinygrad/pull/4995). Note that you should keep your branch up-to-date with master.

tinygrad-0.9.1.dist-info/RECORD

@@ -0,0 +1,63 @@
+tinygrad/__init__.py,sha256=jC-35zswLSXLuRRThG_o6yar6qQjLCqmeaFCj_XKN08,449
+tinygrad/device.py,sha256=oy8vQAn0bLSVnLtu6MfusWF6Axv83xEwrJB3C1sWcPo,18287
+tinygrad/dtype.py,sha256=LGHdreCNcM-FJTwT3uNrsGIMCSnBHi4ZeGf3K0QHXGU,6199
+tinygrad/function.py,sha256=1PS3nZBicJmMAZ-IvE0SJAHWrGyRMkWM9-J_Hr5az0E,9473
+tinygrad/helpers.py,sha256=SQfdz3AxiAT_k7GExSIXKBYqvCugxaF7ftQ9gr8qvFg,16814
+tinygrad/lazy.py,sha256=mrmxsa6dAQhVAm2UMs7yM18Y-IeJyG4rqyW4-LL5apQ,13336
+tinygrad/multi.py,sha256=EK9OTqmt2EaEgJIIebZUe944NaoJLTFbnUGxVYIk--k,11520
+tinygrad/ops.py,sha256=ZOKpgKY_y3BGqu5NpL172me10Lw4SYVkkCUeD3HR7MI,8607
+tinygrad/tensor.py,sha256=-TRbZ1ydsxgsHUlMDfx6AuEY4qtxUbYhYQ8hK7w8mt0,133764
+tinygrad/codegen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/codegen/kernel.py,sha256=Bw53Hb0vfdyyz2JQhE7YV7DqfZh3gVsxtLFABkx8Xps,36961
+tinygrad/codegen/linearizer.py,sha256=oxqSuzBwpbOtZL1vTxVR4qQqjdvvAUwayhVMltkkX1s,32964
+tinygrad/codegen/uops.py,sha256=GzEwfeyATi9uydzHnyFeN_V47hrk7pkz3gHePjyp6jY,25454
+tinygrad/engine/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/engine/graph.py,sha256=2w16XKpAXSkEML8Hl3m1bbPnLDV-MBsaBJLBPx2QzA0,5264
+tinygrad/engine/jit.py,sha256=xoiApaot299mcjKPbxwUlqu2SJ3Tbyfs_rm5zKuHNMA,11355
+tinygrad/engine/realize.py,sha256=LsfZIHQ2Zz_79a4LKBU1pMUwbleM1x_oBTyn4D1qRc8,11214
+tinygrad/engine/schedule.py,sha256=ltOshvCgUU5wHqTDVLhxPVo7DGq7VRA6ECRX6C_09oU,18984
+tinygrad/engine/search.py,sha256=rLS1Aj_r42L0_uUJRyytqFTpJvXxPdx0LRD7BzvF2dQ,11359
+tinygrad/nn/__init__.py,sha256=-CH1lNAoxJQjFeGFpp78RYLaD2sD4Q4biJA5uvWwm88,12942
+tinygrad/nn/datasets.py,sha256=3tau-L-HiopP1leFDXZV-VamCR0DxOhAu6wVOkAVSZY,493
+tinygrad/nn/optim.py,sha256=zf85kwumpS17fk1NBjUfe7tOUE7-XH37SL-LjgAfva8,6799
+tinygrad/nn/state.py,sha256=cYxsy6IiMCl1wITZcM04GirsftFOutp-2X0PIow35aY,10111
+tinygrad/renderer/__init__.py,sha256=OVEBlNy-eaREIn9jrY0tBhVsmMYi1eZK3wW_P_f26-M,2935
+tinygrad/renderer/assembly.py,sha256=syUM3dMOrmIEkmVWzOS47E2sZXDUWbVXmkofl0RnnWw,18011
+tinygrad/renderer/cstyle.py,sha256=O2MNxsxZzMCPY7JAC_5w5aUqosJSi4LGdH8PZoA7s2A,24387
+tinygrad/renderer/llvmir.py,sha256=43iw-cfwukwffibuZQOm3jI3fPPcRT8rs4-JVglVIHU,10192
+tinygrad/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/ops_amd.py,sha256=h4ZwP64C82aNuHTUtjRykiEzsdhUwGKtvfNdsXOUXyI,31393
+tinygrad/runtime/ops_clang.py,sha256=XWqwobReRdu-Tj-chbWEJFMx6AQfgdGCcpdWcLWUTOQ,1468
+tinygrad/runtime/ops_cuda.py,sha256=LujFWnZvLEaPlebomy9c3C4WAbbaEswdeuJgCNDIiGg,10860
+tinygrad/runtime/ops_disk.py,sha256=5AGIXyGb_x3zuPQ54M6p0xkROjtQo0MMeiqSUPnQ6MQ,6655
+tinygrad/runtime/ops_gpu.py,sha256=z9MPEQawWbJGfLyN3KOgxXaENEbi3ULm8xMWcrln1c0,7532
+tinygrad/runtime/ops_llvm.py,sha256=dODiyVSlPofHyDIZrD-V74h8W1d94VPnr_-A4gNbSO4,2229
+tinygrad/runtime/ops_metal.py,sha256=cb4xIXp_4wyMjk7E2XuSwcCzh1JEwc21p6hJxB8LPgc,5802
+tinygrad/runtime/ops_npy.py,sha256=HSDb47vcbPvDiG3EmSGPC0W-EPO9O5R0TqlQKnYxANU,404
+tinygrad/runtime/ops_nv.py,sha256=ClqazkWrMOheiyRDl1ZY7azMPqFHstnEhlAEKNSkklg,37657
+tinygrad/runtime/ops_python.py,sha256=2B1ty1fr4SpPq-bfEyA-gDxxvjI7ulS1EAqHwHFn-QE,11051
+tinygrad/runtime/autogen/amd_gpu.py,sha256=kqu1ygAUxymdwFznsIrXnsJYDa0UDpVndTlTuj9XK_4,581516
+tinygrad/runtime/autogen/comgr.py,sha256=mhQtuF_vGDrJZD3iyQ_38FrS_2jp3WtImNZPC4TBuAg,39793
+tinygrad/runtime/autogen/cuda.py,sha256=GgRl4AfU54JG0G1XJj2dq2FbrUZ8XG_AnFrPAZJpSSg,250380
+tinygrad/runtime/autogen/hip.py,sha256=1yUHDCwL3KkD15if2Q1Ud3GbJiR7DxsNorKZTCINw54,245532
+tinygrad/runtime/autogen/hsa.py,sha256=7Hsrn17HmChyeFOSX_3Fnzl9c0COtq2Z2ExqGu5FNiU,277716
+tinygrad/runtime/autogen/io_uring.py,sha256=cWYzokworiNV1x4I6eBmZpqV2nCt5B_rVEDx7pFP2W0,57664
+tinygrad/runtime/autogen/kfd.py,sha256=dDmLFL6HL_QXRW2rZOCArY55PRuXuLN9563XCipV2jM,29935
+tinygrad/runtime/autogen/nv_gpu.py,sha256=Hd8p-ay8dbdqFwvGyEtXLX3Ymjr3DP1XSRMAe0zKrq4,1686859
+tinygrad/runtime/autogen/opencl.py,sha256=aW-luGFF5PXFEZ6SgrGANhA9qpkI-fZyEsmDfpez2Ss,82654
+tinygrad/runtime/driver/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/driver/hip_comgr.py,sha256=rCS6DaUjkYflI1OvP5XGCOvcy67XCAgeSvQMUb1K9iY,3840
+tinygrad/runtime/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/runtime/graph/clang.py,sha256=NjobP8AVzNC-CS25FjO5gl2-B8KGJb8K8kLzNXEycPg,2005
+tinygrad/runtime/graph/cuda.py,sha256=ZXOhUhTzz3L2ut-E0tCTrqTEcDFExx_FaZzBm8HcQIk,5268
+tinygrad/runtime/graph/hcq.py,sha256=iJcAw-95PXwjkNgQbwwhFmq1m9Zprq8wwPNFkil_SEQ,11877
+tinygrad/runtime/graph/metal.py,sha256=bwB6uAsqjEbwv5ML5ziWduBtmTpseJboo6J9ssVa4v4,4579
+tinygrad/shape/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+tinygrad/shape/shapetracker.py,sha256=oJqr3q9HoDwNN8S2VV9H_IpmG_ovGY10CTSaWWZjUz4,6387
+tinygrad/shape/symbolic.py,sha256=ORyH9bkb94u3ak0XvM7TvZj2n-_RlR9FcYlQqmIUc2M,16617
+tinygrad/shape/view.py,sha256=H9tdtNWD1xKci_Qhf3EJsWKPD8x4M-e7MAxgoUDqG9k,17999
+tinygrad-0.9.1.dist-info/LICENSE,sha256=ABRhUPEILzINYIukgazD-_rPipkUNUwslrb0RxnV6Xc,1058
+tinygrad-0.9.1.dist-info/METADATA,sha256=YApgO_KsHfOTp9m10ytDj_VIVghIq5GP9PTIiIhre2c,10990
+tinygrad-0.9.1.dist-info/WHEEL,sha256=mguMlWGMX-VHnMpKOjjQidIo1ssRlCFu4a4mBpz1s2M,91
+tinygrad-0.9.1.dist-info/top_level.txt,sha256=vDABMCWBFQnx2kn9Azueu88FP-1klQdePoHikQhHymc,9
+tinygrad-0.9.1.dist-info/RECORD,,

{tinygrad-0.9.0.dist-info → tinygrad-0.9.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: setuptools (70.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tinygrad/runtime/driver/hsa.py

@@ -1,143 +0,0 @@
-import ctypes, collections
-import tinygrad.runtime.autogen.hsa as hsa
-from tinygrad.helpers import init_c_var
-
-def check(status):
-  if status != 0:
-    hsa.hsa_status_string(status, ctypes.byref(status_str := ctypes.POINTER(ctypes.c_char)()))
-    raise RuntimeError(f"HSA Error {status}: {ctypes.string_at(status_str).decode()}")
-
-# Precalulated AQL info
-AQL_PACKET_SIZE = ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t)
-EMPTY_SIGNAL = hsa.hsa_signal_t()
-
-DISPATCH_KERNEL_SETUP = 3 << hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS
-DISPATCH_KERNEL_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
-DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
-DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
-DISPATCH_KERNEL_HEADER |= hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE
-
-BARRIER_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
-BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
-BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
-BARRIER_HEADER |= hsa.HSA_PACKET_TYPE_BARRIER_AND << hsa.HSA_PACKET_HEADER_TYPE
-
-class AQLQueue:
-  def __init__(self, device, sz=-1):
-    self.device = device
-
-    check(hsa.hsa_agent_get_info(self.device.agent, hsa.HSA_AGENT_INFO_QUEUE_MAX_SIZE, ctypes.byref(max_queue_size := ctypes.c_uint32())))
-    queue_size = min(max_queue_size.value, sz) if sz != -1 else max_queue_size.value
-
-    null_func = ctypes.CFUNCTYPE(None, hsa.hsa_status_t, ctypes.POINTER(hsa.struct_hsa_queue_s), ctypes.c_void_p)()
-    self.hw_queue = init_c_var(ctypes.POINTER(hsa.hsa_queue_t)(), lambda x: check(
-      hsa.hsa_queue_create(self.device.agent, queue_size, hsa.HSA_QUEUE_TYPE_SINGLE, null_func, None, (1<<32)-1, (1<<32)-1, ctypes.byref(x))))
-
-    self.next_doorbell_index = 0
-    self.queue_base = self.hw_queue.contents.base_address
-    self.queue_size = self.hw_queue.contents.size * AQL_PACKET_SIZE # in bytes
-    self.write_addr = self.queue_base
-    self.write_addr_end = self.queue_base + self.queue_size - 1 # precalc saves some time
-    self.available_packet_slots = self.hw_queue.contents.size
-
-    check(hsa.hsa_amd_queue_set_priority(self.hw_queue, hsa.HSA_AMD_QUEUE_PRIORITY_HIGH))
-    check(hsa.hsa_amd_profiling_set_profiler_enabled(self.hw_queue, 1))
-
-  def __del__(self):
-    if hasattr(self, 'hw_queue'): check(hsa.hsa_queue_destroy(self.hw_queue))
-
-  def submit_kernel(self, prg, global_size, local_size, kernargs, completion_signal=None):
-    if self.available_packet_slots == 0: self._wait_queue()
-
-    packet = hsa.hsa_kernel_dispatch_packet_t.from_address(self.write_addr)
-    packet.workgroup_size_x = local_size[0]
-    packet.workgroup_size_y = local_size[1]
-    packet.workgroup_size_z = local_size[2]
-    packet.reserved0 = 0
-    packet.grid_size_x = global_size[0] * local_size[0]
-    packet.grid_size_y = global_size[1] * local_size[1]
-    packet.grid_size_z = global_size[2] * local_size[2]
-    packet.private_segment_size = prg.private_segment_size
-    packet.group_segment_size = prg.group_segment_size
-    packet.kernel_object = prg.handle
-    packet.kernarg_address = kernargs
-    packet.reserved2 = 0
-    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
-    packet.setup = DISPATCH_KERNEL_SETUP
-    packet.header = DISPATCH_KERNEL_HEADER
-    self._submit_packet()
-
-  def submit_barrier(self, wait_signals=None, completion_signal=None):
-    assert wait_signals is None or len(wait_signals) <= 5
-    if self.available_packet_slots == 0: self._wait_queue()
-
-    packet = hsa.hsa_barrier_and_packet_t.from_address(self.write_addr)
-    packet.reserved0 = 0
-    packet.reserved1 = 0
-    for i in range(5):
-      packet.dep_signal[i] = wait_signals[i] if wait_signals and len(wait_signals) > i else EMPTY_SIGNAL
-    packet.reserved2 = 0
-    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
-    packet.header = BARRIER_HEADER
-    self._submit_packet()
-
-  def blit_packets(self, packet_addr, packet_cnt):
-    if self.available_packet_slots < packet_cnt: self._wait_queue(packet_cnt)
-
-    tail_blit_packets = min((self.queue_base + self.queue_size - self.write_addr) // AQL_PACKET_SIZE, packet_cnt)
-    rem_packet_cnt = packet_cnt - tail_blit_packets
-    ctypes.memmove(self.write_addr, packet_addr, AQL_PACKET_SIZE * tail_blit_packets)
-    if rem_packet_cnt > 0: ctypes.memmove(self.queue_base, packet_addr + AQL_PACKET_SIZE * tail_blit_packets, AQL_PACKET_SIZE * rem_packet_cnt)
-
-    self._submit_packet(packet_cnt)
-
-  def wait(self):
-    self.submit_barrier([], finish_signal := self.device.alloc_signal(reusable=True))
-    hsa.hsa_signal_wait_scacquire(finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-    self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE
-
-  def _wait_queue(self, need_packets=1):
-    while self.available_packet_slots < need_packets:
-      rindex = hsa.hsa_queue_load_read_index_relaxed(self.hw_queue)
-      self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE - (self.next_doorbell_index - rindex)
-
-  def _submit_packet(self, cnt=1):
-    self.available_packet_slots -= cnt
-    self.next_doorbell_index += cnt
-    hsa.hsa_queue_store_write_index_relaxed(self.hw_queue, self.next_doorbell_index)
-    hsa.hsa_signal_store_screlease(self.hw_queue.contents.doorbell_signal, self.next_doorbell_index-1)
-
-    self.write_addr += AQL_PACKET_SIZE * cnt
-    if self.write_addr > self.write_addr_end:
-      self.write_addr = self.queue_base + (self.write_addr - self.queue_base) % self.queue_size
-
-def scan_agents():
-  agents = collections.defaultdict(list)
-
-  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_agent_t, ctypes.c_void_p)
-  def __scan_agents(agent, data):
-    status = hsa.hsa_agent_get_info(agent, hsa.HSA_AGENT_INFO_DEVICE, ctypes.byref(device_type := hsa.hsa_device_type_t()))
-    if status == 0: agents[device_type.value].append(agent)
-    return hsa.HSA_STATUS_SUCCESS
-
-  hsa.hsa_iterate_agents(__scan_agents, None)
-  return agents
-
-def find_memory_pool(agent, segtyp=-1, location=-1):
-  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_amd_memory_pool_t, ctypes.c_void_p)
-  def __filter_amd_memory_pools(mem_pool, data):
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SEGMENT, ctypes.byref(segment := hsa.hsa_amd_segment_t())))
-    if segtyp >= 0 and segment.value != segtyp: return hsa.HSA_STATUS_SUCCESS
-
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_LOCATION, ctypes.byref(loc:=hsa.hsa_amd_memory_pool_location_t())))
-    if location >= 0 and loc.value != location: return hsa.HSA_STATUS_SUCCESS
-
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SIZE, ctypes.byref(sz := ctypes.c_size_t())))
-    if sz.value == 0: return hsa.HSA_STATUS_SUCCESS
-
-    ret = ctypes.cast(data, ctypes.POINTER(hsa.hsa_amd_memory_pool_t))
-    ret[0] = mem_pool
-    return hsa.HSA_STATUS_INFO_BREAK
-
-  hsa.hsa_amd_agent_iterate_memory_pools(agent, __filter_amd_memory_pools, ctypes.byref(region := hsa.hsa_amd_memory_pool_t()))
-  return region

tinygrad/runtime/graph/hsa.py

@@ -1,171 +0,0 @@
-import ctypes, collections, time, itertools
-from typing import List, Any, Dict, cast, Optional, Tuple
-from tinygrad.helpers import GraphException, init_c_var, round_up
-from tinygrad.device import Buffer, BufferOptions
-from tinygrad.device import Compiled, Device
-from tinygrad.shape.symbolic import Variable
-from tinygrad.runtime.ops_hsa import HSADevice, PROFILE, Profiler
-from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
-from tinygrad.engine.jit import MultiGraphRunner
-import tinygrad.runtime.autogen.hsa as hsa
-from tinygrad.runtime.driver.hsa import check, AQLQueue, AQL_PACKET_SIZE, EMPTY_SIGNAL
-
-def dedup_signals(signals): return [hsa.hsa_signal_t(hndl) for hndl in set([x.handle for x in signals if isinstance(x, hsa.hsa_signal_t)])]
-
-class VirtAQLQueue(AQLQueue):
-  def __init__(self, device, sz):
-    self.device = device
-    self.virt_queue = (hsa.hsa_kernel_dispatch_packet_t * sz)()
-    self.queue_base = self.write_addr = ctypes.addressof(self.virt_queue)
-    self.packets_count = 0
-    self.available_packet_slots = sz
-  def _wait_queue(self, need_packets=1): assert False, f"VirtQueue is too small to handle {self.packets_count+need_packets} packets!"
-  def _submit_packet(self):
-    self.write_addr += AQL_PACKET_SIZE
-    self.packets_count += 1
-    self.available_packet_slots -= 1
-
-class HSAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
-    super().__init__(jit_cache, input_rawbuffers, var_vals)
-
-    # Check all jit items are compatible.
-    compiled_devices = set()
-    for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): compiled_devices.add(ji.prg.device)
-      elif isinstance(ji.prg, BufferXfer):
-        for x in ji.bufs[0:2]: compiled_devices.add(Device[cast(Buffer, x).device])
-      else: raise GraphException
-    if any(not isinstance(d, HSADevice) for d in compiled_devices): raise GraphException
-
-    self.devices: List[HSADevice] = list(compiled_devices) #type:ignore
-
-    # Allocate kernel args.
-    kernargs_size: Dict[Compiled, int] = collections.defaultdict(int)
-    for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): kernargs_size[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
-    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferOptions()) for dev,sz in kernargs_size.items()}
-
-    # Fill initial arguments.
-    self.ji_kargs_structs: Dict[int, ctypes.Structure] = {}
-    for j,ji in enumerate(self.jit_cache):
-      if not isinstance(ji.prg, CompiledRunner): continue
-      self.ji_kargs_structs[j] = ji.prg.clprg.args_struct_t.from_address(kernargs_ptrs[ji.prg.device])
-      kernargs_ptrs[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
-      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
-      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])
-
-    # Build queues.
-    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
-    self.packets = {}
-    self.transfers = []
-    self.ji_to_transfer: Dict[int, int] = {} # faster to store transfers as list and update using this mapping table.
-    self.signals_to_reset: List[hsa.hsa_signal_t] = []
-    self.signals_to_devices: Dict[ctypes.c_uint64, List[HSADevice]] = {}
-    self.profile_info: Dict[Compiled, List[Tuple[Any, ...]]] = collections.defaultdict(list)
-
-    # Special packet to wait for the world.
-    self.kickoff_signals: Dict[HSADevice, hsa.hsa_signal_t] = {dev:self.alloc_signal(reset_on_start=True) for dev in self.devices}
-    for dev in self.devices: self.virt_aql_queues[dev].submit_barrier([], self.kickoff_signals[dev])
-
-    for j,ji in enumerate(self.jit_cache):
-      if isinstance(ji.prg, CompiledRunner):
-        wait_signals = self.access_resources(ji.bufs[(outs:=ji.prg.p.outcount):], ji.bufs[:outs], new_dependency=j, sync_with_aql_packets=False)
-        for i in range(0, len(wait_signals), 5):
-          self.virt_aql_queues[ji.prg.device].submit_barrier(wait_signals[i:i+5])
-        self.packets[j] = hsa.hsa_kernel_dispatch_packet_t.from_address(self.virt_aql_queues[ji.prg.device].write_addr)
-
-        sync_signal = self.alloc_signal(reset_on_start=True) if PROFILE else None
-        self.virt_aql_queues[ji.prg.device].submit_kernel(ji.prg.clprg, *ji.prg.p.launch_dims(var_vals), #type:ignore
-                                                          ctypes.addressof(self.ji_kargs_structs[j]), completion_signal=sync_signal)
-        if PROFILE: self.profile_info[ji.prg.device].append((sync_signal, ji.prg.clprg.name, False))
-      elif isinstance(ji.prg, BufferXfer):
-        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
-        dest_dev, src_dev = cast(HSADevice, Device[dest.device]), cast(HSADevice, Device[src.device])
-        sync_signal = self.alloc_signal(reset_on_start=True, wait_on=[dest_dev, src_dev])
-
-        wait_signals = self.access_resources(read=[src], write=[dest], new_dependency=sync_signal, sync_with_aql_packets=True)
-        self.transfers.append([dest._buf, dest_dev.agent, src._buf, src_dev.agent, dest.nbytes, len(wait_signals),
-                              (hsa.hsa_signal_t*len(wait_signals))(*wait_signals), sync_signal, hsa.HSA_AMD_SDMA_ENGINE_0, True])
-        self.ji_to_transfer[j] = len(self.transfers) - 1
-        if PROFILE: self.profile_info[src_dev].append((sync_signal, f"transfer: HSA:{src_dev.device_id} -> HSA:{dest_dev.device_id}", True))
-
-    # Wait for all active signals to finish the graph
-    wait_signals_to_finish: Dict[HSADevice, List[hsa.hsa_signal_t]] = collections.defaultdict(list)
-    for v in dedup_signals(list(self.w_dependency_map.values()) + list(itertools.chain.from_iterable(self.r_dependency_map.values()))):
-      for dev in self.signals_to_devices[v.handle]:
-        wait_signals_to_finish[dev].append(v)
-
-    self.finish_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
-    for dev in self.devices:
-      wait_signals = wait_signals_to_finish[dev]
-      for i in range(0, max(1, len(wait_signals)), 5):
-        self.virt_aql_queues[dev].submit_barrier(wait_signals[i:i+5], completion_signal=self.finish_signal if i+5>=len(wait_signals) else None)
-
-    # Zero signals to allow graph to start and execute.
-    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
-    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)
-
-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
-    # Wait and restore signals
-    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
-    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, len(self.devices))
-
-    # Update rawbuffers
-    for (j,i),input_idx in self.input_replace.items():
-      if j in self.ji_kargs_structs:
-        self.ji_kargs_structs[j].__setattr__(f'f{i}', input_rawbuffers[input_idx]._buf)
-      else:
-        if i == 0: self.transfers[self.ji_to_transfer[j]][0] = input_rawbuffers[input_idx]._buf # dest
-        elif i == 1: self.transfers[self.ji_to_transfer[j]][2] = input_rawbuffers[input_idx]._buf # src
-
-    # Update var_vals
-    for j in self.jc_idx_with_updatable_var_vals:
-      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
-        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])
-
-    # Update launch dims
-    for j in self.jc_idx_with_updatable_launch_dims:
-      gl, lc = cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
-      self.packets[j].workgroup_size_x = lc[0]
-      self.packets[j].workgroup_size_y = lc[1]
-      self.packets[j].workgroup_size_z = lc[2]
-      self.packets[j].grid_size_x = gl[0] * lc[0]
-      self.packets[j].grid_size_y = gl[1] * lc[1]
-      self.packets[j].grid_size_z = gl[2] * lc[2]
-
-    for dev in self.devices:
-      dev.flush_hdp()
-      dev.hw_queue.blit_packets(self.virt_aql_queues[dev].queue_base, self.virt_aql_queues[dev].packets_count)
-
-    for transfer_data in self.transfers:
-      check(hsa.hsa_amd_memory_async_copy_on_engine(*transfer_data))
-
-    et = None
-    if wait:
-      st = time.perf_counter()
-      hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-      et = time.perf_counter() - st
-
-    for profdev,profdata in self.profile_info.items(): Profiler.tracked_signals[profdev] += profdata
-    return et
-
-  def alloc_signal(self, reset_on_start=False, wait_on=None):
-    sync_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
-    if reset_on_start: self.signals_to_reset.append(sync_signal)
-    if wait_on is not None: self.signals_to_devices[sync_signal.handle] = wait_on
-    return sync_signal
-
-  def dependency_as_signal(self, dep, sync_with_aql_packets) -> Optional[hsa.hsa_signal_t]:
-    if isinstance(dep, hsa.hsa_signal_t): return dep
-    elif sync_with_aql_packets and isinstance(packet := self.packets.get(dep), hsa.hsa_kernel_dispatch_packet_t):
-      if packet.completion_signal.handle == EMPTY_SIGNAL.handle: packet.completion_signal = self.alloc_signal(reset_on_start=True)
-      return packet.completion_signal
-    return None
-
-  def access_resources(self, read, write, new_dependency, sync_with_aql_packets=False):
-    rdeps = self._access_resources(read, write, new_dependency)
-    wait_signals = [self.dependency_as_signal(dep, sync_with_aql_packets=sync_with_aql_packets) for dep in rdeps]
-    if sync_with_aql_packets: wait_signals += [self.kickoff_signals[cast(HSADevice, Device[rawbuf.device])] for rawbuf in read+write]
-    return dedup_signals(wait_signals)