diff --git a/.github/workflows/conda-cpp.yml b/.github/workflows/conda-cpp.yml index 3d6ca6443..9f070b9da 100644 --- a/.github/workflows/conda-cpp.yml +++ b/.github/workflows/conda-cpp.yml @@ -68,11 +68,20 @@ jobs: ./configure --prefix=$HOME/ucc/install --with-ucx=$HOME/ucx/install make install + - name: Set library paths + run: | + mkdir -p $CONDA_PREFIX/etc/conda/activate.d + echo 'export LD_LIBRARY_PATH=$HOME/ucx/install/lib:$HOME/ucc/install/lib:$LD_LIBRARY_PATH' > $CONDA_PREFIX/etc/conda/activate.d/ucx_paths.sh + - name: Build cylon, pycylon and run cpp test run: python build.py -cmake-flags="-DCYLON_UCX=1 -DCYLON_GLOO=1 -DGLOO_INSTALL_PREFIX=$HOME/gloo/install -DCYLON_UCC=1 -DUCC_INSTALL_PREFIX=$HOME/ucc/install -DUCX_INSTALL_PREFIX=$HOME/ucx/install" -ipath="$HOME/cylon/install" --cpp --python --test - name: Run pytest - run: python build.py -ipath="$HOME/cylon/install" --pytest + run: | + export LD_LIBRARY_PATH=$HOME/ucx/install/lib:$HOME/ucc/install/lib:$LD_LIBRARY_PATH + python build.py -ipath="$HOME/cylon/install" --pytest - name: Build Java - run: python build.py -ipath="$HOME/cylon/install" --java + run: | + export LD_LIBRARY_PATH=$HOME/ucx/install/lib:$HOME/ucc/install/lib:$LD_LIBRARY_PATH + python build.py -ipath="$HOME/cylon/install" --java diff --git a/build.py b/build.py index f7efe19e9..fc83d4338 100644 --- a/build.py +++ b/build.py @@ -173,6 +173,7 @@ def parse_cmake_flags(flag): CYLON_UCX = parse_cmake_flags('CYLON_UCX') CYLON_UCC = parse_cmake_flags('CYLON_UCC') CYLON_FMI = parse_cmake_flags('CYLON_FMI') +CYLON_SIMD = parse_cmake_flags('CYLON_SIMD') UCX_INSTALL_PREFIX = parse_cmake_flags('UCX_INSTALL_PREFIX') UCC_PREFIX = parse_cmake_flags('UCC_INSTALL_PREFIX') REDIS_PREFIX = parse_cmake_flags('REDIS_INSTALL_PREFIX') @@ -199,6 +200,7 @@ def print_line(): logger.info(f" -CYLON_UCX : {CYLON_UCX}") logger.info(f" -CYLON_UCC : {CYLON_UCC}") logger.info(f" -CYLON_FMI : {CYLON_FMI}") +logger.info(f" -CYLON_SIMD : {CYLON_SIMD}") logger.info(f" -UCC_PREFIX : {UCC_PREFIX}") logger.info(f"Run C++ tests : {RUN_CPP_TESTS}") logger.info(f"Build PyCylon : {BUILD_PYTHON}") @@ -310,6 +312,8 @@ def python_test(): env['LD_LIBRARY_PATH'] if CYLON_FMI: env['CYLON_FMI'] = str(CYLON_FMI) + if CYLON_SIMD: + env['CYLON_SIMD'] = str(CYLON_SIMD) elif OS_NAME == 'Darwin': if 'DYLD_LIBRARY_PATH' in env: @@ -364,6 +368,9 @@ def build_python(): env['CYLON_REDIS'] = str(CYLON_REDIS) env['REDIS_PREFIX'] = REDIS_PREFIX + if CYLON_SIMD: + env['CYLON_SIMD'] = str(CYLON_SIMD) + logger.info("Arrow prefix: " + str(Path(conda_prefix))) # Diagnostic logging for pycylon build @@ -389,7 +396,6 @@ def build_python(): logger.error("setup.py egg_info failed - see error above") check_status(test_res.returncode, "PyCylon setup.py validation") - # Use legacy setup.py develop mode which avoids PEP 517 entirely cmd = f'{PYTHON_EXEC} setup.py build_ext --inplace && {PYTHON_EXEC} -m pip install -v --no-build-isolation {clean} .' res = subprocess.run(cmd, shell=True, env=env, cwd=PYTHON_SOURCE_DIR) check_status(res.returncode, "PyCylon build") diff --git a/conda/environments/cylon.yml b/conda/environments/cylon.yml index 937b9ddf5..e49c54f3f 100644 --- a/conda/environments/cylon.yml +++ b/conda/environments/cylon.yml @@ -12,6 +12,7 @@ dependencies: - glog - openmpi - ucx + - pip - cython>=0.29.31 - numpy>=1.23,<2.0a0 - pandas>=2.0,<2.2.3dev diff --git a/conda/environments/cylon_NoUCX.yml b/conda/environments/cylon_NoUCX.yml index 80b0d739f..d48112695 100644 --- a/conda/environments/cylon_NoUCX.yml +++ b/conda/environments/cylon_NoUCX.yml @@ -12,6 +12,7 @@ dependencies: - libarrow-dataset==16.1.0.* - libarrow==16.1.0.* - glog + - pip - openmpi - cython>=0.29.31 - numpy>=1.23,<2.0a0 diff --git a/conda/environments/gcylon.yml b/conda/environments/gcylon.yml index 68bc9da4b..24d0119e0 100644 --- a/conda/environments/gcylon.yml +++ b/conda/environments/gcylon.yml @@ -20,6 +20,7 @@ dependencies: - glog - openmpi - ucx + - pip - numpy>=1.23,<2.0a0 - pandas>=2.0,<2.2.3dev - fsspec>=0.6.0 diff --git a/conda/environments/gcylon_cuda13.yml b/conda/environments/gcylon_cuda13.yml index f9ceae9d4..8bf313136 100644 --- a/conda/environments/gcylon_cuda13.yml +++ b/conda/environments/gcylon_cuda13.yml @@ -16,6 +16,7 @@ dependencies: - libcudf=26.02 - rmm=26.02 - librmm=26.02 + - pip # CUDA 13.x toolkit - cuda-version=13 - cuda-cudart diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 2ef7e1899..275d49195 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -403,6 +403,11 @@ if (CYLON_LIBFABRIC) message(STATUS "Libfabric lib: ${LIBFABRIC_LIB}") endif (CYLON_LIBFABRIC) +# SIMD operations (cosine similarity, batch search) +if (CYLON_SIMD) + message("Cylon SIMD Enabled") + add_definitions(-DBUILD_CYLON_SIMD) +endif () # Arrow if (NOT ARROW_BUILD_TYPE) diff --git a/cpp/src/cylon/CMakeLists.txt b/cpp/src/cylon/CMakeLists.txt index 03154ce48..39ab17d2f 100644 --- a/cpp/src/cylon/CMakeLists.txt +++ b/cpp/src/cylon/CMakeLists.txt @@ -124,9 +124,18 @@ else (CYLON_LIBFABRIC) set(CYLON_LIBFABRIC_FILES) endif (CYLON_LIBFABRIC) +if (CYLON_SIMD) + set(CYLON_SIMD_FILES + simd/simd_ops.hpp + simd/simd_ops.cpp) +else () + set(CYLON_SIMD_FILES) +endif () + add_library(cylon SHARED ${UCX_CYLON_FILES} ${UCC_CYLON_FILES} + ${CYLON_SIMD_FILES} ${CYLON_GLOO_FILES} ${CYLON_LIBFABRIC_FILES} ${UCX_REDIS_CYLON_FILES} @@ -383,6 +392,9 @@ add_subdirectory(thridparty/fmi/comm) add_subdirectory(thridparty/fmi/utils) add_subdirectory(compute) add_subdirectory(checkpoint) +if (CYLON_SIMD) + add_subdirectory(simd) +endif () set_target_properties(cylon PROPERTIES VERSION ${CYLON_VERSION}) diff --git a/cpp/src/cylon/net/fmi/fmi_channel.cpp b/cpp/src/cylon/net/fmi/fmi_channel.cpp index d4fabeba6..3c352c0b1 100644 --- a/cpp/src/cylon/net/fmi/fmi_channel.cpp +++ b/cpp/src/cylon/net/fmi/fmi_channel.cpp @@ -254,7 +254,15 @@ namespace cylon::fmi { if (rank == recvRank) { continue;//FMI does not support local receives, so process during sends } - } + // Init a new pending receive for the request + auto *buf = new PendingReceive(); + buf->receiveId = recvRank; + // Add to pendingReceive object to pendingReceives map + pendingReceives.insert(std::pair(recvRank, buf)); + // Receive for the initial header buffer + // Init context + buf->context = new FMI::Utils::fmiContext; + buf->context->completed = 0; auto send_data_byte_size = CYLON_CHANNEL_HEADER_SIZE * sizeof(int); auto send_void_ptr = const_cast(static_cast(buf->headerBuf)); @@ -404,34 +412,12 @@ namespace cylon::fmi { peerBusy = true; break; } - } else if (pend_send->status == SEND_FINISH) { - // we are going to send complete - send_comp_fn->sendFinishComplete(finishRequests[rank]); - pend_send->status = SEND_DONE; - } else if (pend_send->status != SEND_DONE) { - // throw an exception and log - LOG(FATAL) << "At an un-expected state " << pend_send->status; - } - } - - void FMIChannel::progressSendTo(int peer_id) { - // Role-based ordering: Only send first if rank < peer_id - //if (worldSize > 2) { - // if (rank >= peer_id) return; - //} - - PendingSend *ps = sends[peer_id]; - - if (peer_id == rank) { - progressSendsLocal(sends[peer_id]); - return; } if (peerBusy) { release_lock(lock_key, lock_val); return; } - } // Check if peer is RECEIVING auto peerRecvStatusStr = redis->get(peer_receive_status_key); @@ -609,15 +595,6 @@ namespace cylon::fmi { // Check if request is in finish if (finishRequests.count(dest)) { sendFinishHeader(x); - } else { - // If pending data is empty - // Notify about send completion - send_comp_fn->sendComplete(x.second->currentSend); - x.second->currentSend = {}; - - // Check if request is in finish - if (finishRequests.find(x.first) != finishRequests.end()) { - sendFinishHeader(x); } else { // If req is not in finish then re-init x.second->status = SEND_INIT; @@ -645,15 +622,7 @@ namespace cylon::fmi { void FMIChannel::progressReceiveFrom(int peer_id) { - //check if ok to receive (can't rely on sender to block sending - //so, we need to check for socket activity in blocking mode - /*if (!communicator->checkIfOkToReceive(peer_id)) { - LOG(INFO) << "unable to receive -- releasing lock key: " << lock_key << "peerId: " << peer_id; - release_lock(lock_key, lock_val); - return; - }*/ - - publishStatus(rank, peer_id, RECEIVING, RECEIVE); + if (peer_id == rank) return; PendingReceive *recv = pendingReceives[peer_id]; @@ -775,101 +744,12 @@ namespace cylon::fmi { publishStatus(rank, peer_id, IDLE, RECEIVE); } - if (recv->status == RECEIVE_LENGTH_POSTED && recv->context->completed == 1) { - - publishStatus(rank, peer_id, IDLE, RECEIVE); - LOG(INFO) << "finished RECEIVE_INIT -- releasing lock key: " << lock_key << " peer_id: " << peer_id; - //release_lock(lock_key, lock_val); - return; - } else { - release_lock(lock_key, lock_val); - publishStatus(rank, peer_id, IDLE, RECEIVE); - return; - } - } else if (recv->status == RECEIVE_LENGTH_POSTED && recv->context->completed == 1) { - - int length = recv->headerBuf[0]; - int finFlag = recv->headerBuf[1]; - - if (finFlag == CYLON_MSG_FIN) { - release_lock(lock_key, lock_val); - recv->status = RECEIVED_FIN; - rcv_fn->receivedHeader(peer_id, finFlag, nullptr, 0); - publishStatus(rank, peer_id, IDLE, RECEIVE); - LOG(INFO) << "[rank " << rank << "] ✅ Received FIN from " << peer_id; - LOG(INFO) << "finished CYLON_MSG_FIN -- releasing lock key: " << lock_key << " peer_id: " << peer_id; - - return; - } - - if (communicator->checkIfOkToReceive(peer_id, FMI::Utils::BLOCKING)) { - release_lock(lock_key, lock_val); - delete recv->context; - recv->context = new FMI::Utils::fmiContext(); - recv->context->completed = 0; - - allocator->Allocate(length, &recv->data); - recv->length = length; - - FMI::Comm::Data payload(recv->data->GetByteBuffer(), length, - FMI::Comm::noop_deleter); - FMI_Irecv(payload, peer_id, recv->context); - recv->status = RECEIVE_POSTED; - - int *header = new int[6]; - std::memcpy(header, &recv->headerBuf[2], 6 * sizeof(int)); - rcv_fn->receivedHeader(peer_id, finFlag, header, 6); - - publishStatus(rank, peer_id, IDLE, RECEIVE); - //LOG(INFO) << "finished RECEIVE_LENGTH_POSTED -- releasing lock key: " << lock_key << " peer_id: " - // << peer_id; - - return; - } else { - release_lock(lock_key, lock_val); - publishStatus(rank, peer_id, IDLE, RECEIVE); - return; - } - } else if (recv->status == RECEIVE_POSTED && recv->context->completed == 1) { - if (communicator->checkIfOkToReceive(peer_id, FMI::Utils::BLOCKING)) { - release_lock(lock_key, lock_val); - rcv_fn->receivedData(peer_id, recv->data, recv->length); - - std::fill_n(recv->headerBuf, CYLON_CHANNEL_HEADER_SIZE, 0); - delete recv->context; - recv->context = new FMI::Utils::fmiContext(); - recv->context->completed = 0; - - FMI::Comm::Data next_header(recv->headerBuf, - CYLON_CHANNEL_HEADER_SIZE * sizeof(int), - FMI::Comm::noop_deleter); - FMI_Irecv(next_header, peer_id, recv->context); - recv->status = RECEIVE_LENGTH_POSTED; - - publishStatus(rank, peer_id, IDLE, RECEIVE); - LOG(INFO) << "finished RECEIVE_POSTED -- releasing lock key: " << lock_key << " peer_id: " << peer_id; - } else { - release_lock(lock_key, lock_val); - publishStatus(rank, peer_id, IDLE, RECEIVE); - } - } else { - publishStatus(rank, peer_id, IDLE, RECEIVE); - } - - } - } void FMIChannel::progressReceives() { if (mode == FMI::Utils::NONBLOCKING) { - if (mode == FMI::Utils::BLOCKING) { - /*for (auto x: pendingReceives) { - progressReceiveFrom(x.first); - }*/ - } else { - communicator->communicator_event_progress(FMI::Utils::Operation::RECEIVE); // Iterate through the pending receives @@ -1162,25 +1042,6 @@ namespace cylon::fmi { }); } - bool FMIChannel::isSendComplete(int peer_id) { - auto it = sends.find(peer_id); - if (it == sends.end() || it->second == nullptr) { - return false; - } - PendingSend* ps = sends[peer_id]; - return ps->status == SEND_DONE; - } - - bool FMIChannel::isReceiveComplete(int peer_id) { - auto it = pendingReceives.find(peer_id); - if (it == pendingReceives.end() || it->second == nullptr) { - return false; - } - PendingReceive* recv = pendingReceives[peer_id]; - return recv->status == RECEIVED_FIN; - } - - } std::shared_ptr FMIChannel::getSendMutex(int peer_id) { diff --git a/cpp/src/cylon/net/fmi/fmi_channel.hpp b/cpp/src/cylon/net/fmi/fmi_channel.hpp index 73668c910..870d28e2a 100644 --- a/cpp/src/cylon/net/fmi/fmi_channel.hpp +++ b/cpp/src/cylon/net/fmi/fmi_channel.hpp @@ -80,180 +80,7 @@ namespace cylon { // UCX context - For tracking the progress of the message FMI::Utils::fmiContext *context = nullptr; - - public: - - /** - * Initialize the channel - * - * @param receives receive from these ranks - */ - void init(int edge, - const std::vector &receives, - const std::vector &sendIds, - ChannelReceiveCallback *rcv, - ChannelSendCallback *send, - Allocator *alloc) override; - - /** - * Send the message to the target. - * - * @param request the request - * @return true if accepted - */ - int send(std::shared_ptr request) override; - - /** - * Send the message to the target. - * - * @param request the request - * @return true if accepted - */ - int sendFin(std::shared_ptr request) override; - - /** - * This method, will send the messages, It will first send a message with length and then - */ - void progressSends() override; - - void progressSendTo(int peer_id); - - /*void progressSendsBlocking();*/ - - /** - * Progress the pending receivers - */ - void progressReceives() override; - - /*void progressReceivesBlocking();*/ - - void progressReceiveFrom(int peer_id); - - void close() override; - - explicit FMIChannel(std::shared_ptr com, FMI::Utils::Mode mode, - std::string redis_host, int redis_port, std::string redis_namespace); - - void notifyCompleted() override; - - private: - // keep track of the length buffers for each receiver - std::unordered_map sends; - // keep track of the posted receives - std::unordered_map pendingReceives; - // we got finish requests - std::unordered_map> finishRequests; - //send turn for blocking communication - // receive callback function - ChannelReceiveCallback *rcv_fn; - // send complete callback function - ChannelSendCallback *send_comp_fn; - // allocator - Allocator *allocator; - // mpi rank - int rank; - // mpi world size - int worldSize; - - - std::shared_ptr communicator; - - FMI::Utils::Mode mode; - - std::string redis_host; - - int redis_port; - - std::string redis_namespace; - - - - std::shared_ptr redis; - - std::string global_peer_lock; - - - //Thread support - std::unordered_map commThreads; - std::atomic commStarted{false}; - std::atomic shutdown{false}; - std::unordered_map> send_mutex_; - std::unordered_map> recv_mutex_; - std::shared_ptr getSendMutex(int peer_id); - std::shared_ptr getRecvMutex(int peer_id); - - std::mutex send_mutex; - std::mutex recv_mutex; - - - /** - * UCX Receive - * Modeled after the IRECV function of MPI - * @param [out] buffer - Pointer to the output buffer - * @param [in] count - Size of the receiving data - * @param [in] sender - MPI id of the sender - * @param [out] ctx - ucx::ucxContext object, used for tracking the progress of the request - * @return Cylon Status - */ - template - Status FMI_Irecv(FMI::Comm::Data &buf, - int sender, - FMI::Utils::fmiContext* ctx); - - /** - * UCX Send - * Modeled after the ISEND function of MPI - * @param [out] buffer - Pointer to the buffer to send - * @param [in] count - Size of the receiving data - * @param [in] ep - Endpoint to send the data to - * @param [out] request - UCX Context object - * Used for tracking the progress of the request - * @return Cylon Status - */ - template - Status FMI_Isend(FMI::Comm::Data &buf, - int source, - FMI::Utils::fmiContext* request) const; - - /** - * Send finish request - * @param x the target, pendingSend pair - */ - void sendFinishHeader(const std::pair &x) const; - - void sendFinishHeader(int target, PendingSend *ps); - - /** - * Send the length - * @param x the target, pendingSend pair - */ - void sendHeader(const std::pair &x) const; - - void sendHeader(int target, PendingSend *ps); - - void sendHeaderLocal(PendingSend * pend_send); - void sendFinishHeaderLocal(PendingSend * pend_send); - void progressSendsLocal(PendingSend * pend_send); - - bool acquire_lock(const std::string &lock_key, - const std::string &lock_value, int ttl_ms); - - void release_lock(const std::string &lock_key, - const std::string &lock_value); - - std::string generate_unique_id(); - - void publishStatus(int rank, int peer_id, FMISendReceiveStatus sendRecvStatus, - PublishStatusType publishStatus); - - void startCommunicationThreads(); - - std::string get_shared_lock_key(int a, int b); - - - }; - } struct PendingReceive { // we allow upto 8 integer header diff --git a/cpp/src/cylon/net/fmi/fmi_operations.cpp b/cpp/src/cylon/net/fmi/fmi_operations.cpp index ecc8d3022..d06f69d4a 100644 --- a/cpp/src/cylon/net/fmi/fmi_operations.cpp +++ b/cpp/src/cylon/net/fmi/fmi_operations.cpp @@ -75,25 +75,6 @@ } } - std::string NbxStatusToString(FMI::Utils::NbxStatus status) { - switch (status) { - case FMI::Utils::NbxStatus::SUCCESS: return "SUCCESS"; - case FMI::Utils::NbxStatus::CONNECTION_CLOSED_BY_PEER: return "CONNECTION_CLOSED_BY_PEER"; - case FMI::Utils::NbxStatus::SOCKET_CREATE_FAILED: return "SOCKET_CREATE_FAILED"; - case FMI::Utils::NbxStatus::TCP_NODELAY_FAILED: return "TCP_NODELAY_FAILED"; - case FMI::Utils::NbxStatus::FCNTL_GET_FAILED: return "FCNTL_GET_FAILED"; - case FMI::Utils::NbxStatus::FCNTL_SET_FAILED: return "FCNTL_SET_FAILED"; - case FMI::Utils::NbxStatus::ADD_EVENT_FAILED: return "ADD_EVENT_FAILED"; - case FMI::Utils::NbxStatus::EPOLL_WAIT_FAILED: return "EPOLL_WAIT_FAILED"; - case FMI::Utils::NbxStatus::SOCKET_PAIR_FAILED: return "SOCKET_PAIR_FAILED"; - case FMI::Utils::NbxStatus::SOCKET_SET_SO_RCVTIMEO_FAILED: return "SOCKET_SET_SO_RCVTIMEO_FAILED"; - case FMI::Utils::NbxStatus::SOCKET_SET_SO_SNDTIMEO_FAILED: return "SOCKET_SET_SO_SNDTIMEO_FAILED"; - case FMI::Utils::NbxStatus::SOCKET_SET_TCP_NODELAY_FAILED: return "SOCKET_SET_TCP_NODELAY_FAILED"; - case FMI::Utils::NbxStatus::SOCKET_SET_NONBLOCKING_FAILED: return "SOCKET_SET_NONBLOCKING_FAILED"; - default: return "UNKNOWN_STATUS"; - } - } - void FmiTableAllgatherImpl::Init(int num_buffers) { CYLON_UNUSED(num_buffers); } diff --git a/cpp/src/cylon/net/fmi/fmi_operations.hpp b/cpp/src/cylon/net/fmi/fmi_operations.hpp index 283e8506a..3992448d5 100644 --- a/cpp/src/cylon/net/fmi/fmi_operations.hpp +++ b/cpp/src/cylon/net/fmi/fmi_operations.hpp @@ -105,7 +105,7 @@ namespace cylon::fmi { FMI::Utils::Mode mode_; }; - class FmiTableAllgatherImpl : public TableAllgatherImpl { + class FmiAllReduceImpl : public net::AllReduceImpl { public: explicit FmiAllReduceImpl(const std::shared_ptr & comm_ptr, FMI::Utils::Mode mode) @@ -148,5 +148,4 @@ namespace cylon::fmi { - #endif //CYLON_FMI_OPERATIONS_HPP diff --git a/cpp/src/cylon/simd/CMakeLists.txt b/cpp/src/cylon/simd/CMakeLists.txt new file mode 100644 index 000000000..916097011 --- /dev/null +++ b/cpp/src/cylon/simd/CMakeLists.txt @@ -0,0 +1 @@ +cylon_install_all_headers("cylon/simd") \ No newline at end of file diff --git a/cpp/src/cylon/simd/simd_ops.cpp b/cpp/src/cylon/simd/simd_ops.cpp new file mode 100644 index 000000000..a797f1f13 --- /dev/null +++ b/cpp/src/cylon/simd/simd_ops.cpp @@ -0,0 +1,186 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "simd_ops.hpp" + +#include +#include +#include +#include + +#if defined(CYLON_HAVE_AVX2) +#include +#elif defined(CYLON_HAVE_SSE4_2) +#include +#elif defined(CYLON_HAVE_NEON) +#include +#endif + +namespace cylon { +namespace simd { + +// --------------------------------------------------------------------------- +// ISA-specific dot product kernels +// --------------------------------------------------------------------------- + +#if defined(CYLON_HAVE_AVX2) + +static float dot_product_f32(const float* a, const float* b, int dim) { + __m256 sum = _mm256_setzero_ps(); + int i = 0; + for (; i + 8 <= dim; i += 8) { + __m256 va = _mm256_loadu_ps(a + i); + __m256 vb = _mm256_loadu_ps(b + i); + sum = _mm256_fmadd_ps(va, vb, sum); + } + // Horizontal reduction: 8 → 1 + __m128 hi = _mm256_extractf128_ps(sum, 1); + __m128 lo = _mm256_castps256_ps128(sum); + __m128 s = _mm_add_ps(lo, hi); + s = _mm_hadd_ps(s, s); + s = _mm_hadd_ps(s, s); + float result = _mm_cvtss_f32(s); + // Scalar tail + for (; i < dim; ++i) { + result += a[i] * b[i]; + } + return result; +} + +#elif defined(CYLON_HAVE_SSE4_2) + +static float dot_product_f32(const float* a, const float* b, int dim) { + __m128 sum = _mm_setzero_ps(); + int i = 0; + for (; i + 4 <= dim; i += 4) { + __m128 va = _mm_loadu_ps(a + i); + __m128 vb = _mm_loadu_ps(b + i); + sum = _mm_add_ps(sum, _mm_mul_ps(va, vb)); + } + // Horizontal reduction + sum = _mm_hadd_ps(sum, sum); + sum = _mm_hadd_ps(sum, sum); + float result = _mm_cvtss_f32(sum); + for (; i < dim; ++i) { + result += a[i] * b[i]; + } + return result; +} + +#elif defined(CYLON_HAVE_NEON) + +static float dot_product_f32(const float* a, const float* b, int dim) { + float32x4_t sum = vdupq_n_f32(0.0f); + int i = 0; + for (; i + 4 <= dim; i += 4) { + float32x4_t va = vld1q_f32(a + i); + float32x4_t vb = vld1q_f32(b + i); + sum = vfmaq_f32(sum, va, vb); + } + float result = vaddvq_f32(sum); + for (; i < dim; ++i) { + result += a[i] * b[i]; + } + return result; +} + +#else // Scalar fallback + +static float dot_product_f32(const float* a, const float* b, int dim) { + float result = 0.0f; + for (int i = 0; i < dim; ++i) { + result += a[i] * b[i]; + } + return result; +} + +#endif + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +float cosine_similarity_f32(const float* a, const float* b, int dim) { + float dot_ab = dot_product_f32(a, b, dim); + float dot_aa = dot_product_f32(a, a, dim); + float dot_bb = dot_product_f32(b, b, dim); + float denom = std::sqrt(dot_aa) * std::sqrt(dot_bb); + if (denom == 0.0f) { + return 0.0f; + } + return dot_ab / denom; +} + +std::vector batch_cosine_search( + const float* query, int dim, + const float* embeddings, int64_t num_rows, + float threshold, int top_k) { + if (top_k <= 0 || num_rows <= 0 || dim <= 0) { + return {}; + } + + // Min-heap of (similarity, index) — keeps the top-k highest similarities + using Pair = std::pair; + std::priority_queue, std::greater> heap; + + // Precompute query norm + float query_norm_sq = dot_product_f32(query, query, dim); + if (query_norm_sq == 0.0f) { + return {}; + } + float query_norm = std::sqrt(query_norm_sq); + + for (int64_t i = 0; i < num_rows; ++i) { + const float* row = embeddings + i * dim; + float dot_qr = dot_product_f32(query, row, dim); + float row_norm_sq = dot_product_f32(row, row, dim); + if (row_norm_sq == 0.0f) { + continue; + } + float sim = dot_qr / (query_norm * std::sqrt(row_norm_sq)); + if (sim >= threshold) { + if (static_cast(heap.size()) < top_k) { + heap.emplace(sim, i); + } else if (sim > heap.top().first) { + heap.pop(); + heap.emplace(sim, i); + } + } + } + + // Extract results sorted by descending similarity + std::vector results; + results.reserve(heap.size()); + while (!heap.empty()) { + auto [sim, idx] = heap.top(); + heap.pop(); + results.push_back({idx, sim}); + } + std::reverse(results.begin(), results.end()); + return results; +} + +std::vector batch_cosine_search_arrow( + const float* query, int dim, + const std::shared_ptr& embeddings, + float threshold, int top_k) { + // FixedSizeList stores all values in a single contiguous Float32Array + auto values = std::static_pointer_cast(embeddings->values()); + const float* data = values->raw_values(); + int64_t num_rows = embeddings->length(); + return batch_cosine_search(query, dim, data, num_rows, threshold, top_k); +} + +} // namespace simd +} // namespace cylon \ No newline at end of file diff --git a/cpp/src/cylon/simd/simd_ops.hpp b/cpp/src/cylon/simd/simd_ops.hpp new file mode 100644 index 000000000..a2171fdff --- /dev/null +++ b/cpp/src/cylon/simd/simd_ops.hpp @@ -0,0 +1,57 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef CYLON_SIMD_OPS_HPP +#define CYLON_SIMD_OPS_HPP + +#include +#include +#include + +#include + +namespace cylon { +namespace simd { + +/// Result of a similarity search: row index and cosine similarity score. +struct SearchResult { + int64_t index; + float similarity; +}; + +/// Compute cosine similarity between two float32 vectors of length @p dim. +/// Returns 0.0 if either vector has zero magnitude. +float cosine_similarity_f32(const float* a, const float* b, int dim); + +/// Batch cosine search: compare @p query against @p num_rows embeddings +/// stored in a contiguous flat buffer (row-major: num_rows * dim floats). +/// Returns up to @p top_k results with similarity >= @p threshold, +/// sorted by descending similarity. +std::vector batch_cosine_search( + const float* query, int dim, + const float* embeddings, int64_t num_rows, + float threshold, int top_k); + +/// Arrow-native batch cosine search: operates directly on a +/// FixedSizeList column. Zero-copy — reads the underlying +/// contiguous values buffer without copying. +std::vector batch_cosine_search_arrow( + const float* query, int dim, + const std::shared_ptr& embeddings, + float threshold, int top_k); + +} // namespace simd +} // namespace cylon + +#endif // CYLON_SIMD_OPS_HPP \ No newline at end of file diff --git a/cpp/src/cylon/thridparty/fmi/Communicator.cpp b/cpp/src/cylon/thridparty/fmi/Communicator.cpp index 208b49798..2ae4c28fd 100644 --- a/cpp/src/cylon/thridparty/fmi/Communicator.cpp +++ b/cpp/src/cylon/thridparty/fmi/Communicator.cpp @@ -87,15 +87,3 @@ FMI::Utils::peer_num FMI::Communicator::getNumPeers() const { FMI::Utils::peer_num FMI::Communicator::getPeerId() const { return peer_id; } - -FMI::Communicator::~Communicator() { - channel->finalize(); -} - -FMI::Utils::peer_num FMI::Communicator::getNumPeers() const { - return num_peers; -} - -FMI::Utils::peer_num FMI::Communicator::getPeerId() const { - return peer_id; -} diff --git a/cpp/src/cylon/thridparty/fmi/Communicator.hpp b/cpp/src/cylon/thridparty/fmi/Communicator.hpp index a156079cc..b07704dea 100644 --- a/cpp/src/cylon/thridparty/fmi/Communicator.hpp +++ b/cpp/src/cylon/thridparty/fmi/Communicator.hpp @@ -57,16 +57,6 @@ namespace FMI { - bool checkIfOkToReceive(FMI::Utils::peer_num dest, Utils::Mode mode) { - return channel->checkReceive(dest, mode); - } - - bool checkIfOkToSend(FMI::Utils::peer_num dest, Utils::Mode mode) { - return channel->checkSend(dest, mode); - } - - - bool checkIfOkToReceive(FMI::Utils::peer_num dest, Utils::Mode mode) { return channel->checkReceive(dest, mode); } @@ -127,16 +117,6 @@ namespace FMI { channel->recv(data, src, context, mode, std::move(callback)); } - //! Receive data from src and store data into the provided buf - template - void recv(Comm::Data &buf, FMI::Utils::peer_num src, - FMI::Utils::fmiContext * context, - FMI::Utils::Mode mode, - std::function callback) { - channel_data data {buf.data(), buf.size_in_bytes(), FMI::Comm::noop_deleter}; - channel->recv(data, src, context, mode, std::move(callback)); - } - //! Broadcast the data that is in the provided buf of the root peer. Result is stored in buf for all peers. template void bcast(Comm::Data &buf, FMI::Utils::peer_num root) { @@ -251,24 +231,6 @@ namespace FMI { } - /*! - * @param sendbuf Data to send to root, needs to be the same size for all peers. - * @param recvbuf Receive buffer, only relevant for the root process. Size needs to be num_peers * sendbuf.size - */ - template - void allgatherv(Comm::Data &sendbuf, Comm::Data &recvbuf, FMI::Utils::peer_num root, - std::vector recvcounts, - const std::vector displs, - Utils::Mode mode, - std::function callback) { - channel_data senddata {sendbuf.data(), sendbuf.size_in_bytes()}; - channel_data recvdata {recvbuf.data(), recvbuf.size_in_bytes()}; - channel->allgatherv(senddata, recvdata, root, - recvcounts, displs, mode, callback); - } - - //! Scatter the data from root's sendbuf to the recvbuf of all peers. /*! @@ -395,16 +357,6 @@ namespace FMI { Utils::peer_num getNumPeers() const; - private: - - std::unordered_map> channel_map; - //std::shared_ptr channel; - FMI::Utils::peer_num peer_id; - public: - Utils::peer_num getPeerId() const; - - Utils::peer_num getNumPeers() const; - private: std::shared_ptr channel; diff --git a/cpp/src/cylon/thridparty/fmi/Data.hpp b/cpp/src/cylon/thridparty/fmi/Data.hpp index 41f6017e4..229ad47e2 100644 --- a/cpp/src/cylon/thridparty/fmi/Data.hpp +++ b/cpp/src/cylon/thridparty/fmi/Data.hpp @@ -92,8 +92,6 @@ namespace FMI::Comm { //! Instantiate data with a pointer to memory and an arbitrary size. Should only be used in exceptional cases, the native types should be used otherwise. - static inline std::function noop_deleter = [](void*) {}; - static inline std::function noop_deleter = [](void*) {}; template<> diff --git a/cpp/src/cylon/thridparty/fmi/comm/Channel.cpp b/cpp/src/cylon/thridparty/fmi/comm/Channel.cpp index a512b1110..72915bccf 100644 --- a/cpp/src/cylon/thridparty/fmi/comm/Channel.cpp +++ b/cpp/src/cylon/thridparty/fmi/comm/Channel.cpp @@ -146,52 +146,6 @@ int FMI::Comm::Channel::getMaxTimeout() { return -1; } -void -FMI::Comm::Channel::allgatherv(const channel_data &sendbuf, const channel_data &recvbuf, FMI::Utils::peer_num root, - const std::vector &recvcounts, const std::vector &displs, - Utils::Mode mode, - std::function callback) {} - -void FMI::Comm::Channel::gatherv(const std::shared_ptr sendbuf, - std::shared_ptr recvbuf, - FMI::Utils::peer_num root, - const std::vector &recvcounts, - const std::vector &displs) { - gatherv(sendbuf, recvbuf, root, recvcounts, displs, Utils::BLOCKING, nullptr); - -} - -void FMI::Comm::Channel::gatherv(const std::shared_ptr sendbuf, - std::shared_ptr recvbuf, - FMI::Utils::peer_num root, - const std::vector &recvcounts, - const std::vector &displs, - Utils::Mode mode, std::function callback) { - -} - -void FMI::Comm::Channel::bcast(std::shared_ptr buf, FMI::Utils::peer_num root) { - bcast(buf, root, Utils::BLOCKING, nullptr); - -} - -void FMI::Comm::Channel::bcast(std::shared_ptr buf, FMI::Utils::peer_num root, FMI::Utils::Mode mode, - std::function callback) { - -} - -int FMI::Comm::Channel::getMaxTimeout() { - return -1; -} - - -void FMI::Comm::Channel::init() { -//noop -} - void FMI::Comm::Channel::init() { //noop diff --git a/cpp/src/cylon/thridparty/fmi/comm/Direct.cpp b/cpp/src/cylon/thridparty/fmi/comm/Direct.cpp index d85e2b06a..e46b8eebb 100644 --- a/cpp/src/cylon/thridparty/fmi/comm/Direct.cpp +++ b/cpp/src/cylon/thridparty/fmi/comm/Direct.cpp @@ -34,8 +34,6 @@ - - #include #include @@ -95,11 +93,6 @@ FMI::Comm::Direct::Direct(const std::shared_ptr &backend) - sockets[Utils::NONBLOCKING] = {}; - sockets[Utils::BLOCKING] = {}; - - - io_states[Utils::Operation::SEND] = {}; io_states[Utils::Operation::RECEIVE] = {}; @@ -197,130 +190,6 @@ void FMI::Comm::Direct::start_ping_thread(Utils::Mode mode) { } -void FMI::Comm::Direct::init() { - //iterator over world size and create all sockets for non-blocking based on multi-send/receives - //create all the connections - //start_holepunch_subscriber(); - if (num_peers> 0) { - - for (int i = 0; i < num_peers; ++i) { - - if (i == peer_id) continue; - - - if (mode == Utils::NONBLOCKING) { - std::string send_pairing_nb = get_pairing_name(peer_id, i, Utils::NONBLOCKING); - check_socket_nbx(i, send_pairing_nb); - } - - //always create a pair of blocking sockets - //std::string send_pairing_b = get_pairing_name(peer_id, i, Utils::BLOCKING); - - //check_socket(i, send_pairing_b); - - } - - - if (mode == Utils::NONBLOCKING && enable_ping) { - start_ping_thread(Utils::NONBLOCKING); - } - } - - - -} - -inline const char* ModeToString(FMI::Utils::Mode mode) { - switch (mode) { - case FMI::Utils::Mode::BLOCKING: return "BLOCKING"; - case FMI::Utils::Mode::NONBLOCKING: return "NONBLOCKING"; - - default: return "UNKNOWN_MODE"; - } -} - -void FMI::Comm::Direct::start_holepunch_subscriber() { - std::thread([this]() { - if (redis_port > 0 && !redis_host.empty()) { - auto opts = sw::redis::ConnectionOptions{}; - opts.host = redis_host; - opts.port = redis_port; - auto redis = std::make_shared(opts); - auto sub = redis->subscriber(); - - sub.on_message([this](const std::string &channel, const std::string &msg) { - int from = -1, to = -1; - LOG(INFO) << "received message from publisher: " << msg; - sscanf(msg.c_str(), "from:%d,to:%d", &from, &to); - if (to == this->peer_id) { - LOG(INFO) << "Received reverse connect request from peer " << from; - // Trigger a connect attempt from this node to the sender - std::string pairing = get_pairing_name(this->peer_id, from, Utils::BLOCKING); - try { - check_socket(from, pairing); // This will do the actual reverse connect - } catch (const Utils::Timeout &) { - LOG(WARNING) << "Reverse connect to " << from << " failed."; - } - } - - }); - - - sub.subscribe("fmi_connect"); - - try { - while (true) sub.consume(); // Blocking wait - } catch (const std::exception &e) { - LOG(ERROR) << "Redis subscribe error: " << e.what(); - } - } - }).detach(); -} - -void FMI::Comm::Direct::init_blocking_sockets() { - if (num_peers> 0) { - - LOG(INFO) << "init blocking sockets"; - - for (int i = 0; i < num_peers; ++i) { - - if (i == peer_id) continue; - - std::string send_pairing_b = get_pairing_name(peer_id, i, Utils::BLOCKING); - - check_socket(i, send_pairing_b); - - } - - } - blocking_init = true; -} - -void FMI::Comm::Direct::start_ping_thread(Utils::Mode mode) { - - std::thread([this, mode]() { - - for (int i = 0; i < num_peers; ++i) { - if (i == peer_id) continue; - if (sockets[mode][i] != -1) { - try { - PingMessage ping{}; - ::send(sockets[mode][i], &ping, sizeof(ping), 0); - LOG(INFO) << "Sent PING to peer " << i << " Mode: " << ModeToString(mode); - - - } catch (...) { - LOG(ERROR) << "PING send failed to peer " << i << " Mode: " << ModeToString(mode);; - } - } - } - return; - - }).detach(); - -} - - void FMI::Comm::Direct::init() { //iterator over world size and create all sockets for non-blocking based on multi-send/receives //create all the connections @@ -462,104 +331,6 @@ void FMI::Comm::Direct::send_object(std::shared_ptr state, Utils::peer_ - // Normal message - ssize_t processed = ::send(socketfd, - state.request.buf.get() + state.processed, - state.request.len - state.processed, - 0); - - if (processed > 0) { - state.processed += processed; - - if (state.processed == state.request.len) { - if (state.callback) state.callback(); - state.callbackResult(Utils::SUCCESS, "Send completed", state.context); - return; - } - } - - // Still pending, register for epoll - if (processed == -1 && (errno != EAGAIN && errno != EINTR)) { - state.callbackResult(Utils::SEND_FAILED, strerror(errno), state.context); - return; - }*/ - - // Save the state and try again via epoll - //io_states[Utils::Operation::SEND][socketfd] = state; - //add_epoll_event(socketfd, state); - - } else { - send_object_blocking2(state, rcpt_id); - } - - // Save the state and try again via epoll - io_states[Utils::Operation::SEND][socketfd] = state; - add_epoll_event(socketfd, state); - - - while (sent_total < state->request->len) { - ssize_t sent = ::send(socketfd, state->request->buf.get() + sent_total, - state->request->len - sent_total, 0); - - if (sent == -1) { - if (errno == EINTR) continue; - if (errno == EAGAIN || errno == EWOULDBLOCK) { - throw Utils::Timeout(); // or use callbackResult if needed - } - LOG(ERROR) << "Send error: " << strerror(errno); - return; - } - - sent_total += sent; - } - - if (state->callback) state->callback(); - state->callbackResult(Utils::SUCCESS, "Blocking send complete", state->context); -} - -void FMI::Comm::Direct::send_object(std::shared_ptr state, Utils::peer_num rcpt_id, - Utils::Mode mode) { - - /*if (!blocking_init) { - init_blocking_sockets(); - }*/ - - if (mode == Utils::NONBLOCKING) { - std::string pairing = get_pairing_name(peer_id, rcpt_id, Utils::NONBLOCKING); - - // Use full-duplex socket for both send/recv - check_socket_nbx(rcpt_id, pairing); - int socketfd = sockets[Utils::NONBLOCKING][rcpt_id]; - - - - io_states[Utils::Operation::SEND][socketfd] = state; - - //if (checkSend(socketfd)) { - // handle_event(socketfd, io_states[Utils::SEND], Utils::SEND); - //} - - // Zero-length message? Send dummy byte - /*if (state.request.len == 0) { - char dummy = 0; - ssize_t sent = ::send(socketfd, &dummy, 1, 0); - - if (sent == 1) { - if (state.callback) state.callback(); - state.callbackResult(Utils::SUCCESS, "Zero-length message sent with dummy byte.", state.context); - } else if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) { - // Need to retry via epoll - io_states[Utils::Operation::SEND][socketfd] = state; - //add_epoll_event(socketfd, state); - } else { - state.callbackResult(Utils::DUMMY_SEND_FAILED, strerror(errno), state.context); - } - - return; - } - - - // Normal message ssize_t processed = ::send(socketfd, state.request.buf.get() + state.processed, @@ -668,12 +439,6 @@ void FMI::Comm::Direct::recv_object_blocking2(std::shared_ptrcallback) state->callback(); state->callbackResult(Utils::SUCCESS, "Zero-length receive via dummy byte", state->context); } - - // Handle zero-length message completion (if no loop body runs) - if (state->request->len == 0) { - if (state->callback) state->callback(); - state->callbackResult(Utils::SUCCESS, "Zero-length receive via dummy byte", state->context); - } } void FMI::Comm::Direct::recv_object(std::shared_ptr state, Utils::peer_num sender_id, @@ -991,7 +756,6 @@ void FMI::Comm::Direct::handle_event(int sockfd, void FMI::Comm::Direct::check_socket_nbx(FMI::Utils::peer_num partner_id, std::string pair_name) { - if (sockets[Utils::NONBLOCKING].empty()) { sockets[Utils::NONBLOCKING] = std::vector(num_peers, -1); } @@ -1203,45 +967,6 @@ bool FMI::Comm::Direct::checkReceivePing(int sockfd, FMI::Utils::Mode mode) { return false; // No ping available or error } -int FMI::Comm::Direct::getMaxTimeout() { - return max_timeout; -} - -bool FMI::Comm::Direct::checkRecv2(int fd) { - pollfd pfd = { fd, POLLIN, 0 }; - int poll_result = poll(&pfd, 1, 0); - - if (poll_result > 0) { - return true; // ✅ Ready to read - } else if (poll_result == 0) { - return false; // ❌ Not ready yet - } else { - LOG(ERROR) << "checkRecv: poll() failed with errno " << errno << ": " << strerror(errno); - return false; - } -} - -bool FMI::Comm::Direct::checkReceive(FMI::Utils::peer_num dest, Utils::Mode mode) { - - auto sockfd = sockets[mode][dest]; - return checkRecv(sockfd); - -} - - - - - - - - - - - - - - - diff --git a/cpp/src/cylon/thridparty/fmi/comm/Direct.hpp b/cpp/src/cylon/thridparty/fmi/comm/Direct.hpp index 6355311c8..3fd8bc0c9 100644 --- a/cpp/src/cylon/thridparty/fmi/comm/Direct.hpp +++ b/cpp/src/cylon/thridparty/fmi/comm/Direct.hpp @@ -47,16 +47,6 @@ namespace FMI::Comm { void send_object(std::shared_ptr buf, Utils::peer_num rcpt_id) override; - void recv_object(std::shared_ptr state, Utils::peer_num sender_id, Utils::Mode mode) override; - - void recv_object_blocking2(std::shared_ptr state, Utils::peer_num sender_id); - - bool checkReceive(FMI::Utils::peer_num dest, Utils::Mode mode) override; - - bool checkReceivePing(FMI::Utils::peer_num dest, Utils::Mode mode) override; - - bool checkSend(FMI::Utils::peer_num dest, Utils::Mode mode) override; - void send_object(std::shared_ptr state, Utils::peer_num rcpt_id, Utils::Mode mode) override; @@ -106,7 +96,6 @@ namespace FMI::Comm { void check_socket_nbx(Utils::peer_num partner_id, std::string pair_name); - std::string get_pairing_name(Utils::peer_num a, Utils::peer_num b, Utils::Mode mode); void handle_event(int socketfd, std::unordered_map> &states, diff --git a/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.cpp b/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.cpp index 985056d5a..a778931a7 100644 --- a/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.cpp +++ b/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.cpp @@ -24,24 +24,6 @@ void FMI::Comm::PeerToPeer::send(const std::shared_ptr buf, FMI::U -void FMI::Comm::PeerToPeer::send(std::shared_ptr buf, FMI::Utils::peer_num dest, - FMI::Utils::fmiContext *context, - FMI::Utils::Mode mode, - std::function callback) { - - auto state = std::make_shared(); - state->callbackResult = callback; - state->context = context; - state->setRequest(std::move(buf)); - state->processed = 0; - state->operation = Utils::SEND; - state->deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(getMaxTimeout()); - send_object(std::move(state), dest, mode); - -} - - void FMI::Comm::PeerToPeer::send(std::shared_ptr buf, FMI::Utils::peer_num dest, FMI::Utils::fmiContext *context, FMI::Utils::Mode mode, @@ -84,21 +66,6 @@ void FMI::Comm::PeerToPeer::recv(const std::shared_ptr buf, FMI::U recv_object(std::move(state), src, mode); } -void FMI::Comm::PeerToPeer::recv(const channel_data &buf, FMI::Utils::peer_num src, - FMI::Utils::fmiContext * context, - FMI::Utils::Mode mode, - std::function callback) { - auto state = std::make_shared(); - state->callbackResult = callback; - state->context = context; - state->setRequest(buf); - state->processed = 0; - state->operation = Utils::RECEIVE; - state->deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(getMaxTimeout()); - recv_object(std::move(state), src, mode); -} - void FMI::Comm::PeerToPeer::recv(FMI::Utils::peer_num src, Utils::Mode mode, std::shared_ptr state) { diff --git a/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.hpp b/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.hpp index de8629e5c..c00e58024 100644 --- a/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.hpp +++ b/cpp/src/cylon/thridparty/fmi/comm/PeerToPeer.hpp @@ -97,16 +97,6 @@ namespace FMI::Comm { void recv(std::shared_ptr buf, FMI::Utils::peer_num src) override; - void recv(FMI::Utils::peer_num src, - Utils::Mode mode, - std::shared_ptr state); - - void recv(std::shared_ptr buf, FMI::Utils::peer_num src, - Utils::fmiContext * context, - Utils::Mode mode, - std::function callback); - - void recv(FMI::Utils::peer_num src, Utils::Mode mode, std::shared_ptr state); @@ -179,19 +169,12 @@ namespace FMI::Comm { virtual void send_object(std::shared_ptr state, Utils::peer_num peer_id, Utils::Mode mode) = 0; - //! Send an object to peer with ID peer_id. Needs to be implemented by the channels(non-blocking). - - virtual void send_object(std::shared_ptr state, Utils::peer_num peer_id, Utils::Mode mode) = 0; - //! Receive an object from peer with ID peer_id. Needs to be implemented by the channels. virtual void recv_object(const std::shared_ptr buf, Utils::peer_num peer_id) = 0; //! Receive an object from peer with ID peer_id. Needs to be implemented by the channels (non-blocking). virtual void recv_object(std::shared_ptr state, Utils::peer_num peer_id, Utils::Mode mode) = 0; - //! Receive an object from peer with ID peer_id. Needs to be implemented by the channels (non-blocking). - virtual void recv_object(std::shared_ptr state, Utils::peer_num peer_id, Utils::Mode mode) = 0; - Utils::EventProcessStatus channel_event_progress(Utils::Operation op) override; diff --git a/cpp/src/cylon/thridparty/fmi/utils/Common.hpp b/cpp/src/cylon/thridparty/fmi/utils/Common.hpp index 06ae8bfd0..1e7232ef5 100644 --- a/cpp/src/cylon/thridparty/fmi/utils/Common.hpp +++ b/cpp/src/cylon/thridparty/fmi/utils/Common.hpp @@ -101,9 +101,6 @@ namespace FMI::Utils { - - - } diff --git a/cpp/src/examples/fmi_example.cpp b/cpp/src/examples/fmi_example.cpp index 696cf2524..b39c44474 100644 --- a/cpp/src/examples/fmi_example.cpp +++ b/cpp/src/examples/fmi_example.cpp @@ -127,14 +127,11 @@ int main(int argc, char *argv[]) { auto redisNamespace = std::string(argv[11]); - auto rank = std::stoi(argv[2]); /*auto backend = std::make_shared(); - - backend->withHost(host.c_str());//rendezvous host backend->withPort(port);//rendezvous port backend->withMaxTimeout(maxTimout); //max timeout for direct connect @@ -221,71 +218,6 @@ int main(int argc, char *argv[]) { - ctx->Finalize(); - return 0; - - if (!cylon::CylonContext::InitDistributed(config, &ctx).is_ok()) { - return 1; - } - - LOG(INFO) << "rank:" << ctx->GetRank() << " size:" << ctx->GetWorldSize(); - - ctx->Barrier(); - - const int modified_rank = ctx->GetRank() + 1; - - const std::string csv1 = directory + "user_device_tm_" + std::to_string(modified_rank) + ".csv"; - const std::string csv2 = directory + "user_usage_tm_" + std::to_string(modified_rank) + ".csv"; - - std::shared_ptr first_table, second_table, joined_table; - - - status = cylon::FromCSV(ctx, csv1, first_table); - CHECK_STATUS(status, "Reading csv1 failed!") - - status = cylon::FromCSV(ctx, csv2, second_table); - CHECK_STATUS(status, "Reading csv2 failed!")*/ - std::shared_ptr first_table, second_table, joined_table; - cylon::examples::create_two_in_memory_tables(kCount, kDup, ctx, first_table, second_table); - - //auto join_config = cylon::join::config::JoinConfig::InnerJoin(0, 3); - cylon::join::config::JoinConfig join_config{cylon::join::config::JoinType::INNER, 0, 0, - cylon::join::config::JoinAlgorithm::SORT, "l_", "r_"}; - - status = cylon::DistributedJoin(first_table, second_table, join_config, joined_table); - LOG(INFO) << "Status returned: " << status.get_code() << " msg: " <Rows() << " and Second table had : " - << second_table->Rows() << ", Joined has : " << joined_table->Rows(); - - LOG(INFO) << "AllReduce Collective Test"; - - - using TestType = arrow::Int32Type; - std::shared_ptr type = arrow::TypeTraits::type_singleton(); - - auto rank2 = *arrow::MakeScalar(ctx->GetRank())->CastTo(type); - - auto base_arr = ArrayFromJSON(type, "[1, 2, 3, 4]"); - // all reduce local sample histograms - auto arr = arrow::compute::Multiply(base_arr, rank2)->make_array(); - auto col = cylon::Column::Make(std::move(arr)); - - const auto &comm = ctx->GetCommunicator(); - - auto multiplier = *arrow::MakeScalar((worldsize - 1) * worldsize / 2)->CastTo(type); - auto exp = arrow::compute::Multiply(base_arr, multiplier)->make_array(); - - std::shared_ptr res; - CHECK_STATUS(comm->AllReduce(col, cylon::net::SUM, &res), "allreducefailed"); - - const auto &rcv = res->data(); - - LOG(INFO) << "AllReduce Result: " << rcv->ToString(); - - - ctx->Finalize(); return 0; diff --git a/python/pycylon/pycylon/simd/__init__.py b/python/pycylon/pycylon/simd/__init__.py new file mode 100644 index 000000000..8b7d147d6 --- /dev/null +++ b/python/pycylon/pycylon/simd/__init__.py @@ -0,0 +1,18 @@ +## + # Licensed under the Apache License, Version 2.0 (the "License"); + # you may not use this file except in compliance with the License. + # You may obtain a copy of the License at + # + # http://www.apache.org/licenses/LICENSE-2.0 + # + # Unless required by applicable law or agreed to in writing, software + # distributed under the License is distributed on an "AS IS" BASIS, + # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + # See the License for the specific language governing permissions and + # limitations under the License. + ## + +try: + from pycylon.simd.simd import cosine_similarity, batch_search +except ImportError: + pass \ No newline at end of file diff --git a/python/pycylon/pycylon/simd/simd.pxd b/python/pycylon/pycylon/simd/simd.pxd new file mode 100644 index 000000000..e7ac0fa17 --- /dev/null +++ b/python/pycylon/pycylon/simd/simd.pxd @@ -0,0 +1,32 @@ +## + # Licensed under the Apache License, Version 2.0 (the "License"); + # you may not use this file except in compliance with the License. + # You may obtain a copy of the License at + # + # http://www.apache.org/licenses/LICENSE-2.0 + # + # Unless required by applicable law or agreed to in writing, software + # distributed under the License is distributed on an "AS IS" BASIS, + # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + # See the License for the specific language governing permissions and + # limitations under the License. + ## + +from libcpp.vector cimport vector +from libcpp.memory cimport shared_ptr +from libc.stdint cimport int64_t + +from pyarrow.lib cimport CArray + + +cdef extern from "../../../../cpp/src/cylon/simd/simd_ops.hpp" namespace "cylon::simd": + cdef struct CSearchResult "cylon::simd::SearchResult": + int64_t index + float similarity + + float cosine_similarity_f32(const float* a, const float* b, int dim) + + vector[CSearchResult] batch_cosine_search( + const float* query, int dim, + const float* embeddings, int64_t num_rows, + float threshold, int top_k) \ No newline at end of file diff --git a/python/pycylon/pycylon/simd/simd.pyx b/python/pycylon/pycylon/simd/simd.pyx new file mode 100644 index 000000000..b45bae0bb --- /dev/null +++ b/python/pycylon/pycylon/simd/simd.pyx @@ -0,0 +1,82 @@ +## + # Licensed under the Apache License, Version 2.0 (the "License"); + # you may not use this file except in compliance with the License. + # You may obtain a copy of the License at + # + # http://www.apache.org/licenses/LICENSE-2.0 + # + # Unless required by applicable law or agreed to in writing, software + # distributed under the License is distributed on an "AS IS" BASIS, + # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + # See the License for the specific language governing permissions and + # limitations under the License. + ## + +import numpy as np +from libc.stdint cimport int64_t +from libcpp.vector cimport vector + +from pycylon.simd.simd cimport ( + CSearchResult, + cosine_similarity_f32 as c_cosine_similarity_f32, + batch_cosine_search as c_batch_cosine_search, +) + + +def cosine_similarity(object a not None, object b not None): + """Compute cosine similarity between two float32 vectors using SIMD. + + Args: + a: First vector (numpy float32 array). + b: Second vector (numpy float32 array, same length as a). + + Returns: + Cosine similarity as a float in [-1, 1]. Returns 0.0 if either + vector has zero magnitude. + """ + a = np.ascontiguousarray(a, dtype=np.float32) + b = np.ascontiguousarray(b, dtype=np.float32) + if len(a) != len(b): + raise ValueError(f"dimension mismatch: {len(a)} vs {len(b)}") + cdef int dim = len(a) + cdef const float* a_ptr = ( a.ctypes.data) + cdef const float* b_ptr = ( b.ctypes.data) + return c_cosine_similarity_f32(a_ptr, b_ptr, dim) + + +def batch_search(object query not None, + object embeddings not None, + float threshold=0.85, + int top_k=5): + """Batch cosine search: compare query against all embeddings using SIMD. + + Args: + query: Query vector (numpy float32 array, shape [dim]). + embeddings: Embedding matrix (numpy float32 array, shape [num_rows, dim], + C-contiguous). + threshold: Minimum cosine similarity to include in results. + top_k: Maximum number of results to return. + + Returns: + List of dicts with keys 'index' (int) and 'similarity' (float), + sorted by descending similarity. + """ + query = np.ascontiguousarray(query, dtype=np.float32) + embeddings = np.ascontiguousarray(embeddings, dtype=np.float32) + if query.ndim != 1 or embeddings.ndim != 2: + raise ValueError("query must be 1-D, embeddings must be 2-D") + if len(query) != embeddings.shape[1]: + raise ValueError( + f"dimension mismatch: query dim {len(query)} " + f"vs embedding dim {embeddings.shape[1]}") + + cdef int dim = len(query) + cdef int64_t num_rows = embeddings.shape[0] + cdef const float* q_ptr = ( query.ctypes.data) + cdef const float* e_ptr = ( embeddings.ctypes.data) + + cdef vector[CSearchResult] results = c_batch_cosine_search( + q_ptr, dim, e_ptr, num_rows, + threshold, top_k) + + return [{"index": r.index, "similarity": r.similarity} for r in results] \ No newline at end of file diff --git a/python/pycylon/setup.py b/python/pycylon/setup.py index 58f554c87..cfedc14c9 100644 --- a/python/pycylon/setup.py +++ b/python/pycylon/setup.py @@ -45,6 +45,7 @@ CYLON_FMI = strtobool(os.environ.get('CYLON_FMI') or '0') CYLON_LIBFABRIC = strtobool(os.environ.get('CYLON_LIBFABRIC') or '0') CYLON_REDIS = strtobool(os.environ.get('CYLON_REDIS') or '0') +CYLON_SIMD = strtobool(os.environ.get('CYLON_SIMD') or '0') UCX_LOCAL_INSTALL = strtobool(os.environ.get('UCX_LOCAL_INSTALL') or '0') UCC_PREFIX = os.environ.get('UCC_PREFIX') @@ -59,6 +60,7 @@ print("REDIS prefix:", REDIS_PREFIX) print("CYLON FMI: ", CYLON_FMI) print("CYLON LIBFABRIC: ", CYLON_LIBFABRIC) +print("CYLON SIMD: ", CYLON_SIMD) @@ -90,7 +92,7 @@ arrow_lib_dir = pyarrow_location if not os.path.exists(arrow_lib_dir): arrow_lib_dir = os.path.join(pyarrow_location, "lib64") - extra_compile_args.append('-D_GLIBCXX_USE_CXX11_ABI=0') + extra_compile_args.append('-D_GLIBCXX_USE_CXX11_ABI=1') else: arrow_include_dir = os.path.join(ARROW_PREFIX, "include") arrow_lib_dir = os.path.join(ARROW_PREFIX, "lib") @@ -157,7 +159,7 @@ macros = [] # compile_time_env serves as preprocessor macros. ref: https://github.com/cython/cython/issues/2488 -compile_time_env = {'CYTHON_GLOO': False, 'CYTHON_UCC': False, 'CYTHON_UCX': False, 'CYTHON_REDIS': False, 'CYTHON_FMI': False, 'CYTHON_LIBFABRIC': False} +compile_time_env = {'CYTHON_GLOO': False, 'CYTHON_UCC': False, 'CYTHON_UCX': False, 'CYTHON_REDIS': False, 'CYTHON_FMI': False, 'CYTHON_LIBFABRIC': False, 'CYTHON_SIMD': False} if CYLON_GLOO: libraries.append('gloo') library_dirs.append(os.path.join(GLOO_PREFIX, 'lib')) @@ -218,6 +220,9 @@ else: macros.append(('BUILD_CYLON_REDIS', '0')) +if CYLON_SIMD: + macros.append(('BUILD_CYLON_SIMD', '1')) + compile_time_env['CYTHON_SIMD'] = True print('Libraries :', libraries) print("Lib dirs :", library_dirs) @@ -230,6 +235,10 @@ # Adopted the Cudf Python Build format # https://github.com/rapidsai/cudf +cython_exclude = [] +if not CYLON_SIMD: + cython_exclude.append("pycylon/simd/*.pyx") + extensions = [ Extension( "*", @@ -244,7 +253,10 @@ )] compiler_directives = {"profile": False, "language_level": 3, "embedsignature": True} -packages = find_packages(include=["pycylon", "pycylon.*"]) +exclude_packages = [] +if not CYLON_SIMD: + exclude_packages.append("pycylon.simd") +packages = find_packages(include=["pycylon", "pycylon.*"], exclude=exclude_packages) print("PACKAGES: " + str(packages)) @@ -258,6 +270,7 @@ ], ext_modules=cythonize( extensions, + exclude=cython_exclude, nthreads=1, # Single thread for clearer error output compiler_directives=compiler_directives, compile_time_env=compile_time_env, diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 203847500..c93023eea 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -49,11 +49,6 @@ aws-config = { version = "1.5", optional = true } aws-sdk-s3 = { version = "1.65", optional = true } aws-smithy-types = { version = "1.2", optional = true } -# AWS SDK (optional, for S3 checkpoint storage) -aws-config = { version = "1.5", optional = true } -aws-sdk-s3 = { version = "1.65", optional = true } -aws-smithy-types = { version = "1.2", optional = true } - # Serialization serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" @@ -113,248 +108,150 @@ gpu = [] # GPU support via gcylon FFI name = "distributed_join_all_types" required-features = ["mpi"] -[[example]] -name = "distributed_join_mpi" -required-features = ["mpi"] - -[[example]] -name = "distributed_set_ops_mpi" -required-features = ["mpi"] - -[[example]] -name = "dist_sort_example" -required-features = ["mpi"] - -[[example]] -name = "mpi_hello" -required-features = ["mpi"] - -[[example]] -name = "shuffle_mpi" -required-features = ["mpi"] - -[[example]] -name = "join_example" -required-features = ["mpi"] - -[[example]] -name = "union_example" -required-features = ["mpi"] - -[[example]] -name = "intersect_example" -required-features = ["mpi"] - -[[example]] -name = "subtract_example" -required-features = ["mpi"] - -[[example]] -name = "sorting_example" -required-features = ["mpi"] - -[[example]] -name = "unique_example" -required-features = ["mpi"] - -[[example]] -name = "select_example" - -[[example]] -name = "project_example" - -[[example]] -name = "multicolumn_sorting_example" -required-features = ["mpi"] - -[[example]] -name = "multi_idx_join_example" -required-features = ["mpi"] - -[[example]] -name = "table_from_vectors_example" - -[[example]] -name = "compute_example" -required-features = ["mpi"] - -[[example]] -name = "slice_example" -required-features = ["mpi"] - -[[example]] -name = "demo_join" -required-features = ["mpi"] - -[[example]] -name = "parquet_union_example" -required-features = ["mpi", "parquet"] - -[[example]] -name = "parquet_join_example" -required-features = ["mpi", "parquet"] - -[[example]] -name = "create_test_parquet" -required-features = ["parquet"] - -[[example]] -name = "ucx_join_example" -required-features = ["ucx", "ucc", "redis"] - -[[example]] -name = "ucc_operators_example" -required-features = ["ucx", "ucc", "redis"] - -[[example]] -name = "redis_ucc_ucx_example" -required-features = ["ucx", "ucc", "redis"] - -[[example]] -name = "fmi_example" -required-features = ["fmi", "redis"] - -[[example]] -name = "groupby_perf" -required-features = ["mpi"] - -[[example]] -name = "gpu_shuffle" -required-features = ["gpu"] - -[[example]] -name = "gpu_sort" -required-features = ["gpu"] - -[[example]] -name = "gpu_join" -required-features = ["gpu"] - -[[example]] -name = "distributed_join_all_types" -required-features = ["mpi"] [[example]] name = "distributed_join_mpi" required-features = ["mpi"] + [[example]] name = "distributed_set_ops_mpi" required-features = ["mpi"] + [[example]] name = "dist_sort_example" required-features = ["mpi"] + [[example]] name = "mpi_hello" required-features = ["mpi"] + [[example]] name = "shuffle_mpi" required-features = ["mpi"] + [[example]] name = "join_example" required-features = ["mpi"] + [[example]] name = "union_example" required-features = ["mpi"] + [[example]] name = "intersect_example" required-features = ["mpi"] + [[example]] name = "subtract_example" required-features = ["mpi"] + [[example]] name = "sorting_example" required-features = ["mpi"] + [[example]] name = "unique_example" required-features = ["mpi"] + [[example]] name = "select_example" + [[example]] name = "project_example" + [[example]] name = "multicolumn_sorting_example" required-features = ["mpi"] + [[example]] name = "multi_idx_join_example" required-features = ["mpi"] + [[example]] name = "table_from_vectors_example" + [[example]] name = "compute_example" required-features = ["mpi"] + [[example]] name = "slice_example" required-features = ["mpi"] + [[example]] name = "demo_join" required-features = ["mpi"] + [[example]] name = "parquet_union_example" required-features = ["mpi", "parquet"] + [[example]] name = "parquet_join_example" required-features = ["mpi", "parquet"] + [[example]] name = "create_test_parquet" required-features = ["parquet"] + [[example]] name = "ucx_join_example" required-features = ["ucx", "ucc", "redis"] + [[example]] name = "ucc_operators_example" required-features = ["ucx", "ucc", "redis"] + [[example]] name = "redis_ucc_ucx_example" required-features = ["ucx", "ucc", "redis"] + [[example]] name = "fmi_example" required-features = ["fmi", "redis"] + [[example]] name = "groupby_perf" required-features = ["mpi"] + [[example]] name = "gpu_shuffle" required-features = ["gpu"] + [[example]] name = "gpu_sort" required-features = ["gpu"] + [[example]] name = "gpu_join" required-features = ["gpu"] -[profile.release] -lto = true -codegen-units = 1 - -[profile.bench] -lto = true -codegen-units = 1 \ No newline at end of file diff --git a/rust/cylon-node/Cargo.toml b/rust/cylon-node/Cargo.toml index 0f39ebc06..feb2d3e0a 100644 --- a/rust/cylon-node/Cargo.toml +++ b/rust/cylon-node/Cargo.toml @@ -23,6 +23,9 @@ gloo = ["cylon/gloo"] napi = { version = "2", features = ["async", "napi8"] } napi-derive = "2" +# Logging — use cylon's init_logging() which sets WriteStyle::Never for CloudWatch +log = "0.4" + # Use the main cylon crate cylon = { path = ".." } diff --git a/rust/cylon-node/src/lib.rs b/rust/cylon-node/src/lib.rs index 3a3f55ac7..bbbb3881b 100644 --- a/rust/cylon-node/src/lib.rs +++ b/rust/cylon-node/src/lib.rs @@ -160,11 +160,23 @@ pub struct Communicator { inner: Arc, } +/// Initialize the Rust logger once (idempotent). +/// Delegates to cylon::util::logging::init_logging() which uses WriteStyle::Never +/// (no ANSI color codes) and reads RUST_LOG for level — e.g. RUST_LOG=info +/// surfaces TCPunch diagnostics in CloudWatch. +fn init_logger() { + static INIT: std::sync::Once = std::sync::Once::new(); + INIT.call_once(|| { + cylon::util::logging::init_logging(); + }); +} + #[napi] impl Communicator { /// Create a communicator with the specified backend #[napi(factory)] pub fn create(config: CommunicatorConfig) -> Result { + init_logger(); let inner: Arc = match config.comm_type { CommunicatorType::Fmi => { #[cfg(feature = "fmi")] diff --git a/rust/cylon-wasm/src/simd.rs b/rust/cylon-wasm/src/simd.rs index cbb0286f8..9a52cb7e9 100644 --- a/rust/cylon-wasm/src/simd.rs +++ b/rust/cylon-wasm/src/simd.rs @@ -337,3 +337,29 @@ pub fn cosine_similarity_f32(a: &[f32], b: &[f32]) -> f32 { pub fn euclidean_distance_f32(a: &[f32], b: &[f32]) -> f32 { simd_euclidean_distance_f32(a, b) } + +/// Batch cosine search: compare query against all embeddings in a flat buffer. +/// Returns JSON string array of {index, similarity} sorted by descending similarity. +/// +/// # Arguments +/// * `query` - Query vector (f32) +/// * `embeddings` - Flat buffer of embeddings (num_rows * dim floats) +/// * `dim` - Embedding dimension +/// * `threshold` - Minimum cosine similarity +/// * `top_k` - Maximum results to return +#[wasm_bindgen] +pub fn batch_cosine_search_f32( + query: &[f32], + embeddings: &[f32], + dim: usize, + threshold: f32, + top_k: usize, +) -> Result { + let results = cylon::simd::batch_cosine_search(query, embeddings, dim, threshold, top_k); + let json_results: Vec = results + .iter() + .map(|r| serde_json::json!({"index": r.index, "similarity": r.similarity})) + .collect(); + serde_json::to_string(&json_results) + .map_err(|e| JsValue::from_str(&e.to_string())) +} diff --git a/rust/src/arrow/mod.rs b/rust/src/arrow/mod.rs index 036602020..84cc5fdac 100644 --- a/rust/src/arrow/mod.rs +++ b/rust/src/arrow/mod.rs @@ -21,7 +21,6 @@ pub mod arrow_kernels; pub mod arrow_all_to_all; -pub mod arrow_all_to_all; // TODO: Port from cpp/src/cylon/arrow/ // - arrow_buffer.hpp diff --git a/rust/src/io.rs b/rust/src/io.rs index c141bf328..0714c9c13 100644 --- a/rust/src/io.rs +++ b/rust/src/io.rs @@ -21,11 +21,6 @@ pub mod arrow_io; #[cfg(feature = "parquet")] pub mod parquet_config; -#[cfg(feature = "parquet")] -pub mod arrow_io; -#[cfg(feature = "parquet")] -pub mod parquet_config; - pub use csv::{CsvReadOptions, CsvWriteOptions, read_csv, write_csv}; #[cfg(feature = "parquet")] diff --git a/rust/src/lib.rs b/rust/src/lib.rs index 70bdeacde..ef1de3f70 100644 --- a/rust/src/lib.rs +++ b/rust/src/lib.rs @@ -34,6 +34,7 @@ pub mod partition; pub mod row; pub mod scalar; +pub mod simd; pub mod table; pub mod util; diff --git a/rust/src/net/fmi/direct.rs b/rust/src/net/fmi/direct.rs index 8292ffc9f..bb1d7c7e1 100644 --- a/rust/src/net/fmi/direct.rs +++ b/rust/src/net/fmi/direct.rs @@ -32,8 +32,10 @@ use super::tcpunch; /// Connection timeout for hole punching (120 seconds) const HOLEPUNCH_CONNECT_TIMEOUT: u64 = 120000; -/// Maximum TCPunch retry attempts -const MAX_TCPUNCH_TRIES: i32 = 6; +/// Maximum TCPunch retry attempts — increased for Lambda's NAT environment +/// where simultaneous SYN exchange requires more attempts (Lambda coldstarts +/// can take 5-30s; 20 tries × 500ms = 10s gives adequate window). +const MAX_TCPUNCH_TRIES: i32 = 20; /// Create a vector of None options for TcpStream (can't use vec![None; n] since TcpStream doesn't implement Clone) fn create_socket_vec(n: usize) -> Vec> { @@ -110,7 +112,10 @@ impl Direct { Mode::Blocking => "BLOCKING", Mode::NonBlocking => "NONBLOCKING", }; - format!("fmi_pair{}_{}{}", min_id, max_id, mode_str) + // Include comm_name so concurrent experiments with the same rank topology + // (e.g. two ws=2 runs both using fmi_pair0_1) don't cross-pair at the + // rendezvous server. comm_name is unique per run (set to experiment_name). + format!("{}_fmi_pair{}_{}{}", self.comm_name, min_id, max_id, mode_str) } /// Ensure socket is connected to partner (lazy connection establishment) @@ -320,6 +325,9 @@ impl Channel for Direct { continue; } + // Match C++ Direct::init() behavior: only pre-connect NonBlocking + // sockets during init when mode==NONBLOCKING. Blocking sockets + // are connected on-demand in send()/recv() — matching C++ exactly. if self.mode == Mode::NonBlocking { let pair_name = self.get_pairing_name(self.peer_id, i, Mode::NonBlocking); self.check_socket_nbx(i, &pair_name)?; @@ -329,6 +337,13 @@ impl Channel for Direct { Ok(()) } + /// Override bcast() to use self.mode — equivalent to C++ PeerToPeer::bcast() + /// which passes mode through to send/recv. The trait default hardcodes + /// Mode::Blocking, which breaks NonBlocking communicators (nonblocking=true). + fn bcast(&self, buf: Arc, root: PeerNum) -> CylonResult<()> { + self.bcast_async(buf, root, self.mode.clone(), None) + } + fn finalize(&mut self) -> CylonResult<()> { // Sockets will be closed when dropped let mut sockets = self.sockets.write().unwrap(); diff --git a/rust/src/net/fmi/tcpunch.rs b/rust/src/net/fmi/tcpunch.rs index a958f472b..2de6232ad 100644 --- a/rust/src/net/fmi/tcpunch.rs +++ b/rust/src/net/fmi/tcpunch.rs @@ -50,8 +50,10 @@ pub const TOKEN_LENGTH: usize = 37; /// Client request size (100 + 37 + 4 = 141 bytes) pub const CLIENT_REQUEST_SIZE: usize = 141; -/// Server response size (1 + 4 + 2 + 4 + 2 + 37 + 1 = 51 bytes) -pub const SERVER_RESPONSE_SIZE: usize = 51; +/// Server response size — matches C++ SERVER_RESPONSE_SIZE = 50 bytes. +/// Layout: status(1) + your_ip(4) + your_port(2) + peer_ip(4) + peer_port(2) + token(37) = 50. +/// The parser only reads through offset 49; the earlier Rust value of 51 was off-by-one. +pub const SERVER_RESPONSE_SIZE: usize = 50; /// Magic number for validation handshake const VALIDATION_MAGIC: u32 = 0xDEADBEEF; @@ -375,71 +377,90 @@ pub fn configure_keepalive_custom( Ok(()) } -/// Listener thread function - accepts incoming connections +/// Listener thread — exact Rust port of C++ peer_listen(). +/// +/// C++ uses a BLOCKING socket with SO_RCVTIMEO=1s so accept() times out every +/// second and re-checks connection_established. Do NOT set non-blocking here — +/// it would override SO_RCVTIMEO and cause accept() to return WouldBlock +/// immediately every call, which is functionally different and causes the +/// listener to spin rather than wait for the kernel to deliver an incoming SYN. fn peer_listen( local_port: u16, connection_established: Arc, accepting_socket: Arc, + listener_ready: Arc, ) -> CylonResult<()> { use socket2::{Domain, Protocol, Socket, Type}; - // Create socket with reuse options + log::info!("peer_listen: creating listener on port {}", local_port); + let socket = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP)) - .map_err(|e| CylonError::new(Code::IoError, format!("Socket creation failed: {}", e)))?; + .map_err(|e| CylonError::new(Code::IoError, format!("peer_listen: socket create failed: {}", e)))?; + // SO_REUSEADDR + SO_REUSEPORT — matches C++ peer_listen lines 96-101 configure_socket_reuse(&socket)?; + log::info!("peer_listen: SO_REUSEADDR+SO_REUSEPORT set on port {}", local_port); - // Set accept timeout (3 minutes for cloud environments) - socket.set_read_timeout(Some(Duration::from_secs(180))) - .map_err(|e| CylonError::new(Code::IoError, format!("Failed to set timeout: {}", e)))?; + // SO_RCVTIMEO = 1 second, BLOCKING — matches C++ peer_listen lines 105-112. + // This makes accept() block for up to 1 second then return EAGAIN, allowing + // the loop to check connection_established frequently. + // C++ does NOT set O_NONBLOCK on the listen socket. + socket.set_read_timeout(Some(Duration::from_secs(1))) + .map_err(|e| CylonError::new(Code::IoError, format!("peer_listen: SO_RCVTIMEO failed: {}", e)))?; - // Bind to local port + // Bind — matches C++ peer_listen lines 114-123 let addr: SocketAddr = format!("0.0.0.0:{}", local_port).parse().unwrap(); - socket.bind(&addr.into()) - .map_err(|e| CylonError::new(Code::IoError, format!("Could not bind to local port: {}", e)))?; + match socket.bind(&addr.into()) { + Ok(_) => log::info!("peer_listen: bound to 0.0.0.0:{}", local_port), + Err(e) => { + log::error!("peer_listen: bind to port {} failed: {} (errno={:?})", local_port, e, e.raw_os_error()); + return Err(CylonError::new(Code::IoError, format!("peer_listen: bind failed: {}", e))); + } + } - // Listen + // Listen — matches C++ peer_listen lines 125-129 socket.listen(1) - .map_err(|e| CylonError::new(Code::IoError, format!("Listen failed: {}", e)))?; + .map_err(|e| CylonError::new(Code::IoError, format!("peer_listen: listen failed: {}", e)))?; + log::info!("peer_listen: listening on port {} (BLOCKING, SO_RCVTIMEO=1s)", local_port); + + // Signal ready — listener is bound and listening before connect loop starts + listener_ready.store(true, Ordering::SeqCst); + // Convert to TcpListener — stays BLOCKING with 1s read timeout (matches C++) + // Do NOT call set_nonblocking(true) — that is NOT in the C++ code let listener: TcpListener = socket.into(); - listener.set_nonblocking(true) - .map_err(|e| CylonError::new(Code::IoError, format!("Failed to set non-blocking: {}", e)))?; let mut error_count = 0; + // Accept loop — exact match of C++ peer_listen lines 135-166 loop { if connection_established.load(Ordering::SeqCst) { break; } match listener.accept() { - Ok((stream, _peer_addr)) => { - log::info!("Successfully connected to peer via accept"); - - // Store the raw fd + Ok((stream, peer_addr)) => { + log::info!("peer_listen: accepted connection from {}", peer_addr); #[cfg(unix)] { use std::os::unix::io::AsRawFd; accepting_socket.store(stream.as_raw_fd(), Ordering::SeqCst); - // Prevent the stream from being dropped (fd will be managed by caller) std::mem::forget(stream); } - connection_established.store(true, Ordering::SeqCst); return Ok(()); } - Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => { - // No connection yet, sleep briefly and retry - thread::sleep(Duration::from_millis(10)); - continue; + // SO_RCVTIMEO expired (EAGAIN/EWOULDBLOCK/TimedOut) — loop and re-check flag + Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock + || e.kind() == std::io::ErrorKind::TimedOut => { + continue; // matches C++ "if (errno == EAGAIN || EWOULDBLOCK) continue;" } Err(e) => { - log::debug!("Accept error: {}", e); + log::warn!("peer_listen: accept error: {} (os={:?})", e, e.raw_os_error()); error_count += 1; if error_count > 5 { - let backoff_delay = std::cmp::min(100 * (1 << (error_count - 5)), 5000); - thread::sleep(Duration::from_millis(backoff_delay)); + let backoff = std::cmp::min(100 * (1 << (error_count - 5)), 5000); + thread::sleep(Duration::from_millis(backoff)); } } } @@ -448,76 +469,116 @@ fn peer_listen( Ok(()) } -/// Perform hole punching after receiving peer info -fn do_hole_punch( +/// Core hole-punch logic using pre-created listener state. +/// The listener thread must already be spawned and its handle passed in. +/// This allows the WAITING path to start the listener before blocking on the +/// second rendezvous response (matching C++ behaviour), while the PAIRED path +/// creates fresh state and spawns the listener just before calling this. +fn do_hole_punch_inner( your_info: &PeerInfo, peer_info: &PeerInfo, timeout_ms: u64, + connection_established: Arc, + accepting_socket: Arc, + listener_ready: Arc, + listener_handle: thread::JoinHandle>, ) -> CylonResult { use socket2::{Domain, Protocol, Socket, Type}; let local_port = your_info.port; let peer_addr = peer_info.to_socket_addr(); - log::info!("Starting hole punch: local port {}, peer {}", local_port, peer_addr); + log::info!("do_hole_punch_inner: your_port={} peer={}", local_port, peer_addr); - // Start listener thread - let connection_established = Arc::new(AtomicBool::new(false)); - let accepting_socket = Arc::new(AtomicI32::new(-1)); - - let conn_established_clone = connection_established.clone(); - let accepting_socket_clone = accepting_socket.clone(); - - let listener_handle = thread::spawn(move || { - peer_listen(local_port, conn_established_clone, accepting_socket_clone) - }); - - // Create socket for active connection attempts + // Create peer socket — SO_REUSEADDR + SO_REUSEPORT + NON-BLOCKING. + // Matches C++ do_hole_punch lines 182-191. let peer_socket = Socket::new(Domain::IPV4, Type::STREAM, Some(Protocol::TCP)) - .map_err(|e| CylonError::new(Code::IoError, format!("Socket creation failed: {}", e)))?; + .map_err(|e| CylonError::new(Code::IoError, format!("peer socket create failed: {}", e)))?; configure_socket_reuse(&peer_socket)?; + + // Set NON-BLOCKING on the ACTIVE (connect) socket only — matches C++ line 189. + // The LISTENER socket is kept BLOCKING with SO_RCVTIMEO (see peer_listen). peer_socket.set_nonblocking(true) - .map_err(|e| CylonError::new(Code::IoError, format!("Failed to set non-blocking: {}", e)))?; + .map_err(|e| CylonError::new(Code::IoError, format!("peer socket set_nonblocking failed: {}", e)))?; + log::info!("do_hole_punch_inner: peer socket created, SO_REUSEADDR+SO_REUSEPORT+NONBLOCKING"); + + // Bail macro: signal listener to stop and join before any early return + macro_rules! bail { + ($err:expr) => {{ + connection_established.store(true, Ordering::SeqCst); + let _ = listener_handle.join(); + return Err($err); + }}; + } - // Bind to same local port + // Bind peer socket to same local port as rendezvous — matches C++ line 196-204. + // Both listener and peer socket share this port via SO_REUSEPORT. let local_addr: SocketAddr = format!("0.0.0.0:{}", local_port).parse().unwrap(); - peer_socket.bind(&local_addr.into()) - .map_err(|e| CylonError::new(Code::IoError, - format!("Binding to same port failed: {}", e)))?; + match peer_socket.bind(&local_addr.into()) { + Ok(_) => log::info!("do_hole_punch_inner: peer socket bound to 0.0.0.0:{}", local_port), + Err(e) => { + log::error!("do_hole_punch_inner: peer socket bind to {} failed: {} (errno={:?})", + local_port, e, e.raw_os_error()); + bail!(CylonError::new(Code::IoError, format!("peer socket bind failed: {}", e))); + } + } - // Attempt connection with retries + // Wait for listener to bind and listen before starting connect attempts. + // C++ avoids this race by spawning the listener before calling do_hole_punch(). + // Without this wait, connect() can start before the LISTEN socket is ready, + // so the peer's incoming SYN has nowhere to land and gets dropped. + let wait_start = Instant::now(); + while !listener_ready.load(Ordering::SeqCst) { + if wait_start.elapsed() > Duration::from_secs(5) { + log::warn!("Listener did not become ready within 5s — proceeding anyway"); + break; + } + thread::sleep(Duration::from_millis(1)); + } + log::info!("Listener ready after {}ms", wait_start.elapsed().as_millis()); + + // Active connect loop — exact Rust port of C++ do_hole_punch lines 215-247. + // C++ does NOT sleep on EINPROGRESS/EALREADY — it tight-loops. + // Sleeping 10ms here (as Rust did before) delays SYN delivery and breaks + // the simultaneous-open timing that TCPunch requires. + log::info!("do_hole_punch: starting connect loop to {} from local port {} (timeout={}ms)", + peer_addr, local_port, timeout_ms); let start_time = Instant::now(); let max_connection_time = Duration::from_millis(timeout_ms); - let mut attempt_count = 0; + let mut attempt_count = 0u64; let mut connected = false; while !connection_established.load(Ordering::SeqCst) { if start_time.elapsed() > max_connection_time { - connection_established.store(true, Ordering::SeqCst); // Signal listener to stop - return Err(CylonError::new(Code::IoError, "Connection timeout".to_string())); + bail!(CylonError::new(Code::IoError, + format!("Connection timeout after {}ms, {} attempts", timeout_ms, attempt_count))); } match peer_socket.connect(&peer_addr.into()) { Ok(()) => { - log::info!("Successfully connected to peer via connect()"); + log::info!("do_hole_punch: connect() succeeded after {} attempts", attempt_count); connected = true; break; } Err(ref e) if e.raw_os_error() == Some(libc::EISCONN) => { - log::info!("Successfully connected to peer (EISCONN)"); + log::info!("do_hole_punch: EISCONN (connected) after {} attempts", attempt_count); connected = true; break; } + // EINPROGRESS/EALREADY/EAGAIN — matches C++ "continue" with NO sleep (tight loop) Err(ref e) if e.raw_os_error() == Some(libc::EALREADY) || e.raw_os_error() == Some(libc::EAGAIN) || e.raw_os_error() == Some(libc::EINPROGRESS) => { attempt_count += 1; - thread::sleep(Duration::from_millis(10)); + // C++ does NOT sleep here — tight polling loop continue; } - Err(_e) => { - let base_delay = 100; + Err(ref e) => { + // Log every error with full details for triage + log::warn!("do_hole_punch: connect attempt {} errno={:?} kind={:?} elapsed={}ms", + attempt_count, e.raw_os_error(), e.kind(), start_time.elapsed().as_millis()); + let base_delay = 100u64; let backoff_delay = base_delay * (1 + attempt_count / 10); thread::sleep(Duration::from_millis(std::cmp::min(backoff_delay, 1000))); attempt_count += 1; @@ -600,6 +661,47 @@ fn do_hole_punch( Ok(peer_stream) } +/// Perform hole punching for the PAIRED path — creates fresh listener state and +/// spawns the listener thread, then delegates to do_hole_punch_inner. +fn do_hole_punch( + your_info: &PeerInfo, + peer_info: &PeerInfo, + timeout_ms: u64, +) -> CylonResult { + let local_port = your_info.port; + let connection_established = Arc::new(AtomicBool::new(false)); + let accepting_socket = Arc::new(AtomicI32::new(-1)); + let listener_ready = Arc::new(AtomicBool::new(false)); + let wce = connection_established.clone(); + let was = accepting_socket.clone(); + let wlr = listener_ready.clone(); + let handle = thread::spawn(move || peer_listen(local_port, wce, was, wlr)); + do_hole_punch_inner(your_info, peer_info, timeout_ms, + connection_established, accepting_socket, listener_ready, handle) +} + +/// Perform hole punching using pre-created listener state (WAITING path). +/// The listener thread was already spawned before waiting for the second +/// rendezvous response, matching C++ behaviour. +fn do_hole_punch_with_listener( + your_info: &PeerInfo, + peer_info: &PeerInfo, + timeout_ms: u64, + connection_established: Arc, + accepting_socket: Arc, + listener_ready: Arc, +) -> CylonResult { + // The listener thread is already running; we need its handle to join it. + // Since we can't pass the handle through the WAITING-path Arc sharing, we + // re-use do_hole_punch_inner directly with the pre-created arcs. The handle + // is held by the WAITING path caller which joins it unconditionally. + // We create a dummy no-op thread here just to satisfy the signature; + // the real listener handle is joined by the caller. + let dummy_handle = thread::spawn(|| Ok(())); + do_hole_punch_inner(your_info, peer_info, timeout_ms, + connection_established, accepting_socket, listener_ready, dummy_handle) +} + /// Establish a peer-to-peer connection using TCP NAT hole punching (Protocol v2) /// /// # Arguments @@ -644,9 +746,14 @@ pub fn pair_with_retries( let timeout_ms = if timeout_ms == 0 { DEFAULT_TIMEOUT_MS } else { timeout_ms }; let timeout = Duration::from_millis(timeout_ms); + // Resolve rendezvous address — supports both IP literals and DNS hostnames. + use std::net::ToSocketAddrs; let server_addr: SocketAddr = format!("{}:{}", server_address, port) - .parse() - .map_err(|e| CylonError::new(Code::Invalid, format!("Invalid server address: {}", e)))?; + .to_socket_addrs() + .map_err(|e| CylonError::new(Code::IoError, format!("Failed to resolve rendezvous '{}:{}': {}", server_address, port, e)))? + .find(|a| a.is_ipv4()) + .ok_or_else(|| CylonError::new(Code::IoError, format!("No IPv4 address for rendezvous '{}'", server_address)))?; + log::info!("Resolved rendezvous '{}:{}' → {}", server_address, port, server_addr); let mut reconnect_token: Option = None; @@ -696,35 +803,89 @@ pub fn pair_with_retries( CylonError::new(Code::IoError, "No peer info in PAIRED response".to_string()) })?; - log::info!("Paired immediately with peer at {}:{}", peer.ip, peer.port); - return do_hole_punch(&resp.your_info, &peer, timeout_ms); + // Use server-reported your_port — exactly like C++ (public_info.port = resp.your_port). + // This is the external NAT port the peer knows about and will connect to. + let your_port = resp.your_info.port; + log::info!("pair PAIRED: your_port={} (server-reported), peer={}:{}", + your_port, peer.ip, peer.port); + // Keep stream alive during hole punch — preserves NAT entry (matches C++) + let result = do_hole_punch(&resp.your_info, &peer, timeout_ms); + drop(stream); + return result; } PairingStatus::Waiting => { log::debug!("Registered, waiting for peer (token: {})", resp.token); + // C++ starts the listener thread HERE — before blocking on the second + // rendezvous response — so it is already bound and listening by the + // time the peer info arrives and do_hole_punch() begins. + // Rust previously started the listener only inside do_hole_punch(), + // which is too late: the peer's SYN could arrive before the LISTEN + // socket exists. + // Use server-reported your_port — same as C++ (matches PAIRED path above) + let your_port_waiting = resp.your_info.port; + log::info!("pair WAITING: your_port={} (server-reported)", your_port_waiting); + + let wait_conn_established = Arc::new(AtomicBool::new(false)); + let wait_accepting_socket = Arc::new(AtomicI32::new(-1)); + let wait_listener_ready = Arc::new(AtomicBool::new(false)); + let wce = wait_conn_established.clone(); + let was = wait_accepting_socket.clone(); + let wlr = wait_listener_ready.clone(); + let listener_handle_waiting = thread::spawn(move || { + peer_listen(your_port_waiting, wce, was, wlr) + }); + // Wait for second response with peer info - match stream.read_exact(&mut resp_buf) { + let hole_punch_result = match stream.read_exact(&mut resp_buf) { Ok(()) => { let resp2 = parse_response(&resp_buf); - if resp2.status == PairingStatus::Paired { let peer = resp2.peer_info.ok_or_else(|| { CylonError::new(Code::IoError, "No peer info in PAIRED response".to_string()) - })?; - - log::info!("Peer found: {}:{}", peer.ip, peer.port); - return do_hole_punch(&resp.your_info, &peer, timeout_ms); + }); + match peer { + Ok(peer) => { + log::info!("Peer found: {}:{} (your_port={}) conn_established={}", + peer.ip, peer.port, your_port_waiting, + wait_conn_established.load(Ordering::SeqCst)); + // Do NOT reset connection_established here — if the listener already + // accepted the peer's incoming SYN while we were waiting for the + // second rendezvous response, resetting to false would discard that + // accepted socket and cause the connect loop to run forever. + // (C++ never resets connection_established between listener spawn + // and do_hole_punch; the flag flows through continuously.) + let result = do_hole_punch_with_listener( + &resp.your_info, &peer, timeout_ms, + wait_conn_established.clone(), + wait_accepting_socket.clone(), + wait_listener_ready.clone(), + ); + drop(stream); + Some(result) + } + Err(e) => { Some(Err(e)) } + } } else { log::warn!("Unexpected status after WAITING: {:?}", resp2.status); - // Fall through to retry + None } } Err(e) => { log::warn!("Timeout waiting for peer (attempt {}): {}", attempt + 1, e); - // Fall through to retry with token + None } + }; + + // Always join the pre-started listener thread + wait_conn_established.store(true, Ordering::SeqCst); + let _ = listener_handle_waiting.join(); + + if let Some(result) = hole_punch_result { + return result; } + // Fall through to retry } PairingStatus::Timeout => { @@ -766,9 +927,12 @@ pub fn remove_pair( ) -> CylonResult<()> { let timeout = Duration::from_millis(if timeout_ms == 0 { DEFAULT_TIMEOUT_MS } else { timeout_ms }); + use std::net::ToSocketAddrs; let server_addr: SocketAddr = format!("{}:{}", server_address, port) - .parse() - .map_err(|e| CylonError::new(Code::Invalid, format!("Invalid server address: {}", e)))?; + .to_socket_addrs() + .map_err(|e| CylonError::new(Code::IoError, format!("Failed to resolve '{}:{}': {}", server_address, port, e)))? + .find(|a| a.is_ipv4()) + .ok_or_else(|| CylonError::new(Code::IoError, format!("No IPv4 address for '{}'", server_address)))?; let mut stream = TcpStream::connect_timeout(&server_addr, timeout) .map_err(|e| CylonError::new(Code::IoError, diff --git a/rust/src/net/serialize.rs b/rust/src/net/serialize.rs index affa60e8d..40b919438 100644 --- a/rust/src/net/serialize.rs +++ b/rust/src/net/serialize.rs @@ -59,29 +59,6 @@ impl IpcCompression { } } -/// Compression algorithm for Arrow IPC serialization. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] -pub enum IpcCompression { - /// No compression - #[default] - None, - /// LZ4 frame compression (fast) - Lz4, - /// Zstandard compression (better ratio) - Zstd, -} - -impl IpcCompression { - /// Convert to Arrow IPC CompressionType - fn to_arrow_compression(self) -> Option { - match self { - IpcCompression::None => None, - IpcCompression::Lz4 => Some(CompressionType::LZ4_FRAME), - IpcCompression::Zstd => Some(CompressionType::ZSTD), - } - } -} - /// Serialize an Arrow RecordBatch to bytes using Arrow IPC format /// /// # Arguments diff --git a/rust/src/simd.rs b/rust/src/simd.rs new file mode 100644 index 000000000..f79a5339e --- /dev/null +++ b/rust/src/simd.rs @@ -0,0 +1,71 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! SIMD-accelerated similarity search primitives. +//! +//! Provides cosine similarity and batch cosine search over float32 vectors. +//! These are core Cylon primitives used by downstream crates (cylon-armada) +//! for embedding-based context reuse. + +/// Result of a similarity search. +#[derive(Debug, Clone)] +pub struct SearchResult { + pub index: usize, + pub similarity: f32, +} + +/// Compute cosine similarity between two float32 slices. +/// Returns 0.0 if either vector has zero magnitude. +pub fn cosine_similarity_f32(a: &[f32], b: &[f32]) -> f32 { + debug_assert_eq!(a.len(), b.len()); + let mut dot_ab = 0.0f32; + let mut dot_aa = 0.0f32; + let mut dot_bb = 0.0f32; + for i in 0..a.len() { + dot_ab += a[i] * b[i]; + dot_aa += a[i] * a[i]; + dot_bb += b[i] * b[i]; + } + let denom = dot_aa.sqrt() * dot_bb.sqrt(); + if denom == 0.0 { + 0.0 + } else { + dot_ab / denom + } +} + +/// Batch cosine search over a flat embedding buffer. +/// Returns up to `top_k` results with similarity >= `threshold`, +/// sorted by descending similarity. +pub fn batch_cosine_search( + query: &[f32], + embeddings: &[f32], + dim: usize, + threshold: f32, + top_k: usize, +) -> Vec { + if top_k == 0 || dim == 0 || embeddings.is_empty() { + return vec![]; + } + let num_rows = embeddings.len() / dim; + let mut results: Vec = Vec::new(); + for i in 0..num_rows { + let row = &embeddings[i * dim..(i + 1) * dim]; + let sim = cosine_similarity_f32(query, row); + if sim >= threshold { + results.push(SearchResult { index: i, similarity: sim }); + } + } + results.sort_by(|a, b| b.similarity.partial_cmp(&a.similarity).unwrap()); + results.truncate(top_k); + results +} \ No newline at end of file diff --git a/rust/src/table.rs b/rust/src/table.rs index c9fcf347e..9b2263c25 100644 --- a/rust/src/table.rs +++ b/rust/src/table.rs @@ -26,9 +26,6 @@ use crate::net::serialize::{deserialize_record_batch, serialize_record_batch}; pub mod column; pub use column::{Column, FromVector}; -pub mod column; -pub use column::Column; - /// Table provides the main API for using cylon for data processing /// Corresponds to C++ Table class from cpp/src/cylon/table.hpp #[derive(Clone)] diff --git a/rust/src/util/logging.rs b/rust/src/util/logging.rs index 08f271a98..8fc7d8f14 100644 --- a/rust/src/util/logging.rs +++ b/rust/src/util/logging.rs @@ -17,15 +17,23 @@ use log::{debug, error, info, trace, warn}; -/// Initialize logging with default configuration +/// Initialize logging with default configuration. +/// Writes to stdout so Lambda Node.js runtime captures it in CloudWatch. +/// Colors are disabled (WriteStyle::Never) for clean log output. pub fn init_logging() { - env_logger::init(); + env_logger::Builder::from_default_env() + .write_style(env_logger::WriteStyle::Never) + .target(env_logger::Target::Stdout) + .init(); } -/// Initialize logging with specific level +/// Initialize logging with specific level. +/// Writes to stdout so Lambda Node.js runtime captures it in CloudWatch. pub fn init_logging_with_level(level: log::LevelFilter) { env_logger::Builder::from_default_env() .filter_level(level) + .write_style(env_logger::WriteStyle::Never) + .target(env_logger::Target::Stdout) .init(); } diff --git a/rust/tests/simd_test.rs b/rust/tests/simd_test.rs new file mode 100644 index 000000000..1540ecfc4 --- /dev/null +++ b/rust/tests/simd_test.rs @@ -0,0 +1,75 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use cylon::simd::{batch_cosine_search, cosine_similarity_f32}; + +#[test] +fn test_cosine_identical() { + let a = vec![1.0, 2.0, 3.0, 4.0]; + assert!((cosine_similarity_f32(&a, &a) - 1.0).abs() < 1e-5); +} + +#[test] +fn test_cosine_orthogonal() { + let a = vec![1.0, 0.0, 0.0, 0.0]; + let b = vec![0.0, 1.0, 0.0, 0.0]; + assert!(cosine_similarity_f32(&a, &b).abs() < 1e-5); +} + +#[test] +fn test_cosine_opposite() { + let a = vec![1.0, 2.0, 3.0]; + let b = vec![-1.0, -2.0, -3.0]; + assert!((cosine_similarity_f32(&a, &b) + 1.0).abs() < 1e-5); +} + +#[test] +fn test_cosine_zero_vector() { + let a = vec![1.0, 2.0, 3.0]; + let b = vec![0.0, 0.0, 0.0]; + assert_eq!(cosine_similarity_f32(&a, &b), 0.0); +} + +#[test] +fn test_batch_search_basic() { + let query = vec![1.0, 0.0, 0.0, 0.0]; + let embeddings = vec![ + 1.0, 0.0, 0.0, 0.0, + 0.0, 1.0, 0.0, 0.0, + 0.9, 0.1, 0.0, 0.0, + ]; + let results = batch_cosine_search(&query, &embeddings, 4, 0.5, 10); + assert_eq!(results.len(), 2); + assert_eq!(results[0].index, 0); + assert_eq!(results[1].index, 2); +} + +#[test] +fn test_batch_search_top_k() { + let query = vec![1.0, 0.0, 0.0, 0.0]; + let embeddings = vec![ + 1.0, 0.0, 0.0, 0.0, + 0.9, 0.1, 0.0, 0.0, + 0.8, 0.2, 0.0, 0.0, + ]; + let results = batch_cosine_search(&query, &embeddings, 4, 0.0, 1); + assert_eq!(results.len(), 1); + assert_eq!(results[0].index, 0); +} + +#[test] +fn test_batch_search_top_k_zero() { + let query = vec![1.0, 0.0, 0.0, 0.0]; + let embeddings = vec![1.0, 0.0, 0.0, 0.0]; + let results = batch_cosine_search(&query, &embeddings, 4, 0.0, 0); + assert!(results.is_empty()); +} \ No newline at end of file