Class ImageDType

tinygrad/dtype.py:128–147

@dataclass(frozen=True, eq=False)
class ImageDType(PtrDType):
  shape: tuple[int, ...] = ()   # shape of the Image
  def ptr(self, size=-1, addrspace=AddrSpace.GLOBAL) -> PtrDType:
    assert addrspace == AddrSpace.GLOBAL, "images can't be local"
    return self
  def __repr__(self): return f"dtypes.{self.name}({self.shape})" + (f'.vec({self.v})' if self.v != 1 else '')

  # for 1d images on macos, we need to round pitch up to 256 pixels to make CL happy
  @property
  def pitch(self): return (round_up(self.shape[1], 256) if OSX else self.shape[1]) * 4 * self.itemsize

  # get list of (height, width) that do not require pitch padding
  @staticmethod
  def valid_dims(ptr:PtrDType, arch:str) -> list[tuple[int,int]]:
    if (ALIGN:=next((int(p.split('=')[1]) for p in arch.split(',') if p.startswith("IMAGE_PITCH_ALIGNMENT=")), 0)) == 0: return []
    MAXW, pxls = 16384, ptr.size // 4
    if ptr.base not in (dtypes.half, dtypes.float) or ptr.size > 4*MAXW*MAXW: return []
    # height=1 images just need to abide by alignment requirements in bytes, not pixels!
    if ptr.size % (ALIGN * 4) != 0: return [] if ptr.nbytes() % (64 if OSX else ALIGN) != 0 or pxls > MAXW else [(1, pxls)]
    return [(pxls//ALIGN//k, ALIGN*k) for k in range(ceildiv(pxls//ALIGN, MAXW), min(pxls//ALIGN, MAXW//ALIGN)+1) if (pxls//ALIGN)%k == 0]

class dtypes:
  @staticmethod