From 16c27ae400250b16a495187e933f352452a751af Mon Sep 17 00:00:00 2001 From: chenyu Date: Tue, 23 Jul 2024 16:58:12 -0400 Subject: [PATCH] update UOp.SPECIAL arg spec [run_process_replay] (#5661) * update UOp.SPECIAL arg spec [run_process_replay] from `(0, "gid0", 4)` to just `("gid0", 4)`. closer to a Variable * fix ptx --- test/test_linearizer.py | 8 ++++---- test/test_uop_graph.py | 2 +- test/test_uops.py | 6 +++--- tinygrad/codegen/kernel.py | 8 ++++---- tinygrad/codegen/lowerer.py | 2 +- tinygrad/codegen/uopgraph.py | 4 ++-- tinygrad/codegen/uops.py | 2 +- tinygrad/renderer/assembly.py | 8 ++++---- tinygrad/renderer/cstyle.py | 12 ++++++------ tinygrad/runtime/ops_python.py | 4 ++-- 10 files changed, 28 insertions(+), 28 deletions(-) diff --git a/test/test_linearizer.py b/test/test_linearizer.py index 25fbfe14ce..ef1f6150eb 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -783,7 +783,7 @@ class TestLinearizer(unittest.TestCase): idxs = get_grouped_dims(prefix, dims, max_sizes, reverse_dims) loop_idxs = dedup(flatten([[y for y in sorted(list(x.sparents)) if y.op is UOps.SPECIAL] for x in idxs])) loop_idxs = sorted(loop_idxs, key=lambda uop: uop.arg[0]) - sizes = [x.arg[2] for x in loop_idxs] + sizes = [x.arg[1] for x in loop_idxs] assert len(idxs) == len(dims), f"expected idxs to have same length as dims {len(dims)}, got {len(idxs)}" assert len(loop_idxs) == min(len(sizes), len(dims)), f"expected idxs to have length {min(len(sizes), len(dims))}, got {len(loop_idxs)}" assert sizes == expected_sizes, f"expected sizes={expected_sizes}, got {sizes=}" @@ -840,9 +840,9 @@ class TestLinearizer(unittest.TestCase): k = helper_linearizer_opt(t+1)[0] idxs = dedup([uop for uop in k.uops if uop.op is UOps.SPECIAL]) idxs = sorted(idxs, key=lambda uop: uop.arg[0]) - assert idxs[0].arg == (0, 'gidx0', 6), idxs[0].arg - assert idxs[1].arg == (1, 'gidx1', 5), idxs[1].arg - assert idxs[2].arg == (2, 'gidx2', 4), idxs[2].arg + assert idxs[0].arg == ('gidx0', 6), idxs[0].arg + assert idxs[1].arg == ('gidx1', 5), idxs[1].arg + assert idxs[2].arg == ('gidx2', 4), idxs[2].arg def test_div_collapse(self): def helper(t, msg, max_ops=0): diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 323f50f9ed..433b6585f8 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -222,7 +222,7 @@ class TestUOpGraph(TestUOps): def test_fold_gated_load_local(self): glbl0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.int), (), (0, True)) smem = UOp(UOps.DEFINE_LOCAL, PtrDType(dtypes.int), (), ("temp", 1)) - lidx = UOp(UOps.SPECIAL, dtypes.int, (), (0, "lidx1", 16)) + lidx = UOp(UOps.SPECIAL, dtypes.int, (), ("lidx0", 16)) st = UOp(UOps.STORE, None, (smem, lidx, UOp.load(glbl0, lidx, dtype=dtypes.int))) barrier = UOp(UOps.BARRIER, None, (st, )) ld0 = UOp(UOps.LOAD, dtypes.int, (smem, lidx+1, UOp.const(dtypes.bool, False), UOp.const(dtypes.int, 2), barrier)) diff --git a/test/test_uops.py b/test/test_uops.py index 569368939d..dd0a9bcafb 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -240,7 +240,7 @@ class TestGatedStoreRewrite(unittest.TestCase): @unittest.expectedFailure def test_tiny_gate_store(self): gmem = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), (0, True)) - gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), (0, 'gidx0', 4)) + gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), ('gidx0', 4)) idx = gidx0 * UOp.const(dtypes.int, 2) val = UOp.const(dtypes.float, 42.0) gate = gidx0.lt(UOp.const(dtypes.int, 1)) @@ -258,7 +258,7 @@ class TestGatedStoreRewrite(unittest.TestCase): def test_gate_some_stores(self): gmem0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), (0, True)) gmem1 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), (1, True)) - gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), (0, 'gidx0', 4)) + gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), ('gidx0', 4)) idx = gidx0*UOp.const(dtypes.int, 2) val = UOp.const(dtypes.float, 42.0) gate = gidx0.lt(UOp.const(dtypes.int, 1)) @@ -277,7 +277,7 @@ class TestGatedStoreRewrite(unittest.TestCase): def test_merge_ifs_alt(self): gmem0 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), (0, True)) gmem1 = UOp(UOps.DEFINE_GLOBAL, PtrDType(dtypes.float), (), (1, True)) - gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), (0, 'gidx0', 4)) + gidx0 = UOp(UOps.SPECIAL, dtypes.int, (), ('gidx0', 4)) idx = gidx0*UOp.const(dtypes.int, 2) val = UOp.const(dtypes.float, 42.0) gate = gidx0.lt(UOp.const(dtypes.int, 1)) diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 04b3184a72..8a67e973d5 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -740,12 +740,12 @@ class Kernel: self.local_size: Optional[List[int]] = [1,1,1] for u in uop_sink.parents: if u.op is UOps.SPECIAL: - if u.arg[1][0] == 'i': self.local_size = None - if u.arg[1][0] == 'l': + if u.arg[0][0] == 'i': self.local_size = None + if u.arg[0][0] == 'l': assert self.local_size is not None - self.local_size[u.arg[0]] = u.arg[2] + self.local_size[int(u.arg[0][-1])] = u.arg[1] else: - self.global_size[u.arg[0]] = u.arg[2] + self.global_size[int(u.arg[0][-1])] = u.arg[1] else: self.global_size, self.local_size = None, None diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index d4987c6a77..5cbb738070 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -75,7 +75,7 @@ def _limit_dims(dims:Tuple[sint, ...], max_sizes:Tuple[int, ...]): def get_grouped_dims(prefix, dims:Tuple[sint, ...], max_sizes:Optional[Tuple[int, ...]], reverse=False) -> List[UOp]: if reverse: dims = dims[::-1] limited = _limit_dims(dims, max_sizes) if max_sizes is not None else dims - ret = raw_idxs = [UOp(UOps.SPECIAL, dtypes.bigint, (), (i, f"{prefix}{i}", s)) for i,s in enumerate(limited)] + ret = raw_idxs = [UOp(UOps.SPECIAL, dtypes.bigint, (), (f"{prefix}{i}", s)) for i,s in enumerate(limited)] if limited != dims: ret = [] # cast for mypy, get_contraction won't be None diff --git a/tinygrad/codegen/uopgraph.py b/tinygrad/codegen/uopgraph.py index 0116eae5c6..5478ee8cf2 100644 --- a/tinygrad/codegen/uopgraph.py +++ b/tinygrad/codegen/uopgraph.py @@ -192,9 +192,9 @@ constant_folder = PatternMatcher([ # other arange folders (UOp.cvar("c1") - (UOp.var("x") + UOp.cvar("c2")), lambda c1, c2, x: (c1-c2)-x), # c1 - (x + c2) -> (c1-c2) - x # max on special can go away (TODO: special should be variable, same thing applies) - (UOp.max(UOp.cvar('c'), UOp(UOps.SPECIAL).name('s')), lambda c,s: c if (s.arg[2]-1) <= c.arg else None), + (UOp.max(UOp.cvar('c'), UOp(UOps.SPECIAL).name('s')), lambda c,s: c if (s.arg[1]-1) <= c.arg else None), (UOp.max(UOp.cvar('c'), UOp(UOps.SPECIAL).name('s')+UOp.cvar('c2')), lambda c,s,c2: (s+c2) if 0 >= c.arg else None), # TODO: generic - (UOp.max(UOp.cvar('c'), -(UOp(UOps.SPECIAL).name('s')+UOp.cvar('c2'))), lambda c,s,c2: -(s+c2) if -(s.arg[2]-1+c2.arg) >= c.arg else None), + (UOp.max(UOp.cvar('c'), -(UOp(UOps.SPECIAL).name('s')+UOp.cvar('c2'))), lambda c,s,c2: -(s+c2) if -(s.arg[1]-1+c2.arg) >= c.arg else None), # max on range can go away (ugh: copy of SPECIAL, and with/without const) (UOp.max(UOp.cvar('c'), UOp(UOps.RANGE).name('s')), lambda c,s: s if s.src[0].arg >= c.arg else None), # TODO: generic (UOp.max(UOp.cvar('c'), UOp(UOps.RANGE).name('s')+UOp.cvar('c2')), lambda c,s,c2: (s+c2) if s.src[0].arg >= c.arg else None), # TODO: generic diff --git a/tinygrad/codegen/uops.py b/tinygrad/codegen/uops.py index 4cc90ffc74..a0690c7881 100644 --- a/tinygrad/codegen/uops.py +++ b/tinygrad/codegen/uops.py @@ -205,7 +205,7 @@ def type_verify(uops): assert dtype == src[1].dtype == src[2].dtype, f"{arg} choice dtype mismatch {dtype=} != {src[1].dtype=} != {src[2].dtype=}" def uop_alu_resolve(u:UOp) -> sint: - if u.op is UOps.SPECIAL: return u.arg[2]-1 + if u.op is UOps.SPECIAL: return u.arg[1]-1 if u.op in {UOps.CONST, UOps.DEFINE_VAR}: return u.arg if u.op is UOps.ALU: return exec_alu(u.arg, cast(DType,u.dtype), tuple(map(uop_alu_resolve, u.src))) raise RuntimeError(f"ALU resolve fail @ {u.op}") diff --git a/tinygrad/renderer/assembly.py b/tinygrad/renderer/assembly.py index 8cc06185c9..0262da6c81 100644 --- a/tinygrad/renderer/assembly.py +++ b/tinygrad/renderer/assembly.py @@ -168,10 +168,10 @@ class PTXRenderer(Renderer): for uu in r[u]: kk(f"mov.b{self.types[dtype.scalar()][1:]} {uu}, {const(src[0].arg, dtype.scalar())};") else: kk(f"mov.{f'b{self.types[dtype][1:]}' if dtype != dtypes.bool else 'pred'} {ssa('acc', u)}, {const(src[0].arg, dtype)};") elif uop is UOps.SPECIAL: - assert args[1][0] != "i", "idx not supported" - kk(f"mov.u32 %{args[1]}, {(self.gid if args[1][0] == 'g' else self.lid)[args[0]]};") - r[u] = "%" + args[1] - kernel = [f".reg .u32 %{args[1]};"] + kernel + assert args[0][0] != "i", "idx not supported" + kk(f"mov.u32 %{args[0]}, {(self.gid if args[0][0] == 'g' else self.lid)[int(args[0][-1])]};") + r[u] = "%" + args[0] + kernel = [f".reg .u32 %{args[0]};"] + kernel elif uop is UOps.DEFINE_VAR: bufs.append((args.expr, dtype)) r[u] = f"%{args.expr}" diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index 3eb2969014..5dd07ea7cd 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -139,8 +139,8 @@ class CStyleLanguage(Renderer): if child_count[u] <= 1 and args is not BinaryOps.MAX and not getenv("EXPAND_SSA"): r[u] = val else: kk(f"{self.render_dtype(dtype)} {ssa('alu',u)} = {val};") elif uop is UOps.SPECIAL: - kk(f"int {args[1]} = {self.code_for_workitem[args[1][0]](args[0])}; /* {args[2]} */") - r[u] = args[1] + kk(f"int {args[0]} = {self.code_for_workitem[args[0][0]](args[0][-1])}; /* {args[1]} */") + r[u] = args[0] elif uop is UOps.DEFINE_VAR: assert args.expr not in seen_vars, f"duplicate variable {args.expr}" seen_vars.add(args.expr) @@ -226,7 +226,7 @@ class MetalRenderer(CStyleLanguage): barrier = "threadgroup_barrier(mem_flags::mem_threadgroup);" float4 = "float4" uses_ptr_arithmetic = True - code_for_workitem = {"g": lambda x: f"gid.{chr(120+x)}", "l": lambda x: f"lid.{chr(120+x)}"} + code_for_workitem = {"g": lambda x: f"gid.{chr(120+int(x))}", "l": lambda x: f"lid.{chr(120+int(x))}"} # uint3 used for gid/lid - TODO: this should probably be `ushort3 lid [[thread_position_in_threadgroup]]` extra_args = ['uint3 gid [[threadgroup_position_in_grid]]', 'uint3 lid [[thread_position_in_threadgroup]]'] type_map = {dtypes.bfloat16: "bfloat"} @@ -274,8 +274,8 @@ class CUDARenderer(CStyleLanguage): smem_prefix_for_cast = False barrier = "__syncthreads();" float4 = "make_float4" - code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+x)}", "l": lambda x: f"threadIdx.{chr(120+x)}", - "i": lambda x: f"(blockIdx.{chr(120+x)}*blockDim.{chr(120+x)}+threadIdx.{chr(120+x)})"} + code_for_workitem = {"g": lambda x: f"blockIdx.{chr(120+int(x))}", "l": lambda x: f"threadIdx.{chr(120+int(x))}", + "i": lambda x: f"(blockIdx.{chr(120+int(x))}*blockDim.{chr(120+x)}+threadIdx.{chr(120+int(x))})"} code_for_op = {**CStyleLanguage().code_for_op, **code_for_op_half} type_map = {dtypes.bfloat16: "nv_bfloat16"} @@ -385,7 +385,7 @@ static inline __attribute__((device)) bool operator==(hip_bfloat16 a, hip_bfloat return super().render_kernel(function_name, kernel, bufs, uops, prefix) def get_kernel_modifier(self, uops:UOpGraph) -> str: - requiredMaxThreadsPerBlock = prod(u.arg[2] for u in uops if u.op is UOps.SPECIAL and u.arg[1][0] == "l") + requiredMaxThreadsPerBlock = prod(u.arg[1] for u in uops if u.op is UOps.SPECIAL and u.arg[0][0] == "l") # https://clang.llvm.org/docs/AttributeReference.html#amdgpu-flat-work-group-size # NOTE: this makes hlb_cifar10 twice as fast, there may be more gains in tweaking these parameters return f"__attribute__((amdgpu_flat_work_group_size(1, {requiredMaxThreadsPerBlock})))" diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index f649bb44b6..f523a3ae6c 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -84,8 +84,8 @@ class PythonProgram: elif uop is UOps.DEFINE_VAR: ul[i] = [pvals.pop(0)] * warp_size elif uop is UOps.SPECIAL: - if arg[1][0] == 'g': ul[i] = [idxs[2-arg[0]]] * warp_size - elif arg[1][0] == 'l': ul[i] = [x[2-arg[0]] for x in warp] + if arg[0][0] == 'g': ul[i] = [idxs[2-int(arg[0][-1])]] * warp_size + elif arg[0][0] == 'l': ul[i] = [x[2-int(arg[0][-1])] for x in warp] elif uop is UOps.CONST: ul[i] = [[arg] * warp_size for _ in range(dtype.count)] if dtype.count > 1 else [arg] * warp_size elif uop is UOps.DEFINE_ACC: