improve opencl, why is it OOMing
MODIFY
accel/opencl/ops_opencl.py
accel/opencl/ops_opencl.py
@@ -10,7 +10,7 @@ import numpy as np
import pyopencl as cl
UNSAFE_FLOAT4 = int(os.getenv("UNSAFE_FLOAT4", 0))
-NATIVE_EXPLOG = int(os.getenv("NATIVE_EXPLOG", 0))
+FLOAT16 = int(os.getenv("FLOAT16", 0))
import pathlib
def load(x):
@@ -20,15 +20,15 @@ def load(x):
CONV_SRC = load(pathlib.Path(__file__).resolve().parent.parent.parent / 'accel/opencl/conv.cl')
MATMUL_SRC = load(pathlib.Path(__file__).resolve().parent.parent.parent / 'accel/opencl/matmul.cl')
-class ECL(CL):
- @staticmethod
- def image(shape):
- if int(os.getenv("FLOAT16", 0)):
+class CLImage:
+ def __init__(self, shape):
# HALF_FLOAT breaks tests
- fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.HALF_FLOAT)
- else:
- fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.FLOAT)
- return cl.Image(CL().cl_ctx, cl.mem_flags.READ_WRITE, fmt, shape=shape)
+ fmt = cl.ImageFormat(cl.channel_order.RGBA, cl.channel_type.HALF_FLOAT if FLOAT16 else cl.channel_type.FLOAT)
+ self.cl = cl.Image(CL().cl_ctx, cl.mem_flags.READ_WRITE, fmt, shape=shape)
+ CL.mem_used += self.cl.row_pitch * self.cl.height
+
+ def __del__(self):
+ CL.mem_used -= self.cl.row_pitch * self.cl.height
def get_replacements(prg_src:str, opencl_type:List[str]) -> Dict[str, str]:
middle_code = []
@@ -102,9 +102,9 @@ def roundup(x, n=4): return (x+(n-1))//n * n
class OpenCLBuffer(GPUBuffer):
code_for_op = {
UnaryOps.NOOP: "(A)", UnaryOps.NEG: "(-(A))", UnaryOps.RELU: "max(A, (float)0.)", UnaryOps.SIGN: "sign(A)",
- UnaryOps.EXP: "native_exp(A)" if NATIVE_EXPLOG else "exp(A)",
- UnaryOps.LOG: "native_log(A)" if NATIVE_EXPLOG else "log(A)",
- UnaryOps.RECIPROCAL: "native_recip(A)" if NATIVE_EXPLOG else "(1.0/A)",
+ UnaryOps.EXP: "native_exp(A)",
+ UnaryOps.LOG: "native_log(A)",
+ UnaryOps.RECIPROCAL: "native_recip(A)",
BinaryOps.ADD: "(A+B)", BinaryOps.SUB: "(A-B)", BinaryOps.MUL: "(A*B)", BinaryOps.DIV: "(A/B)", BinaryOps.POW: "pow(A,B)", BinaryOps.CMPEQ: "(A==B)",
ReduceOps.SUM: "(acc + A)", ReduceOps.MAX: "max(A, acc)"
}
@@ -127,7 +127,7 @@ class OpenCLBuffer(GPUBuffer):
self._buf = CLBuffer(4*roundup(prod(self.shape)))
if self._image is not None:
- self._buf = CLBuffer(4*roundup(prod(self._image.shape)*4))
+ self._buf = CLBuffer(4*roundup(prod(self._image.cl.shape)*4))
if self._backing is not None:
CL.enqueue_copy(self._buf.cl, self._backing, is_blocking=False)
#self._backing = None
@@ -143,7 +143,7 @@ class OpenCLBuffer(GPUBuffer):
int W = get_image_width(in);
out[l.y*W + l.x] = read_imagef(in, smp, l);
}
- """)(self._image.shape, None, self._buf.cl, self._image)
+ """)(self._image.cl.shape, None, self._buf.cl, self._image.cl)
self._image = None
return self._buf.cl
@@ -153,9 +153,9 @@ class OpenCLBuffer(GPUBuffer):
def image(self):
if self._image is None:
assert len(self.shape) == 3 and self.shape[2] == 4, f"bad shape for image {self.shape}"
- self._image = ECL.image(shape=(self.shape[1], self.shape[0]))
+ self._image = CLImage(shape=(self.shape[1], self.shape[0]))
if self._buf is not None:
- assert prod(self.shape) == prod(self._image.shape)*4
+ assert prod(self.shape) == prod(self._image.cl.shape)*4
#print(f"converting {self.shape} to image with shape {self._image.shape}")
CLProgram("to_image", """
__kernel void to_image(
@@ -167,9 +167,9 @@ class OpenCLBuffer(GPUBuffer):
int W = get_image_width(out);
write_imagef(out, l, in[l.y*W + l.x]);
}
- """)(self._image.shape, None, self._image, self._buf.cl)
+ """)(self._image.cl.shape, None, self._image.cl, self._buf.cl)
self._buf = None
- return self._image
+ return self._image.cl
SUPPORTS_PADDING = True
def processing_op(x, op:ProcessingOps, w:GPUBuffer, C:ConvArgs):

MODIFY
examples/stable_diffusion.py
examples/stable_diffusion.py
@@ -10,6 +10,7 @@ from collections import namedtuple
from extra.utils import fake_torch_load_zipped, get_child
from tinygrad.nn import Conv2d
from tinygrad.tensor import Tensor
+from tinygrad.helpers import prod
# TODO: rename to GroupNorm and put in nn.py
class Normalize:
@@ -116,6 +117,7 @@ class Decoder:
bs,c,py,px = x.shape
x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
x = l['upsample']['conv'](x)
+ x.realize()
return self.conv_out(self.norm_out(x).swish())
@@ -361,17 +363,19 @@ class UNetModel:
saved_inputs = []
for i,b in enumerate(self.input_blocks):
- #print("input block", i)
+ print("input block", i)
for bb in b:
x = run(x, bb)
saved_inputs.append(x)
+ x.realize()
for bb in self.middle_block:
x = run(x, bb)
for i,b in enumerate(self.output_blocks):
- #print("output block", i)
+ print("output block", i)
x = x.cat(saved_inputs.pop(), dim=1)
for bb in b:
x = run(x, bb)
+ x.realize()
return x.sequential(self.out)
class CLIPMLP:
@@ -516,7 +520,7 @@ FILENAME = "/home/kafka/model.ckpt"
if __name__ == "__main__":
Tensor.no_init = True
- # WTF!! no_grad breaks it
+ # WTF!! no_grad breaks it (only with OPENCL, now fixed)
Tensor.no_grad = True
model = StableDiffusion()
@@ -541,24 +545,32 @@ if __name__ == "__main__":
# "penguin with fire extinguisher"
#phrase = [49406, 14952, 593, 1769, 38567, 4510, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407]
- context = model.cond_stage_model.transformer.text_model(phrase)
+ context = model.cond_stage_model.transformer.text_model(phrase).realize()
print("got CLIP context", context.shape)
phrase = [49406, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407, 49407]
- unconditional_context = model.cond_stage_model.transformer.text_model(phrase)
+ unconditional_context = model.cond_stage_model.transformer.text_model(phrase).realize()
print("got unconditional CLIP context", unconditional_context.shape)
+ # done with clip model
+ del model.cond_stage_model
+
def get_model_output(latent, t):
# put into diffuser
timesteps = Tensor([t])
- unconditional_latent = model.model.diffusion_model(latent, timesteps, unconditional_context)
-
- latent = model.model.diffusion_model(latent, timesteps, context)
+ from tinygrad.llops.ops_gpu import CL
+ import gc
+
+ print(CL.mem_used/1e9, sum([prod(x.shape)*4 for x in gc.get_objects() if isinstance(x, Tensor)])/1e9)
+ unconditional_latent = model.model.diffusion_model(latent, timesteps, unconditional_context).realize()
+ print(CL.mem_used/1e9, sum([prod(x.shape)*4 for x in gc.get_objects() if isinstance(x, Tensor)])/1e9)
+ latent = model.model.diffusion_model(latent, timesteps, context).realize()
+ print(CL.mem_used/1e9, sum([prod(x.shape)*4 for x in gc.get_objects() if isinstance(x, Tensor)])/1e9)
unconditional_guidance_scale = 7.5
e_t = unconditional_latent + unconditional_guidance_scale * (latent - unconditional_latent)
return e_t
- TIMESTEPS = 50
+ TIMESTEPS = 5
timesteps = list(np.arange(1, 1000, 1000//TIMESTEPS))
print(f"running for {timesteps} timesteps")
alphas = [model.alphas_cumprod.numpy()[t] for t in timesteps]

MODIFY
openpilot/compile.py
openpilot/compile.py
@@ -83,9 +83,7 @@ def compile(input, output_fn):
# realize all non GCed tensors (fix for batchnorm folding)
import gc
gc.collect()
- def tensors_allocated():
- return [x for x in gc.get_objects() if isinstance(x, Tensor)]
- for x in tensors_allocated():
+ for x in [x for x in gc.get_objects() if isinstance(x, Tensor)]:
x.realize()
# real run

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -13,9 +13,9 @@ class CLBuffer:
def __init__(self, size):
if len(CL.BUFFER_CACHE[size]) > 0: self.cl = CL.BUFFER_CACHE[size].pop()
else:
- CL.mem_used += size
# TODO: on GPU OOM, clear the cache
self.cl = cl.Buffer(CL().cl_ctx, cl.mem_flags.READ_WRITE, size)
+ CL.mem_used += self.cl.size
def __del__(self):
if CLCACHE: CL.BUFFER_CACHE[self.cl.size].append(self.cl)
@@ -58,7 +58,7 @@ class CLProgram:
if DEBUG >= 1:
CL.time_sum += 0 if DEBUG <= 1 or CL.CACHE is not None else (e.profile.end - e.profile.start)
CL.ops_sum += op_estimate
- print(f"**CL** {CL.kernel_count:6d} {self.name:20s} args {len(args[2:]):5d} kernels {str(args[0]):18s} {str(args[1]):12s} OPs {op_estimate/1e6:6.1f}M/{CL.ops_sum/1e9:7.2f}G " +
+ print(f"**CL** {CL.kernel_count:6d} {self.name:20s} args {len(args[2:]):5d} kernels {str(args[0]):18s} {str(args[1]):12s} OPs {op_estimate/1e6:6.1f}M/{CL.ops_sum/1e9:7.2f}G mem {CL.mem_used/1e9:5.2f} GB " +
("" if DEBUG <= 1 or CL.CACHE is not None else f"tm {(e.profile.end - e.profile.start)/1e3:9.2f}us/{CL.time_sum/1e6:9.2f}ms ({op_estimate/(e.profile.end - e.profile.start):8.2f} GFLOPS)"))
if DEBUG >= 4: print(self.prg)

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -144,7 +144,7 @@ def _realize_binaryops(self:LazyBuffer) -> Tuple[DeviceBuffer, List[DeviceBuffer
# if there's *one* processing or reduce op in here, we can corealize it. we can corealize binary op siblings as well
# NOTE: if it references the same conv multiple times, they should already be merged by the dictionary
psrcs : List[Tuple[LazyBuffer, LazyBuffer]] = [(k,x) for k,x in zip(real_srcs.keys(), map(get_movementroot_contiguous, real_srcs.keys())) if x.optype in [ProcessingOps,ReduceOps] and x.realized is None and len(x.children) <= 1 and len(k.children) <= 1]
- if len(psrcs) == 1 and MERGE_ONE_REDUCE_INTO_ELEMENTWISE:
+ if len(psrcs) == 1 and MERGE_ONE_REDUCE_INTO_ELEMENTWISE and (self.device != "OPENCL" or self.shape[-1] == 4):
if psrcs[0][1].optype == ProcessingOps:
# TODO: do something similar to what i did with reduceop to use the ast engine?
# it's hard because conv also has convargs

improve opencl, why is it OOMing