| 388 | } |
| 389 | |
| 390 | void CompiledTransformation::compile() { |
| 391 | // these ops require seq order, so we link them to an mm_io_link to ensure order |
| 392 | static std::unordered_set<Typeinfo*> mm_io_ops = { |
| 393 | CollectiveComm::typeinfo(), RemoteSend::typeinfo(), RemoteRecv::typeinfo()}; |
| 394 | mgb_assert(!m_executable, "already compiled"); |
| 395 | // FIXME: mm_io_link and io_links should be merged |
| 396 | SymbolVarArray io_links; |
| 397 | SymbolVar mm_io_link; |
| 398 | auto make_input = [&](VarInfo* var_info) { |
| 399 | mgb_assert( |
| 400 | var_info->kind == VarKind::External, "input node should be external"); |
| 401 | VarAccessor accessor; |
| 402 | auto box = make_box<DeviceTensorND>(); |
| 403 | // TODO: attach ref count, release early |
| 404 | auto outputs = opr::InputCallback::make( |
| 405 | *m_graph, [box] { return box->take_value(); }, *var_info->device, |
| 406 | *var_info->dtype, var_info->shape, io_links, m_input_shape_static); |
| 407 | // attach input_callback to io_links |
| 408 | accessor.node = outputs[0].node(); |
| 409 | io_links = {outputs[1]}; |
| 410 | accessor.data_setter = [box](DeviceTensorND data) { box->try_set_value(data); }; |
| 411 | return accessor; |
| 412 | }; |
| 413 | auto make_output = [&](TraceResult::VarInfo* var_info, SymbolVar node) { |
| 414 | VarAccessor accessor; |
| 415 | accessor.node = node.node(); |
| 416 | if (auto bound_data = var_info->bound_data) { |
| 417 | accessor.shape_getter = [bound_data]() -> TensorShape { |
| 418 | return bound_data.shape()->as_tensor_shape(); |
| 419 | }; |
| 420 | accessor.data_getter = [bound_data]() -> DeviceTensorND { |
| 421 | return bound_data.dev_tensor()->as_nd(); |
| 422 | }; |
| 423 | accessor.value_getter = [bound_data]() -> HostTensorND { |
| 424 | return bound_data.numpy()->as_nd(); |
| 425 | }; |
| 426 | return accessor; |
| 427 | } |
| 428 | if (var_info->data_required) { |
| 429 | // reduce d2h when data is available |
| 430 | // FIXME: compile should not change var_info in-place |
| 431 | var_info->shape_required = false; |
| 432 | } |
| 433 | if (var_info->shape_required) { |
| 434 | // TODO: use static infer manager for some vars? |
| 435 | auto box = make_box<TensorShape>(); |
| 436 | auto callback = [box](DeviceTensorND data) { |
| 437 | box->try_set_value(data.shape()); |
| 438 | }; |
| 439 | SymbolVarArray inputs = io_links; |
| 440 | inputs.insert(inputs.begin(), node); |
| 441 | auto output = opr::OutputCallback::make({callback, true, false}, inputs); |
| 442 | io_links = {output}; |
| 443 | accessor.shape_getter = [box]() -> TensorShape { return box->get_value(); }; |
| 444 | } |
| 445 | if (var_info->data_required) { |
| 446 | auto box = make_box<DeviceTensorND>(); |
| 447 | auto callback = [box](DeviceTensorND data) { box->try_set_value(data); }; |