arrayfire · umar456 · Aug 19, 2023 · Aug 19, 2023 · Aug 19, 2023
diff --git a/src/api/c/reorder.cpp b/src/api/c/reorder.cpp
@@ -33,12 +33,14 @@ using std::swap;
 
 template<typename T>
 static inline af_array reorder(const af_array in, const af::dim4 &rdims0) {
-    Array<T> In = getArray<T>(in);
+    Array<T> In = detail::createEmptyArray<T>(af::dim4(0));
     dim4 rdims  = rdims0;
 
     if (rdims[0] == 1 && rdims[1] == 0) {
-        In = transpose(In, false);
+        In = transpose(getArray<T>(in), false);
         std::swap(rdims[0], rdims[1]);
+    } else {
+        In = getArray<T>(in);
     }
     const dim4 idims    = In.dims();
     const dim4 istrides = In.strides();
@@ -48,8 +50,7 @@ static inline af_array reorder(const af_array in, const af::dim4 &rdims0) {
 
     af_array out;
     if (rdims[0] == 0 && rdims[1] == 1 && rdims[2] == 2 && rdims[3] == 3) {
-        const Array<T> &Out = In;
-        out                 = getHandle(Out);
+        out = getHandle(In);
     } else if (rdims[0] == 0) {
         dim4 odims    = dim4(1, 1, 1, 1);
         dim4 ostrides = dim4(1, 1, 1, 1);

diff --git a/src/backend/common/jit/Node.cpp b/src/backend/common/jit/Node.cpp
@@ -76,6 +76,12 @@ auto isScalar(const Node &ptr) -> bool { return ptr.isScalar(); }
 
 bool Node::isLinear(const dim_t dims[4]) const { return true; }
 
+/// This function returns true if the \p node is a Shift node or a Buffer node
+auto isBufferOrShift(const Node_ptr &node) -> bool {
+    return node->getNodeType() == kNodeType::Buffer ||
+           node->getNodeType() == kNodeType::Shift;
+}
+
 }  // namespace common
 }  // namespace arrayfire
 

diff --git a/src/backend/common/jit/Node.hpp b/src/backend/common/jit/Node.hpp
@@ -14,6 +14,7 @@
 #include <types.hpp>
 #include <af/defines.h>
 
+#include <nonstd/span.hpp>
 #include <algorithm>
 #include <array>
 #include <functional>
@@ -330,9 +331,79 @@ std::string getFuncName(const std::vector<Node *> &output_nodes,
                         const bool is_linear, const bool loop0,
                         const bool loop1, const bool loop2, const bool loop3);
 
+/// Returns true if the \p ptr is a Buffer Node
 auto isBuffer(const Node &ptr) -> bool;
 
+/// Returns true if the \p ptr is a Scalar Node
 auto isScalar(const Node &ptr) -> bool;
 
+/// Returns true if \p node is a Buffer or a Shift node
+auto isBufferOrShift(const Node_ptr &node) -> bool;
+
+template<typename T>
+inline void applyShifts(std::array<int, 4> &shifts, nonstd::span<T> dims) {
+    std::array<T, 4> out;
+    for (size_t i = 0; i < shifts.size(); i++) { out[i] = dims[shifts[i]]; }
+    std::copy(begin(out), std::end(out), std::begin(dims));
+}
+
+template<typename ArrayT>
+inline std::array<int, 4> compressArray(ArrayT dims) {
+    std::array<int, 4> shifts{0, 1, 2, 3};
+    bool changed;
+    do {
+        changed = false;
+        for (int i = 0; i < AF_MAX_DIMS - 1; i++) {
+            if (dims[i] == 1 && dims[i + 1] != 1) {
+                std::swap(dims[i], dims[i + 1]);
+                std::swap(shifts[i], shifts[i + 1]);
+                changed = true;
+            }
+        }
+    } while (changed);
+    return shifts;
+}
+
+/// Removes empty columns from output and the other node pointers in \p nodes
+template<typename ParamT, typename BufferNodeT, typename ShiftNodeT>
+void removeEmptyDimensions(nonstd::span<ParamT> outputs,
+                           nonstd::span<Node_ptr> nodes) {
+    dim_t *outDims{outputs[0].dims_ptr()};
+    dim_t *outStrides{outputs[0].strides_ptr()};
+    auto shifts = compressArray(outDims);
+    applyShifts<dim_t>(shifts, {outStrides, AF_MAX_DIMS});
+    for (auto nodeIt{begin(nodes)}, endIt{end(nodes)};
+         (nodeIt = find_if(nodeIt, endIt, isBufferOrShift)) != endIt;
+         ++nodeIt) {
+        switch ((*nodeIt)->getNodeType()) {
+            case kNodeType::Buffer: {
+                BufferNodeT *buf{static_cast<BufferNodeT *>(nodeIt->get())};
+                applyShifts<dim_t>(shifts,
+                                   {buf->m_param.dims_ptr(), AF_MAX_DIMS});
+                applyShifts<dim_t>(shifts,
+                                   {buf->m_param.strides_ptr(), AF_MAX_DIMS});
+            } break;
+            case kNodeType::Shift: {
+                ShiftNodeT &shiftNode{
+                    *static_cast<ShiftNodeT *>(nodeIt->get())};
+                BufferNodeT &buf{shiftNode.getBufferNode()};
+                applyShifts<dim_t>(shifts,
+                                   {buf.m_param.dims_ptr(), AF_MAX_DIMS});
+                applyShifts<dim_t>(shifts,
+                                   {buf.m_param.strides_ptr(), AF_MAX_DIMS});
+
+                auto &node_shifts = shiftNode.getShifts();
+                applyShifts<int>(shifts, node_shifts);
+            } break;
+            default: break;
+        }
+    }
+    std::for_each(
+        std::begin(outputs) + 1, std::end(outputs), [&shifts](ParamT &output) {
+            applyShifts<dim_t>(shifts, {output.dims_ptr(), AF_MAX_DIMS});
+            applyShifts<dim_t>(shifts, {output.strides_ptr(), AF_MAX_DIMS});
+        });
+}
+
 }  // namespace common
 }  // namespace arrayfire
diff --git a/src/backend/common/jit/ShiftNodeBase.hpp b/src/backend/common/jit/ShiftNodeBase.hpp
@@ -53,6 +53,8 @@ class ShiftNodeBase : public Node {
         return *this;
     }
 
+    std::array<int, 4> &getShifts() { return m_shifts; }
+
     std::unique_ptr<Node> clone() final {
         return std::make_unique<ShiftNodeBase>(*this);
     }
@@ -65,6 +67,7 @@ class ShiftNodeBase : public Node {
         swap(m_shifts, other.m_shifts);
     }
 
+    BufferNode &getBufferNode() { return *m_buffer_node; }
     const BufferNode &getBufferNode() const { return *m_buffer_node; }
 
     bool isLinear(const dim_t dims[4]) const final {

diff --git a/src/backend/cuda/CMakeLists.txt b/src/backend/cuda/CMakeLists.txt
@@ -553,6 +553,7 @@ add_library(afcuda
     wrap.hpp
 
     jit/BufferNode.hpp
+    jit/ShiftNode.hpp
     jit/kernel_generators.hpp
 
     ${scan_by_key_sources}

diff --git a/src/backend/cuda/Param.hpp b/src/backend/cuda/Param.hpp
@@ -35,6 +35,9 @@ class Param {
         return dims[0] * dims[1] * dims[2] * dims[3];
     }
 
+    dim_t *dims_ptr() { return dims; }
+    dim_t *strides_ptr() { return strides; }
+
     Param(const Param<T> &other) noexcept               = default;
     Param(Param<T> &&other) noexcept                    = default;
     Param<T> &operator=(const Param<T> &other) noexcept = default;

diff --git a/src/backend/cuda/jit.cpp b/src/backend/cuda/jit.cpp
@@ -20,6 +20,7 @@
 #include <debug_cuda.hpp>
 #include <device_manager.hpp>
 #include <err_cuda.hpp>
+#include <jit/ShiftNode.hpp>
 #include <kernel_headers/jit_cuh.hpp>
 #include <math.hpp>
 #include <platform.hpp>
@@ -38,13 +39,17 @@ using arrayfire::common::findModule;
 using arrayfire::common::getEnvVar;
 using arrayfire::common::getFuncName;
 using arrayfire::common::half;
+using arrayfire::common::isBufferOrShift;
+using arrayfire::common::kNodeType;
 using arrayfire::common::ModdimNode;
 using arrayfire::common::Node;
 using arrayfire::common::Node_ids;
 using arrayfire::common::Node_map_t;
 using arrayfire::common::Node_ptr;
 using arrayfire::common::NodeIterator;
 using arrayfire::common::saveKernel;
+using arrayfire::cuda::jit::BufferNode;
+using arrayfire::cuda::jit::ShiftNode;
 
 using std::array;
 using std::equal;
@@ -58,7 +63,6 @@ using std::vector;
 
 namespace arrayfire {
 namespace cuda {
-using jit::BufferNode;
 
 static string getKernelString(const string& funcName,
                               const vector<Node*>& full_nodes,
@@ -474,22 +478,9 @@ void evalNodes(vector<Param<T>>& outputs, const vector<Node*>& output_nodes) {
                 }
             }
             if (emptyColumnsFound) {
-                const auto isBuffer{
-                    [](const Node_ptr& node) { return node->isBuffer(); }};
-                for (auto nodeIt{begin(node_clones)}, endIt{end(node_clones)};
-                     (nodeIt = find_if(nodeIt, endIt, isBuffer)) != endIt;
-                     ++nodeIt) {
-                    BufferNode<T>* buf{
-                        static_cast<BufferNode<T>*>(nodeIt->get())};
-                    removeEmptyColumns(outDims, ndims, buf->m_param.dims,
-                                       buf->m_param.strides);
-                }
-                for_each(++begin(outputs), end(outputs),
-                         [outDims, ndims](Param<T>& output) {
-                             removeEmptyColumns(outDims, ndims, output.dims,
-                                                output.strides);
-                         });
-                ndims = removeEmptyColumns(outDims, ndims, outDims, outStrides);
+                common::removeEmptyDimensions<Param<T>, BufferNode<T>,
+                                              ShiftNode<T>>(outputs,
+                                                            node_clones);
             }
 
             full_nodes.clear();

diff --git a/src/backend/cuda/jit/ShiftNode.hpp b/src/backend/cuda/jit/ShiftNode.hpp
@@ -0,0 +1,22 @@
+/*******************************************************
+ * Copyright (c) 2023, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <common/jit/ShiftNodeBase.hpp>
+#include <jit/BufferNode.hpp>
+
+namespace arrayfire {
+namespace cuda {
+namespace jit {
+
+template<typename T>
+using ShiftNode = common::ShiftNodeBase<BufferNode<T>>;
+
+}  // namespace jit
+}  // namespace cuda
+}  // namespace arrayfire
diff --git a/src/backend/cuda/shift.cpp b/src/backend/cuda/shift.cpp
@@ -11,16 +11,16 @@
 #include <common/jit/ShiftNodeBase.hpp>
 #include <err_cuda.hpp>
 #include <jit/BufferNode.hpp>
+#include <jit/ShiftNode.hpp>
 #include <shift.hpp>
 
 #include <memory>
 
 using af::dim4;
 
 using arrayfire::common::Node_ptr;
-using arrayfire::common::ShiftNodeBase;
-
 using arrayfire::cuda::jit::BufferNode;
+using arrayfire::cuda::jit::ShiftNode;
 
 using std::array;
 using std::make_shared;
@@ -29,8 +29,6 @@ using std::string;
 
 namespace arrayfire {
 namespace cuda {
-template<typename T>
-using ShiftNode = ShiftNodeBase<BufferNode<T>>;
 
 template<typename T>
 Array<T> shift(const Array<T> &in, const int sdims[4]) {

diff --git a/src/backend/oneapi/CMakeLists.txt b/src/backend/oneapi/CMakeLists.txt
@@ -99,6 +99,7 @@ add_library(afoneapi
   ireduce.hpp
   jit.cpp
   jit/BufferNode.hpp
+  jit/ShiftNode.hpp
   jit/kernel_generators.hpp
   join.cpp
   join.hpp

diff --git a/src/backend/oneapi/Param.hpp b/src/backend/oneapi/Param.hpp
@@ -27,6 +27,9 @@ struct Param {
     Param(const Param& other)            = default;
     Param(Param&& other)                 = default;
 
+    dim_t* dims_ptr() { return info.dims; }
+    dim_t* strides_ptr() { return info.strides; }
+
     // AF_DEPRECATED("Use Array<T>")
     Param() : data(nullptr), info{{0, 0, 0, 0}, {0, 0, 0, 0}, 0} {}
 
@@ -54,6 +57,9 @@ struct AParam {
     AParam(const AParam& other)            = default;
     AParam(AParam&& other)                 = default;
 
+    dim_t* dims_ptr() { return dims.get(); }
+    dim_t* strides_ptr() { return strides.get(); }
+
     // AF_DEPRECATED("Use Array<T>")
     AParam() : data(), dims{0, 0, 0, 0}, strides{0, 0, 0, 0}, offset(0) {}
 

diff --git a/src/backend/oneapi/jit.cpp b/src/backend/oneapi/jit.cpp
@@ -8,6 +8,7 @@
  ********************************************************/
 
 #include <CL/cl.h>
+#include <jit/ShiftNode.hpp>
 #include <jit/kernel_generators.hpp>
 
 #include <kernel_headers/KParam.hpp>
@@ -55,6 +56,7 @@ using arrayfire::common::NodeIterator;
 using arrayfire::common::ShiftNodeBase;
 using arrayfire::oneapi::getActiveDeviceBaseBuildFlags;
 using arrayfire::oneapi::jit::BufferNode;
+using arrayfire::oneapi::jit::ShiftNode;
 
 using std::array;
 using std::begin;
@@ -468,21 +470,8 @@ void evalNodes(vector<Param<T>>& outputs, const vector<Node*>& output_nodes) {
             }
         }
         if (emptyColumnsFound) {
-            const auto isBuffer{
-                [](const Node_ptr& ptr) { return ptr->isBuffer(); }};
-            for (auto nodeIt{begin(node_clones)}, endIt{end(node_clones)};
-                 (nodeIt = find_if(nodeIt, endIt, isBuffer)) != endIt;
-                 ++nodeIt) {
-                BufferNode<T>* buf{static_cast<BufferNode<T>*>(nodeIt->get())};
-                removeEmptyColumns(outDims, ndims, buf->m_param.dims.get(),
-                                   buf->m_param.strides.get());
-            }
-            for_each(++begin(outputs), end(outputs),
-                     [outDims, ndims](Param<T>& output) {
-                         removeEmptyColumns(outDims, ndims, output.info.dims,
-                                            output.info.strides);
-                     });
-            ndims = removeEmptyColumns(outDims, ndims, outDims, outStrides);
+            common::removeEmptyDimensions<Param<T>, BufferNode<T>,
+                                          ShiftNode<T>>(outputs, node_clones);
         }
     }
 

diff --git a/src/backend/oneapi/jit/ShiftNode.hpp b/src/backend/oneapi/jit/ShiftNode.hpp
@@ -0,0 +1,22 @@
+/*******************************************************
+ * Copyright (c) 2023, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <common/jit/ShiftNodeBase.hpp>
+#include <jit/BufferNode.hpp>
+
+namespace arrayfire {
+namespace oneapi {
+namespace jit {
+
+template<typename T>
+using ShiftNode = common::ShiftNodeBase<BufferNode<T>>;
+
+}  // namespace jit
+}  // namespace oneapi
+}  // namespace arrayfire
diff --git a/src/backend/opencl/CMakeLists.txt b/src/backend/opencl/CMakeLists.txt
@@ -468,6 +468,7 @@ target_sources(afopencl
 target_sources(afopencl
   PRIVATE
     jit/BufferNode.hpp
+    jit/ShiftNode.hpp
     jit/kernel_generators.hpp
   )
 

diff --git a/src/backend/opencl/Param.hpp b/src/backend/opencl/Param.hpp
@@ -22,6 +22,9 @@ struct Param {
     Param(const Param& other)            = default;
     Param(Param&& other)                 = default;
 
+    dim_t* dims_ptr() { return info.dims; }
+    dim_t* strides_ptr() { return info.strides; }
+
     // AF_DEPRECATED("Use Array<T>")
     Param();
     // AF_DEPRECATED("Use Array<T>")