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

{tinygrad-0.8.0.dist-info → tinygrad-0.9.1.dist-info}/LICENSE

@@ -1,4 +1,4 @@
-Copyright (c) 2023 George Hotz
+Copyright (c) 2024, the tiny corp
 
 Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinygrad
-Version: 0.8.0
+Version: 0.9.1
 Summary: You like pytorch? You like micrograd? You love tinygrad! <3
 Author: George Hotz
 License: MIT
@@ -10,12 +10,16 @@ Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: numpy
-Requires-Dist: tqdm
-Requires-Dist: gpuctypes
 Requires-Dist: pyobjc-framework-Metal ; platform_system == "Darwin"
 Requires-Dist: pyobjc-framework-libdispatch ; platform_system == "Darwin"
 Provides-Extra: arm
 Requires-Dist: unicorn ; extra == 'arm'
+Provides-Extra: docs
+Requires-Dist: mkdocs-material ; extra == 'docs'
+Requires-Dist: mkdocstrings[python] ; extra == 'docs'
+Requires-Dist: markdown-callouts ; extra == 'docs'
+Requires-Dist: markdown-exec[ansi] ; extra == 'docs'
+Requires-Dist: black ; extra == 'docs'
 Provides-Extra: linting
 Requires-Dist: pylint ; extra == 'linting'
 Requires-Dist: mypy ; extra == 'linting'
@@ -30,10 +34,11 @@ Requires-Dist: torch ; extra == 'testing'
 Requires-Dist: pillow ; extra == 'testing'
 Requires-Dist: pytest ; extra == 'testing'
 Requires-Dist: pytest-xdist ; extra == 'testing'
-Requires-Dist: onnx ==1.15.0 ; extra == 'testing'
+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'
@@ -41,18 +46,26 @@ Requires-Dist: tiktoken ; extra == 'testing'
 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'
 Provides-Extra: triton
 Requires-Dist: triton-nightly >=2.1.0.dev20231014192330 ; extra == 'triton'
 
 <div align="center">
 
-[![logo](https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png)](https://tinygrad.org)
+<picture>
+  <source media="(prefers-color-scheme: light)" srcset="/docs/logo_tiny_light.svg">
+  <img alt="tiny corp logo" src="/docs/logo_tiny_dark.svg" width="50%" height="50%">
+</picture>
 
 tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) and [karpathy/micrograd](https://github.com/karpathy/micrograd). Maintained by [tiny corp](https://tinygrad.org).
 
 <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>
 
@@ -122,17 +135,15 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers
 
 tinygrad already supports numerous accelerators, including:
 
-- [x] [CPU](tinygrad/runtime/ops_cpu.py)
 - [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
-- [x] [C Code (Clang)](tinygrad/runtime/ops_clang.py)
+- [x] [CLANG (C Code)](tinygrad/runtime/ops_clang.py)
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
-- [x] [PyTorch](tinygrad/runtime/ops_torch.py)
-- [x] [HIP](tinygrad/runtime/ops_hip.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.
-More information can be found in the [documentation for adding new accelerators](/docs/adding_new_accelerators.md).
 
 ## Installation
 
@@ -154,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
 
@@ -193,13 +204,14 @@ We'll start with what will get your PR closed with a pointer to this section:
 - All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
 - Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
 - In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
+- If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.
 
 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.
 
@@ -215,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.8.0.dist-info → tinygrad-0.9.1.dist-info}/WHEEL

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

tinygrad/features/image.py

@@ -1,93 +0,0 @@
-from typing import Tuple
-from tinygrad.helpers import prod, IMAGE, getenv, DEBUG
-from tinygrad.dtype import dtypes
-
-# *** image Tensor function replacements ***
-
-def image_dot(self, w):
-  # NOTE: we use a 1x1 conv2d to do the matmul. mxk @ kxn = (1,k,m,1).conv2d(n,k,1,1)
-  n1, n2 = len(self.shape), len(w.shape)
-  assert n1 != 0 and n2 != 0, f"both arguments to matmul need to be at least 1D, but they are {n1}D and {n2}D"
-  assert self.shape[-1] == w.shape[-min(n2, 2)], f"Input Tensor shapes {self.shape} and {w.shape} cannot be multiplied ({self.shape[-1]} != {w.shape[-min(n2, 2)]})"  # noqa: E501
-  bs, groups, cin, cout = prod(self.shape[0:-2]), prod(w.shape[0:-2]), w.shape[-2], w.shape[-1]
-  out_shape_t = self.shape[0:-2] + (cout,-1) if len(self.shape) > 1 else (cout, )
-
-  # NOTE: with NHWC we can remove the transposes
-  # bs x groups*cin x H x W
-  cx = self.transpose(self.ndim-1, self.ndim-2).reshape((bs//groups, groups*cin, -1, 1))
-  # groups*cout x cin x H, W
-  cw = w.transpose(w.ndim-1, w.ndim-2).reshape((groups*cout, cin, 1, 1))
-  return image_conv2d(cx, cw, groups=groups).reshape(out_shape_t).transpose(self.ndim-1, self.ndim-2)
-
-def image_conv2d(self, weight, bias=None, groups=1, stride=1, dilation=1, padding=0):
-  base_image_type = dtypes.imageh if getenv("FLOAT16", 0) else dtypes.imagef
-
-  (bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape
-  x, w = self, weight.reshape(groups, (rcout := cout//groups), cin, H, W)
-
-  # hack for non multiples of 4 on cin
-  if cin % 4 != 0 and not (cin == 1 and groups%4 == 0):
-    x = x.reshape(bs, groups, cin, iy, ix)   # do this always?
-    added_input_channels = 4 - (cin % 4)
-    w = w.pad(tuple((0, added_input_channels) if i == 2 else None for i in range(w.ndim)))
-    x = x.pad(tuple((0, added_input_channels) if i == 2 else None for i in range(x.ndim)))
-    cin = cin + added_input_channels
-    x = x.reshape(bs, groups*cin, iy, ix)
-
-  # hack for non multiples of 4 on rcout
-  added_output_channels = 0
-  if rcout % 4 != 0 and not (rcout == 1 and groups%4 == 0):
-    added_output_channels = 4 - (rcout % 4)
-    rcout += added_output_channels
-    cout = groups * rcout
-    w = w.pad(tuple((0, added_output_channels) if i == 1 else None for i in range(w.ndim)))
-
-  # packed (note: flipping bs and iy would make the auto-padding work)
-  x = x.permute(0,2,3,1)
-  cin_last = iy == 1 and ix == 1
-  if cin == 1: w = w.reshape(cout//4,4,H,W).permute(0,2,3,1)
-  elif cin_last: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,1,3)
-  else: w = w.reshape(cout//4,4,cin//4,4,H,W).permute(0,4,2,5,3,1)
-
-  # contiguous creates the image, and early realize static weights (TODO: test for the static weight)
-  if IMAGE >= 2: x,w = x.cast(base_image_type((bs*iy, ix*groups*cin//4, 4))), w.cast(base_image_type((cout//4, H*W*cin, 4)))
-  x, w = x.contiguous(), w.contiguous()
-
-  # expand out
-  rcin_hi, rcin_lo = cin//4 if cin >= 4 else 1, 4 if cin >= 4 else 1
-  cout_expand = [groups//4 if cin == 1 else groups, 4 if cin == 1 else 1, rcout//4 if rcout >= 4 else 1, 4 if rcout >= 4 else 1]
-  x = x.reshape(bs, iy, ix, groups, rcin_hi, rcin_lo)
-  if cin_last: w = w.reshape(cout//4, H, rcin_hi, W, 4, rcin_lo)
-  else: w = w.reshape(cout//4, H, rcin_hi, W, rcin_lo, 4).permute(0,1,2,3,5,4)
-
-  # padding
-  padding_ = [padding]*4 if isinstance(padding, int) else (padding if len(padding) == 4 else [padding[1], padding[1], padding[0], padding[0]])
-  x = x.slice((None, (-padding_[2], x.shape[1]+padding_[3]), (-padding_[0], x.shape[2]+padding_[1]), None, None, None))
-
-  # prepare input
-  x = x.permute(0,3,4,5,1,2)._pool((H, W), stride, dilation) # -> (bs, groups, rcin_hi, rcin_lo, oy, ox, H, W)
-  x = x.permute(0,4,5,1,2,3,6,7).reshape(bs, (oy := x.shape[4]), (ox := x.shape[5]), *cout_expand[0:2], 1, 1, rcin_hi, rcin_lo, H, W)
-
-  # prepare weights
-  w = w.permute(0,4,2,5,1,3).reshape((1, 1, 1, *cout_expand, rcin_hi, rcin_lo, H, W))
-
-  # the conv!
-  ret = (x*w).cast(base_image_type((bs*oy, ox*cout//4, 4)) if IMAGE >= 2 else dtypes.float32).sum((-4, -3, -2, -1))
-
-  # undo hack for non multiples of 4 on C.rcout
-  if added_output_channels != 0:
-    ret = ret.reshape(bs, oy, ox, groups, rcout)[:, :, :, :, :-added_output_channels]
-    cout = groups * (rcout - added_output_channels)
-
-  # NCHW output
-  ret = ret.reshape(bs, oy, ox, cout).permute(0,3,1,2)
-  return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1))
-
-# *** images have weird indexing requirements ***
-
-from tinygrad.shape.symbolic import Node
-def to_image_idx(base_shape:Tuple[int, ...], idxy:Node, valid:Node) -> Tuple[Tuple[Node, Node], Node]:
-  idx, idy = (idxy // 4) % base_shape[1], (idxy // (4 * base_shape[1]))
-  # TODO: bring back the valid removal logic (correct!)
-  if DEBUG>=5: print("to_image_idx", base_shape, idx.min, idx.max, idy.min, idy.max, idx, idy, valid)
-  return (idx, idy), valid

tinygrad/features/multi.py

@@ -1,103 +0,0 @@
-from __future__ import annotations
-from typing import Optional, Union, Any, Tuple, List
-import functools
-from tinygrad.helpers import all_same, dedup
-from tinygrad.dtype import DType
-from tinygrad.ops import BinaryOps, LoadOps, UnaryOps, TernaryOps, ReduceOps
-from tinygrad.lazy import LazyBuffer, create_schedule
-from tinygrad.shape.shapetracker import ShapeTracker, sint
-
-def all_reduce(lbs):
-  # TODO: replace this with ring reduce
-  return [functools.reduce(lambda x,y: x.e(BinaryOps.ADD, y), [x.copy_to_device(lb.device) for x in lbs]) for lb in lbs]
-
-def to_sharded(lbs:List[LazyBuffer], axis:int) -> List[LazyBuffer]:
-  assert lbs[0].shape[axis] % len(lbs) == 0, f"{lbs[0].shape=} {axis=} {len(lbs)=}"
-  sz = lbs[0].shape[axis] // len(lbs)
-  return [lb.shrink(tuple((0,s) if a != axis else (sz*i,sz*(i+1)) for a,s in enumerate(lb.shape))) for i,lb in enumerate(lbs)]
-
-class MultiLazyBuffer:
-  def __init__(self, lbs:List[LazyBuffer], axis:Optional[int]):
-    assert all(isinstance(x, LazyBuffer) for x in lbs) and len(lbs) >= 2, "all lbs must be LazyBuffers, and we need at least two of them"
-    assert all_same([(x.shape, x.dtype, x.st) for x in lbs]), "all multilazybuffer needs same shape, dtype, and st"
-    self.lbs, self.axis, self.dtype, self.device = lbs, axis, lbs[0].dtype, tuple(x.device for x in lbs)
-    self.shape = tuple(s*len(self.lbs) if a == self.axis else s for a,s in enumerate(lbs[0].shape))
-
-  def __repr__(self):
-    return f"<MLB{chr(10)}{chr(10).join([f'{x.device} {x.st}' for x in self.lbs])}>"
-
-  @staticmethod
-  def from_sharded(lb:LazyBuffer, devices:Tuple[str, ...], axis:Optional[int]=None):
-    lbs = [lb.contiguous() if lb.base != lb else lb] * len(devices)
-    return MultiLazyBuffer([lb.copy_to_device(d).contiguous() for lb,d in zip(to_sharded(lbs, axis) if axis is not None else lbs, devices)], axis)
-
-  def copy_to_device(self, device:str) -> LazyBuffer:
-    if self.axis is None: return self.lbs[0].copy_to_device(device)
-    sz = self.lbs[0].shape[self.axis]
-    llbs = []
-    for i,lb in enumerate([lb.copy_to_device(device) for lb in self.lbs]):
-      pad_arg = tuple((0,0) if a != self.axis else (sz*i,(s*len(self.lbs))-sz*(i+1)) for a,s in enumerate(lb.shape))
-      llbs.append(lb.pad(pad_arg))
-    return functools.reduce(lambda x,y: x.e(BinaryOps.ADD, y), llbs)
-
-  # TODO: fix this
-  def is_unrealized_contiguous_const(self): return False
-
-  # passthroughs
-  def schedule(self, seen=None): return create_schedule(self.lbs, seen)
-  def cast(self, dtype:DType, bitcast:bool=False): return MultiLazyBuffer([x.cast(dtype, bitcast) for x in self.lbs], self.axis)
-  def const(self, val:Union[float, int]) -> MultiLazyBuffer: return MultiLazyBuffer([x.const(val) for x in self.lbs], self.axis)
-  def contiguous(self): return MultiLazyBuffer([x.contiguous() for x in self.lbs], self.axis)
-
-  # elementwise is simple
-  def e(self, op:Union[LoadOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:MultiLazyBuffer, arg:Optional[Any]=None) -> MultiLazyBuffer:
-    msrcs = (self,)+in_srcs
-    assert all(isinstance(x, MultiLazyBuffer) for x in msrcs), f"all buffers must be MultiLazyBuffer {msrcs}"
-    assert all_same([x.device for x in msrcs]), f"all buffers must have the same device {[x.device for x in msrcs]}"
-
-    # NOTE: they all have to share an axis, we always choose [-1]
-    axis = axes[-1] if len(axes := dedup([x.axis for x in msrcs if x.axis is not None])) else None
-    srcs = []
-    for mlb in msrcs:
-      if mlb.axis == axis: srcs.append(mlb.lbs)
-      elif mlb.axis is None and axis is not None: srcs.append(to_sharded(mlb.lbs, axis))
-      else: srcs.append(to_sharded([mlb.copy_to_device(lb.device) for lb in mlb.lbs], axis))
-    return MultiLazyBuffer([lsrcs[0].e(op, *lsrcs[1:], arg=arg) for lsrcs in zip(*srcs)], axis)
-
-  def _shape_to_single_shard(self, shape): return tuple(s//len(self.lbs) if a == self.axis else s for a,s in enumerate(shape))
-
-  def r(self, op:ReduceOps, new_shape:Tuple[sint, ...]) -> MultiLazyBuffer:
-    if self.axis is not None and new_shape[self.axis] == 1:
-      # all-reduce on sharded axes
-      return MultiLazyBuffer(all_reduce([x.r(op, new_shape) for x in self.lbs]), None)
-    # reduce on non sharded axes, piecewise is fine. if axis is None this is also correct
-    return MultiLazyBuffer([x.r(op, self._shape_to_single_shard(new_shape)) for x in self.lbs], self.axis)
-
-  # *** movement ops ***
-
-  def reshape(self, arg:Tuple[sint, ...]):
-    if self.axis is None: return MultiLazyBuffer([x.reshape(arg) for x in self.lbs], None)
-    # TODO: this can be wrong
-    st = ShapeTracker.from_shape(self.shape)
-    rs = st.real_strides()[self.axis]
-    new_axis = st.reshape(arg).real_strides().index(rs)
-    narg = tuple(s//len(self.lbs) if a == new_axis else s for a,s in enumerate(arg))
-    return MultiLazyBuffer([x.reshape(narg) for x in self.lbs], new_axis)
-
-  def pad(self, arg:Tuple[Tuple[sint, sint], ...]):
-    assert self.axis is None or arg[self.axis] == (0,0), "padding not supported on sharded axis"
-    return MultiLazyBuffer([x.pad(arg) for x in self.lbs], self.axis)
-  def expand(self, arg:Tuple[sint, ...]):
-    # NOTE: this assert isn't needed, sharded axis can have dim 1
-    assert self.axis is None or arg[self.axis] == self.lbs[0].shape[self.axis] * len(self.lbs), "expand not supported on sharded axis"
-    return MultiLazyBuffer([x.expand(self._shape_to_single_shard(arg)) for x in self.lbs], self.axis)
-  def permute(self, arg:Tuple[int, ...]):
-    # all permutes supported!
-    return MultiLazyBuffer([x.permute(arg) for x in self.lbs], arg.index(self.axis) if self.axis is not None else None)
-  def shrink(self, arg:Tuple[Tuple[sint, sint], ...]):
-    assert self.axis is None or arg[self.axis] == (0, self.lbs[0].shape[self.axis] * len(self.lbs)), "shrinking not supported on sharded axis"
-    narg = tuple((s1//len(self.lbs), s2//len(self.lbs)) if a == self.axis else (s1,s2) for a,(s1,s2) in enumerate(arg))
-    return MultiLazyBuffer([x.shrink(narg) for x in self.lbs], self.axis)
-  def stride(self, arg:Tuple[int, ...]):
-    assert self.axis is None or arg[self.axis] == 1, "flipping not supported on sharded axis"
-    return MultiLazyBuffer([x.stride(arg) for x in self.lbs], self.axis)

tinygrad/features/search.py

@@ -1,160 +0,0 @@
-from typing import Dict, List, cast, DefaultDict, Optional, Tuple, Callable
-import itertools, random, math, time, multiprocessing, traceback, signal
-from tinygrad.device import Device, Compiled, Buffer, CompiledASTRunner
-from tinygrad.ops import MemBuffer, LazyOp
-from tinygrad.helpers import prod, flatten, DEBUG, CACHELEVEL, diskcache_get, diskcache_put, getenv, Context, colored, to_function_name
-from tinygrad.dtype import ImageDType
-from tinygrad.codegen.linearizer import Linearizer
-from collections import defaultdict
-from tinygrad.tensor import Tensor
-from tinygrad.shape.symbolic import sym_infer
-
-from tinygrad.codegen.kernel import Opt, OptOps
-actions = [Opt(op=OptOps.UPCAST, axis=axis, amt=amt) for amt in [0,2,3,4,7] for axis in range(6)]
-actions += [Opt(op=OptOps.UNROLL, axis=axis, amt=amt) for amt in [0,4] for axis in range(4)]
-actions += [Opt(op=OptOps.LOCAL, axis=axis, amt=amt) for amt in [2,3,4,8,13,16,29] for axis in range(5)]
-actions += [Opt(op=OptOps.GROUPTOP, axis=axis, amt=amt) for amt in [13,16,29,32,256] for axis in range(3)]
-actions += [Opt(op=OptOps.PADTO, axis=axis, amt=amt) for amt in [32] for axis in range(7)]
-actions += [Opt(op=OptOps.LOCAL, axis=0, amt=32), Opt(op=OptOps.UPCASTMID, axis=1, amt=4),
-            Opt(op=OptOps.GROUP, axis=0, amt=4), Opt(op=OptOps.GROUP, axis=0, amt=8), Opt(op=OptOps.GROUP, axis=1, amt=8),]
-if getenv("NOLOCALS"): actions += [Opt(op=OptOps.NOLOCALS)]
-
-def _get_test_global_size(global_size, max_global_size, var_vals):
-  test_global_size, factor = [sym_infer(sz, var_vals) for sz in global_size], 1
-  while prod(test_global_size) > max_global_size:
-    for j in range(len(global_size)-1,-1,-1):
-      if test_global_size[j] > 16:
-        test_global_size[j] //= 2
-        factor *= 2
-        break
-  return test_global_size, factor
-
-def _time_program(ast:LazyOp, rdev:Compiled, lib:bytes, global_size, local_size, var_vals, rawbufs, early_stop=None, max_global_size=65536, clear_l2=False, cnt=3, name="test"):  # noqa: E501
-  factor = 1
-  if global_size is not None and max_global_size is not None:
-    global_size, factor = _get_test_global_size(global_size, max_global_size, var_vals)
-  car = CompiledASTRunner(ast, name, "", lib, global_size, local_size).build(rdev.runtime)
-  tms = []
-  for _ in range(cnt):
-    if clear_l2:
-      with Context(DEBUG=0): Tensor.rand(1024,1024).realize()
-    tms.append(car(rawbufs, var_vals, wait=True, do_update_stats=False)*factor)
-    if early_stop is not None and early_stop < tms[-1]: break
-  return tms
-
-def _compile_linearizer(rdev:Compiled, lin:Linearizer, name:Optional[str]=None) -> Tuple[bytes, Optional[List[int]], Optional[List[int]]]:
-  lin.linearize()
-  src = rdev.renderer(name if name is not None else to_function_name(lin.name), lin.uops)   # NOTE: these all have the same name for deduping
-  return rdev.compiler(src), lin.global_size, lin.local_size
-
-def _try_compile_linearized_w_idx(x):
-  try: return (x[0], _compile_linearizer(cast(Compiled, Device[Device.DEFAULT]), x[1], "test"))
-  except Exception:
-    if DEBUG >= 4: traceback.print_exc()
-    return (x[0], None)
-
-# workers should ignore ctrl c
-def _init_worker(): signal.signal(signal.SIGINT, signal.SIG_IGN)
-
-# *** external API ***
-
-# get (scrap) buffers for timing the linearizer
-def bufs_from_lin(lin:Linearizer) -> List[Buffer]:
-  bufsts:DefaultDict[int, List[MemBuffer]] = defaultdict(list)
-  for x in lin.membufs: bufsts[x.idx].append(x)
-  rawbufs:List[Optional[Buffer]] = [None]*len(bufsts)
-  for k,lx in bufsts.items():
-    buf_size = prod(lx[0].dtype.shape) if isinstance(lx[0].dtype, ImageDType) else max(y.st.real_size() for y in lx)
-    rawbufs[k] = Buffer(Device.DEFAULT, buf_size, lx[0].dtype)
-  assert all(r is not None for r in rawbufs)
-  return cast(List[Buffer], rawbufs)
-
-# get dictionary of all possible actions
-def get_linearizer_actions(lin:Linearizer, include_0=True) -> Dict[int, Linearizer]:
-  acted_lins = {0:lin} if include_0 else {}
-  for i,a in enumerate(actions):
-    if a.axis is not None and a.axis >= lin.shape_len: continue
-    if a.axis is not None and lin.full_shape[a.axis] == a.amt and Opt(a.op, a.axis, 0) in actions: continue
-    lin2 = lin.copy()
-    try:
-      lin2.apply_opt(a)
-      up, lcl = 1, 1
-      for s,c in zip(lin2.full_shape, lin2.colors()):
-        if c in {"magenta", "yellow"}: up *= s
-        if c in {"cyan", "green", "white"}: lcl *= s
-      if up > 256 or lcl > 256: continue
-      acted_lins[i+1] = lin2
-    except Exception:
-      pass
-  return acted_lins
-
-def beam_search(lin:Linearizer, rawbufs, amt:int, allow_test_size=True) -> Linearizer:
-  key = {"ast": str(lin.ast), "amt": amt, "allow_test_size": allow_test_size, "device": Device.DEFAULT}
-  if (val:=diskcache_get("beam_search", key)) is not None and not getenv("IGNORE_BEAM_CACHE") and CACHELEVEL >= 1:
-    ret = lin.copy()
-    for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
-    return ret
-
-  beam: List[Tuple[Linearizer, float]] = []
-  seen_libs = set()
-
-  default_parallel = 1 if Device.DEFAULT in {"CUDA", "HIP"} else 0
-  pool = multiprocessing.Pool(multiprocessing.cpu_count(), _init_worker) if getenv("PARALLEL", default_parallel) else None
-
-  try:
-    var_vals = {k:(k.max+k.min)//2 for k in lin.ast.vars()}
-    exiting, st = False, time.perf_counter()
-    dev = Device[Device.DEFAULT]
-    assert isinstance(dev, Compiled)
-    while not exiting:
-      acted_lins = flatten([get_linearizer_actions(lin, include_0=False).values() for lin,_ in beam]) if len(beam) else [lin]
-      timed_lins: List[Tuple[Linearizer, float]] = []
-      for i,proc in (pool.imap_unordered(_try_compile_linearized_w_idx, enumerate(acted_lins)) if pool is not None else map(_try_compile_linearized_w_idx, enumerate(acted_lins))):  # noqa: E501
-        if proc is None: continue
-        lib, global_size, local_size = proc
-        if lib in seen_libs: continue
-        seen_libs.add(lib)
-        tms = _time_program(lin.ast, dev, lib, global_size, local_size, var_vals, rawbufs, early_stop=beam[0][1]*3 if len(beam) else 1.0)
-        timed_lins.append((acted_lins[i], min(tms)))
-        if DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s: {timed_lins[-1][1]*1e6:12.2f} us       {len(timed_lins):4d}/{len(acted_lins):4d}         {timed_lins[-1][0].colored_shape()}\033[K", end="")  # noqa: E501
-
-      # done
-      opts = sorted(timed_lins, key=lambda x: x[1])
-      exiting = len(opts) == 0 or (len(beam) > 0 and beam[0][1] <= opts[0][1])
-      if not exiting: beam = opts[:amt]
-      assert len(beam) > 0, "no BEAM items succeeded?!?"
-      if DEBUG >= 2: print(f"\r{time.perf_counter() - st:7.2f}s:", colored(f"{beam[0][1]*1e6:12.2f} us", "green" if exiting else None), f"from {len(acted_lins):3d} -> {len(opts):3d} actions\033[K", beam[0][0].colored_shape())  # noqa: E501
-    if pool is not None: pool.close()    # the pool is closed
-  except KeyboardInterrupt as e:
-    if pool is not None: pool.terminate()
-    raise e
-
-  if CACHELEVEL >= 1: diskcache_put("beam_search", key, beam[0][0].applied_opts)
-  if DEBUG >= 3: print(beam[0][0].applied_opts)
-  return beam[0][0]
-
-def optimize_local_size(clprg:Callable, global_size:List[int], rawbufs:List[Buffer]) -> List[int]:
-  test_rawbuffers = [Buffer(rawbufs[0].device, rawbufs[0].size, rawbufs[0].dtype), *rawbufs[1:]] if rawbufs[0] in rawbufs[1:] else rawbufs
-  MAX_WORKGROUP = 1024
-  local_dims = [[x for x in set([sz, 1, 2, 4, 8, 16, 32, 64, 128, 256, MAX_WORKGROUP]) if x<=sz] for sz in global_size]
-  local_sizes = [list(x) for x in itertools.product(*local_dims) if prod(x) <= MAX_WORKGROUP] * 2  # try each valid size twice
-  def try_exec(local_size):
-    try: return clprg(*[x._buf for x in test_rawbuffers], global_size=[g//l if g%l == 0 else g/l for g,l in zip(global_size, local_size)], local_size=local_size, wait=True)  # noqa: E501
-    except Exception: return float('inf')
-  ret = min([(try_exec(local_size), local_size) for local_size in random.sample(local_sizes, len(local_sizes))])
-  assert not math.isinf(ret[0]), "all optimize_local_size exec failed"
-  return ret[1]
-
-def time_linearizer(lin:Linearizer, rawbufs:List[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
-  key = {"ast": str(lin.ast), "opts": str(lin.applied_opts), "allow_test_size": allow_test_size, "max_global_size": max_global_size, "clear_l2": clear_l2, "device": Device.DEFAULT}  # noqa: E501
-  if not disable_cache and CACHELEVEL >= 2 and (val:=diskcache_get("time_linearizer", key)) is not None: return min(val)
-
-  dev = Device[Device.DEFAULT]
-  assert isinstance(dev, Compiled)
-
-  var_vals = {k:(k.max+k.min)//2 for k in lin.ast.vars()}
-  lib, global_size, local_size = _compile_linearizer(dev, lin)
-  tms = _time_program(lin.ast, dev, lib, global_size, local_size, var_vals, rawbufs, max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))  # noqa: E501
-
-  if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms)
-  return min(tms)