Use getArray instead of castArray if types are the same in arithOp

umar456 · Aug 17, 2021 · da6518e · da6518e
1 parent a1de6ac
commit da6518e
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 5 deletions.
diff --git a/src/api/c/binary.cpp b/src/api/c/binary.cpp
@@ -27,6 +27,7 @@
 #include <common/half.hpp>
 
 using af::dim4;
+using af::dtype;
 using common::half;
 using detail::arithOp;
 using detail::arithOpD;
@@ -41,9 +42,17 @@ using detail::ushort;
 template<typename T, af_op_t op>
 static inline af_array arithOp(const af_array lhs, const af_array rhs,
                                const dim4 &odims) {
-    af_array res =
-        getHandle(arithOp<T, op>(castArray<T>(lhs), castArray<T>(rhs), odims));
-    return res;
+    const ArrayInfo &linfo = getInfo(lhs);
+    const ArrayInfo &rinfo = getInfo(rhs);
+
+    dtype type = static_cast<af::dtype>(af::dtype_traits<T>::af_type);
+
+    const detail::Array<T> &l =
+        linfo.getType() == type ? getArray<T>(lhs) : castArray<T>(lhs);
+    const detail::Array<T> &r =
+        rinfo.getType() == type ? getArray<T>(rhs) : castArray<T>(rhs);
+
+    return getHandle(arithOp<T, op>(l, r, odims));
 }
 
 template<typename T, af_op_t op>

diff --git a/src/backend/common/jit/NaryNode.hpp b/src/backend/common/jit/NaryNode.hpp
@@ -94,7 +94,9 @@ common::Node_ptr createNaryNode(
     const af::dim4 &odims, FUNC createNode,
     std::array<const detail::Array<Ti> *, N> &&children) {
     std::array<common::Node_ptr, N> childNodes;
-    for (int i = 0; i < N; i++) { childNodes[i] = children[i]->getNode(); }
+    for (int i = 0; i < N; i++) {
+        childNodes[i] = move(children[i]->getNode());
+    }
 
     common::Node_ptr ptr = createNode(childNodes);
 

diff --git a/src/backend/cpu/Array.cpp b/src/backend/cpu/Array.cpp
@@ -273,7 +273,7 @@ kJITHeuristics passesJitHeuristics(Node *root_node) {
 
 template<typename T>
 Array<T> createNodeArray(const dim4 &dims, Node_ptr node) {
-    Array<T> out = Array<T>(dims, node);
+    Array<T> out(dims, node);
     return out;
 }
 

diff --git a/src/backend/cpu/arith.hpp b/src/backend/cpu/arith.hpp
@@ -15,6 +15,12 @@
 
 namespace cpu {
 
+template<typename T, af_op_t op>
+Array<T> arithOp(const Array<T> &&lhs, const Array<T> &&rhs,
+                 const af::dim4 &odims) {
+    return common::createBinaryNode<T, T, op>(lhs, rhs, odims);
+}
+
 template<typename T, af_op_t op>
 Array<T> arithOp(const Array<T> &lhs, const Array<T> &rhs,
                  const af::dim4 &odims) {

diff --git a/src/backend/cuda/arith.hpp b/src/backend/cuda/arith.hpp
@@ -14,6 +14,13 @@
 #include <af/dim4.hpp>
 
 namespace cuda {
+
+template<typename T, af_op_t op>
+Array<T> arithOp(const Array<T> &&lhs, const Array<T> &&rhs,
+                 const af::dim4 &odims) {
+    return common::createBinaryNode<T, T, op>(lhs, rhs, odims);
+}
+
 template<typename T, af_op_t op>
 Array<T> arithOp(const Array<T> &lhs, const Array<T> &rhs,
                  const af::dim4 &odims) {

diff --git a/src/backend/opencl/arith.hpp b/src/backend/opencl/arith.hpp
@@ -15,6 +15,13 @@
 #include <af/dim4.hpp>
 
 namespace opencl {
+
+template<typename T, af_op_t op>
+Array<T> arithOp(const Array<T> &&lhs, const Array<T> &&rhs,
+                 const af::dim4 &odims) {
+    return common::createBinaryNode<T, T, op>(lhs, rhs, odims);
+}
+
 template<typename T, af_op_t op>
 Array<T> arithOp(const Array<T> &lhs, const Array<T> &rhs,
                  const af::dim4 &odims) {