Function run_dimshuffle

src/jit/test/codegen.cpp:515–541

/* ===================== TestJITMlirDimshuffle ===================== */

void run_dimshuffle(CompNode cn, TensorShape ishape, const std::vector<int>& pattern) {
    set_backend(Backend::MLIR);
    auto graph = ComputingGraph::make();
    HostTensorGenerator<> gen;

    auto host_x = gen(ishape, cn);
    auto x = opr::Host2DeviceCopy::make(*graph, host_x);
    auto y = opr::Dimshuffle::make(x, pattern);

    auto ig_gen = std::make_unique<InternalGraphGenerator>(y.node()->owner_opr());

    for (auto i : get_rev_topo_order(y)) {
        if (!i->template same_type<opr::Host2DeviceCopy>()) {
            ig_gen->add_opr(i);
        }
    }

    auto igraph = ig_gen->generate();
    auto y_jit = JITExecutor::make(igraph, ig_gen->orig_inps());

    HostTensorND host_y, host_y_jit;
    auto func = graph->compile(
            {make_callback_copy(y, host_y), make_callback_copy(y_jit, host_y_jit)});
    func->execute();

    MGB_ASSERT_TENSOR_EQ(host_y, host_y_jit);
}

void run_dimshuffle_cases(CompNode cn) {
    run_dimshuffle(cn, {3, 4, 5}, {2, 0, 1});