code: def save(self, output_fn):
      # this is the struct that will be saved
      jdat = {
          "binaries": [],
          "programs": {},
          "kernels": [],
          "objects": [],
      }
  
      # build the pieces of this struct
      weights = []
      binaries = []
      saved_objs = set()
      saved_binaries = set()
      for prg, args in self.cl_cache:
          # get binaries for saving
          if prg.name not in saved_binaries:
              binary = prg.clprogram.get_info(cl.program_info.BINARIES)
              assert len(binary) == 1
              jdat["binaries"].append(
                  {"name": prg.name, "length": len(binary[0])}
              )
              binaries.append(binary[0])
              saved_binaries.add(prg.name)
  
          # get the args from the kernel, some need the data saved
          targs, args_size = [], []
          argdtypes = [None] * (len(args) - 2)
          for a, d in zip(args[2:], argdtypes):
              if d == np.int16:
                  targs.append(struct.pack("H", a).decode("latin_1"))
                  args_size.append(2)
              elif d == np.int32:
                  targs.append(struct.pack("I", a).decode("latin_1"))
                  args_size.append(4)
              elif isinstance(a, cl.LocalMemory):
                  targs.append("")
                  args_size.append(a.size)
              elif d is None:
                  if getattr(a, "global_id", None) is None:
                      setattr(a, "global_id", self.gobj)
                      self.gobj += 1
                  ptr = struct.pack("Q", a.global_id).decode("latin_1")
                  if ptr not in saved_objs:
                      if isinstance(a, cl.Buffer):
                          needs_load = a in self.buffers_to_save
                          jdat["objects"].append(
                              {
                                  "id": ptr,
                                  "arg_type": "float*",
                                  "needs_load": needs_load,
                                  "size": a.size,
                              }
                          )
                          if needs_load:
                              data = np.empty(
                                  a.size // 4, dtype=np.float32
                              )
                              cl.enqueue_copy(
                                  CL.queue, data, a, is_blocking=True
                              )
                              weights.append(data.tobytes())
                      elif isinstance(a, cl.Image):
                          assert a.format == cl.ImageFormat(
                              cl.channel_order.RGBA,
                              cl.channel_type.HALF_FLOAT
                              if FLOAT16
                              else cl.channel_type.FLOAT,
                          ), "wrong type"
                          needs_load = a in self.buffers_to_save
                          row_pitch = (
                              (
                                  a.shape[0] * 4 * (2 if FLOAT16 else 4)
                                  + 63
                              )
                              // 64
                              * 64
                          )
                          size = row_pitch * a.shape[1]
                          # this is *2 if float16 and *4 if float32
                          buf = cl.Buffer(
                              CL.ctx,
                              cl.mem_flags.READ_WRITE,
                              size=size * (2 if FLOAT16 else 1),
                          )
  
                          # zero out the buffer
                          cl.enqueue_copy(
                              CL.queue,
                              buf,
                              b"\x00" * buf.size,
                              is_blocking=True,
                          )
  
                          CLProgram(
                              CL,
                              "from_image_strided",
                              compile_gpu(
                                  """
                  __kernel void from_image_strided(read_only image2d_t in, __global float4 *out, int row_pitch) {
                    const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;
                    int2 l;
                    l.y = get_global_id(1);
                    l.x = get_global_id(0);
                    out[l.y*row_pitch + l.x] = read_imagef(in, smp, l);
                  }
                """
                              ),
                              bufs=2,
                              vars=1,
                          )(
                              a,
                              buf,
                              row_pitch // (4 * (2 if FLOAT16 else 4)),
                              global_size=a.shape,
                          )
  
                          # multiple of 32 isn't enough
                          jdat["objects"].append(
                              {
                                  "id": ptr,
                                  "needs_load": needs_load,
                                  "size": size,
                                  "arg_type": "image2d_t",
                                  "width": a.shape[0],
                                  "height": a.shape[1],
                                  "row_pitch": row_pitch,
                                  "float32": not FLOAT16,
                              }
                          )
  
                          if needs_load:
                              data = np.empty(
                                  size // (2 if FLOAT16 else 4),
                                  dtype=np.float32,
                              )
                              cl.enqueue_copy(
                                  CL.queue, data, buf, is_blocking=True
                              )
                              if FLOAT16:
                                  data = data.astype(np.float16)
                              weights.append(data.tobytes())
                      else:
                          raise Exception("unknown object", a)
                      # print(jdat['objects'][-1])
                      saved_objs.add(ptr)
                  targs.append(ptr)
                  args_size.append(8)
              else:
                  raise Exception("idk this type")
  
          # save the kernel itself
          jdat["kernels"].append(
              {
                  "name": prg.name,
                  "work_dim": len(args[0]),
                  "global_work_size": args[0],
                  # TODO: C++ thneed requires a local_work_size, so we fill it with ones
                  "local_work_size": [1 for _ in args[0]]
                  if args[1] is None
                  else args[1],
                  "num_args": len(args) - 2,
                  "args": targs,
                  "args_size": args_size,
              }
          )
  
      jdat["outputs"] = [
          {
              "buffer_id": struct.pack("Q", x.global_id).decode(
                  "latin_1"
              ),
              "size": x.size,
          }
          for x in self.outputs
      ]
  
      jdat["inputs"] = [
          {
              "buffer_id": struct.pack("Q", v.global_id).decode(
                  "latin_1"
              ),
              "size": v.size,
              "name": k,
          }
          for k, v in self.inputs.items()
      ][::-1]
  
      print(f"saving thneed to {output_fn}")
      with open(output_fn, "wb") as f:
          j = json.dumps(jdat, ensure_ascii=False).encode("latin_1")
          f.write(struct.pack("I", len(j)))
          f.write(j)
          f.write(b"".join(weights))
          f.write(b"".join(binaries))
  
file_path: tinygrad/extra/thneed.py
imports: import time
  import struct
  import json
  import traceback
  import numpy
  from tinygrad.runtime.ops_gpu import CLProgram, compile_gpu
  from tinygrad.device import Device
  from tinygrad.helpers import DEBUG, getenv
  from collections import defaultdict
  import pyopencl
  from tinygrad.runtime.ops_gpu import OSX_TIMING_RATIO
label_desc: Use image node_id 3 for calling the Thneed obj's underlying member method code with example usage: obj.save(output_fn) without return types