(reducer, inp, world_size)
| 40 | |
| 41 | |
| 42 | def _all_reduce(reducer, inp, world_size): |
| 43 | def _replica_groups_hlo(replica_groups: Sequence[Sequence[int]]): |
| 44 | groups = np.array( |
| 45 | list(itertools.zip_longest(*replica_groups, fillvalue=-1)), dtype=np.int64 |
| 46 | ).T |
| 47 | return ir.DenseIntElementsAttr.get(np.ascontiguousarray(groups)) |
| 48 | |
| 49 | replica_groups = _replica_groups_hlo([[i for i in range(world_size)]]) |
| 50 | hlo_cfgs = {} |
| 51 | |
| 52 | all_reduce_op = hlo.AllReduceOp( |
| 53 | inp.tensor.type, inp.tensor, replica_groups=replica_groups, **hlo_cfgs |
| 54 | ) |
| 55 | scalar_type = ir_utils.make_ir_type_according_meta(tuple(), inp.dtype) |
| 56 | reducer_region = all_reduce_op.regions[0].blocks.append(scalar_type, scalar_type) |
| 57 | with ir.InsertionPoint(reducer_region): |
| 58 | reducer_ret = reducer(*reducer_region.arguments) |
| 59 | hlo.ReturnOp(reducer_ret.results) |
| 60 | return HLOTensor(all_reduce_op.results) |
| 61 | |
| 62 | |
| 63 | all_reduce_sum = partial(_all_reduce, hlo.AddOp) |
nothing calls this directly
no test coverage detected