diff --git a/src/backend/common/jit/BufferNodeBase.hpp b/src/backend/common/jit/BufferNodeBase.hpp index fd63e89932..85576304ad 100644 --- a/src/backend/common/jit/BufferNodeBase.hpp +++ b/src/backend/common/jit/BufferNodeBase.hpp @@ -119,6 +119,18 @@ class BufferNodeBase : public common::Node { } return false; } + + virtual void modDims(const af::dim4 &newDim) override { + af::dim4 strides(1, 1, 1, 1); + for(dim_t i = 1; i < 4; ++i) { + strides[i] = strides[i - 1] * newDim[i - 1]; + } + + for(dim_t i = 0; i < 4; ++i) { + m_param.dims[i] = newDim[i]; + m_param.strides[i] = strides[i]; + } + } }; } // namespace common diff --git a/src/backend/common/jit/Node.hpp b/src/backend/common/jit/Node.hpp index 4641ff182c..794c10c14c 100644 --- a/src/backend/common/jit/Node.hpp +++ b/src/backend/common/jit/Node.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include @@ -311,6 +312,10 @@ class Node { } virtual std::unique_ptr clone() = 0; + virtual void modDims(const af::dim4 &newDim) { + UNUSED(newDim); + } + #ifdef AF_CPU template friend void arrayfire::cpu::kernel::evalMultiple( diff --git a/src/backend/common/moddims.cpp b/src/backend/common/moddims.cpp index cf9d8d6bb9..25edfa5b0a 100644 --- a/src/backend/common/moddims.cpp +++ b/src/backend/common/moddims.cpp @@ -20,21 +20,27 @@ using detail::createNodeArray; using std::make_shared; using std::shared_ptr; +using std::array; +using arrayfire::common::Node; +using arrayfire::common::Node_ptr; using std::vector; namespace arrayfire { namespace common { + +Node_ptr copyModdims(const Node_ptr &in, const af::dim4 &newDim) { + + Node_ptr out = in->clone(); + for(int i = 0; i < in->kMaxChildren && in->m_children[i] != nullptr; ++i) { + out->m_children[i] = copyModdims(in->m_children[i], newDim); + } + if(out->isBuffer()) out->modDims(newDim); + + return out; +} + template Array moddimOp(const Array &in, af::dim4 outDim) { - using arrayfire::common::Node; - using arrayfire::common::Node_ptr; - using std::array; - - auto createModdim = [outDim](array &operands) { - return make_shared( - outDim, static_cast(af::dtype_traits::af_type), - operands[0]); - }; const auto &node = in.getNode(); @@ -49,8 +55,9 @@ Array moddimOp(const Array &in, af::dim4 outDim) { } if (all_linear == false) in.eval(); - Node_ptr out = createNaryNode(outDim, createModdim, {&in}); - return createNodeArray(outDim, out); + Array out = createNodeArray(outDim, copyModdims(in.getNode(), outDim)); + + return out; } template diff --git a/src/backend/cpu/jit/BufferNode.hpp b/src/backend/cpu/jit/BufferNode.hpp index 32a94b2a74..ca3cfe7bb5 100644 --- a/src/backend/cpu/jit/BufferNode.hpp +++ b/src/backend/cpu/jit/BufferNode.hpp @@ -175,6 +175,19 @@ class BufferNode : public TNode { } return false; } + + virtual void modDims(const af::dim4 &newDim) override { + af::dim4 strides(1, 1, 1, 1); + for(dim_t i = 1; i < 4; ++i) { + strides[i] = strides[i - 1] * newDim[i - 1]; + } + + for(dim_t i = 0; i < 4; ++i) { + m_dims[i] = newDim[i]; + m_strides[i] = strides[i]; + } + } + }; } // namespace jit diff --git a/src/backend/cuda/jit/BufferNode.hpp b/src/backend/cuda/jit/BufferNode.hpp index 195353fdd8..8692b72515 100644 --- a/src/backend/cuda/jit/BufferNode.hpp +++ b/src/backend/cuda/jit/BufferNode.hpp @@ -27,7 +27,16 @@ bool BufferNodeBase::operator==( // clang-format off return m_data.get() == other.m_data.get() && m_bytes == other.m_bytes && - m_param.ptr == other.m_param.ptr; + m_param.ptr == other.m_param.ptr && + m_linear_buffer == other.m_linear_buffer && + m_param.dims[0] == other.m_param.dims[0] && + m_param.dims[1] == other.m_param.dims[1] && + m_param.dims[2] == other.m_param.dims[2] && + m_param.dims[3] == other.m_param.dims[3] && + m_param.strides[0] == other.m_param.strides[0] && + m_param.strides[1] == other.m_param.strides[1] && + m_param.strides[2] == other.m_param.strides[2] && + m_param.strides[3] == other.m_param.strides[3]; // clang-format on } diff --git a/src/backend/oneapi/jit/BufferNode.hpp b/src/backend/oneapi/jit/BufferNode.hpp index 94655f23e7..d10ca24cc3 100644 --- a/src/backend/oneapi/jit/BufferNode.hpp +++ b/src/backend/oneapi/jit/BufferNode.hpp @@ -31,7 +31,16 @@ bool BufferNodeBase::operator==( // clang-format off return m_data.get() == other.m_data.get() && m_bytes == other.m_bytes && - m_param.offset == other.m_param.offset; + m_param.offset == other.m_param.offset && + m_linear_buffer == other.m_linear_buffer && + m_param.dims[0] == other.m_param.dims[0] && + m_param.dims[1] == other.m_param.dims[1] && + m_param.dims[2] == other.m_param.dims[2] && + m_param.dims[3] == other.m_param.dims[3] && + m_param.strides[0] == other.m_param.strides[0] && + m_param.strides[1] == other.m_param.strides[1] && + m_param.strides[2] == other.m_param.strides[2] && + m_param.strides[3] == other.m_param.strides[3]; // clang-format on } diff --git a/src/backend/opencl/jit/BufferNode.hpp b/src/backend/opencl/jit/BufferNode.hpp index e188fb429f..14521030f7 100644 --- a/src/backend/opencl/jit/BufferNode.hpp +++ b/src/backend/opencl/jit/BufferNode.hpp @@ -28,7 +28,16 @@ bool BufferNodeBase::operator==( // clang-format off return m_data.get() == other.m_data.get() && m_bytes == other.m_bytes && - m_param.offset == other.m_param.offset; + m_param.offset == other.m_param.offset && + m_linear_buffer == other.m_linear_buffer && + m_param.dims[0] == other.m_param.dims[0] && + m_param.dims[1] == other.m_param.dims[1] && + m_param.dims[2] == other.m_param.dims[2] && + m_param.dims[3] == other.m_param.dims[3] && + m_param.strides[0] == other.m_param.strides[0] && + m_param.strides[1] == other.m_param.strides[1] && + m_param.strides[2] == other.m_param.strides[2] && + m_param.strides[3] == other.m_param.strides[3]; // clang-format on } diff --git a/test/jit.cpp b/test/jit.cpp index 3848a22242..487fdcb6e2 100644 --- a/test/jit.cpp +++ b/test/jit.cpp @@ -814,3 +814,43 @@ TEST(JIT, setKernelCacheDirectory) { // Reset to the old path ASSERT_SUCCESS(af_set_kernel_cache_directory(old_path.c_str(), false)); } + +// Ensure that a correct result is obtained when evaluating an expression +// that contains both an array and its transpose - see ISSUE 3660 +TEST(JIT, evaluateBothArrayAndItsTranspose) { + float X2_ptr[25] = { -1., -1., -1., -1., -1., + -0.5, -0.5, -0.5, -0.5, -0.5, + 0., 0., 0., 0., 0., + 0.5, 0.5, 0.5, 0.5, 0.5, + 1., 1., 1., 1., 1. }; + array X2_gold(5, 5, X2_ptr); + + float Y2_ptr[25] = { -1., -0.5, 0., 0.5, 1., + -1., -0.5, 0., 0.5, 1., + -1., -0.5, 0., 0.5, 1., + -1., -0.5, 0., 0.5, 1., + -1., -0.5, 0., 0.5, 1. }; + array Y2_gold(5, 5, Y2_ptr); + + float X2Y2_ptr[25] = { -2., -1.5, -1., -0.5, 0., + -1.5, -1., -0.5, 0., 0.5, + -1., -0.5, 0., 0.5, 1., + -0.5, 0., 0.5, 1., 1.5, + 0., 0.5, 1., 1.5, 2. }; + array X2Y2_gold(5, 5, X2Y2_ptr); + + int n = 5; + int half = (n - 1) / 2; + double delta = 1.0 / half; + + array coord = delta * (af::range(n) - half); + + array X2 = tile(coord.T(), n, 1); + array Y2 = tile(coord, 1, n); + + array X2Y2 = X2 + Y2; + + ASSERT_ARRAYS_EQ(X2_gold, X2); + ASSERT_ARRAYS_EQ(Y2_gold, Y2); + ASSERT_ARRAYS_EQ(X2Y2_gold, X2Y2); +}