From 9847108c12c27c130f8323a1893ddc2104daaca7 Mon Sep 17 00:00:00 2001 From: Christophe Murphy Date: Mon, 23 Jun 2025 15:06:15 -0700 Subject: [PATCH] Fix for evaluation containing array and its inverse. A bug was found in the jit calculation tree where under certain circumstances the incorrect result would be found when both an array and a transpose of the same array are used in the same evaluation. When the transpose method is used on an array the treatment depends on the structure of the jit tree of that array. In the particular case that the array is linear (e.g. not a sub-array) and is not simply a buffer (i.e. a combination of buffers and/or scalars with some operators) a moddim node will be added to the root of the tree with the new dimensions. When the tree is evaluated the dimensions of the child buffer(s) of a moddim node are changed and the moddim node is deleted. If an evaluation happens to include both an array and the transpose of that same array then when the moddim node is applied it changes the dimensions of the children for the remaining lifetime of the evaluation including any subsequent use of these nodes. As a result if the transpose of an array is used in an expression followed by the array itself the second instance will also be transposed. In this fix instead of creating a moddims node a deep copy of the calculation tree is made at that point with the transpose applied to the child buffer nodes. These buffer node copies will have the transposed dimensions and strides but will still point to the same memory location for the data so no copy of the underlying data needs to be made. --- src/backend/common/jit/BufferNodeBase.hpp | 12 +++++++ src/backend/common/jit/Node.hpp | 5 +++ src/backend/common/moddims.cpp | 29 +++++++++------- src/backend/cpu/jit/BufferNode.hpp | 13 ++++++++ src/backend/cuda/jit/BufferNode.hpp | 11 ++++++- src/backend/oneapi/jit/BufferNode.hpp | 11 ++++++- src/backend/opencl/jit/BufferNode.hpp | 11 ++++++- test/jit.cpp | 40 +++++++++++++++++++++++ 8 files changed, 118 insertions(+), 14 deletions(-) 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); +}