| 520 | } |
| 521 | |
| 522 | void CustomOp::compute( |
| 523 | const std::vector<Tensor>& inputs, const Param& param, |
| 524 | std::vector<Tensor>& outputs) const { |
| 525 | assert_inputs_size_right(inputs); |
| 526 | assert_outputs_size_right(outputs); |
| 527 | if (outputs.size() == 0) { |
| 528 | return; |
| 529 | } |
| 530 | |
| 531 | Device device = outputs[0].device(); |
| 532 | std::string device_str = device.str(); |
| 533 | for (size_t i = 1; i < outputs.size(); ++i) { |
| 534 | mgb_assert( |
| 535 | outputs[i].device().str() == device_str, |
| 536 | "all output tensors should have the same device attribute"); |
| 537 | } |
| 538 | |
| 539 | // need to add other input/output check |
| 540 | mgb_assert( |
| 541 | Device::is_legal(device_str), "unsupported device type: %s", |
| 542 | device_str.c_str()); |
| 543 | |
| 544 | auto preprocess_func = OpImplRef(m_impl.get())->preprocess_funcs[device_str]; |
| 545 | auto forward_func = OpImplRef(m_impl.get())->compute_funcs[device_str]; |
| 546 | auto postprocess_func = OpImplRef(m_impl.get())->postprocess_funcs[device_str]; |
| 547 | |
| 548 | RuntimeArgs rt_args(device); |
| 549 | |
| 550 | preprocess_func(inputs, param, outputs, rt_args); |
| 551 | forward_func(inputs, param, outputs, rt_args); |
| 552 | postprocess_func(outputs, param, outputs, rt_args); |
| 553 | assert_outputs_size_right(outputs); |
| 554 | } |
| 555 | |
| 556 | void compute_impl( |
| 557 | std::shared_ptr<const CustomOp> op, const Param& param, |
no test coverage detected