binop also

* cache apply_movement_op

* pyling and clear cache

* fix types

* ignore

* cleanup

.github/workflows/benchmark.yml

@@ -625,11 +625,11 @@ jobs:
     - name: benchmark openpilot 0.9.9 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.9.9 driving_policy
       run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.9.9 dmonitoring
-      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: openpilot compile3 Space Lab policy + vision
       run: |
         PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29

autogen_stubs.sh

@@ -435,8 +435,8 @@ generate_sqtt() {
     -o extra/sqtt/rocprof/rocprof.py
   fixup extra/sqtt/rocprof/rocprof.py
   sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes, tinygrad.helpers.fetch as tgfetch/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(tgfetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes\nfrom tinygrad.helpers import fetch/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
 }
 
 generate_webgpu() {

extra/sqtt/rgptool.py

@@ -156,6 +156,9 @@ class RGP:
     sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
     if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
     device_props = sqtt_events[0].props
+    gfx_ver = device_props['gfx_target_version'] // 10000
+    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
+                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
     sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
     sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
     sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -193,7 +196,7 @@ class RGP:
         flags=0,
         trace_shader_core_clock=0x93f05080,
         trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
+        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550}[device_props['gfx_target_version']],
         device_revision_id=0xc8,
         vgprs_per_simd=1536,
         sgprs_per_simd=128*16,
@@ -207,7 +210,7 @@ class RGP:
         sgpr_alloc_granularity=128,
         hardware_contexts=8,
         gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
-        gfxip_level=sqtt.SQTT_GFXIP_LEVEL_GFXIP_11_0,
+        gfxip_level=gfx_iplvl,
         gpu_index=0,
         gds_size=0,
         gds_per_shader_engine=0,
@@ -258,7 +261,7 @@ class RGP:
             major_version=0, minor_version=2,
           ),
           shader_engine_index=sqtt_event.se,
-          sqtt_version=sqtt.SQTT_VERSION_3_2,
+          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
           _0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
             v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
               instrumentation_spec_version=1,

extra/sqtt/roc.py

@@ -20,14 +20,15 @@ class InstInfo:
 class _ROCParseCtx:
   def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
     self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
-    self.wave_events = {}
+    self.wave_events, self.disasms, self.addr2prg = {}, {}, {}
+
+    for prog in prog_evs:
+      for addr, info in comgr_get_address_table(prog.lib).items():
+        self.disasms[prog.base + addr] = info
+        self.addr2prg[prog.base + addr] = prog
 
   def next_sqtt(self): return next(self.sqtt_evs, None)
-  def find_program(self, idx): return self.prog_evs[idx]
-  def get_instr_info(self, idx, exec_addr): return self.disasm_program(idx)[exec_addr - self.find_program(idx).base]
-
-  @functools.lru_cache(None)
-  def disasm_program(self, idx): return comgr_get_address_table(self.find_program(idx).lib)
+  def find_program(self, addr): return self.addr2prg[addr]
 
   def on_occupancy_ev(self, ev):
     if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
@@ -39,10 +40,10 @@ class _ROCParseCtx:
     for j in range(ev.instructions_size):
       inst_ev = ev.instructions_array[j]
       inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.get_instr_info(inst_ev.pc.code_object_id, inst_ev.pc.address)[0]))
+      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.disasms[inst_ev.pc.address][0]))
       asm[inst_ev.pc.address].on_ev(inst_ev)
 
-    self.wave_events[(self.find_program(ev.instructions_array[0].pc.code_object_id).name, ev.wave_id, ev.cu, ev.simd)] = asm
+    self.wave_events[(self.find_program(ev.instructions_array[0].pc.address).name, ev.wave_id, ev.cu, ev.simd)] = asm
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
@@ -78,7 +79,7 @@ if __name__ == "__main__":
 
   @rocprof.rocprof_trace_decoder_isa_callback_t
   def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.get_instr_info(pc.code_object_id, pc.address)
+    instr, mem_size_ptr[0] = ROCParseCtx.disasms[pc.address]
 
     # this is the number of bytes to next instruction, set to 0 for end_pgm
     if instr == "s_endpgm": mem_size_ptr[0] = 0

extra/sqtt/rocprof/rocprof.py

@@ -7,7 +7,8 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes, tinygrad.helpers.fetch as tgfetch
+import ctypes
+from tinygrad.helpers import fetch
 
 
 class AsDictMixin:
@@ -155,7 +156,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['FIXME_STUB'] = ctypes.CDLL(str(tgfetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
 
 
 

extra/sqtt/sqtt.h

@@ -43,6 +43,7 @@ enum sqtt_version
    SQTT_VERSION_2_3 = 0x6, /* GFX9 */
    SQTT_VERSION_2_4 = 0x7, /* GFX10+ */
    SQTT_VERSION_3_2 = 0xb, /* GFX11+ */
+   SQTT_VERSION_3_3 = 0xc, /* GFX12+ */
 };
 
 enum sqtt_file_chunk_type
@@ -144,6 +145,8 @@ enum sqtt_gfxip_level
    SQTT_GFXIP_LEVEL_GFXIP_10_1 = 0x7,
    SQTT_GFXIP_LEVEL_GFXIP_10_3 = 0x9,
    SQTT_GFXIP_LEVEL_GFXIP_11_0 = 0xc,
+   SQTT_GFXIP_LEVEL_GFXIP_11_5 = 0xd,
+   SQTT_GFXIP_LEVEL_GFXIP_12 = 0x10,
 };
 
 enum sqtt_memory_type
@@ -427,6 +430,8 @@ enum elf_gfxip_level
    EF_AMDGPU_MACH_AMDGCN_GFX1010 = 0x033,
    EF_AMDGPU_MACH_AMDGCN_GFX1030 = 0x036,
    EF_AMDGPU_MACH_AMDGCN_GFX1100 = 0x041,
+   EF_AMDGPU_MACH_AMDGCN_GFX1150 = 0x043,
+   EF_AMDGPU_MACH_AMDGCN_GFX1200 = 0x04e,
 };
 
 struct sqtt_file_chunk_spm_db {

ruff.toml

@@ -50,4 +50,7 @@ exclude = [
   "E303", "E304", "E501", "E702", "E703", "E731", "W191",
   "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
   "FURB110", "RUF018", "F541", "F841"
-]
+]
+
+[format]
+exclude = ["*"]

test/external/external_uop_gc.py

@@ -2,6 +2,7 @@ import gc
 from tinygrad import Tensor, UOp, Device, nn
 from tinygrad.shape.shapetracker import views_to_valid_uop
 from tinygrad.engine.realize import method_cache, get_program
+from tinygrad.schedule.indexing import apply_movement_op
 from test.test_tiny import TestTiny
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
@@ -69,6 +70,7 @@ if __name__ == "__main__":
     # these caches will keep uops alive
     method_cache.clear()
     views_to_valid_uop.cache_clear()
+    apply_movement_op.cache_clear()
     Tensor._device_seeds.clear()
     Tensor._device_rng_counters.clear()
 

test/test_ops.py

@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -901,7 +901,6 @@ class TestOps(unittest.TestCase):
   def test_abs_exact(self):
     helper_test_op(None, torch.abs, Tensor.abs, vals=[[-1.,0,1]])
 
-  @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes")
   def test_log(self):
     helper_test_op([(45,65)], torch.log, Tensor.log)
     helper_test_op(None, torch.log, Tensor.log, vals=[[math.inf, -math.inf, math.nan]])
@@ -911,7 +910,6 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, torch.log2, Tensor.log2, vals=[[math.inf, -math.inf, math.nan]])
     helper_test_op([()], torch.log2, Tensor.log2)
 
-  @unittest.skipIf(TRANSCENDENTAL and Device.DEFAULT=="AMD", "TODO: remu crashes")
   def test_exp(self):
     helper_test_op([(45,65)], torch.exp, Tensor.exp)
     helper_test_op(None, torch.exp, Tensor.exp, vals=[[math.inf, -math.inf, math.nan]])
@@ -1549,7 +1547,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,4,5,6)], lambda x: torch.stack(torch.std_mean(x, axis=(1,2))),
                                 lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))
 
-  @unittest.skip("TODO: this fails because of loaded nan in mul folding")
   def test_std_mean_loaded_nan(self):
     helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
                                   lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))

test/test_outerworld.py

@@ -1,5 +1,15 @@
 import unittest
-from tinygrad import Tensor, UOp, GlobalCounters, Context
+from tinygrad import Tensor, UOp
+from tinygrad.uop.ops import AxisType, Ops
+
+class TestOuterworldReduce(unittest.TestCase):
+  def test_reduce(self):
+    x = Tensor.ones(5, 5).contiguous()
+    a = UOp.range(5, -1, AxisType.REDUCE)
+    out = x[a]
+    # TODO: syntax for this
+    t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, a), arg=Ops.ADD))
+    self.assertListEqual(t.tolist(), [5.,5.,5.,5.,5.])
 
 class TestOuterworld(unittest.TestCase):
   def test_range_plus_1(self):
@@ -13,6 +23,17 @@ class TestOuterworld(unittest.TestCase):
 
     self.assertTrue((t+1==cpy).all().item())
 
+  def test_range_plus_1_transpose(self):
+    t = Tensor.arange(100).reshape(10,10).realize()
+
+    # passthrough ranges
+    a = UOp.range(10, -1)
+    sel = t[a] + 1
+    assert sel.shape == (10,)
+    cpy = sel.reshape(10, 1).expand(10, a).contiguous().realize()
+
+    self.assertTrue(((t+1).T==cpy).all().item())
+
   def test_flip_range(self):
     t = Tensor.rand(10, 10).realize()
 
@@ -37,39 +58,17 @@ class TestOuterworld(unittest.TestCase):
     out.realize()
     self.assertTrue((out==20).all().item())
 
-  @unittest.skip("opts don't work")
-  def test_triple_gemm(self):
-    x = Tensor.rand(1, 16).realize()
-    W = Tensor.rand(3, 16, 16).realize()
+  def test_fancy_vmap(self):
+    def f(x,y): return x+y
 
-    manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
+    x = Tensor.arange(9).reshape(3,3).contiguous()
+    y = Tensor.arange(9).reshape(3,3).contiguous()
 
     a = UOp.range(3, -1)
-    x = x.assign(x @ W[a])
-    out = x.contiguous(a)[-1].contiguous().realize()
-
-    self.assertTrue((manual==out).all().item())
-
-  def test_setitem_pyrange(self):
-    with Context(DEBUG=0):
-      t = Tensor.rand(10).realize()
-      o = Tensor.empty(10)
-    GlobalCounters.reset()
-    for i in range(10):
-      o[i] = t[i]
-    o.realize()
-    self.assertTrue((t==o).all().item())
-
-  @unittest.skip("TODO: fix this")
-  def test_setitem(self):
-    with Context(DEBUG=0):
-      t = Tensor.rand(10).realize()
-      o = Tensor.empty(10)
-    GlobalCounters.reset()
-    i = UOp.range(10, -1)
-    o[i] = t[i]
-    o.contiguous(i).realize()
-    self.assertTrue((t==o).all().item())
+    out = f(x[:,a], y[a,:])
+    # TODO: this should support flatten
+    out = out.reshape(1, 3).expand(a, 3).contiguous().realize()
+    self.assertListEqual([[0,4,8],[4,8,12],[8,12,16]], out.tolist())
 
 if __name__ == '__main__':
   unittest.main()

test/test_schedule.py

@@ -1526,7 +1526,7 @@ class TestSchedule(unittest.TestCase):
     # run_schedule(check_schedule(out, 1))
     run_schedule(check_schedule(out, 4))
     np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
-                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-6)
+                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)
 
   def test_shrink_pad_safe(self):
     a = Tensor.ones((3, )).contiguous().realize()

test/unit/test_helpers.py

@@ -1,9 +1,8 @@
 import ctypes, gzip, unittest, timeit
 from tinygrad import Variable
-from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
+from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address, get_contraction
 from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
 from tinygrad.tensor import Tensor, get_shape
-from tinygrad.shape.view import get_contraction
 import numpy as np
 
 VARIABLE = ContextVar("VARIABLE", 0)

test/unit/test_shapetracker.py

@@ -154,21 +154,6 @@ class TestRealStrides(unittest.TestCase):
     ))
     self.assertEqual(st.is_expanded(), (False, False, False, True, False))
 
-class TestRealSimplifies(unittest.TestCase):
-  def tearDown(self):
-    self.st = self.st.simplify()
-    assert len(self.st.views) == 1
-
-  def test_1(self):
-    self.st = ShapeTracker((
-      View.create((1, 3, 2, 11, 4, 28), (0, 308, 0, 28, 0, 1), 0, None),
-      View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)))
-
-  def test_2(self):
-    self.st = ShapeTracker((
-      View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
-      View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
-
 class TestIndexExpressions2d(unittest.TestCase):
   def setUp(self):
     shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5

test/unit/test_shapetracker_math.py

@@ -62,6 +62,7 @@ class TestShapeTrackerAdd(unittest.TestCase):
     b = ShapeTracker.from_shape((100,))
     assert a+b == b
 
+  @unittest.skip("no longer simplifies")
   def test_simple_add_permute(self):
     a = ShapeTracker.from_shape((10, 10))
     a = a.permute((1,0))

test/unit/test_symbolic_shapetracker.py

@@ -28,7 +28,7 @@ class TestSymbolic(unittest.TestCase):
   def test_merge_view_recursion_err(self):
     vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
     vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    self.assertEqual(vm2+vm1, vm1)
+    self.assertEqual(vm2+vm1, None)
 
   def test_merge_view_recursion_err2(self):
     vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)

test/unit/test_view.py

@@ -69,161 +69,5 @@ class TestMergeDims(unittest.TestCase):
     # print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
     self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
 
-class TestMergeViews(unittest.TestCase):
-  def test_with_mask_0(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(1, 1, 5, 8), strides=(0, 0, 5, 1), offset=-3, mask=((0, 1), (0, 1), (0, 5), (3, 8)), contiguous=False)
-    v1 = View(shape=(1, 1, 2, 2), strides=(0, 0, 8, 1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 2, 2), strides=(0, 0, 5, 1), offset=0, mask=None, contiguous=False))
-
-  def test_with_mask_1(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, 3, 1), offset=-6, mask=((0, 3), (0, 3), (2, 4), (1, 3)), contiguous=False)
-    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, 3, 1), offset=7, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=1, mask=None, contiguous=False))
-
-  def test_with_mask_2(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (0, 2)), contiguous=False)
-    v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, -3, 1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=3, mask=None, contiguous=False))
-
-  def test_with_mask_3(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(3, 3, 4, 4), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False)
-    v1 = View(shape=(3, 3, 4, 2), strides=(48, 16, 4, 1), offset=0, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 4, 2), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False))
-
-  def test_with_mask_4(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (1, 3)), contiguous=False)
-    v1 = View(shape=(3, 3, 3, 3), strides=(45, 15, 3, 1), offset=6, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(3, 3, 3, 3), strides=(0, 0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0), (0, 0)), contiguous=False))
-
-  def test_with_mask_5(self):
-    # from test/test_ops.py::TestOps::test_pad_reflect_mode
-    # has a mask in the final view
-    v0 = View(shape=(1, 1, 6, 5), strides=(0, 0, 5, 1), offset=-5, mask=((0, 1), (0, 1), (1, 6), (0, 5)), contiguous=False)
-    v1 = View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=3, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=-2, mask=((0, 1), (0, 1), (1, 6), (0, 3)), contiguous=False))
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer1(self):
-    v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
-    v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
-    target = View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer2(self):
-    v0 = View(shape=(5, 10, 12), strides=(100, 1, 10), offset=-20, mask=((0, 5), (0, 10), (2, 12)), contiguous=False)
-    v1 = View(shape=(10, 6, 5, 2, 2), strides=(12, 2, 120, 1, 0), offset=0, mask=None, contiguous=False)
-    target = View(shape=(10, 6, 5, 2, 2), strides=(1, 20, 100, 10, 0), offset=-20, mask=((0, 10), (1, 6), (0, 5), (0, 2), (0, 2)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer3(self):
-    v0 = View(shape=(8, 7, 3), strides=(1, 12, -4), offset=6, mask=((2, 6), (0, 7), (0, 3)), contiguous=False)
-    v1 = View(shape=(4, 2, 6, 2, 1), strides=(42, 21, 3, 1, 0), offset=4, mask=None, contiguous=False)
-    target = View(shape=(4, 2, 6, 2, 1), strides=(2, 1, 12, -4, 0), offset=14, mask=((1, 3), (0, 2), (0, 6), (0, 2), (0, 1)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer4(self):
-    v0 = View(shape=(7, 21, 3), strides=(54, 3, 1), offset=-9, mask=((0, 6), (3, 21), (0, 3)), contiguous=False)
-    v1 = View(shape=(5, 3, 3, 7), strides=(63, 1, 3, 9), offset=63, mask=None, contiguous=False)
-    target = View(shape=(5, 3, 3, 7), strides=(54, 1, 3, 9), offset=45, mask=((0, 5), (0, 3), (0, 3), (1, 7)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  @unittest.expectedFailure  # TODO: fix these
-  def test_merges_from_fuzzer5(self):
-    v0 = View(shape=(5, 1, 24), strides=(20, 0, 1), offset=-2, mask=((0, 5), (0, 1), (2, 22)), contiguous=False)
-    v1 = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 24, 0, 0), offset=0, mask=None, contiguous=False)
-    target = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 20, 0, 0), offset=-2, mask=((1, 11), (0, 2), (0, 5), (0, 2), (0, 1)), contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  def test_merge_views_variable(self):
-    from tinygrad import Variable
-    N = 100
-    start_pos = Variable("start_pos", 1, N-1)
-    v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
-    v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
-    target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, target)
-
-  def test_view_padded_area1(self):
-    # test_multinomial
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
-    v1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1,), strides=(0,), offset=0, mask=((0, 0),), contiguous=False))
-
-  def test_view_padded_area2(self):
-    # test_pad_reflect_mode
-    v0 = View(shape=(1, 1, 10, 7), strides=(0, 0, 5, 1), offset=-15, mask=((0, 1), (0, 1), (3, 8), (0, 5)), contiguous=False)
-    v1 = View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True))
-
-  def test_view_padded_area3(self):
-    # test_roll
-    v0 = View(shape=(2, 4), strides=(0, 1), offset=4, mask=((0, 1), (0, 4)), contiguous=False)
-    v1 = View(shape=(1, 4), strides=(0, 1), offset=4, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 4), strides=(0, 0), offset=0, mask=((0, 0), (0, 0)), contiguous=False))
-
-  def test_view_padded_area4(self):
-    # test_std_mean
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((0, 1),), contiguous=False)
-    v1 = View(shape=(1, 1, 1), strides=(0, 0, 0), offset=1, mask=None, contiguous=False)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0)), contiguous=False))
-
-  def test_empty_shape_view1(self):
-    # test_stack_slice
-    v0 = View(shape=(3, 5), strides=(0, 1), offset=0, mask=((0, 1), (0, 5)), contiguous=False)
-    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    self.assertIsNotNone(v)
-    self.assertEqual(v, View(shape=(), strides=(), offset=0, mask=None, contiguous=True))
-
-  def test_empty_shape_view2(self):
-    # test_std_mean
-    v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
-    v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
-    v = v0 + v1
-    # TODO: why is this different?
-    self.assertIsNone(v)
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/gpudims.py

@@ -1,8 +1,7 @@
 import math
 from tinygrad.uop.ops import UOp, Ops, sint, PatternMatcher, UPat, KernelInfo, ssimplify, AxisType, sint_to_uop
-from tinygrad.helpers import all_int, dedup
+from tinygrad.helpers import all_int, dedup, get_contraction
 from tinygrad.dtype import dtypes
-from tinygrad.shape.view import get_contraction
 from tinygrad.renderer import Renderer
 
 def _group_dims(dims:tuple[sint, ...], max_sizes:tuple[int, ...]):

tinygrad/codegen/late/devectorizer.py

@@ -11,7 +11,7 @@ from tinygrad.renderer import Renderer
 # ***** image load valid simplification *****
 
 def simplify_valid_load(buf:UOp, start_idx:UOp, valid:UOp) -> UOp|None:
-  if (idx:=uop_given_valid(valid, start_idx)) is None: return buf.index(UOp.invalid())
+  idx = uop_given_valid(valid, start_idx)
   if not isinstance(buf.dtype, ImageDType): return None if idx is start_idx else buf.index(idx.valid(valid))
 
   # wait for it to be image indexed before running simplification

tinygrad/codegen/opt/__init__.py

@@ -19,7 +19,7 @@ class Opt:
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                 AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
 axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
-               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "green", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
 
 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):

tinygrad/gradient.py

@@ -39,7 +39,7 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.EXPAND, name="ret"), lambda ctx, ret: (ctx.r(Ops.ADD, tuple(i for i,(si,so) in enumerate(zip(ret.src[0].shape, ret.arg)) if si!=so)),)),
   (UPat(Ops.MULTI, name="ret"), lambda ctx, ret: ctx.shard(ret.device, ret.axis).src),
   # there's no gradient for bitcast
-  (UPat(Ops.BITCAST), lambda ctx: (None,)),
+  (UPat(Ops.BITCAST), lambda: (None,)),
 ])
 
 def _deepwalk(root:UOp, targets:set[UOp]) -> list[UOp]:

tinygrad/helpers.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, sqlite3, tempfile, pathlib, string, ctypes, sys, gzip, getpass
-import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal
+import urllib.request, subprocess, shutil, math, types, copyreg, inspect, importlib, decimal, itertools
 from dataclasses import dataclass, field
 from typing import ClassVar, Iterable, Any, TypeVar, Callable, Sequence, TypeGuard, Iterator, Generic, Generator
 
@@ -82,6 +82,13 @@ def word_wrap(x, wrap=80):
   while len(ansistrip(x[:i])) < wrap and i < len(x): i += 1
   return x[:i] + "\n" + word_wrap(x[i:], wrap)
 
+# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
+def get_contraction(old_shape:tuple[T, ...], new_shape:tuple[T, ...]) -> list[list[int]]|None: # T is sint
+  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
+  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
+  except ValueError: return None
+  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
+
 def suppress_finalizing(func):
   def wrapper(*args, **kwargs):
     try: return func(*args, **kwargs)
@@ -89,7 +96,7 @@ def suppress_finalizing(func):
       if not getattr(sys, 'is_finalizing', lambda: True)(): raise # re-raise if not finalizing
   return wrapper
 
-def unwrap_class_type(cls_t:T): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
+def unwrap_class_type(cls_t): return cls_t.func if isinstance(cls_t, functools.partial) else cls_t
 
 def pluralize(st:str, cnt:int): return f"{cnt} {st}"+('' if cnt == 1 else 's')
 

tinygrad/renderer/cstyle.py

@@ -2,7 +2,7 @@ from typing import Literal, Callable, cast
 import os, math, sys
 from collections import defaultdict, Counter
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, sint_to_uop, range_str
+from tinygrad.uop.ops import GroupOp, Ops, UOp, PatternMatcher, UPat, range_str
 from tinygrad.helpers import strip_parens, getenv, prod, dedup, AMX, CPU_COUNT
 from tinygrad.dtype import ImageDType, dtypes, DType, PtrDType, AddrSpace, truncate
 from tinygrad.renderer import Renderer
@@ -112,7 +112,7 @@ class CStyleLanguage(Renderer):
     buftypes = [(name, self.render_dtype(dtype, mutable)+self.buffer_suffix if isinstance(dtype, (ImageDType, PtrDType)) else
                 self.arg_int_prefix if dtype == dtypes.int else None) for name,(dtype,mutable) in bufs]
     local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
-    launch_bounds = sint_to_uop(prod(local_dims)).vmax
+    launch_bounds = prod([d.vmax for d in local_dims])
     prg = ''.join([f"{self.kernel_typedef.format(launch_bounds=launch_bounds)} {function_name}(",] +
     [', '.join([f'{t} {name}' for name,t in buftypes] + self.extra_args)] +
     [") {\n" + tmp] + ['\n'.join(kernel), "\n}"])

tinygrad/renderer/llvmir.py

@@ -4,7 +4,7 @@ from tinygrad.codegen.opt import tc
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import AMDRenderer
 from tinygrad.uop.decompositions import xexp2, xlog2
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, sint_to_uop, range_str
+from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops, GroupOp, range_str
 from tinygrad.dtype import dtypes, DType, PtrDType, truncate
 from tinygrad.helpers import prod, AMX
 
@@ -226,7 +226,7 @@ class AMDLLVMRenderer(LLVMRenderer):
   def _render_footer(self, uops: list[UOp]) -> str:
     # TODO: this is copied from cstyle
     local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
-    requiredMaxThreadsPerBlock = sint_to_uop(prod(local_dims)).vmax
+    requiredMaxThreadsPerBlock = prod([d.vmax for d in local_dims])
     attributes = ["alwaysinline", "nounwind", '"no-builtins"',
                   f'"amdgpu-flat-work-group-size"="1,{requiredMaxThreadsPerBlock}"', '"no-trapping-math"="true"']
     return 'attributes #0 = { ' + ' '.join(attributes) + ' }'

tinygrad/renderer/ptx.py

@@ -2,7 +2,7 @@ from typing import cast, Callable
 import struct
 from collections import defaultdict
 from tinygrad.codegen.opt import tc
-from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp, sint_to_uop
+from tinygrad.uop.ops import Ops, UOp, PatternMatcher, UPat, GroupOp
 from tinygrad.dtype import dtypes, DType, PtrDType, AddrSpace
 from tinygrad.renderer import Renderer
 from tinygrad.renderer.cstyle import CUDARenderer
@@ -157,7 +157,7 @@ class PTXRenderer(Renderer):
     def fmt(line): return line if line[0]=="$" else "\t" + line.replace(" ", "\t" if len(line.split(" ")[0]) > 7 else "\t\t", 1)
     kernel = '\n'.join(map(fmt, [f".reg .{reg.split('_')[-2]} %{reg}<{cnt}>;" for reg,cnt in regs] + kernel + ["ret;"]))
     local_dims = [u.src[0] for u in uops if u.op is Ops.SPECIAL and u.arg[0] == "l"]
-    launch_bounds = sint_to_uop(prod(local_dims)).vmax
+    launch_bounds = prod([d.vmax for d in local_dims])
     params = ',\n\t'.join([f".param .{'u64' if dtype.__class__ == PtrDType else self.types[dtype]} {name}" for name,dtype in bufs])
     return f"{self.kernel_prefix.format(launch_bounds=launch_bounds)} {function_name} (\n\t{params}\n)\n.maxntid {launch_bounds}\n{{\n{kernel}\n}}"
 

tinygrad/runtime/autogen/sqtt.py

@@ -174,12 +174,14 @@ sqtt_version__enumvalues = {
     6: 'SQTT_VERSION_2_3',
     7: 'SQTT_VERSION_2_4',
     11: 'SQTT_VERSION_3_2',
+    12: 'SQTT_VERSION_3_3',
 }
 SQTT_VERSION_NONE = 0
 SQTT_VERSION_2_2 = 5
 SQTT_VERSION_2_3 = 6
 SQTT_VERSION_2_4 = 7
 SQTT_VERSION_3_2 = 11
+SQTT_VERSION_3_3 = 12
 sqtt_version = ctypes.c_uint32 # enum
 
 # values for enumeration 'sqtt_file_chunk_type'
@@ -336,6 +338,8 @@ sqtt_gfxip_level__enumvalues = {
     7: 'SQTT_GFXIP_LEVEL_GFXIP_10_1',
     9: 'SQTT_GFXIP_LEVEL_GFXIP_10_3',
     12: 'SQTT_GFXIP_LEVEL_GFXIP_11_0',
+    13: 'SQTT_GFXIP_LEVEL_GFXIP_11_5',
+    16: 'SQTT_GFXIP_LEVEL_GFXIP_12',
 }
 SQTT_GFXIP_LEVEL_NONE = 0
 SQTT_GFXIP_LEVEL_GFXIP_6 = 1
@@ -346,6 +350,8 @@ SQTT_GFXIP_LEVEL_GFXIP_9 = 5
 SQTT_GFXIP_LEVEL_GFXIP_10_1 = 7
 SQTT_GFXIP_LEVEL_GFXIP_10_3 = 9
 SQTT_GFXIP_LEVEL_GFXIP_11_0 = 12
+SQTT_GFXIP_LEVEL_GFXIP_11_5 = 13
+SQTT_GFXIP_LEVEL_GFXIP_12 = 16
 sqtt_gfxip_level = ctypes.c_uint32 # enum
 
 # values for enumeration 'sqtt_memory_type'
@@ -806,12 +812,16 @@ elf_gfxip_level__enumvalues = {
     51: 'EF_AMDGPU_MACH_AMDGCN_GFX1010',
     54: 'EF_AMDGPU_MACH_AMDGCN_GFX1030',
     65: 'EF_AMDGPU_MACH_AMDGCN_GFX1100',
+    67: 'EF_AMDGPU_MACH_AMDGCN_GFX1150',
+    78: 'EF_AMDGPU_MACH_AMDGCN_GFX1200',
 }
 EF_AMDGPU_MACH_AMDGCN_GFX801 = 40
 EF_AMDGPU_MACH_AMDGCN_GFX900 = 44
 EF_AMDGPU_MACH_AMDGCN_GFX1010 = 51
 EF_AMDGPU_MACH_AMDGCN_GFX1030 = 54
 EF_AMDGPU_MACH_AMDGCN_GFX1100 = 65
+EF_AMDGPU_MACH_AMDGCN_GFX1150 = 67
+EF_AMDGPU_MACH_AMDGCN_GFX1200 = 78
 elf_gfxip_level = ctypes.c_uint32 # enum
 class struct_sqtt_file_chunk_spm_db(Structure):
     pass
@@ -1607,7 +1617,8 @@ __all__ = \
     'ApiCmdUpdateBuffer', 'ApiCmdWaitEvents', 'ApiCmdWriteTimestamp',
     'ApiInvalid', 'ApiRayTracingSeparateCompiled',
     'EF_AMDGPU_MACH_AMDGCN_GFX1010', 'EF_AMDGPU_MACH_AMDGCN_GFX1030',
-    'EF_AMDGPU_MACH_AMDGCN_GFX1100', 'EF_AMDGPU_MACH_AMDGCN_GFX801',
+    'EF_AMDGPU_MACH_AMDGCN_GFX1100', 'EF_AMDGPU_MACH_AMDGCN_GFX1150',
+    'EF_AMDGPU_MACH_AMDGCN_GFX1200', 'EF_AMDGPU_MACH_AMDGCN_GFX801',
     'EF_AMDGPU_MACH_AMDGCN_GFX900', 'EventCmdBlitImage',
     'EventCmdBuildAccelerationStructuresIndirectKHR',
     'EventCmdBuildAccelerationStructuresKHR',
@@ -1671,7 +1682,8 @@ __all__ = \
     'SQTT_FILE_CHUNK_TYPE_SQTT_DESC', 'SQTT_FILE_MAGIC_NUMBER',
     'SQTT_FILE_VERSION_MAJOR', 'SQTT_FILE_VERSION_MINOR',
     'SQTT_GFXIP_LEVEL_GFXIP_10_1', 'SQTT_GFXIP_LEVEL_GFXIP_10_3',
-    'SQTT_GFXIP_LEVEL_GFXIP_11_0', 'SQTT_GFXIP_LEVEL_GFXIP_6',
+    'SQTT_GFXIP_LEVEL_GFXIP_11_0', 'SQTT_GFXIP_LEVEL_GFXIP_11_5',
+    'SQTT_GFXIP_LEVEL_GFXIP_12', 'SQTT_GFXIP_LEVEL_GFXIP_6',
     'SQTT_GFXIP_LEVEL_GFXIP_7', 'SQTT_GFXIP_LEVEL_GFXIP_8',
     'SQTT_GFXIP_LEVEL_GFXIP_8_1', 'SQTT_GFXIP_LEVEL_GFXIP_9',
     'SQTT_GFXIP_LEVEL_NONE', 'SQTT_GPU_NAME_MAX_SIZE',
@@ -1697,9 +1709,9 @@ __all__ = \
     'SQTT_QUEUE_TYPE_COMPUTE', 'SQTT_QUEUE_TYPE_DMA',
     'SQTT_QUEUE_TYPE_UNIVERSAL', 'SQTT_QUEUE_TYPE_UNKNOWN',
     'SQTT_SA_PER_SE', 'SQTT_VERSION_2_2', 'SQTT_VERSION_2_3',
-    'SQTT_VERSION_2_4', 'SQTT_VERSION_3_2', 'SQTT_VERSION_NONE',
-    'UserEventObjectName', 'UserEventPop', 'UserEventPush',
-    'UserEventTrigger', 'elf_gfxip_level',
+    'SQTT_VERSION_2_4', 'SQTT_VERSION_3_2', 'SQTT_VERSION_3_3',
+    'SQTT_VERSION_NONE', 'UserEventObjectName', 'UserEventPop',
+    'UserEventPush', 'UserEventTrigger', 'elf_gfxip_level',
     'rgp_sqtt_marker_event_type', 'rgp_sqtt_marker_general_api_type',
     'rgp_sqtt_marker_identifier', 'rgp_sqtt_marker_user_event_type',
     'sqtt_api_type', 'sqtt_engine_type',

tinygrad/runtime/ops_amd.py

@@ -130,8 +130,9 @@ class AMDComputeQueue(HWQueue):
       self.wreg(self.gc.regSQ_THREAD_TRACE_USERDATA_2, *data_ints[i:i+2])
 
   def sqtt_config(self, tracing:bool):
-    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1,
-              rt_freq=self.soc.SQ_TT_RT_FREQ_4096_CLK, util_timer=self.soc.SQ_TT_UTIL_TIMER_250_CLK, mode=int(tracing))
+    trace_ctrl = {'rt_freq': self.soc.SQ_TT_RT_FREQ_4096_CLK} if self.dev.target < (12,0,0) else {}
+    self.wreg(self.gc.regSQ_THREAD_TRACE_CTRL, draw_event_en=1, spi_stall_en=1, sq_stall_en=1, reg_at_hwm=2, hiwater=1, util_timer=1,
+      mode=int(tracing), **trace_ctrl)
 
   # Magic values from mesa/src/amd/vulkan/radv_sqtt.c:radv_emit_spi_config_cntl and src/amd/common/ac_sqtt.c:ac_sqtt_emit_start
   def sqtt_start(self, buf0s:list[HCQBuffer], se_mask:int):
@@ -140,24 +141,35 @@ class AMDComputeQueue(HWQueue):
     # One buffer for one SE, mesa does it with a single buffer and ac_sqtt_get_data_offset, but this is simpler and should work just as well
     for se in range(len(buf0s)):
       self.wreg(self.gc.regGRBM_GFX_INDEX, se_index=se, instance_broadcast_writes=1)
-      buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr>>12)
-      self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size>>12)
-      self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
+      buf0_lo, buf0_hi = data64_le(buf0s[se].va_addr >> 12)
+      if self.dev.target >= (12,0,0):
+        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, size=buf0s[se].size >> 12)
+        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_LO, base_lo=buf0_lo)
+        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE_HI, base_hi=buf0_hi)
+      else:
+        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_SIZE, base_hi=buf0_hi, size=buf0s[se].size >> 12)
+        self.wreg(self.gc.regSQ_THREAD_TRACE_BUF0_BASE, base_lo=buf0_lo)
       # NOTE: SQTT can only trace instructions on one simd per se, this selects first simd in first wgp in first sa.
       # For RGP to display instruction trace it has to see it on first SE. Howerver ACE/MEC/whatever does the dispatching starting with second se,
       # and on amdgpu/non-AM it also does weird things with dispatch order inside se: around 7 times out of 10 it starts from the last cu, but
       # sometimes not, especially if the kernel has more than one wavefront which means that kernels with small global size might get unlucky and
       # be dispatched on something else and not be seen in instruction tracing tab. You can force the wavefronts of a kernel to be dispatched on the
       # CUs you want to by disabling other CUs via bits in regCOMPUTE_STATIC_THREAD_MGMT_SE<x> and trace even kernels that only have one wavefront.
-      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT, simd_sel=0, wgp_sel=0, sa_sel=0)
+      cs_wtype = (1 << 6) if self.dev.target >= (12,0,0) else self.soc.SQ_TT_WTYPE_INCLUDE_CS_BIT
+      self.wreg(self.gc.regSQ_THREAD_TRACE_MASK, wtype_include=cs_wtype, simd_sel=0, wgp_sel=0, sa_sel=0)
       reg_include = self.soc.SQ_TT_TOKEN_MASK_SQDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_SHDEC_BIT | self.soc.SQ_TT_TOKEN_MASK_GFXUDEC_BIT | \
-                    self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
-      token_exclude = 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT
+                    self.soc.SQ_TT_TOKEN_MASK_COMP_BIT | self.soc.SQ_TT_TOKEN_MASK_CONTEXT_BIT
+      token_exclude = (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_PERF_SHIFT) if self.dev.target < (12,0,0) else 0
+
+      # disable tracing
       if not (se_mask >> se) & 0b1:
-        token_exclude |= 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
+        # gfx12 doesn't have enums with all fields, so it's hardcoded, but it's the same as gfx11.
+        token_exclude |= (1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VMEMEXEC_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_ALUEXEC_SHIFT | \
                          1 << self.soc.SQ_TT_TOKEN_EXCLUDE_VALUINST_SHIFT | 1 << self.soc.SQ_TT_TOKEN_EXCLUDE_IMMEDIATE_SHIFT | \
-                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT
-      self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1)
+                         1 << self.soc.SQ_TT_TOKEN_EXCLUDE_INST_SHIFT) if self.dev.target < (12,0,0) else 0x927
+
+      token_mask = {} if self.dev.target < (12,0,0) else {'exclude_barrier_wait': 1}
+      self.wreg(self.gc.regSQ_THREAD_TRACE_TOKEN_MASK, reg_include=reg_include, token_exclude=token_exclude, bop_events_token_include=1, **token_mask)
       # Enable SQTT
       self.sqtt_config(tracing=True)
     # Restore global broadcasting
@@ -178,9 +190,6 @@ class AMDComputeQueue(HWQueue):
       # Wait for FINISH_PENDING==0
       self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_EQ),
                 self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_pending'), 4)
-      # Wait for FINISH_DONE!=0
-      self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, self.pm4.WAIT_REG_MEM_FUNCTION(WAIT_REG_MEM_FUNCTION_NEQ),
-                self.gc.regSQ_THREAD_TRACE_STATUS.addr[0], 0, 0, self.gc.regSQ_THREAD_TRACE_STATUS.fields_mask('finish_done'), 4)
       # Disable SQTT
       self.sqtt_config(tracing=False)
       # Wait for BUSY==0
@@ -804,7 +813,7 @@ class AMDDevice(HCQCompiled):
     # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
     self.sqtt_enabled = PROFILE and bool(getenv("SQTT", 0))
     if self.sqtt_enabled:
-      if self.target[0] != 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
+      if self.target[0] < 11: raise RuntimeError(f'SQ Thread Tracing is not supported on gc:{self.target}')
       if not self.is_am() and (ppfeaturemask:=int(FileIOInterface('/sys/module/amdgpu/parameters/ppfeaturemask', os.O_RDONLY).read(), 16))&0x8000:
         raise RuntimeError("SQTT can't be enabled because of hardware bug, to workaround either use AMD_IFACE=PCI or add "
                            f"ppfeaturemask={(ppfeaturemask&~0x8000):#x} (current {ppfeaturemask=:#x} & ~PP_GFXOFF_MASK) to amdgpu module parameters\n"
@@ -872,7 +881,7 @@ class AMDDevice(HCQCompiled):
       self.synchronize()
       if DEBUG >= 2: print(f'{self.device}: Saving SQTT in profile...')
       for i,buf0 in enumerate(self.sqtt_buffers):
-        wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - ((buf0.va_addr//32) & 0x1FFFFFFF)) * 32
+        wptr = ((struct.unpack('<I', wptrs[i*4:i*4+4])[0] & 0x1FFFFFFF) - (((buf0.va_addr//32) & 0x1FFFFFFF) if self.target < (12,0,0) else 0)) * 32
         if DEBUG >= 2: print(f'\t{self.device}: SE {i} blob size {wptr:#x}')
         assert wptr >= 0 and wptr <= buf0.size, f"{wptr} > {buf0.size}, should never happen"
         # When sqtt buffer overflows, wptr stops at the last dword

tinygrad/runtime/support/llvm.py

@@ -16,7 +16,7 @@ elif OSX:
 else:
   LLVM_PATH = ctypes.util.find_library('LLVM')
   # use newer LLVM if possible
-  for ver in reversed(range(14, 20+1)):
+  for ver in reversed(range(14, 21+1)):
     if LLVM_PATH is not None: break
     LLVM_PATH = ctypes.util.find_library(f'LLVM-{ver}')
   if LLVM_PATH is None:

tinygrad/schedule/indexing.py

@@ -1,4 +1,4 @@
-from typing import Iterator, Sequence
+from typing import Iterator
 import functools, operator, itertools
 from dataclasses import dataclass, field
 from tinygrad.dtype import dtypes, AddrSpace
@@ -41,7 +41,7 @@ class BufferizeOpts:
 @dataclass
 class IndexingContext:
   realize_map: dict[UOp, None] = field(default_factory=dict)
-  range_map: dict[UOp, tuple[list[UOp], list[UOp]]] = field(default_factory=dict)
+  range_map: dict[UOp, tuple[tuple[UOp, ...], tuple[UOp, ...]]] = field(default_factory=dict)
 
   # create ranges
   range_idx: Iterator[int] = field(default_factory=itertools.count)
@@ -103,30 +103,31 @@ pm_apply_rangeify = PatternMatcher([
 ])
 
 # this is the definition of the movement ops
-def apply_movement_op(x:UOp, rngs:Sequence[UOp]) -> list[UOp]:
-  match x.op:
-    case Ops.SHRINK:  rngs = [a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, x.arg)]
-    case Ops.PERMUTE: rngs = [rngs[p] for p in argsort(x.arg)]
-    case Ops.FLIP:    rngs = [((s-1)-a) if f else a for a,s,f in zip(rngs, x.shape, x.arg)]
-    case Ops.EXPAND:  rngs = [a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, x.src[0].shape, x.shape)]
+@functools.cache
+def apply_movement_op(op:Ops, in_shape:tuple[sint,...], arg:tuple, rngs:tuple[UOp, ...]) -> tuple[UOp, ...]:
+  match op:
+    case Ops.SHRINK:  rngs = tuple(a if ss == 0 else a+ss for a,(ss,_) in zip(rngs, arg))
+    case Ops.PERMUTE: rngs = tuple(rngs[p] for p in argsort(arg))
+    case Ops.FLIP:    rngs = tuple(((s-1)-a) if f else a for a,s,f in zip(rngs, in_shape, arg))
+    case Ops.EXPAND:  rngs = tuple(a if in_sh == out_sh else a.const_like(0) for a,in_sh,out_sh in zip(rngs, in_shape, arg))
     case Ops.PAD:
       # TODO: why is multiple graph_rewrites faster than one here?
-      rngs = [r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh-e))).where(r-s, UOp.invalid()), sym, name="pad")
-              for r,sh,(s,e) in zip(rngs, x.shape, x.arg)]
+      rngs = tuple(r if (s == 0 and e == 0) else graph_rewrite(((r >= s) & (r < (sh+s))).where(r-s, UOp.invalid()), sym, name="pad")
+        for r,sh,(s,e) in zip(rngs, in_shape, arg))
     case Ops.RESHAPE:
       acc = 1
       axes_in:list[UOp] = []
-      for s,src in list(zip(x.shape, rngs))[::-1]:
+      for s,src in list(zip(arg, rngs))[::-1]:
         axes_in.append(acc*src)
         acc *= s
       combined_axes = sum(axes_in, start=UOp.const(dtypes.index, 0))
       axes_out:list[UOp] = []
-      for s in x.src[0].shape[::-1]:
+      for s in in_shape[::-1]:
         axes_out.append(combined_axes % s)
         combined_axes //= s
       # this simplify is doing a lot of heavy lifting. this is the replacement for the reshape view merging code
-      rngs = list(graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src)
-    case _: raise RuntimeError(f"{x.op} is not a MovementOp")
+      rngs = graph_rewrite(UOp.sink(*axes_out[::-1]), symbolic, name="reshape").src
+    case _: raise RuntimeError(f"{op} is not a MovementOp")
   return rngs
 
 @cpu_profile(TracingKey("run_rangeify"), "TINY")
@@ -157,7 +158,7 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     consumer_rngs = [rctx.range_map[c][0] for c in consumer_map[x] if c in rctx.range_map]
     if x in rctx.realize_map:
       # if this is in the realize_map, we create new ranges (at the output)
-      out_rngs = [rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape]
+      out_rngs = tuple(rctx.new_range(s) if not isinstance(s, UOp) or s.op is not Ops.RANGE else s for s in x.shape)
       # all ranges are ended now
       ending_ranges[x] = False
     elif x.op in {Ops.MSTACK, Ops.MSELECT}:
@@ -181,15 +182,16 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
 
       # TODO: in RANGEIFY > 1 all_all_same isn't required
       all_all_same = all(same_rngs for _,_,same_rngs in rngs_valids)
-      out_rngs = []
+      _out_rngs = []
       for i,(local_rngs,valids,same_rngs) in enumerate(rngs_valids):
         # we compare the ranges without their valids
         if all_all_same:
           # the new valid is the OR of all the children valids
           minimum_valid = functools.reduce(operator.or_, valids, UOp.const(dtypes.bool, False))
-          out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
+          _out_rngs.append(graph_rewrite(minimum_valid.where(local_rngs[0], UOp.invalid()), symbolic, name="minimum_valid"))
         else:
-          out_rngs.append(rctx.new_range(x.shape[i]))
+          _out_rngs.append(rctx.new_range(x.shape[i]))
+      out_rngs = tuple(_out_rngs)
 
       # we have to realize here if there's new ranges
       if not all_all_same: rctx.realize_map[x] = None
@@ -203,18 +205,16 @@ def run_rangeify(tsink:UOp, debug:bool=False) -> tuple[UOp, IndexingContext]:
     #  2. newly created for REDUCE_AXIS
     #  3. passed through for everything else
 
-    rngs = out_rngs  # rngs is the input ranges
+    rngs = out_rngs  # rngs is the input ranges  # pylint: disable=possibly-used-before-assignment
 
     # apply movement ops
-    if x.op in GroupOp.Movement: rngs = apply_movement_op(x, rngs)
+    if x.op in GroupOp.Movement: rngs = apply_movement_op(x.op, x.src[0].shape, x.arg, rngs)
     # if the EXPAND is used to inject a range, we don't mark it as ending_ranges. otherwise we do.
     if x.op is Ops.EXPAND and all(isinstance(y, int) or y.op is not Ops.RANGE for y in x.shape): ending_ranges[x] = True
 
     # REDUCE_AXIS creates ranges for the axes it is reducing
     if x.op is Ops.REDUCE_AXIS:
-      rngs = rngs[:]
-      for i,s in enumerate(x.src[0].shape):
-        if i in x.arg[1]: rngs[i] = rctx.new_range(s, axistype=AxisType.REDUCE)
+      rngs = tuple(rctx.new_range(s, axistype=AxisType.REDUCE) if i in x.arg[1] else r for i,(r,s) in enumerate(zip(rngs, x.src[0].shape)))
 
     if debug:
       print("***" if x in rctx.realize_map else "   ", len(consumer_map[x]), f"{str(x.op):20s}",

tinygrad/schedule/rangeify.py

@@ -103,7 +103,7 @@ earliest_rewrites = PatternMatcher([
 # movement op on INDEX as a PatternMatcher
 pm_mops = PatternMatcher([
   (UPat(GroupOp.Movement, name="r").f(Ops.INDEX, allow_any_len=True, name="idx"),
-   lambda r,idx: r.src[0].index(*apply_movement_op(r, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),
+   lambda r,idx: r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.arg, idx.src[1:]), dtype=idx.dtype, arg=idx.arg)),  # type: ignore
 ])
 
 # *****************
@@ -442,7 +442,7 @@ def tag_uop(ctx:list[UOp], x:UOp):
 add_tags = PatternMatcher([
   # don't tag BUFFERs, they are global
   (UPat(GroupOp.All-{Ops.BUFFER, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.DEFINE_VAR, Ops.BIND,
-                     Ops.MSTACK, Ops.MSELECT}.union(GroupOp.Movement), name="x"), tag_uop),
+                     Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.BUFFER for s in x.src) else tag_uop(ctx, x)),
 ])
 

tinygrad/shape/view.py

@@ -4,14 +4,7 @@ from dataclasses import dataclass
 from typing import cast, Sequence
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import resolve, UOp, Variable, sint, smax, smin, sint_to_uop, Ops, ssimplify
-from tinygrad.helpers import prod, all_int, flatten, ceildiv
-
-# returns the axes to create new_shape if new_shape can be created by combining axis from old_shape
-def get_contraction(old_shape:tuple[sint, ...], new_shape:tuple[sint, ...]) -> list[list[int]]|None:
-  acc_old, acc_new = list(itertools.accumulate(old_shape, operator.mul)), list(itertools.accumulate(new_shape, operator.mul))
-  try: split = [acc_old.index(acc)+1 if acc != 1 else 0 for acc in acc_new]
-  except ValueError: return None
-  return [list(range(st,ed)) for st,ed in zip([0]+split[:-1], split[:-1]+[len(old_shape)])]
+from tinygrad.helpers import prod, all_int, flatten
 
 @functools.cache
 def canonicalize_strides(shape:tuple[sint, ...], strides:tuple[sint, ...]) -> tuple[sint, ...]:
@@ -171,7 +164,6 @@ class View:
     if not all_int(vm1.shape):
       # if all strides are 0 and vm2 is unmasked, return vm1
       if all(x == 0 for x in vm2.strides+vm1.strides) and vm2.mask is None: return vm1
-      # TODO: handle more cases
       return None
 
     # Project vm1's offset and strides on to vm2.
@@ -184,47 +176,7 @@ class View:
         if not resolve((s1 := s1 - o)!=0): continue  # if s1 can possibly be 0
         terms[d2].append((d1, s1))
         strides[d1] += ssimplify(s1 * vm2.strides[d2])
-
-    # Merge dimensions in vm2 if required.
-    # NB: Merging too many dimensions can make it difficult to project vm2's mask, hence only combining when required.
-    idxs: list[UOp] = [UOp.variable(f"idx{i}", 0, s-1, dtypes.index) for i,s in enumerate(vm1.shape)]
-    merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
-    extents: list[tuple[sint, UOp]] = []
-    for term, s, o in zip(reversed(terms), reversed(vm2.shape), reversed(origin)):
-      merged_term += (sum([idxs[d1] * s1 for d1, s1 in term]) + o) * merged_size
-      merged_size *= s
-      if resolve(merged_term < merged_size, False) and resolve(0 <= merged_term, False):
-        extents.append((merged_size, merged_term))
-        merged_size, merged_term = 1, UOp.const(dtypes.index, 0)
-    if resolve(merged_term != 0): return None
-    if (vm2_shape := tuple(s for s,_ in reversed(extents))) != vm2.shape:
-      if (reshaped_vm2 := vm2.reshape(vm2_shape)) is None: return None
-      # NOTE: this != to prevent infinite loop
-      if reshaped_vm2.shape != vm2.shape: return reshaped_vm2 + vm1
-
-    if vm2.mask:
-      # Try to project vm2's mask on to vm1.
-      newb, newe, bad = [0] * len(vm1.shape), list(vm1.shape), False
-      for (b, e), o, term, (_, t) in zip(vm2.mask, origin, terms, reversed(extents)):
-        if resolve(b <= (t := t.simplify()).vmin and t.vmax < e, False): continue
-        if len(term) != 1:
-          if not term and newe:
-            # t should be a constant if no terms contribute to this dimension, but it might not be simplified
-            if t.vmin != t.vmax: return None
-            newe[0] = 0
-          else: bad = True
-          continue
-        d1, s1 = term[0]
-        newb[d1] = smax(newb[d1], ceildiv(b - o if s1 > 0 else e - o - 1, s1))
-        newe[d1] = smin(newe[d1], (b - o if s1 < 0 else e - o - 1) // s1 + 1)
-
-      # If any of vm1 was masked off, try again with that mask in place.
-      if any((b, e) != (0, s) for b, e, s in zip(newb, newe, vm1.shape)):
-        return vm2 + View.create(vm1.shape, vm1.strides, vm1.offset, tuple(zip(newb, newe)))
-      # Otherwise if vm2's mask was violated, then cannot merge.
-      if bad: return None
-
-    return View.create(vm1.shape, tuple(strides), ssimplify(sum(o * s for o, s in zip(origin, vm2.strides)) + vm2.offset))
+    return None
 
   def __unsafe_resize(self, arg: tuple[tuple[sint, sint], ...], mask=None) -> View:
     offset = sum([s * x[0] for s, x in zip(self.strides,arg)])
@@ -292,7 +244,6 @@ class View:
 
     r_strides, r_new_shape = [], reversed(new_shape)
     for merged_size, new_stride, real_size in reversed(merge_dims(self.shape, self.strides, self.mask)):
-      # TODO: write with get_contraction
       acc = 1
       # TODO: third resolve shouldn't be needed
       while resolve(acc <= merged_size) and resolve(acc != merged_size) and resolve((new_dim := next(r_new_shape, 0)) > 0):

tinygrad/uop/mathtraits.py

@@ -1,15 +1,17 @@
+from typing import TypeVar
 from tinygrad.uop import Ops
-from tinygrad.helpers import T
-from tinygrad.dtype import dtypes
+from tinygrad.dtype import dtypes, ConstType
 
+TMathTrait = TypeVar("TMathTrait", bound="MathTrait")
 class MathTrait:
   # required to implement
-  def alu(self:T, op:Ops, *src) -> T: raise NotImplementedError
-  def const_like(self:T, b) -> T: raise NotImplementedError
+  def alu(self:TMathTrait, op:Ops, *src:TMathTrait) -> TMathTrait: raise NotImplementedError
+  def const_like(self:TMathTrait, b:ConstType) -> TMathTrait: raise NotImplementedError
 
   # great functions you get!
-  def ufix(self, x): return self.const_like(x) if not isinstance(x, MathTrait) else x
-  def _binop(self, op, x, reverse): return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
+  def ufix(self:TMathTrait, x:ConstType|TMathTrait) -> TMathTrait: return self.const_like(x) if not isinstance(x, MathTrait) else x
+  def _binop(self:TMathTrait, op:Ops, x:TMathTrait|ConstType, reverse:bool) -> TMathTrait:
+    return self.ufix(x).alu(op, self) if reverse else self.alu(op, self.ufix(x))
   def logical_not(self): return self.ne(True)
   def neg(self):
     if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")

tinygrad/uop/spec.py

@@ -108,6 +108,9 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
   (UPat(Ops.COPY, name="copy", src=(UPat.var("x"), UPat(Ops.DEVICE)), arg=None), lambda copy,x: copy.dtype == x.dtype),
   (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
   (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
+
+  # REDUCE with an outerworld range
+  (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
 ])
 
 # ***** uop type spec *****

tinygrad/uop/symbolic.py

@@ -397,8 +397,8 @@ def parse_valid(valid:UOp) -> tuple[UOp, bool, int]:
   if valid.op is Ops.CMPLT and dtypes.is_int(valid.src[0].dtype): return valid.src[0], True, int((valid.src[1]).vmax)-1
   raise ValueError(f"not able to parse {valid=}")
 
-def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
-  # return None if valid is always False, otherwise the simplified uop (might be the same as input)
+def uop_given_valid(valid:UOp, uop:UOp) -> UOp:
+  # return simplified uop (might be the same as input)
 
   # first, parse valid into {expr: (lower_bound, upper_bound)}
   bounds:defaultdict[UOp, list[ConstType|None]] = defaultdict(lambda: [None, None])
@@ -415,18 +415,13 @@ def uop_given_valid(valid:UOp, uop:UOp) -> UOp|None:
     v0, v1 = (expr.vmin if v[0] is None else v[0], expr.vmax if v[1] is None else v[1])
     expr = expr.substitute(load_subs)  # make sure expr appears in same form in the uop
     # some expr has lower bound > upper bound -> valid is an empty set and we return None
-    if v0 > v1: return None
-    # whole node became a const
-    if v0 == v1:
-      uop = uop.substitute({expr:expr.const_like(v0)}).simplify()
-      continue
     # every candidate is a set of constrained UOp based on valid, and if every item in a set simplifies the uop into a same output, we rewrite uop
     candidates = []
     if expr.op is Ops.ADD and v0 == 1 and all(u.op in GroupOp.Irreducible for u in expr.split_uop(Ops.ADD)):
       # if the constraint is a simplex: X0 + X1 + ... > 0, we can check if all Xi > 0 simplify into the same output
       candidates.append([(Xi, UOp.variable("fake", 1, Xi.vmax, Xi.dtype)) for Xi in expr.split_uop(Ops.ADD)])
     # try checking the whole clause
-    if expr in uop.toposort(): candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
+    candidates.append([(expr, UOp.variable("fake", v0, v1, expr.dtype))])
 
     for candidate in candidates:
       # if every branch in candidate gives the same simplified uop, we can rewrite the uop
@@ -451,7 +446,7 @@ def simplify_valid(valid:UOp) -> UOp|None:
   something_changed = False
   valids = list(valid.split_uop(Ops.AND))
   for stmt in sorted(valids, key=lambda v: _valid_priority(v, valids)):
-    ret.append(newstmt if ret and (newstmt:=uop_given_valid(functools.reduce(operator.and_, ret), stmt)) is not None else stmt)
+    ret.append(uop_given_valid(functools.reduce(operator.and_, ret), stmt) if ret else stmt)
     if ret[-1] is not stmt: something_changed = True
   return functools.reduce(operator.and_, ret) if something_changed else None
 

tinygrad/viz/index.html

@@ -315,15 +315,6 @@
     font-size: 0.95em;
     letter-spacing: 0.03em;
   }
-  .legend {
-    display: flex;
-    align-items: center;
-  }
-  .legend > div {
-    width: 0.95em;
-    height: 0.95em;
-    margin-right: 4px;
-  }
   </style>
 </head>
 <body>

tinygrad/viz/js/index.js

@@ -14,8 +14,9 @@ const darkenHex = (h, p = 0) =>
   ).toString(16).padStart(6, '0')}`;
 
 const ANSI_COLORS = ["#b3b3b3", "#ff6666", "#66b366", "#ffff66", "#6666ff", "#ff66ff", "#66ffff", "#ffffff"];
+const ANSI_COLORS_LIGHT = ["#d9d9d9","#ff9999","#99cc99","#ffff99","#9999ff","#ff99ff","#ccffff","#ffffff"];
 const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g),
-  ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? ANSI_COLORS[(parseInt(code)-30+60)%60] : defaultColor }));
+  ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? (code>=90 ? ANSI_COLORS_LIGHT : ANSI_COLORS)[(parseInt(code)-30+60)%60] : defaultColor }));
 
 const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();
 
@@ -164,7 +165,13 @@ const drawLine = (ctx, x, y, opts) => {
   ctx.stroke();
 }
 
-var data, focusedDevice, canvasZoom, zoomLevel = d3.zoomIdentity;
+function tabulate(rows) {
+  const root = d3.create("div").style("display", "grid").style("grid-template-columns", `${Math.max(...rows.map(x => x[0].length), 0)}ch 1fr`).style("gap", "0.2em");
+  for (const [k,v] of rows) { root.append("div").text(k); root.append("div").node().append(v); }
+  return root;
+}
+
+var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
   displayGraph("profiler");
   d3.select(".metadata").html("");
@@ -272,7 +279,10 @@ async function renderProfiler() {
       for (const [num, {dtype, sz, nbytes, y, x:steps}] of buf_shapes) {
         const x = steps.map(s => timestamps[s]);
         const dur = x.at(-1)-x[0];
-        const arg = {tooltipText:`${dtype} len:${formatUnit(sz)}\n${formatUnit(nbytes, "B")}\nnum:${num}\nalive for ${formatTime(dur)}`};
+        const html = document.createElement("div");
+        const rows = [["DType", dtype], ["Len", formatUnit(sz)], ["Size", formatUnit(nbytes, "B")], ["Lifetime", formatTime(dur)]];
+        const info = html.appendChild(tabulate(rows).node());
+        const arg = {tooltipText:info.outerHTML, html, key:`${k}-${num}`};
         shapes.push({ x, y0:y.map(yscale), y1:y.map(y0 => yscale(y0+nbytes)), arg, fillColor:cycleColors(colorScheme.BUFFER, shapes.length) });
       }
       // generic polygon merger
@@ -341,6 +351,7 @@ async function renderProfiler() {
           for (let i=x.length-1; i>=0; i--) ctx.lineTo(x[i], offsetY+e.y1[i]);
           ctx.closePath();
           ctx.fillStyle = e.fillColor; ctx.fill();
+          if (focusedShape && e.arg?.key === focusedShape) { ctx.lineWidth = 1.4; ctx.strokeStyle = "#c9a8ff"; ctx.stroke(); }
           continue;
         }
         // contiguous rect
@@ -434,6 +445,8 @@ async function renderProfiler() {
     e.preventDefault();
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
     if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step);
+    if (foundRect?.key != focusedShape) { focusedShape = foundRect?.key; render(zoomLevel); }
+    return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? "");
   });
 
   canvas.addEventListener("mousemove", e => {
@@ -644,17 +657,10 @@ async function main() {
           }
         }
       }
-      const summary = metadata.appendChild(document.createElement("table"));
-      for (const s of ret.summary) {
-        const tr = summary.appendChild(document.createElement("tr"));
-        tr.className = "main-row";
-        const td = tr.appendChild(document.createElement("td"));
-        const div = td.appendChild(document.createElement("div"));
-        div.className = "legend";
-        div.appendChild(document.createElement("div")).style.background = cycleColors(colorScheme.CATEGORICAL, s.idx);
-        div.appendChild(document.createElement("p")).textContent = s.label;
-        appendTd(tr, s.value);
-      }
+      metadata.appendChild(tabulate(ret.summary.map(s => {
+        const div = d3.create("div").style("background", cycleColors(colorScheme.CATEGORICAL, s.idx)).style("width", "24px").style("height", "100%");
+        return [s.label.trim(), div.node()];
+      })).node());
     } else root.appendChild(codeBlock(ret.src, "x86asm"));
     return document.querySelector(".disasm").replaceChildren(root);
   }

tinygrad/viz/serve.py

@@ -206,7 +206,7 @@ def get_llvm_mca(asm:str, mtriple:str, mcpu:str) -> dict:
   # disassembly output can include headers / metadata, skip if llvm-mca can't parse those lines
   data = json.loads(subprocess.check_output(["llvm-mca","-skip-unsupported-instructions=parse-failure","--json","-"]+target_args, input=asm.encode()))
   cr = data["CodeRegions"][0]
-  resource_labels = data["TargetInfo"]["Resources"]
+  resource_labels = [repr(x)[1:-1] for x in data["TargetInfo"]["Resources"]]
   rows:list = [[instr] for instr in cr["Instructions"]]
   # add scheduler estimates
   for info in cr["InstructionInfoView"]["InstructionList"]: rows[info["Instruction"]].append(info["Latency"])