(
ctx: ModuleContext,
name: str,
trace_result: TraceResult,
public: bool = True,
in_shardings: Optional[Sequence[Optional[xc.OpSharding]]] = None,
out_shardings: Optional[Sequence[Optional[xc.OpSharding]]] = None,
input_output_aliases: Optional[Sequence[Optional[int]]] = None,
)
| 171 | |
| 172 | |
| 173 | def make_xla_graph( |
| 174 | ctx: ModuleContext, |
| 175 | name: str, |
| 176 | trace_result: TraceResult, |
| 177 | public: bool = True, |
| 178 | in_shardings: Optional[Sequence[Optional[xc.OpSharding]]] = None, |
| 179 | out_shardings: Optional[Sequence[Optional[xc.OpSharding]]] = None, |
| 180 | input_output_aliases: Optional[Sequence[Optional[int]]] = None, |
| 181 | ) -> func_dialect.FuncOp: |
| 182 | assert public is True, "do not process the visibitity of function" |
| 183 | assert ( |
| 184 | in_shardings is None and out_shardings is None |
| 185 | ), "sharding when lowering is not supported yet" |
| 186 | assert ( |
| 187 | input_output_aliases is None or input_output_aliases == [] |
| 188 | ), "donated inputs are not supported yet" |
| 189 | |
| 190 | input_types = [ |
| 191 | mge_varinfo_to_ir_type_tuple(trace_result.vars[idx]) |
| 192 | for idx in trace_result.inputs |
| 193 | ] |
| 194 | output_types = [ |
| 195 | mge_varinfo_to_ir_type_tuple(trace_result.vars[idx]) |
| 196 | for idx in trace_result.outputs |
| 197 | ] |
| 198 | |
| 199 | flat_input_types = utils.flatten_list(input_types) |
| 200 | flat_output_types = utils.flatten_list(output_types) |
| 201 | assert len(flat_input_types) == len(trace_result.inputs) |
| 202 | assert len(flat_output_types) == len(trace_result.outputs) |
| 203 | |
| 204 | ftype = ir.FunctionType.get(flat_input_types, flat_output_types) |
| 205 | func_op = func_dialect.FuncOp(name, ftype, ip=ctx.ip) |
| 206 | func_op.attributes["sym_visibility"] = ir.StringAttr.get( |
| 207 | "public" if public else "private" |
| 208 | ) |
| 209 | ctx.symbol_table.insert(func_op) |
| 210 | |
| 211 | entry_block = func_op.add_entry_block() |
| 212 | with ir.InsertionPoint(entry_block): |
| 213 | flat_args = entry_block.arguments |
| 214 | unflattened_args = utils.unflatten_list(flat_args, map(len, input_types)) |
| 215 | outs = lowering_ops(ctx, trace_result, *unflattened_args) |
| 216 | flat_oups = utils.flatten_list(outs) |
| 217 | func_dialect.ReturnOp(flat_oups) |
| 218 | return func_op |
| 219 | |
| 220 | |
| 221 | def lower( |
no test coverage detected