diff --git a/test/amd/test_emu2_pcode.py b/test/amd/test_emu2_pcode.py index dc20ff2fd1..ff6366372a 100644 --- a/test/amd/test_emu2_pcode.py +++ b/test/amd/test_emu2_pcode.py @@ -152,7 +152,7 @@ class TestDSPcodePatterns(unittest.TestCase): def test_global_atomic_add_f32_parsing(self): """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype.""" - vmem = UOp.param(2, dtypes.uint32, (1024,)) + vmem = UOp.param(2, dtypes.uint32, 1024) srcs = { 'ADDR': UOp.const(0, dtypes.uint64), 'DATA': UOp.const(0x3f800000, dtypes.uint32), @@ -183,7 +183,7 @@ class TestDSPcodePatterns(unittest.TestCase): def test_mem_read_parsing(self): """Test MEM[addr].type read expression parsing.""" # Create a mock LDS buffer - lds = UOp.param(3, dtypes.uint32, (16384,)) + lds = UOp.param(3, dtypes.uint32, 16384) addr = UOp.const(0, dtypes.uint32) vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)} @@ -218,7 +218,7 @@ class TestDSPcodePatterns(unittest.TestCase): pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32) self.assertIsNotNone(pcode) assert pcode is not None - lds = UOp.param(3, dtypes.uint32, (16384,)) + lds = UOp.param(3, dtypes.uint32, 16384) srcs = { 'ADDR': UOp.const(0, dtypes.uint32), 'OFFSET0': UOp.const(0, dtypes.uint32), @@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase): self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32)) def test_permlane64_wave64_pcode_indices(self): - vgpr = UOp.param(0, dtypes.uint32, (256,)) + vgpr = UOp.param(0, dtypes.uint32, 256) srcs = { 'SRC0': UOp.const(0, dtypes.uint32), 'VDST': UOp.const(1, dtypes.uint32), @@ -330,7 +330,7 @@ class TestAllPcode(unittest.TestCase): def _make_srcs(self): """Create dummy source variables for pcode parsing.""" u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64) - lds = UOp.param(3, dtypes.uint32, (16384,)) + lds = UOp.param(3, dtypes.uint32, 16384) return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(), 'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(), 'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(), diff --git a/test/backend/test_isel.py b/test/backend/test_isel.py index f219195a52..54033ba393 100644 --- a/test/backend/test_isel.py +++ b/test/backend/test_isel.py @@ -46,7 +46,7 @@ class TestIselX86(unittest.TestCase): # complex address is [base + index*scale + displacement] def test_complex_address(self): a = UOp.variable("a", 0, 0, dtypes.int32) - load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load() + load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load() n = self.isel_rewrite(load) # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32 self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4) diff --git a/test/backend/test_linearizer_dumb.py b/test/backend/test_linearizer_dumb.py index 317f0abd3b..364bcc2578 100644 --- a/test/backend/test_linearizer_dumb.py +++ b/test/backend/test_linearizer_dumb.py @@ -11,16 +11,16 @@ from tinygrad.codegen import to_program class TestLinearizerFailure(unittest.TestCase): @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL") def test_failure_beam_mnist(self): - c0 = UOp.param(0, dtypes.uchar, (4014080,)) + c0 = UOp.param(0, dtypes.uchar, 4014080) c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL) c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL) c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL) - c4 = UOp.param(1, dtypes.int, (512,)) + c4 = UOp.param(1, dtypes.int, 512) c5 = c4.index(c1.valid(UOp.const(True))) c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE) c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE) c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE) - c9 = UOp.param(2, dtypes.uchar, (47040000,)) + c9 = UOp.param(2, dtypes.uchar, 47040000) c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True))) c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD) c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3) diff --git a/test/backend/test_renderer_failures.py b/test/backend/test_renderer_failures.py index 9ec18c1a9b..8f5460d1b9 100644 --- a/test/backend/test_renderer_failures.py +++ b/test/backend/test_renderer_failures.py @@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None): def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): dtype = alu_src_uops[0].dtype - a = UOp.param(0, dtype, (1,)) - b = UOp.param(1, dtype, (1,)) + a = UOp.param(0, dtype, 1) + b = UOp.param(1, dtype, 1) idx = UOp.const(0) ld = b.index(idx).load() alu = ld.alu(alu_op, *alu_src_uops) @@ -34,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp): class TestRendererFailures(unittest.TestCase): @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer") def test_gated_store_with_alu(self): - a = UOp.param(0, dtypes.int, (4,)) + a = UOp.param(0, dtypes.int, 4) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int))) sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo()) @@ -43,7 +43,7 @@ class TestRendererFailures(unittest.TestCase): @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer") def test_gated_store_with_alu_2d(self): - a = UOp.param(0, dtypes.int, (8,)) + a = UOp.param(0, dtypes.int, 8) gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0) gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0) gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int))) @@ -78,7 +78,7 @@ class TestCStyleFailures(unittest.TestCase): @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer") class TestWGSLFailures(unittest.TestCase): def test_folded_packed_store(self): - b = UOp.param(0, dtypes.char, (4,)) + b = UOp.param(0, dtypes.char, 4) idx = b.index(UOp.const(0).cast(dtypes.int)) store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32)) src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort()) @@ -93,9 +93,9 @@ class TestWGSLFailures(unittest.TestCase): # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt def test_gated_load(self): - a = UOp.param(0, dtypes.int, (4,)) - b = UOp.param(1, dtypes.int, (4,)) - c = UOp.param(2, dtypes.int, (4,)) + a = UOp.param(0, dtypes.int, 4) + b = UOp.param(1, dtypes.int, 4) + c = UOp.param(2, dtypes.int, 4) lidx0 = UOp.special(4, "lidx0") gate = lidx0.ne(0) alt = c.index(lidx0).load() @@ -110,7 +110,7 @@ class TestWGSLFailures(unittest.TestCase): class TestPTXFailures(unittest.TestCase): @unittest.skip("INDEX can only have a gate ALU parent, not an IF") def test_gated_store_with_if(self): - a = UOp.param(0, dtypes.int, (4,)) + a = UOp.param(0, dtypes.int, 4) gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0) val = UOp.const(1).cast(dtypes.int) if_uop = UOp(Ops.IF, src=(gate_alu,)) diff --git a/test/backend/test_uops.py b/test/backend/test_uops.py index e1c0c32443..d966f8eb47 100644 --- a/test/backend/test_uops.py +++ b/test/backend/test_uops.py @@ -20,7 +20,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]): def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp: if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype)) - elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,))) + elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1)) else: uops.append(UOp(op, dtype, tuple(src), arg)) return uops[-1] @@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res): class TestBitcastBufferView(unittest.TestCase): @Context(SPEC=2) def test_render(self): - buf = UOp.param(0, dtypes.uint32, (4,)) + buf = UOp.param(0, dtypes.uint32, 4) uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer) idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST) self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK) @@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase): @Context(SPEC=2) def test_load(self): val = 0x1122334455667788 - src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,)) + src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1) ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes()) obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate() run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf]) @@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase): @Context(SPEC=2) def test_store(self): val = 0x1122334455667788 - dst = UOp.param(0, dtypes.uint32, (6,)) + dst = UOp.param(0, dtypes.uint32, 6) buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24)) view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64) # two stores through one view: it must inline, not get a declared vector-pointer run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf]) @@ -249,8 +249,8 @@ class TestLocalAccess(unittest.TestCase): @unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends") class TestAssembly(unittest.TestCase): def test_bitshift_left(self): - g1 = UOp.param(0, dtypes.int32, shape=(3,)) - out = UOp.param(1, dtypes.int32, shape=(2,)) + g1 = UOp.param(0, dtypes.int32, 3) + out = UOp.param(1, dtypes.int32, 2) c1 = UOp.const(2) c2 = UOp.const(3) l1 = g1.index(c1) @@ -278,7 +278,7 @@ class TestAssembly(unittest.TestCase): self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4) def test_mulacc_shl(self): - g1 = UOp.param(0, dtypes.int32, shape=(2,)) + g1 = UOp.param(0, dtypes.int32, 2) c1 = UOp.const(0) c2 = UOp.const(1) expr = g1.index(c1) * UOp.const(4096) + g1.index(c2) @@ -287,7 +287,7 @@ class TestAssembly(unittest.TestCase): self.assertIn(Ops.MULACC, [x.op for x in uops]) def test_use_cmpeq(self): - g = UOp.param(0, dtypes.uint32, shape=(8,)) + g = UOp.param(0, dtypes.uint32, 8) c = UOp.const(7) comp = g.index(c).ne(c).ne(True) uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer) diff --git a/test/backend/test_wait_loop.py b/test/backend/test_wait_loop.py index 83e0543e61..6073c9422b 100644 --- a/test/backend/test_wait_loop.py +++ b/test/backend/test_wait_loop.py @@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp: return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=())) def wait_ext_kernel() -> UOp: - sig = UOp.param(0, dtypes.int, (1,), volatile=True) + sig = UOp.param(0, dtypes.int, 1, volatile=True) l = UOp.loop(0) v = sig.after(l)[0].load() e = v.end(l, v < 1) diff --git a/test/external/external_benchmark_op_conv.py b/test/external/external_benchmark_op_conv.py index 4076facd5d..9280ef4a05 100644 --- a/test/external/external_benchmark_op_conv.py +++ b/test/external/external_benchmark_op_conv.py @@ -28,7 +28,7 @@ def vision_conv_143(): c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float)) c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) - c63 = UOp.param(3, dtypes.float, (128,)) + c63 = UOp.param(3, dtypes.float, 128) c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5) c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5) @@ -54,7 +54,7 @@ def vision_conv_153(): c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float)) c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4)) c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196)) - c63 = UOp.param(3, dtypes.float, (256,)) + c63 = UOp.param(3, dtypes.float, 256) c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5) c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5) @@ -73,7 +73,7 @@ def dm_conv_172(): c18 = UOp.range(8, 2, AxisType.REDUCE) c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4)) c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512)) - c37 = UOp.param(3, dtypes.float, (960,)) + c37 = UOp.param(3, dtypes.float, 960) c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2) c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39 c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666 diff --git a/test/external/external_test_gpu_crash.py b/test/external/external_test_gpu_crash.py index ebe0f1eb2d..238cdc8c2e 100644 --- a/test/external/external_test_gpu_crash.py +++ b/test/external/external_test_gpu_crash.py @@ -36,7 +36,7 @@ class TestGPUCrash(unittest.TestCase): def _run_insts(self, insts: list[Inst]): buf = UOp.new_buffer("AMD", 64, dtypes.uint8) - sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test")) + sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test")) prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts)))) run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True) diff --git a/test/helpers.py b/test/helpers.py index e6d9916d4e..8bdd7d9bbf 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple for buf_dt, data in inputs or []: bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize)) allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data))) - g = UOp.param(0, uop.dtype, (1,)) + g = UOp.param(0, uop.dtype, 1) prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON"))) prog = dev.runtime(prg.to_elf()) prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals) diff --git a/test/mockgpu/amd/emu.py b/test/mockgpu/amd/emu.py index 6c8b3380f4..9d0ac9ef60 100644 --- a/test/mockgpu/amd/emu.py +++ b/test/mockgpu/amd/emu.py @@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None: class _Ctx: """Context for instruction compilation - holds buffers and helpers.""" __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr') - sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,)) - vmem = UOp.param(2, dtypes.uint32, (1 << 46,)) - lds = UOp.param(3, dtypes.uint32, (16384,)) - scratch = UOp.param(4, dtypes.uint8, (1 << 30,)) + sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT) + vmem = UOp.param(2, dtypes.uint32, 1 << 46) + lds = UOp.param(3, dtypes.uint32, 16384) + scratch = UOp.param(4, dtypes.uint8, 1 << 30) # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references _vgpr_cache: dict[int, UOp] = {} _accvgpr_cache: dict[int, UOp] = {} @@ -332,10 +332,10 @@ class _Ctx: def __init__(self, inst_size: int, wave_size: int = 32): self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size self.dyn_fields: list[tuple[int, int]] = [] # (lo, hi) of fields read dynamically - if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,)) + if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size) self.vgpr = _Ctx._vgpr_cache[wave_size] if wave_size == 64: - if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,)) + if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size) self.accvgpr = _Ctx._accvgpr_cache[wave_size] else: self.accvgpr = self.vgpr diff --git a/test/null/test_gpudims.py b/test/null/test_gpudims.py index 8d73d907bc..cde27aadcd 100644 --- a/test/null/test_gpudims.py +++ b/test/null/test_gpudims.py @@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase): def test_global_prod_max(self): g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL) - sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo()) + sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo()) class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128) specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL] self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1) diff --git a/test/null/test_linearizer_failures.py b/test/null/test_linearizer_failures.py index 620cd8d0f3..5c89d4c3d8 100644 --- a/test/null/test_linearizer_failures.py +++ b/test/null/test_linearizer_failures.py @@ -7,14 +7,14 @@ from tinygrad.codegen import to_program class TestLinearizerFailures(unittest.TestCase): def test_fail_1(self): - c0 = UOp.param(0, dtypes.float, (64,)) + c0 = UOp.param(0, dtypes.float, 64) c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK) c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK) c3 = ((c1*UOp.const(32))+c2) - c4 = UOp.param(1, dtypes.float, (163840,)) + c4 = UOp.param(1, dtypes.float, 163840) c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE) c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920)))) - c7 = UOp.param(2, dtypes.float, (64,)) + c7 = UOp.param(2, dtypes.float, 64) c8 = c7.index(c3) c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal() c10 = c0.index(c3).store(c9).end(c1, c2) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 2166e69e8a..05f831fb83 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -15,7 +15,7 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx") def get_gated_load_uop(valid:UOp, idx:UOp): - return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load() + return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load() def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]): return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load() @@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase): from tinygrad.codegen.late.coalesce import indexing_simplify from tinygrad.uop.ops import graph_rewrite from tinygrad.uop.symbolic import sym - buf = UOp.param(0, dtypes.int, (1,)) + buf = UOp.param(0, dtypes.int, 1) idx = UOp.const(0) true_gate = UOp.const(True) index_with_gate = buf.index(idx.valid(true_gate)) @@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase): # one load guards r < 4, but another load uses r without a gate -> no shrink r = Range(0, 204) load1 = get_gated_load_uop(r < UOp.const(4), r) - load2 = UOp.param(1, dtypes.float, (204,)).index(r).load() + load2 = UOp.param(1, dtypes.float, 204).index(r).load() ranges = self.get_ranges(UOp.sink(load1, load2)) self.assert_range_end(ranges, 204) @@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(1.0), Invalid) - ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink()) self.assert_range_end(ranges, 4) def test_range_shrink_store_where_invalid_flipped(self): @@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase): from tinygrad.dtype import Invalid r = Range(0, 204) x = (r < 4).where(UOp.const(1.0), Invalid) - ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink()) + ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink()) self.assert_range_end(ranges, 4) if __name__ == '__main__': diff --git a/test/null/test_transcendental_helpers.py b/test/null/test_transcendental_helpers.py index 1653b176c7..7602184241 100644 --- a/test/null/test_transcendental_helpers.py +++ b/test/null/test_transcendental_helpers.py @@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase): def test_payne_hanek_reduction(self): # TODO: Test constant input when constant folding is fixed (or maybe test both variants) # Load input value from a buffer to prevent constant folding - input_buf = UOp.param(1, dtypes.double, (1,)) + input_buf = UOp.param(1, dtypes.double, 1) loaded_value = input_buf.index(UOp.const(0)).load() def eval_payne_hanek_reduction(v:float) -> tuple[float, int]: return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value)) diff --git a/test/null/test_uop_graph.py b/test/null/test_uop_graph.py index 9cc3f88688..602144f6e3 100644 --- a/test/null/test_uop_graph.py +++ b/test/null/test_uop_graph.py @@ -214,8 +214,8 @@ class TestUOpGraph(unittest.TestCase): for i, const in enumerate(consts): self.assertIs(vec.index(i), const) def test_cast_alu_fold(self): - d0 = UOp.param(0, dtypes.bool, (1,)) - d1 = UOp.param(1, dtypes.int, (1,)) + d0 = UOp.param(0, dtypes.bool, 1) + d1 = UOp.param(1, dtypes.int, 1) idx = UOp.const(0) ld = d1.index(idx) alu = (ld<1).cast(dtypes.bool) @@ -224,8 +224,8 @@ class TestUOpGraph(unittest.TestCase): self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0) def test_double_cast_fold(self): - d0 = UOp.param(0, dtypes.float, (1,)) - d1 = UOp.param(1, dtypes.int, (1,)) + d0 = UOp.param(0, dtypes.float, 1) + d1 = UOp.param(1, dtypes.int, 1) idx = UOp.const(0, dtypes.int) ld = d1.index(idx) alu = ld.cast(dtypes.float).cast(dtypes.float) @@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase): def test_bitcast_to_same_dtype_fold(self): for dt in dtypes.ints + dtypes.floats + (dtypes.bool,): - d0 = UOp.param(0, dt, (1,)) + d0 = UOp.param(0, dt, 1) v = d0.index(UOp.const(0)) uops = to_uops_list([v.bitcast(dt)]) self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}") @@ -255,10 +255,10 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_fold(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.long, (100,)) + d0 = UOp.param(0, dtypes.long, 100) ld = d0.index(ridx0.valid(ridx0<50)) w = (ridx0<50).where(ld, 5) - out = UOp.param(1, dtypes.long, (100,)) + out = UOp.param(1, dtypes.long, 100) uops = to_uops_list([out.index(ridx0).store(w)]) expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0] for u in uops: @@ -267,7 +267,7 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_folds_swapped_branches(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.long, (100,)) + d0 = UOp.param(0, dtypes.long, 100) ld = d0.index(ridx0.valid((ridx0<50).logical_not())) w = (ridx0<50).where(5, ld) uops = to_uops_list([w]) @@ -278,11 +278,11 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_gated_load_with_cast(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.int, (100,)) + d0 = UOp.param(0, dtypes.int, 100) gate_idx = ridx0.valid((ridx0<50)) ld = d0.index(gate_idx).cast(dtypes.float) w = (ridx0<50).where(ld, 5.0) - out = UOp.param(1, dtypes.float, (100,)) + out = UOp.param(1, dtypes.float, 100) uops = to_uops_list([out.index(ridx0).store(w)]) expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0] for u in uops: @@ -291,27 +291,27 @@ class TestUOpGraph(unittest.TestCase): def test_where_on_casted_gated_load_extra_cond(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.float, (100,)) + d0 = UOp.param(0, dtypes.float, 100) ld = d0.index(ridx0.valid(ridx0<50)) w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half) - out = UOp.param(1, dtypes.half, (100,)) + out = UOp.param(1, dtypes.half, 100) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE def test_where_on_casted_gated_load_extra_cond_swapped(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.float, (100,)) + d0 = UOp.param(0, dtypes.float, 100) ld = d0.index(ridx0.valid(ridx0<50)) w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half) - out = UOp.param(1, dtypes.half, (100,)) + out = UOp.param(1, dtypes.half, 100) uops = to_uops_list([out.index(ridx0).store(w)]) for u in uops: assert u.op is not Ops.WHERE def test_where_in_store_becomes_gate(self): ridx0 = UOp.range(100, 0) - d0 = UOp.param(0, dtypes.long, (100,)) + d0 = UOp.param(0, dtypes.long, 100) idx = d0.index(ridx0) ld = idx.load() val = (ridx0<50).where(5, ld) @@ -325,14 +325,14 @@ class TestUOpGraph(unittest.TestCase): def test_load_idx_becomes_int(self): # mnist indexing with split reduceop # Make sure we are not doign math on the loaded index, which would promote it to long - c0 = UOp.param(0, dtypes.uchar, (128000,)) + c0 = UOp.param(0, dtypes.uchar, 128000) c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK) c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK) - c3 = UOp.param(1, dtypes.int, (512,)) + c3 = UOp.param(1, dtypes.int, 512) c4 = c3.index(c1) c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE) c6 = ((c2*UOp.const(240))+c5) - c7 = UOp.param(2, dtypes.uchar, (60000,)) + c7 = UOp.param(2, dtypes.uchar, 60000) c8 = c7.index(c6) c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD) c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2) @@ -342,14 +342,14 @@ class TestUOpGraph(unittest.TestCase): def test_load_idx_no_math_on_loaded(self): # test the (x+y) NOOP rule. This rule matches patterns that EMERGE during simplification.""" def test_store_load_folding(self): # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0) - buf = UOp.param(0, dtypes.int, (1,)) + buf = UOp.param(0, dtypes.int, 1) index = buf.index(UOp.const(0)) # Direct: store(idx, load(idx)) -> NOOP self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP) @@ -1417,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase): class TestMoveWhereOnLoad(unittest.TestCase): def test_bool_index_preserves_dtype(self): - buf = UOp.param(0, dtypes.bool, (8,)) + buf = UOp.param(0, dtypes.bool, 8) a = Variable("a", 0, 7) r = UOp.range(8, 0) # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid @@ -1475,7 +1475,7 @@ class TestRangeSplitting(unittest.TestCase): from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range r0 = UOp.range(uconst(8), 0) # create a simple expression using the range with mod: store range%2 to a buffer - buf = UOp.param(0, dtypes.int, (1,)) + buf = UOp.param(0, dtypes.int, 1) val = (r0 % uconst(2)).cast(dtypes.int) store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val)) sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),)) diff --git a/test/null/test_uop_vmin_vmax.py b/test/null/test_uop_vmin_vmax.py index 6398c28bb2..fcc7834089 100644 --- a/test/null/test_uop_vmin_vmax.py +++ b/test/null/test_uop_vmin_vmax.py @@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase): def test_vmin_vmax_multiplication_0_inf(self): # vmin and vmax for multiplication with a variable x = UOp.const(0.0) - y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0)) + y = UOp.load(UOp.param(0, dtypes.float, 1), UOp.const(0)) uop = x * y # TODO: these should be 0, but definitely should not be nan self.assertEqual(uop.vmin, -math.inf) @@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase): def test_vmin_vmax_vector_with_gep(self): # vmin and vmax for a vector constant of bool values - d1 = UOp.param(1, dtypes.int, (1,)) + d1 = UOp.param(1, dtypes.int, 1) idx = UOp.const(0) val = UOp(Ops.LOAD, src=(d1.index(idx),)) uop = (val // 32) diff --git a/test/null/test_uops.py b/test/null/test_uops.py index f4230bf9e8..636dfb0957 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -58,7 +58,7 @@ class TestDTypeFromUOp(unittest.TestCase): for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid) for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)), UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)), - UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared) + UOp.param(0, dtypes.float32, 4).index(invalid)): type_verify(u, spec_shared) gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ()) self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float)) type_verify(out.sink(), spec_program) @@ -73,7 +73,7 @@ class TestLowerIndexDtype(unittest.TestCase): def test_gated_shrink_lowers_to_selected_width(self): # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the # width the offset bounds select (this one needs long) - buf = UOp.param(0, dtypes.float, (2**31+64,)) + buf = UOp.param(0, dtypes.float, 2**31+64) i = UOp.variable("i", 0, 2**28) shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4))) lowered = graph_rewrite(shrink.sink(), pm_lower_weak) @@ -221,7 +221,7 @@ class TestExecALU(unittest.TestCase): class TestGatedStoreRewrite(unittest.TestCase): def test_tiny_gate_store(self): - gmem = UOp.param(0, dtypes.float, (8,)) + gmem = UOp.param(0, dtypes.float, 8) gidx0 = UOp.special(4, 'gidx0') gate = gidx0= 0) & (v < 16))).load()]) # valid with self.assertRaises(RuntimeError): @@ -48,7 +48,7 @@ class TestValidateOOB(unittest.TestCase): def test_gated_store(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) v = Variable("v", 0, 20) to_uops_list([buf.index(v.valid(v < 16)).store(0)]) # valid with self.assertRaises(RuntimeError): @@ -57,14 +57,14 @@ class TestValidateOOB(unittest.TestCase): # ALU ops in index def test_floordiv(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()]) # 0..16 oob def test_mod(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) r = UOp.range(100, 0, AxisType.GLOBAL) to_uops_list([buf.index(r % 16).load()]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -72,14 +72,14 @@ class TestValidateOOB(unittest.TestCase): def test_shr(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()]) # 0..31 oob def test_shl(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (64,)) + buf = UOp.param(0, dtypes.int, 64) r = UOp.range(8, 0, AxisType.GLOBAL) to_uops_list([buf.index(r << 2).load()]) # 0..28 valid with self.assertRaises(RuntimeError): @@ -87,7 +87,7 @@ class TestValidateOOB(unittest.TestCase): def test_and(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) r = UOp.range(100, 0, AxisType.GLOBAL) to_uops_list([buf.index(r & 15).load()]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -102,14 +102,14 @@ class TestValidateOOB(unittest.TestCase): def test_max(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()]) # 0..15 valid with self.assertRaises(RuntimeError): to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()]) # 0..20 oob def test_xor_in_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) r = UOp.range(32, 0, AxisType.GLOBAL) to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()]) # 0..15 valid with self.assertRaises(RuntimeError): @@ -118,22 +118,22 @@ class TestValidateOOB(unittest.TestCase): # cast patterns def test_float_cast_in_index(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (16,)) + buf = UOp.param(0, dtypes.int, 16) r = UOp.range(20, 0) i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int) to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()]) def test_bool_cast_in_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf = UOp.param(0, dtypes.int, (1,)) + buf = UOp.param(0, dtypes.int, 1) r = UOp.range(20, 0) to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()]) # only r=0 valid # load result as index/mask def test_load_as_index(self): with Context(CHECK_OOB=1, SPEC=2): - buf0 = UOp.param(0, dtypes.int, (16,)) - buf1 = UOp.param(1, dtypes.int, (64,)) + buf0 = UOp.param(0, dtypes.int, 16) + buf1 = UOp.param(1, dtypes.int, 64) r = UOp.range(42, 0, AxisType.GLOBAL) ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint) to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()]) # valid @@ -142,16 +142,16 @@ class TestValidateOOB(unittest.TestCase): def test_load_from_shrink_as_index(self): with Context(CHECK_OOB=1, SPEC=2): - buf0 = UOp.param(0, dtypes.int, (16,)) - buf1 = UOp.param(1, dtypes.int, (64,)) + buf0 = UOp.param(0, dtypes.int, 16) + buf1 = UOp.param(1, dtypes.int, 64) shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4))) ld0 = shrink.load().index(0) to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()]) def test_load_bool_as_mask(self): with Context(CHECK_OOB=1, SPEC=2): - buf_bool = UOp.param(0, dtypes.bool, (16,)) - buf_int = UOp.param(1, dtypes.int, (8,)) + buf_bool = UOp.param(0, dtypes.bool, 16) + buf_int = UOp.param(1, dtypes.int, 8) gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0") ld_bool = buf_bool.index(gidx).load() with self.assertRaises(RuntimeError): @@ -162,7 +162,7 @@ class TestValidateOOB(unittest.TestCase): def test_in_bounds_access_gated_local(self): with Context(CHECK_OOB=1): # Define buffers - gbuf = UOp.param(0, dtypes.uint, (400,)) + gbuf = UOp.param(0, dtypes.uint, 400) sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL) # Define indices, valids and barrier @@ -186,8 +186,8 @@ class TestValidateOOB(unittest.TestCase): @unittest.skip("Bool load is not supported yet") def test_load_mask(self): with Context(CHECK_OOB=1): - glbl0 = UOp.param(0, dtypes.int, (16,)) - mask = UOp.param(0, dtypes.bool, (16,)) + glbl0 = UOp.param(0, dtypes.int, 16) + mask = UOp.param(0, dtypes.bool, 16) ridx = UOp.range(20, 0) ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx)))) to_uops_list([ld0]) diff --git a/test/unit/test_dtype_weak.py b/test/unit/test_dtype_weak.py index 5ed5f379c4..88025ef476 100644 --- a/test/unit/test_dtype_weak.py +++ b/test/unit/test_dtype_weak.py @@ -116,7 +116,7 @@ class TestWeakPromotion(unittest.TestCase): def test_store_weak_value_uses_destination_dtype(self): with Context(DEFAULT_FLOAT=dtypes.float16): - dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32)) + dst = UOp.param(0, dtypes.bfloat16, 1).index(UOp.const(0).cast(dtypes.int32)) gate = UOp.const(True) out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak) # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it @@ -132,7 +132,7 @@ class TestWeakPromotion(unittest.TestCase): def test_derivable_const_rounds_at_the_derived_width(self): # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire - x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load() + x = UOp.param(0, dtypes.float32, 1).index(UOp.const(0).cast(dtypes.int32)).load() mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak) self.assertIs(mul.src[1], UOp.const(-1.0)) self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 1ce6b428ea..012781ca7b 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -26,24 +26,24 @@ MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('cl def signal_prog(): val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU) - return UOp.param(0, dtypes.uint32, (1,))[0].store(val.cast(dtypes.uint32)) + return UOp.param(0, dtypes.uint32, 1)[0].store(val.cast(dtypes.uint32)) def wait_prog(): val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU) - return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32)) + return (v:=UOp.param(0, dtypes.uint32, 1, volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32)) def timestamp_prog(): if WIN: val = UOp.const(0, dtypes.uint64) else: - fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG) + fn, ts = UOp.param(1, dtypes.uint64, 1), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG) call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts) val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load() - return UOp.param(0, dtypes.uint64, (1,))[0].store(val) + return UOp.param(0, dtypes.uint64, 1)[0].store(val) def worker_prog(): - ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True) - wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True) - sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last + ring = UOp.param(0, dtypes.uint64, RING_SLOTS * CMD_SIZE, volatile=True) + wait, done = UOp.param(1, dtypes.uint64, 1, volatile=True), UOp.param(2, dtypes.uint64, 1, volatile=True) + sem, cur = UOp.param(3, dtypes.uint64, 1), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last # spin on windows, sem_wait to sleep on posix if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur) diff --git a/tinygrad/runtime/ops_qcom.py b/tinygrad/runtime/ops_qcom.py index 7b12d8cc82..63fc9badef 100644 --- a/tinygrad/runtime/ops_qcom.py +++ b/tinygrad/runtime/ops_qcom.py @@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2 def dcache_flush(): from tinygrad.uop.ops import UOp, Ops, KernelInfo from tinygrad.codegen import to_program - buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU) + buf, n = UOp.param(0, dtypes.uint8, 1), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU) i = UOp.range(n, 0, dtype=dtypes.int) flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg=('__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");', dtypes.void)) sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg=('__asm__ volatile("dsb sy" ::: "memory");', dtypes.void)), diff --git a/tinygrad/runtime/support/hcq2.py b/tinygrad/runtime/support/hcq2.py index 88c914bf66..3bc3806b29 100644 --- a/tinygrad/runtime/support/hcq2.py +++ b/tinygrad/runtime/support/hcq2.py @@ -139,7 +139,7 @@ def _get_enqueue_devs(call:UOp) -> Any|None: def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None: if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None - d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2)) + d, s = (UOp.param(i, dst.dtype, n:=dst.max_numel(), device=devs) for i in range(2)) ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1) return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 28b7662de5..aafd8147c7 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1166,9 +1166,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # TODO: this should replace placeholder @staticmethod - def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, + def param(slot:int, dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None, multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None, volatile:bool=False): if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}") + if isinstance(shape, (int, UOp)): shape = (shape,) if shape is not None and axis is not None and isinstance(device, tuple): shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape)) src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)