| 150 | } |
| 151 | |
| 152 | void run_mlir_broadcast(CompNode cn) { |
| 153 | set_backend(Backend::MLIR); |
| 154 | auto graph = ComputingGraph::make(); |
| 155 | HostTensorGenerator<dtype::Float32> gen; |
| 156 | |
| 157 | auto host_x0 = gen({10, 20, 5, 6}, cn), host_x1 = gen({1, 20, 1, 1}, cn), |
| 158 | host_x2 = gen({10, 1, 5, 1}, cn), host_x3 = gen({10, 1, 1, 1}, cn); |
| 159 | |
| 160 | auto a = opr::Host2DeviceCopy::make(*graph, host_x0), |
| 161 | b = opr::Host2DeviceCopy::make(*graph, host_x1), |
| 162 | c = opr::Host2DeviceCopy::make(*graph, host_x2), |
| 163 | d = opr::Host2DeviceCopy::make(*graph, host_x3); |
| 164 | |
| 165 | auto y = opr::Elemwise::make({a, b, c}, opr::Elemwise::Mode::FUSE_MUL_ADD3) + |
| 166 | opr::Elemwise::make({d}, opr::Elemwise::Mode::ABS) - 0.3f; |
| 167 | |
| 168 | auto ig_gen = std::make_unique<InternalGraphGenerator>(y.node()->owner_opr()); |
| 169 | |
| 170 | for (auto i : get_rev_topo_order(y)) { |
| 171 | if (!i->same_type<opr::Host2DeviceCopy>()) { |
| 172 | ig_gen->add_opr(i); |
| 173 | } |
| 174 | } |
| 175 | |
| 176 | auto igraph = ig_gen->generate(); |
| 177 | auto y_jit = JITExecutor::make(igraph, ig_gen->orig_inps()); |
| 178 | |
| 179 | HostTensorND host_y, host_y_jit; |
| 180 | auto func = graph->compile( |
| 181 | {make_callback_copy(y, host_y), make_callback_copy(y_jit, host_y_jit)}); |
| 182 | func->execute(); |
| 183 | |
| 184 | MGB_ASSERT_TENSOR_EQ(host_y, host_y_jit); |
| 185 | } |
| 186 | |
| 187 | void run_mlir_different_shape(CompNode cn) { |
| 188 | set_backend(Backend::MLIR); |