(self, shape, dtype, addrspace:AddrSpace, name:str|None=None)
| 61 | return rng |
| 62 | |
| 63 | def alloc(self, shape, dtype, addrspace:AddrSpace, name:str|None=None): |
| 64 | match addrspace: |
| 65 | case AddrSpace.GLOBAL: |
| 66 | slot = self.global_slot |
| 67 | self.global_slot += 1 |
| 68 | case AddrSpace.LOCAL: |
| 69 | slot = self.shared_slot |
| 70 | self.shared_slot += 1 |
| 71 | case AddrSpace.REG: |
| 72 | slot = self.register_slot |
| 73 | self.register_slot += 1 |
| 74 | |
| 75 | uop = UOp.placeholder(shape, dtype, slot=slot, addrspace=addrspace) |
| 76 | |
| 77 | if name: |
| 78 | if (name, shape) in self.allocs: return self.allocs[(name, shape)] |
| 79 | self.allocs[(name, shape)] = uop |
| 80 | |
| 81 | return uop |
| 82 | |
| 83 | def gl(self, shape, dtype): return GL.create(shape, dtype, self) |
| 84 | def st(self, shape, dtype, layout=TileLayout.ROW, base_shape=ST_16X16): return ST.create(shape, dtype, layout, base_shape, self) |
no test coverage detected