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