MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / dispatch_kernel

Method dispatch_kernel

imperative/src/impl/interpreter/interpreter_impl.cpp:352–428  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

350}
351
352void ChannelImpl::dispatch_kernel(
353 std::shared_ptr<OpDef> op, const SmallVector<TensorInfo*>& input_infos,
354 const SmallVector<LogicalTensorDesc>& input_descs,
355 SmallVector<Handle>* outputs) {
356 auto& state = get_channel_state();
357 auto& options = state.options;
358
359 std::optional<StackManager::Guard> guard;
360 if (Profiler::is_profiling()) {
361 guard.emplace(op->trait()->make_name(*op), &state.stack_manager);
362 }
363
364 auto [output_descs, validated] =
365 OpDef::infer_output_attrs_fallible(*op, input_descs);
366 MGB_RECORD_EVENT(ShapeInferEvent, validated);
367
368 SmallVector<TensorInfo*> output_infos;
369 output_infos.reserve(output_descs.size());
370
371 outputs->reserve(output_descs.size());
372 for (int i = 0; i < output_descs.size(); ++i) {
373 auto&& desc = output_descs[i];
374 auto info = alloc();
375 init(info, std::move(desc));
376 // make sure desc's value is consistent with h_value
377 if (!info->desc.value.empty()) {
378 info->h_value = HostTensorND::make_proxy(info->desc.value)
379 .proxy_to_comp_node(info->desc.comp_node);
380 }
381 output_infos.push_back(info);
382 outputs->push_back(reinterpret_cast<Handle>(info));
383 }
384 auto& bt = get_backtrace();
385 ApplyOp cmd{Profiler::next_id(), std::move(op), std::move(input_infos),
386 std::move(output_infos), validated, bt};
387 if (Profiler::is_profiling()) {
388 auto op_info_getter = [op = cmd.op, bt = cmd.bt] {
389 std::unordered_map<std::string, std::string> op_info;
390 auto props = OpDef::props(*op);
391 for (auto&& [key, value] : props) {
392 op_info[key] = value;
393 }
394 if (bt != nullptr) {
395 if (bt->py_stack_info != nullptr)
396 op_info["python_backtrace"] = bt->py_traceback();
397 if (bt->trans_stack_info.size() > 0)
398 op_info["transformation_backtrace"] =
399 bt->transformation_traceback();
400 }
401 return op_info;
402 };
403 MGB_RECORD_EVENT(
404 OpDispatchEvent, cmd.id, guard.value().name(), op_info_getter,
405 tinfo_to_tid(cmd.inputs), tinfo_to_tid(cmd.outputs),
406 state.stack_manager.dump());
407 m_worker.add_task(
408 {Profiler::next_id(), std::move(cmd),
409 get_channel_state().stack_manager.dump()});

Callers

nothing calls this directly

Calls 15

is_profilingFunction · 0.85
allocFunction · 0.85
tinfo_to_tidFunction · 0.85
proxy_to_comp_nodeMethod · 0.80
py_tracebackMethod · 0.80
check_async_errorMethod · 0.80
initFunction · 0.50
propsFunction · 0.50
emplaceMethod · 0.45
make_nameMethod · 0.45

Tested by

no test coverage detected