diff --git a/.pylintrc b/.pylintrc index 7678d2e71..fe6b6ac4d 100644 --- a/.pylintrc +++ b/.pylintrc @@ -64,7 +64,7 @@ disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1 # either give multiple identifier separated by comma (,) or put this option # multiple time (only on the command line, not in the configuration file where # it should appear only once). See also the "--disable" option for examples. -enable=c-extension-no-member,use-a-generator +enable=c-extension-no-member,use-a-generator, no-else-return [REPORTS] diff --git a/models/mask_rcnn.py b/models/mask_rcnn.py index 5f32b8287..33ea5fb18 100644 --- a/models/mask_rcnn.py +++ b/models/mask_rcnn.py @@ -188,7 +188,7 @@ class BoxList: if self.mode == "xyxy": xmin, ymin, xmax, ymax = self.bbox.chunk(4, dim=-1) return xmin, ymin, xmax, ymax - elif self.mode == "xywh": + if self.mode == "xywh": TO_REMOVE = 1 xmin, ymin, w, h = self.bbox.chunk(4, dim=-1) return ( diff --git a/tinygrad/codegen/assembly.py b/tinygrad/codegen/assembly.py index 3075e126b..f128413f1 100644 --- a/tinygrad/codegen/assembly.py +++ b/tinygrad/codegen/assembly.py @@ -96,9 +96,8 @@ class AssemblyCodegen(Linearizer): ins.append(AssemblyInstruction(UOps.ALU, new_reg, [reg], UnaryOps.NOOP)) reg = new_reg return tor[f"buf{args.i}"], reg, off - else: - reg = render_alu(BinaryOps.ADD, render_cast(reg, dtypes.uint64), tor[f"buf{args.i}"], dtype=dtypes.uint64) - return reg, None, off + reg = render_alu(BinaryOps.ADD, render_cast(reg, dtypes.uint64), tor[f"buf{args.i}"], dtype=dtypes.uint64) + return reg, None, off ins = [] ins += [AssemblyInstruction(UOps.SPECIAL, newreg(f"buf{i}", dtype=dtypes.uint64, scalar=True), [], f"buf{i}") for i in range(len(self.bufs))] diff --git a/tinygrad/codegen/cstyle.py b/tinygrad/codegen/cstyle.py index f9dcafb2f..0e5399c95 100644 --- a/tinygrad/codegen/cstyle.py +++ b/tinygrad/codegen/cstyle.py @@ -30,7 +30,7 @@ class CStyleLanguage(NamedTuple): assert len(x) == var_dtype.sz, f"cast is wrong size {len(x)} != {var_dtype.sz}" assert self.float4 is not None, "cast is not supported on this platform" if var_dtype == dtypes._float4: return f"{self.float4}({','.join(x)})" - elif var_dtype == dtypes._float2: return f"{self.float4.replace('float4', 'float2')}({','.join(x)})" + if var_dtype == dtypes._float2: return f"{self.float4.replace('float4', 'float2')}({','.join(x)})" raise NotImplementedError(f"no cast for {var_dtype}") # returns a str expression of the const with the given type @@ -45,24 +45,22 @@ class CStyleLanguage(NamedTuple): if isinstance(buf_dtype, ImageDType): assert output_dtype == dtypes._float4, "images must be float4" return f"read_imagef({buf_name}, smp, (int2)({idx[0].render(render_cl)}, {idx[1].render(render_cl)}))" - elif self.uses_vload and buf_dtype == dtypes.float16: + if self.uses_vload and buf_dtype == dtypes.float16: return f"vload_half{'' if output_dtype.sz == 1 else str(output_dtype.sz)}(0, {buf_name}+{idx.render(render_cl, strip_parens=True)})" - elif output_dtype.sz > 1: + if output_dtype.sz > 1: return f"({output_dtype.name})(*(({self.smem_prefix if local else self.buffer_prefix}{buf_dtype.name}{output_dtype.sz}*)({buf_name}+{idx.render(render_cl, strip_parens=True)})))" - else: - return f"{buf_name}[{idx.render(render_cl)}]" + return f"{buf_name}[{idx.render(render_cl)}]" # returns a str statement that does the store def render_store(self, buf_name, buf_dtype, var_name, var_dtype, idx, local=False) -> str: if isinstance(buf_dtype, ImageDType): assert var_dtype == dtypes._float4, "images must be float4" return f"write_imagef({buf_name}, (int2)({idx[0].render(render_cl)}, {idx[1].render(render_cl)}), {var_name});" - elif self.uses_vload and buf_dtype == dtypes.float16: + if self.uses_vload and buf_dtype == dtypes.float16: return f"vstore_half{'' if var_dtype.sz == 1 else str(var_dtype.sz)}({var_name}, 0, {buf_name}+{idx.render(render_cl, strip_parens=True)});" - elif var_dtype.sz > 1: + if var_dtype.sz > 1: return f"*(({self.smem_prefix if local else self.buffer_prefix}{buf_dtype.name}{var_dtype.sz}*)({buf_name}+{idx.render(render_cl, strip_parens=True)})) = ({buf_dtype.name}{var_dtype.sz}){var_name};" - else: - return f"{buf_name}[{idx.render(render_cl)}] = {var_name};" + return f"{buf_name}[{idx.render(render_cl)}] = {var_name};" code_for_op: Final[Dict[Op, Callable]] = { UnaryOps.EXP2: lambda x: f"exp2({x})", @@ -85,9 +83,8 @@ def add_gl_dimension(args, i, var, local_size, xid): lidx = (lidx//((lidx.max+1)//local_size[-1]))%(var.max+1) assert lidx.max == var.max and lidx.min == var.min return f"{{ int {var.expr} = {lidx.render(render_cl)}; /* {var.max+1} */" - else: - local_size.append(var.max+1) - return f"{{ int {var.expr} = {xid[min(len(xid), len(args[0]))-1-i]}; /* {var.max+1} */" + local_size.append(var.max+1) + return f"{{ int {var.expr} = {xid[min(len(xid), len(args[0]))-1-i]}; /* {var.max+1} */" def uops_to_cstyle(uops:List[UOp], bufs:List[Union[LocalBuffer,LazyBuffer]], lang:CStyleLanguage) -> Tuple[str, List[int], List[int]]: prekernel: Set[str] = set() diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index ca5124c8d..f3bad6287 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -101,10 +101,10 @@ def get_grouped_maybe_float4(*values:List[Token], grouping_allowed=True): # TODO: generic visitor pattern? def expand_node(idx:Node) -> List[Node]: if isinstance(idx, Variable): return [idx] if idx.expr is not None else [Variable.num(j) for j in range(idx.min, idx.max+1)] - elif isinstance(idx, NumNode): return [idx] - elif isinstance(idx, MulNode): return [x*idx.b for x in expand_node(idx.a)] - elif isinstance(idx, SumNode): return [Variable.sum(list(it)) for it in itertools.product(*[expand_node(x) for x in idx.nodes])] - else: raise NotImplementedError(idx) + if isinstance(idx, NumNode): return [idx] + if isinstance(idx, MulNode): return [x*idx.b for x in expand_node(idx.a)] + if isinstance(idx, SumNode): return [Variable.sum(list(it)) for it in itertools.product(*[expand_node(x) for x in idx.nodes])] + raise NotImplementedError(idx) def expand_idxs(idxs:Sequence[Node]) -> Iterator[Tuple[Node, ...]]: for x in itertools.product(*[expand_node(idx) for idx in idxs[::-1]]): diff --git a/tinygrad/ops.py b/tinygrad/ops.py index e3655171c..ab235e849 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -104,8 +104,7 @@ class Interpreted: assert output.output_buffer.size == ret.size, output.output_buffer.dtype == ret.dtype output.output_buffer._buf = ret._buf return output.output_buffer - else: - return ret + return ret class FlopCounter: def __init__(self, tup:Tuple[Tuple[int, ...], DType, int]): self.shape, self.dtype, self.flops, self._buf = *tup, self diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index 3d3fe06de..dd244ac6a 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -70,8 +70,7 @@ class MetalProgram: if wait: command_buffer.waitUntilCompleted() return command_buffer.GPUEndTime() - command_buffer.GPUStartTime() - else: - METAL.mtl_buffers_in_flight.append(command_buffer) + METAL.mtl_buffers_in_flight.append(command_buffer) class MetalCodegen(CStyleCodegen): lang = CStyleLanguage( diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py index 3fb0b2eb8..a56a5f98d 100644 --- a/tinygrad/shape/shapetracker.py +++ b/tinygrad/shape/shapetracker.py @@ -119,11 +119,9 @@ def _reshape(view: View, new_shape:Tuple[int, ...]) -> Tuple[View, bool]: new_view = View(new_shape, strides_for_shape(new_shape)) if view.contiguous: return new_view, False # NOTE: if it's contiguous it can't have an offset - else: - if (merged_view := merge_views(view, new_view)) is not None: return merged_view, False - else: - if DEBUG >= 4: print(f"WARNING: creating new view with reshape {view} -> {new_shape}") - return new_view, True + if (merged_view := merge_views(view, new_view)) is not None: return merged_view, False + if DEBUG >= 4: print(f"WARNING: creating new view with reshape {view} -> {new_shape}") + return new_view, True @functools.lru_cache(maxsize=None) def get_pad_args(shape:Tuple[int,...], arg:Tuple[Tuple[int, int], ...]): diff --git a/tinygrad/shape/symbolic.py b/tinygrad/shape/symbolic.py index 0c2006b84..2b9b45931 100644 --- a/tinygrad/shape/symbolic.py +++ b/tinygrad/shape/symbolic.py @@ -49,7 +49,7 @@ class Node: return create_node(LtNode(lhs, b)) def __mul__(self, b:int): if b == 0: return NumNode(0) - elif b == 1: return self + if b == 1: return self return create_node(MulNode(self, b)) # *** complex ops *** diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 64c82838c..5d0de7235 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -422,17 +422,16 @@ class Tensor: xup = xup.slice(slc_prefix + flatten(((0,k), (0,o), (0,1)) for k,o in zip(k_, o_))) xup = xup.reshape(*prefix, *flatten((k,o) for k,o in zip(k_, o_))) return xup.permute(*range(len(prefix)), *[len(prefix)+i*2+1 for i in range(len(k_))], *[len(prefix)+i*2 for i in range(len(k_))]) - else: - # TODO: once the shapetracker can optimize well, remove this alternative implementation. or not if the CPU implementation doesn't use ShapeTracker - o_ = [(i+(s-k))//s for i,s,k in zip(i_, s_, k_)] - xup = self.slice(slc_prefix + [(0,o*s) for o,s in zip(o_, s_)]) - xup = xup.reshape(*prefix, *([1]*len(_insert_dims)), *flatten(((o, s) for o,s in zip(o_, s_)))) - if len(_insert_dims): - xup = xup.expand(*prefix, *_insert_dims, *flatten(((o, s) for o,s in zip(o_, s_)))) - prefix += _insert_dims - slc_prefix += [(0,x) for x in _insert_dims] - xup = xup.slice(slc_prefix + flatten(((0,o), (0,k)) for o,k in zip(o_, k_))) - return xup.permute(*range(len(prefix)), *[len(prefix)+i*2 for i in range(len(k_))], *[len(prefix)+i*2+1 for i in range(len(k_))]) + # TODO: once the shapetracker can optimize well, remove this alternative implementation. or not if the CPU implementation doesn't use ShapeTracker + o_ = [(i+(s-k))//s for i,s,k in zip(i_, s_, k_)] + xup = self.slice(slc_prefix + [(0,o*s) for o,s in zip(o_, s_)]) + xup = xup.reshape(*prefix, *([1]*len(_insert_dims)), *flatten(((o, s) for o,s in zip(o_, s_)))) + if len(_insert_dims): + xup = xup.expand(*prefix, *_insert_dims, *flatten(((o, s) for o,s in zip(o_, s_)))) + prefix += _insert_dims + slc_prefix += [(0,x) for x in _insert_dims] + xup = xup.slice(slc_prefix + flatten(((0,o), (0,k)) for o,k in zip(o_, k_))) + return xup.permute(*range(len(prefix)), *[len(prefix)+i*2 for i in range(len(k_))], *[len(prefix)+i*2+1 for i in range(len(k_))]) # NOTE: these work for more than 2D def avg_pool2d(self, kernel_size=(2,2), stride=None): return self._pool(make_pair(kernel_size), stride if stride is not None else kernel_size).mean(axis=tuple(range(0-len(make_pair(kernel_size)), 0)))