(
ctx: ModuleContext, trace_result: TraceResult, *args: Sequence[ir.Value],
)
| 101 | |
| 102 | |
| 103 | def lowering_ops( |
| 104 | ctx: ModuleContext, trace_result: TraceResult, *args: Sequence[ir.Value], |
| 105 | ): |
| 106 | # var_id -> ir.Value |
| 107 | env: Dict[int, Tuple[ir.Value, ...]] = {} |
| 108 | consts = list(map(ir_constant_tuple, trace_result._var_consts)) |
| 109 | |
| 110 | # read ir.Values from env according to var_ids |
| 111 | def read(var_ids): |
| 112 | assert isinstance(var_ids, (list, tuple)) |
| 113 | ret = [] |
| 114 | for vid in var_ids: |
| 115 | assert isinstance(vid, int) |
| 116 | ret.append(env[vid]) |
| 117 | return ret |
| 118 | |
| 119 | # update env with var_ids and ir.Values |
| 120 | def write(var_ids, hlo_nodes): |
| 121 | assert isinstance(var_ids, (list, tuple)) |
| 122 | assert isinstance(hlo_nodes, (map, list, tuple)) |
| 123 | hlo_nodes = list(hlo_nodes) |
| 124 | assert len(var_ids) == len(hlo_nodes), (len(var_ids), len(hlo_nodes)) |
| 125 | for vid, node in zip(var_ids, hlo_nodes): |
| 126 | assert vid not in env |
| 127 | env[vid] = node |
| 128 | |
| 129 | assert len(args) == len(trace_result.inputs) |
| 130 | assert len(consts) == len(trace_result.consts) |
| 131 | assert all(isinstance(v, ir.Value) for vs in consts for v in vs) |
| 132 | |
| 133 | # initialize env with inputs and consts |
| 134 | write(trace_result.inputs, args) |
| 135 | write(trace_result.consts, consts) |
| 136 | |
| 137 | for eqn in trace_result.eqns: |
| 138 | rule_ctx = LoweringRuleContext( |
| 139 | module_context=ctx, |
| 140 | op=eqn.op, |
| 141 | vars_in=[trace_result.vars[inp] for inp in eqn.inputs], |
| 142 | vars_out=[trace_result.vars[oup] for oup in eqn.outputs], |
| 143 | param=eqn.param, |
| 144 | ) |
| 145 | rule = get_rule(eqn.op, use_fake_rule_for_debug=False) |
| 146 | |
| 147 | in_nodes = read(eqn.inputs) |
| 148 | hinps = [ |
| 149 | HLOTensor(irval, var.shape, var.dtype) |
| 150 | for var, irval in zip( |
| 151 | rule_ctx.vars_in, map(_unwrap_singleton_ir_values, in_nodes) |
| 152 | ) |
| 153 | ] |
| 154 | houps = rule(rule_ctx, *hinps) |
| 155 | if isinstance(houps, HLOTensor): |
| 156 | houps = [houps] |
| 157 | |
| 158 | out_nodes = [] |
| 159 | for out_id, hlo_out in zip(eqn.outputs, houps): |
| 160 | var_out = trace_result.vars[out_id] |
no test coverage detected