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