| 629 | } |
| 630 | |
| 631 | bool CompNode::contain_flag(DeviceType device_type, Flag flag) { |
| 632 | Flag cn_flag{}; |
| 633 | switch (resolve_device_type(device_type)) { |
| 634 | case DeviceType::CUDA: |
| 635 | cn_flag = CudaCompNode::sm_flag; |
| 636 | break; |
| 637 | case DeviceType::MULTITHREAD: |
| 638 | case DeviceType::CPU: |
| 639 | cn_flag = CpuCompNode::sm_flag; |
| 640 | break; |
| 641 | case DeviceType::ROCM: |
| 642 | cn_flag = ROCmCompNode::sm_flag; |
| 643 | break; |
| 644 | case DeviceType::CAMBRICON: |
| 645 | cn_flag = CambriconCompNode::sm_flag; |
| 646 | break; |
| 647 | case DeviceType::ATLAS: |
| 648 | cn_flag = AtlasCompNode::sm_flag; |
| 649 | break; |
| 650 | default: |
| 651 | mgb_throw(MegBrainError, "unexpected device type"); |
| 652 | } |
| 653 | return static_cast<bool>(cn_flag & flag); |
| 654 | } |
| 655 | |
| 656 | CompNode CompNode::change_stream(int dest_stream) const { |
| 657 | mgb_assert(m_impl); |