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

Function exec

imperative/src/impl/ops/rng.cpp:495–526  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

493
494template <typename Op>
495void exec(
496 const OpDef& op, const SmallVector<TensorPtr>& inputs,
497 const SmallVector<TensorPtr>& outputs,
498 const SmallVector<TensorPtr>& workspace) {
499 auto&& rng = op.cast_final_safe<Op>();
500
501 auto dest = outputs[0];
502 if (dest->layout().is_empty())
503 return;
504 auto cn = dest->comp_node();
505 auto handle = rng.handle;
506 if (!handle) {
507 handle = RNGDnnOpManager::get_default_handle(cn);
508 }
509
510 // retrieve dnn_op from glob cache
511 auto dnn_op_thread_safe =
512 RNGDnnOpManager::inst().get_dnn_op<typename OpMeth<Op>::DnnOp>(
513 handle, reinterpret_cast<size_t>(op.dyn_typeinfo()), cn);
514 auto initialized = std::get<0>(dnn_op_thread_safe);
515 auto dnn_op = std::get<1>(dnn_op_thread_safe);
516 if (initialized) {
517 auto handle_seed = RNGDnnOpManager::get_seed(handle);
518 mgb_assert(
519 dnn_op->param().seed == handle_seed,
520 "inconsistent rng seed: handle: %lu, dnn_op: %lu", handle_seed,
521 dnn_op->param().seed);
522 }
523 dnn_op->param() = OpMeth<Op>::make_param(rng);
524 _RNGOprInvoker<OpMeth<Op>::DnnOp::NR_INPUTS, OpMeth<Op>::DnnOp::NR_OUTPUTS>::exec(
525 dnn_op, inputs, outputs);
526}
527
528template <typename Op>
529SmallVector<LogicalTensorDesc> infer_output_attrs(

Callers 3

setup.pyFile · 0.50
get_mge_versionFunction · 0.50
PYBIND11_MODULEFunction · 0.50

Calls 5

instFunction · 0.85
is_emptyMethod · 0.45
layoutMethod · 0.45
comp_nodeMethod · 0.45
paramMethod · 0.45

Tested by

no test coverage detected