| 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) |
| 157 | if not self.device.startswith("DISK") and (self.options is None or self.options.external_ptr is None): |
| 158 | GlobalCounters.mem_used += self.nbytes |
| 159 | GlobalCounters.mem_used_per_device[self.device] += self.nbytes |
| 160 | if PROFILE: Buffer.profile_events.append(ProfilePointEvent(self.device, "alloc", self.trace_num, {"dtype":self.dtype, "sz":self.size})) |
| 161 | return self |
| 162 | def deallocate(self): |
| 163 | assert self.device in self._bufs, "buffer must be allocated to deallocate" |
| 164 | if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}") |