| 230 | |
| 231 | # Layout complement |
| 232 | def complement(layout, max_idx=1): |
| 233 | if is_int(layout): |
| 234 | return complement(Layout(layout)) |
| 235 | |
| 236 | result_shape = [] |
| 237 | result_stride = [] |
| 238 | current_idx = 1 |
| 239 | |
| 240 | sorted_DS = sorted(zip(flatten(layout.stride), flatten(layout.shape))) |
| 241 | for (stride, shape) in sorted_DS: |
| 242 | if stride == 0 or shape == 1: |
| 243 | continue |
| 244 | |
| 245 | in_bound = current_idx <= shape * stride |
| 246 | # To support symbolic value which can't be evaluated now |
| 247 | assert (type(in_bound) is not bool) or in_bound |
| 248 | |
| 249 | result_shape.append(stride // current_idx) |
| 250 | result_stride.append(current_idx) |
| 251 | current_idx = shape * stride |
| 252 | |
| 253 | result_shape.append((max_idx + current_idx - 1) // current_idx) # ceil_div |
| 254 | result_stride.append(current_idx) |
| 255 | |
| 256 | return coalesce(Layout(tuple(result_shape), tuple(result_stride))) |
| 257 | |
| 258 | |
| 259 | # Layout right inverse |