(
mge_traced,
func_name=None,
device=None,
keep_unused=True,
donate_invars=None,
verbose=int(os.environ.get("MGE_VERBOSE_XLA_IR", "0")),
return_with_io=False,
return_device_array=False,
ip: str = None,
port: int = None,
)
| 16 | |
| 17 | |
| 18 | def build_xla( |
| 19 | mge_traced, |
| 20 | func_name=None, |
| 21 | device=None, |
| 22 | keep_unused=True, |
| 23 | donate_invars=None, |
| 24 | verbose=int(os.environ.get("MGE_VERBOSE_XLA_IR", "0")), |
| 25 | return_with_io=False, |
| 26 | return_device_array=False, |
| 27 | ip: str = None, |
| 28 | port: int = None, |
| 29 | ): |
| 30 | assert device == None, "cannot specify device now" |
| 31 | assert keep_unused == True, "keep_unused error" |
| 32 | assert donate_invars == None, "donate_invars error" |
| 33 | |
| 34 | # normalize megengine trace result for lowering |
| 35 | tr = TraceResult(mge_traced, func_name) |
| 36 | tr = RngKeyAdder()(tr) |
| 37 | tr = DropoutMaskCanonicalizer()(tr) |
| 38 | |
| 39 | if verbose > 0 and get_rank() == 0: |
| 40 | print("================ Mge Trace Result ================") |
| 41 | print(tr) |
| 42 | |
| 43 | in_is_global = (True,) * len(tr.inputs) |
| 44 | kept_var_idx = set(range(len(tr.inputs))) if keep_unused else set() |
| 45 | |
| 46 | # init for xla distributed and setup device |
| 47 | if is_distributed(): |
| 48 | initialize(ip, port, get_world_size(), get_rank(), [get_local_device_id()]) |
| 49 | backend, device_assignment, platform = get_xla_backend_and_device(device) |
| 50 | |
| 51 | module, keepalive, host_callbacks = lower( |
| 52 | tr, backend, platform, None, None, donate_invars, |
| 53 | ) |
| 54 | |
| 55 | if not is_distributed(): |
| 56 | # setup sharding information |
| 57 | in_shardings = make_unspec_sharding(tr.inputs) |
| 58 | out_shardings = make_unspec_sharding(tr.outputs) |
| 59 | |
| 60 | in_shardings = tuple( |
| 61 | OpShardingSharding.get_replicated(device_assignment) |
| 62 | if _is_unspecified(i) |
| 63 | else i |
| 64 | for i in in_shardings |
| 65 | ) |
| 66 | |
| 67 | computation = MeshComputation( |
| 68 | tr.func_name, |
| 69 | module, |
| 70 | donated_invars=donate_invars, |
| 71 | trace_result=tr, |
| 72 | mesh=None, |
| 73 | in_shardings=in_shardings, |
| 74 | out_shardings=out_shardings, |
| 75 | spmd_lowering=False, |
no test coverage detected