accurate flop estimation
MODIFY
accel/llvm/ops_llvm.py
accel/llvm/ops_llvm.py
@@ -9,7 +9,7 @@ from tinygrad.ops import LazyOp
import ctypes
import numpy as np
from ctypes import CFUNCTYPE
-from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, get_buffers, get_lazyops, ExplicitExecAST, get_lazyop_shape
+from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, get_buffers, get_lazyops, ExplicitExecAST, get_lazyop_info
from llvmlite import ir # type: ignore
import llvmlite.binding as llvm # type: ignore
@@ -178,7 +178,7 @@ class LLVMBuffer(ExplicitExecAST):
reduceops = [x for x in get_lazyops(ast) if isinstance(x.op, ReduceOps)]
assert len(reduceops) <= 1, "max one reduce op in an ast"
earlybufs = get_buffers(reduceops[0]) if len(reduceops) > 0 else []
- ret = cls(get_lazyop_shape(ast))
+ ret = cls(get_lazyop_info(ast).shape)
module = ir.Module(name=__file__)
func = ir.Function(module, ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.FloatType())]*(1+len(bufs))), name='exec')

MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
@@ -231,10 +231,10 @@ class OpenCLBuffer(GPUBuffer):
#ewtypes.append(f"read_only image2d_t {name}_g")
return super().contiguous_view_constant_fold(name, reduce)
- def _processing_op(ret, bufs: List[Tuple[str, OpenCLBuffer]]=[], code:str="acc", C=None, op=ReduceOps.SUM, reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc"):
+ def _processing_op(ret, bufs: List[Tuple[str, OpenCLBuffer]]=[], code:str="acc", C=None, op=ReduceOps.SUM, reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc", op_estimate=0):
if C is None or earlycode != "acc":
# TODO: handle an opencl conv without the conv part
- return super()._processing_op(bufs, code, C, op, reduce_shape, earlybufs, earlycode)
+ return super()._processing_op(bufs, code, C, op, reduce_shape, earlybufs, earlycode, op_estimate)
assert earlycode == "acc"
x = [x for x in bufs if x[0] == "input"][0][1]

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -5,7 +5,7 @@ import pyopencl as cl # type: ignore
from collections import defaultdict
from typing import List, Tuple, Optional, Dict, Union, Set
from tinygrad.helpers import prod, ConvArgs, dedup
-from tinygrad.ops import DEBUG, ProcessingOps, UnaryOps, BinaryOps, ReduceOps, LazyOp, get_buffers, get_lazyops, Op, get_lazyop_shape, ExplicitExecAST
+from tinygrad.ops import DEBUG, ProcessingOps, UnaryOps, BinaryOps, ReduceOps, LazyOp, get_buffers, get_lazyops, Op, get_lazyop_info, ExplicitExecAST
from tinygrad.shapetracker import ShapeTracker
CLCACHE = int(os.getenv("CLCACHE", "1"))
@@ -130,7 +130,8 @@ class GPUBuffer(ExplicitExecAST):
assert len(reduceops) <= 1, f"max one reduce op in an ast, {reduceops}"
earlybufs = dedup(get_buffers(reduceops[0])) if len(reduceops) > 0 else []
reduce_shape = (earlybufs[0].shape, reduceops[0].arg) if len(reduceops) > 0 and isinstance(reduceops[0].op, ReduceOps) else None
- ret = cls(get_lazyop_shape(ast))
+ info = get_lazyop_info(ast)
+ ret = cls(info.shape)
buf_names : Dict[GPUBuffer, str] = {x:f"arg_{i}" for i,x in enumerate(bufs)}
@@ -157,9 +158,9 @@ class GPUBuffer(ExplicitExecAST):
C = reduceops[0].arg if len(reduceops) > 0 and isinstance(reduceops[0].op, ProcessingOps) else None
reduce_op = reduceops[0].op if len(reduceops) > 0 and isinstance(reduceops[0].op, ReduceOps) else ReduceOps.SUM
- return ret._processing_op([(buf_names[x], x) for x in bufs], code, C, reduce_op, reduce_shape, set(buf_names[x] for x in earlybufs), earlycode)
+ return ret._processing_op([(buf_names[x], x) for x in bufs], code, C, reduce_op, reduce_shape, set(buf_names[x] for x in earlybufs), earlycode, info.flops)
- def _processing_op(ret, bufs: List[Tuple[str, GPUBuffer]]=[], code:str="acc", C:Optional[ConvArgs]=None, op=ReduceOps.SUM, reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc") -> GPUBuffer:
+ def _processing_op(ret, bufs: List[Tuple[str, GPUBuffer]]=[], code:str="acc", C:Optional[ConvArgs]=None, op=ReduceOps.SUM, reduce_shape=None, earlybufs:Set[str]=set(), earlycode:str="acc", op_estimate=0) -> GPUBuffer:
assert C is None, f"conv isn't handled by GPU anymore {C}"
# get the input/output shape and the reduce amount
@@ -204,5 +205,5 @@ class GPUBuffer(ExplicitExecAST):
}}
}}""")
- conv_prg([prod(ret.shape), inter_red, 1], [1, inter_red, 1] if inter_red > 1 else None, ret.cl, *buf_cl, op_estimate=prod(reduce_shape[0])*len(earlybufs) + prod(reduce_shape[1])*len(bufs))
+ conv_prg([prod(ret.shape), inter_red, 1], [1, inter_red, 1] if inter_red > 1 else None, ret.cl, *buf_cl, op_estimate=op_estimate)
return ret

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -2,6 +2,7 @@ import os
from enum import Enum
from typing import Union, Type, NamedTuple, Tuple, Any, List
import functools, operator
+from tinygrad.helpers import prod
from tinygrad.shapetracker import ShapeTracker
DEBUG = int(os.getenv("DEBUG", "0"))
@@ -55,13 +56,13 @@ class GenericExecAST(DeviceBuffer):
return ret
class GenericShape(GenericExecAST):
- def __init__(self, x): self.shape = x if isinstance(x, tuple) else x.shape
- def unary_op(self, op:UnaryOps): return self
- def binary_op(self, op:BinaryOps, y): return self
- def reduce_op(self, op:ReduceOps, new_shape:Tuple[int, ...]): return type(self)(new_shape)
- def movement_op(self, op:MovementOps, arg): return type(self)(ShapeTracker(self.shape).movement_op(op, arg))
- def processing_op(self, op:ProcessingOps, w, C): return type(self)(C.out_shape)
-def get_lazyop_shape(ast:LazyOp): return GenericShape.exec_ast(ast, GenericShape).shape
+ def __init__(self, shape, flops=0): self.shape, self.flops = shape, flops
+ def unary_op(self, op:UnaryOps): return GenericShape(self.shape, self.flops + prod(self.shape))
+ def binary_op(self, op:BinaryOps, y): return GenericShape(self.shape, self.flops + y.flops + prod(self.shape))
+ def reduce_op(self, op:ReduceOps, new_shape:Tuple[int, ...]): return GenericShape(new_shape, self.flops + prod(self.shape))
+ def movement_op(self, op:MovementOps, arg): return GenericShape(ShapeTracker(self.shape).movement_op(op, arg).shape, self.flops)
+ def processing_op(self, op:ProcessingOps, w, C): return GenericShape(C.out_shape, float("nan")) # TODO: add flops for this
+def get_lazyop_info(ast:LazyOp): return GenericShape.exec_ast(ast, lambda x: GenericShape(x.shape))
# assumes you are using ShapeTracker
# used in GPUBuffer, OpenCLBuffer, and LLVMBuffer

accurate flop estimation