| 117 | |
| 118 | #if MGB_JIT_MLIR |
| 119 | void run_mlir(CompNode cn) { |
| 120 | set_backend(Backend::MLIR); |
| 121 | auto graph = ComputingGraph::make(); |
| 122 | HostTensorGenerator<dtype::Float32> gen; |
| 123 | |
| 124 | auto host_x0 = gen({23, 42}, cn), host_x1 = gen({23, 1}, cn), |
| 125 | host_x2 = gen({23, 42}, cn); |
| 126 | |
| 127 | auto a = opr::Host2DeviceCopy::make(*graph, host_x0), |
| 128 | b = opr::Host2DeviceCopy::make(*graph, host_x1), |
| 129 | c = opr::Host2DeviceCopy::make(*graph, host_x2); |
| 130 | |
| 131 | auto y = a + b * c + 0.3f; |
| 132 | |
| 133 | auto ig_gen = std::make_unique<InternalGraphGenerator>(y.node()->owner_opr()); |
| 134 | |
| 135 | for (auto i : get_rev_topo_order(y)) { |
| 136 | if (!i->same_type<opr::Host2DeviceCopy>()) { |
| 137 | ig_gen->add_opr(i); |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | auto igraph = ig_gen->generate(); |
| 142 | auto y_jit = JITExecutor::make(igraph, ig_gen->orig_inps()); |
| 143 | |
| 144 | HostTensorND host_y, host_y_jit; |
| 145 | auto func = graph->compile( |
| 146 | {make_callback_copy(y, host_y), make_callback_copy(y_jit, host_y_jit)}); |
| 147 | func->execute(); |
| 148 | |
| 149 | MGB_ASSERT_TENSOR_EQ(host_y, host_y_jit); |
| 150 | } |
| 151 | |
| 152 | void run_mlir_broadcast(CompNode cn) { |
| 153 | set_backend(Backend::MLIR); |