(self, buffer: Buffer)
| 54 | self.var_map[old_buffer.data] = new_buffer.data |
| 55 | |
| 56 | def mutate_buffer(self, buffer: Buffer): |
| 57 | if buffer in self.buffer_map: |
| 58 | return self.buffer_map[buffer] |
| 59 | |
| 60 | # Track mutations for this specific buffer only. Without this reset, |
| 61 | # unrelated buffers can be spuriously cloned and introduce alias buffers. |
| 62 | prev_mutated = self.buffer_attr_var_mutated |
| 63 | self.buffer_attr_var_mutated = False |
| 64 | new_data = self.visit_expr(buffer.data) |
| 65 | new_shape = [self.visit_expr(expr) for expr in buffer.shape] |
| 66 | if isinstance(buffer.layout, TileLayout): |
| 67 | new_shard = [] |
| 68 | new_replicate = [] |
| 69 | for iter in buffer.layout.shard: |
| 70 | new_iter = Iter( |
| 71 | self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis |
| 72 | ) |
| 73 | new_shard.append(new_iter) |
| 74 | for iter in buffer.layout.replica: |
| 75 | new_iter = Iter( |
| 76 | self.visit_expr(iter.extent), self.visit_expr(iter.stride), iter.axis |
| 77 | ) |
| 78 | new_replicate.append(new_iter) |
| 79 | new_layout = TileLayout.from_iters( |
| 80 | new_shard, new_replicate, offset=buffer.layout.offset |
| 81 | ) |
| 82 | else: |
| 83 | new_layout = buffer.layout |
| 84 | buffer_attr_mutated = self.buffer_attr_var_mutated |
| 85 | self.buffer_attr_var_mutated = prev_mutated or buffer_attr_mutated |
| 86 | if not buffer_attr_mutated: |
| 87 | return None |
| 88 | new_buffer = decl_buffer( |
| 89 | new_shape, |
| 90 | buffer.dtype, |
| 91 | buffer.name, |
| 92 | new_data, |
| 93 | buffer.strides, |
| 94 | buffer.elem_offset, |
| 95 | buffer.scope(), |
| 96 | buffer.data_alignment, |
| 97 | buffer.offset_factor, |
| 98 | layout=new_layout, |
| 99 | ) |
| 100 | self.buffer_map[buffer] = new_buffer |
| 101 | return new_buffer |
| 102 | |
| 103 | def visit_var_(self, op: Var): |
| 104 | op = super().visit_var_(op) |
no test coverage detected