| 97 | def __repr__(self): return f"<multibuf real:{self.is_allocated()} device:{tuple(x.device for x in self.bufs)} size:{self.size} dtype:{self.dtype}>" |
| 98 | |
| 99 | class Buffer: |
| 100 | profile_events:list[ProfileEvent] = [] |
| 101 | def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:BufferSpec|None=None, initial_value:bytes|None=None, |
| 102 | uop_refcount=0, base:Buffer|None=None, offset:int=0, preallocate=False): |
| 103 | assert isinstance(dtype, DType) and not isinstance(dtype, PtrDType) |
| 104 | self.device, self.size, self.dtype, self.options, self.offset, self.allocated_views = device, size, dtype, options, offset, 0 |
| 105 | self._bufs: dict[str, Any] = {} |
| 106 | if base is None: |
| 107 | assert offset == 0, "base buffers can't have offset" |
| 108 | self._base = None |
| 109 | self._uop_refcount = uop_refcount |
| 110 | if opaque is not None: self.allocate(opaque) |
| 111 | if initial_value is not None: |
| 112 | self.allocate() |
| 113 | self.copyin(memoryview(initial_value)) |
| 114 | else: |
| 115 | assert base._base is None, "base can't have a base" |
| 116 | assert device == base.device, "base must have the same device" |
| 117 | self._base = base |
| 118 | if preallocate: self.allocate() |
| 119 | @property |
| 120 | def base(self) -> Buffer: return self._base if self._base is not None else self |
| 121 | @property |
| 122 | def uop_refcount(self): return self.base._uop_refcount |
| 123 | @property |
| 124 | def _buf(self) -> Any: return self._bufs[self.device] |
| 125 | def ref(self, cnt): |
| 126 | self.base._uop_refcount += cnt |
| 127 | return self |
| 128 | # check if the underlying buffer is allocated and the current buffer/view is initialized |
| 129 | def is_initialized(self) -> bool: return self.is_allocated() and self.device in self._bufs |
| 130 | # check if the underlying buffer is allocated, possibly from the base object |
| 131 | def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else self.device in self._bufs |
| 132 | def get_buf(self, device: str) -> Any: |
| 133 | if device not in self._bufs: |
| 134 | allocator = Device[device].allocator |
| 135 | if device == self.device: self.ensure_allocated() |
| 136 | elif self._base is not None: |
| 137 | assert hasattr(allocator, "_offset"), "offset function required for view" |
| 138 | self._bufs[device] = allocator._offset(self._base.get_buf(device), self.nbytes, self.offset) |
| 139 | else: self._bufs[device] = allocator._map(self.ensure_allocated()._buf) |
| 140 | return self._bufs[device] |
| 141 | def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self |
| 142 | def allocate(self, opaque=None, external_ptr=None) -> Buffer: |
| 143 | assert not self.is_initialized(), "can't allocate already allocated buffer" |
| 144 | if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}") |
| 145 | if not self.device.startswith("NULL") and self.size > MAX_BUFFER_SIZE > 0 and (self.options is None or self.options.external_ptr is None): |
| 146 | raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large") |
| 147 | self.allocator:Allocator = Device[self.device].allocator |
| 148 | if external_ptr is not None: |
| 149 | self.options = replace(self.options, external_ptr=external_ptr) if self.options else BufferSpec(external_ptr=external_ptr) |
| 150 | if self._base is not None: |
| 151 | self._base.ensure_allocated() |
| 152 | self._base.allocated_views += 1 |
| 153 | assert hasattr(self.allocator, "_offset"), "offset function required for view" |
| 154 | self._bufs[self.device] = self.allocator._offset(self.base._buf, self.nbytes, self.offset) |
| 155 | else: |
| 156 | self._bufs[self.device] = opaque if opaque is not None else self.allocator.alloc(self.nbytes, self.options) |
no outgoing calls
searching dependent graphs…