| 832 | MGB_DYN_TYPE_OBJ_FINAL_IMPL(CondExecMerge); |
| 833 | |
| 834 | CondExecMerge::CondExecMerge( |
| 835 | const VarNodeArrayView& inputs, const VarNodeArrayView& out_shapes, |
| 836 | const Param& param, const OperatorNodeConfig& config) |
| 837 | : Super(inputs[0]->owner_graph(), config, "cond_merge", {}), m_param{param} { |
| 838 | mgb_throw_if( |
| 839 | inputs.size() % param.nr_output, GraphError, |
| 840 | "input size can not divide nr_output: %zu %u", inputs.size(), |
| 841 | param.nr_output); |
| 842 | auto global_registry = CondExecPred::GlobalRegistry::get(*owner_graph()); |
| 843 | auto nr_branch = inputs.size() / param.nr_output; |
| 844 | mgb_assert(param.nr_output); |
| 845 | for (size_t i = 0; i < param.nr_output; ++i) { |
| 846 | auto ovar = add_output(ssprintf("out%zu", i)); |
| 847 | ovar->dtype(inputs[i]->dtype()); |
| 848 | // disable system memory allocation because: |
| 849 | // 1. we can directly forward input storage to output |
| 850 | // 2. dynamic allocator would wait for all inputs to become ready (see |
| 851 | // VarNodeMemManager::DynamicAllocOprInfo::host_wait_input_ready), |
| 852 | // which would cause infinite waiting for unselected inputs. |
| 853 | ovar->add_flag(VarNode::Flag::NO_SYS_MEM_ALLOC) |
| 854 | .add_flag(VarNode::Flag::ALLOW_EMPTY_SHAPE); |
| 855 | } |
| 856 | |
| 857 | MGB_MARK_USED_VAR(mask2str); |
| 858 | m_branch_masks.resize(nr_branch, nullptr); |
| 859 | for (size_t i = 0; i < nr_branch; ++i) { |
| 860 | ExecutionMask* br_mask = nullptr; |
| 861 | for (size_t j = 0; j < param.nr_output; ++j) { |
| 862 | auto ivar = inputs[i * param.nr_output + j]; |
| 863 | auto mask = global_registry->require_mask_from_var(ivar); |
| 864 | mgb_throw_if( |
| 865 | output(j)->dtype() != ivar->dtype(), GraphError, |
| 866 | "CondExecMerge input dtypes mismatch: branch=%zu %s vs %s", i, |
| 867 | output(j)->dtype().name(), ivar->dtype().name()); |
| 868 | if (!j) { |
| 869 | br_mask = mask; |
| 870 | } else { |
| 871 | mgb_throw_if( |
| 872 | br_mask != mask, GraphError, |
| 873 | "CondExecMerge branch %zu have different masks: " |
| 874 | "%s vs %s", |
| 875 | i, mask2str(br_mask).c_str(), mask2str(mask).c_str()); |
| 876 | } |
| 877 | // this flag is added by ExecutionMask; we require flag because |
| 878 | // output var might forward input var storage |
| 879 | mgb_assert(ivar->contain_flag(VarNode::Flag::NO_SYS_STATIC_MEM_ALLOC)); |
| 880 | add_input({ivar}); |
| 881 | } |
| 882 | m_branch_masks[i] = br_mask; |
| 883 | } |
| 884 | add_equivalence_component<PODHash<Param>>(&m_param); |
| 885 | |
| 886 | // handle extra inputs for special modes |
| 887 | if (param.mode == Mode::SUM) { |
| 888 | mgb_assert(out_shapes.size() == param.nr_output); |
| 889 | for (auto i : out_shapes) { |
| 890 | add_input({i}); |
| 891 | } |
nothing calls this directly
no test coverage detected