| 73 | } |
| 74 | |
| 75 | detail::Array<To> operator()(const detail::Array<Ti> &in) { |
| 76 | using detail::jit::UnaryNode; |
| 77 | |
| 78 | common::Node_ptr in_node = in.getNode(); |
| 79 | constexpr af::dtype to_dtype = |
| 80 | static_cast<af::dtype>(af::dtype_traits<To>::af_type); |
| 81 | constexpr af::dtype in_dtype = |
| 82 | static_cast<af::dtype>(af::dtype_traits<Ti>::af_type); |
| 83 | |
| 84 | if (canOptimizeCast(to_dtype, in_dtype)) { |
| 85 | // JIT optimization in the cast of multiple sequential casts that |
| 86 | // become idempotent - check to see if the previous operation was |
| 87 | // also a cast |
| 88 | // TODO: handle arbitrarily long chains of casts |
| 89 | auto in_node_unary = |
| 90 | std::dynamic_pointer_cast<UnaryNode<To, Ti, af_cast_t>>( |
| 91 | in_node); |
| 92 | |
| 93 | if (in_node_unary && in_node_unary->getOp() == af_cast_t) { |
| 94 | // child child's output type is the input type of the child |
| 95 | AF_TRACE("Cast optimiztion performed by removing cast to {}", |
| 96 | af::dtype_traits<Ti>::getName()); |
| 97 | auto in_child_node = in_node_unary->getChildren()[0]; |
| 98 | if (in_child_node->getType() == to_dtype) { |
| 99 | // ignore the input node and simply connect a noop node from |
| 100 | // the child's child to produce this op's output |
| 101 | return detail::createNodeArray<To>(in.dims(), |
| 102 | in_child_node); |
| 103 | } |
| 104 | } |
| 105 | } |
| 106 | |
| 107 | auto node = std::make_shared<UnaryNode<To, Ti, af_cast_t>>(in_node); |
| 108 | |
| 109 | return detail::createNodeArray<To>(in.dims(), move(node)); |
| 110 | } |
| 111 | }; |
| 112 | #else |
| 113 | |