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

Function do_test_dynamic_split

src/opr/test/tensor_manip.cpp:1650–1778  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

1648namespace {
1649
1650void do_test_dynamic_split(bool multiple_cn, bool force_dynamic) {
1651 auto cns = load_multiple_xpus(3);
1652 constexpr size_t N = 2, C = 51;
1653 HostTensorGenerator<> gen;
1654 auto host_x = gen({N, C}, cns[0]), host_sub_begin = gen({1}, cns[0]),
1655 host_sub_end = gen({1}, cns[0]);
1656 host_sub_begin->ptr<float>()[0] = 0;
1657 host_sub_end->ptr<float>()[0] = 2;
1658 auto graph = ComputingGraph::make();
1659
1660 SymbolVar x = opr::Host2DeviceCopy::make_no_fwd(*graph, host_x, {"x"}),
1661 sub_begin = opr::Host2DeviceCopy::make_no_fwd(
1662 *graph, host_sub_begin, {"sub_begin"}),
1663 sub_end = opr::Host2DeviceCopy::make_no_fwd(
1664 *graph, host_sub_end, {"sub_end"}),
1665 xsub = opr::Subtensor::make(
1666 x, {opr::Subtensor::AxisIndexer::make_interval(
1667 1, sub_begin, sub_end, None)})
1668 .rename("xsub");
1669
1670 OperatorNodeConfig split_config("split");
1671 if (multiple_cn) {
1672 split_config.comp_node_arr({cns[1], cns[2]});
1673 }
1674
1675 if (force_dynamic)
1676 xsub = opr::MarkDynamicVar::make(xsub);
1677
1678 auto spl = opr::Split::make(
1679 xsub,
1680 Split::Options::make_callback(
1681 1, 2,
1682 [](size_t s) {
1683 return std::vector<size_t>{s / 2, s - s / 2};
1684 }),
1685 split_config);
1686
1687 if (multiple_cn) {
1688 spl[0] = opr::Sleep::make(spl[0], 0.1);
1689 spl[1] = opr::Sleep::make(spl[1], 0.2);
1690 }
1691 auto cost0 = opr::Dot::make(spl[0].flatten(), spl[0].flatten()),
1692 cost1_ = opr::Dot::make(spl[1].flatten(), spl[1].flatten()),
1693 cost1 = opr::Copy::make(cost1_, OperatorNodeConfig().follow_comp_node(cost0)),
1694 cost = opr::Copy::make(
1695 cost0 + cost1, OperatorNodeConfig().follow_comp_node(x)) *
1696 0.5f,
1697 grad = cg::grad(cost, x);
1698
1699 HostTensorND host_spl0, host_spl1, host_grad;
1700 auto func = graph->compile(
1701 {make_callback_copy(spl[0], host_spl0),
1702 make_callback_copy(spl[1], host_spl1),
1703 make_callback_copy(grad, host_grad)});
1704
1705 if (force_dynamic)
1706 ASSERT_TRUE(!cg::is_static_var_shape(spl[0].node()));
1707 else {

Callers 1

TESTFunction · 0.85

Calls 15

make_no_fwdFunction · 0.85
OperatorNodeConfigClass · 0.85
gradFunction · 0.85
make_callback_copyFunction · 0.85
is_static_var_shapeFunction · 0.85
TensorShapeClass · 0.85
prev_dev_ptrFunction · 0.85
renameMethod · 0.80
makeFunction · 0.70
genFunction · 0.50
flattenMethod · 0.45
compileMethod · 0.45

Tested by

no test coverage detected