| 394 | std::unique_ptr<cg::static_infer::StaticInferUpdater> updater; |
| 395 | |
| 396 | void initialize( |
| 397 | const CompiledOp& op, const SmallVector<LogicalTensorDesc>& input_descs) { |
| 398 | allocator = std::make_shared<DeviceMemoryAllocatorImpl>(); |
| 399 | graph = ComputingGraph::make(); |
| 400 | graph->options().force_dynamic_alloc = true; |
| 401 | graph->options().async_exec_level = 0; |
| 402 | graph->options().graph_opt_level = op.gopt_level; |
| 403 | graph->options().enable_var_mem_defragment = false; |
| 404 | graph->options().comp_seq_sync_device = false; |
| 405 | // set allocator for DTR support |
| 406 | graph->set_device_memory_allocator(allocator); |
| 407 | if constexpr (Kind == HolderKind::ShapeInfer) { |
| 408 | updater = cg::static_infer::StaticInferUpdater::make(); |
| 409 | } |
| 410 | for (auto&& desc : input_descs) { |
| 411 | Input input; |
| 412 | VarNode* input_var = nullptr; |
| 413 | if constexpr (Kind == HolderKind::Execute) { |
| 414 | input.device_value = std::make_shared<DeviceTensorND>(); |
| 415 | input.device_value->dtype(desc.layout.dtype); |
| 416 | input.device_value->comp_node(desc.comp_node); |
| 417 | input.device_value->resize(desc.layout); |
| 418 | auto callback = [value = input.device_value] { return *value; }; |
| 419 | if (!desc.value.empty()) { |
| 420 | input.host_value = std::make_shared<HostTensorND>(); |
| 421 | input.host_value->dtype(desc.layout.dtype); |
| 422 | input.host_value->comp_node(desc.comp_node); |
| 423 | } |
| 424 | input_var = opr::MutableTensor::make( |
| 425 | *graph, input.device_value, input.host_value, {}) |
| 426 | .node(); |
| 427 | // input_var = opr::VolatileSharedDeviceTensor::make(*graph, |
| 428 | // input.device_value).node(); |
| 429 | } else if constexpr (Kind == HolderKind::ShapeInfer) { |
| 430 | if (desc.value.empty()) { |
| 431 | input.host_shape = std::make_shared<HostTensorND>(); |
| 432 | input.host_shape->dtype(dtype::Int32()); |
| 433 | input.host_shape->comp_node(desc.comp_node); |
| 434 | auto input_shape_var = |
| 435 | opr::Host2DeviceCopy::make(*graph, input.host_shape); |
| 436 | input_var = |
| 437 | opr::Alloc::make(input_shape_var, desc.layout.dtype).node(); |
| 438 | } else { |
| 439 | input.host_value = std::make_shared<HostTensorND>(); |
| 440 | input.host_value->dtype(desc.layout.dtype); |
| 441 | input.host_value->comp_node(desc.comp_node); |
| 442 | input_var = |
| 443 | opr::Host2DeviceCopy::make(*graph, input.host_value).node(); |
| 444 | } |
| 445 | } else { |
| 446 | static_assert((Kind != Kind), "unknown holder kind"); |
| 447 | } |
| 448 | input_vars.push_back(input_var); |
| 449 | inputs.push_back(input); |
| 450 | } |
| 451 | // forward to inner op |
| 452 | output_vars = OpDef::apply_on_var_node(*op.op, input_vars); |
| 453 | ComputingGraph::OutputSpec output_spec; |
no test coverage detected