| 41 | |
| 42 | template<typename T> |
| 43 | Array<T> createSelectNode(const Array<char> &cond, const Array<T> &a, |
| 44 | const Array<T> &b, const af::dim4 &odims) { |
| 45 | auto cond_node = cond.getNode(); |
| 46 | auto a_node = a.getNode(); |
| 47 | auto b_node = b.getNode(); |
| 48 | auto a_height = a_node->getHeight(); |
| 49 | auto b_height = b_node->getHeight(); |
| 50 | auto cond_height = cond_node->getHeight(); |
| 51 | const int height = max(max(a_height, b_height), cond_height) + 1; |
| 52 | |
| 53 | auto node = make_shared<NaryNode>( |
| 54 | NaryNode(static_cast<af::dtype>(dtype_traits<T>::af_type), "__select", |
| 55 | 3, {{cond_node, a_node, b_node}}, af_select_t, height)); |
| 56 | |
| 57 | std::array<common::Node *, 1> nodes{node.get()}; |
| 58 | if (detail::passesJitHeuristics<T>(nodes) != kJITHeuristics::Pass) { |
| 59 | if (a_height > max(b_height, cond_height)) { |
| 60 | a.eval(); |
| 61 | } else if (b_height > cond_height) { |
| 62 | b.eval(); |
| 63 | } else { |
| 64 | cond.eval(); |
| 65 | } |
| 66 | return createSelectNode<T>(cond, a, b, odims); |
| 67 | } |
| 68 | return createNodeArray<T>(odims, node); |
| 69 | } |
| 70 | |
| 71 | template<typename T, bool flip> |
| 72 | Array<T> createSelectNode(const Array<char> &cond, const Array<T> &a, |