chore: import upstream snapshot with attribution

2026-07-13 12:47:05 +08:00
commit 4f3b7da785
7394 changed files with 2005594 additions and 0 deletions
@@ -0,0 +1,231 @@
+/* ******************************************************************************
+ *
+ *
+ * This program and the accompanying materials are made available under the
+ * terms of the Apache License, Version 2.0 which is available at
+ * https://www.apache.org/licenses/LICENSE-2.0.
+ *
+ *  See the NOTICE file distributed with this work for additional
+ *  information regarding copyright ownership.
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ ******************************************************************************/
+
+//
+// @author raver119@gmail.com
+//
+#include <array/DataTypeUtils.h>
+#include <array/InteropDataBuffer.h>
+#include <execution/AffinityManager.h>
+#include <helpers/logger.h>
+
+namespace sd {
+InteropDataBuffer::InteropDataBuffer(InteropDataBuffer* dataBuffer, uint64_t length) {
+ if(dataBuffer == nullptr) {
+        THROW_EXCEPTION("InteropDataBuffer::InteropDataBuffer(InteropDataBuffer& dataBuffer, uint64_t length, uint64_t offset) - dataBuffer is nullptr");
+ }
+  if(dataBuffer->_dataBuffer->getDataType() == DataType::UNKNOWN)
+    THROW_EXCEPTION("InteropDataBuffer::InteropDataBuffer(InteropDataBuffer& dataBuffer, uint64_t length, uint64_t offset) - dataBuffer has unknown data type");
+  _dataBuffer = dataBuffer->dataBuffer();
+  _dataType = dataBuffer->_dataType;
+
+  _cachedLenInBytes = length * DataTypeUtils::sizeOf(_dataType);
+  // Cache pointers for deallocation tracking
+  if (_dataBuffer != nullptr) {
+    _cachedPrimaryPtr = _dataBuffer->primary();
+    _cachedSpecialPtr = _dataBuffer->special();
+  }
+
+  owner = false;
+}
+
+InteropDataBuffer::InteropDataBuffer(DataBuffer * databuffer) {
+  _dataBuffer = databuffer;
+  _dataType = databuffer->getDataType();
+  if(_dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION(
+        "InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) - data type is unknown");
+  }
+  // Cache the size to avoid accessing freed memory later
+  _cachedLenInBytes = databuffer != nullptr ? databuffer->getLenInBytes() : 0;
+  // Cache pointers for deallocation tracking
+  if (databuffer != nullptr) {
+    _cachedPrimaryPtr = databuffer->primary();
+    _cachedSpecialPtr = databuffer->special();
+  }
+  // When wrapping an existing DataBuffer, we don't own it by default
+  owner = false;
+}
+
+InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) {
+  if(dtype == DataType::UNKNOWN) {
+    THROW_EXCEPTION(
+        "InteropDataBuffer::InteropDataBuffer(size_t lenInBytes, DataType dtype, bool allocateBoth) - data type is unknown");
+  }
+
+  _cachedLenInBytes = lenInBytes;
+
+  if (lenInBytes == 0) {
+    _dataBuffer = nullptr;
+    this->_dataType = dtype;
+
+  } else {
+    //note this should be size in bytes hence why we multiply the number of elements by the size of the data type
+    _dataBuffer = new DataBuffer(lenInBytes, dtype, nullptr, allocateBoth);
+    this->_dataType = dtype;
+    this->markOwner(true);
+    // Cache pointers for deallocation tracking
+    _cachedPrimaryPtr = _dataBuffer->primary();
+    _cachedSpecialPtr = _dataBuffer->special();
+  }
+}
+
+
+void InteropDataBuffer::printDbAllocationTrace() {
+  if(_dataBuffer == nullptr)
+    return;
+  _dataBuffer->printAllocationTrace();
+}
+
+void InteropDataBuffer::markOwner(bool owner) {
+  this->owner = owner;
+  if(_dataBuffer != nullptr && !_closed) {
+    this->_dataBuffer->_isOwnerPrimary = owner;
+    this->_dataBuffer->_isOwnerSpecial = owner;
+  }
+}
+
+DataBuffer * InteropDataBuffer::getDataBuffer() const {
+  //this can effect size of calculations among others
+  if(_dataType == DataType::UNKNOWN) {
+    THROW_EXCEPTION("All interop buffers must have a known data type.");
+  }
+  // Don't access _dataBuffer if it's been closed/freed
+  if(_dataBuffer != nullptr && !_closed && _dataBuffer->_dataType == DataType::UNKNOWN) {
+    _dataBuffer->_dataType = _dataType;
+  }
+  // Return nullptr if closed to prevent use-after-free
+  return _closed ? nullptr : _dataBuffer;
+}
+
+DataBuffer * InteropDataBuffer::dataBuffer() {
+  if(_dataBuffer == nullptr) {
+    return nullptr;
+  }
+  return _dataBuffer;
+}
+
+
+
+void* InteropDataBuffer::primary() const {
+  if(_dataBuffer == nullptr || _dataBuffer->primary() == nullptr) {
+    return nullptr;
+  }
+  return reinterpret_cast<int8_t*>(_dataBuffer->primary());
+}
+
+void* InteropDataBuffer::special() const {
+  if(_dataBuffer == nullptr)
+    return nullptr;
+  if(_dataBuffer->special() == nullptr) {
+    return nullptr;
+  }
+  return reinterpret_cast<int8_t*>(_dataBuffer->special());
+}
+
+void InteropDataBuffer::setSpecial(void* ptr, size_t length) {
+  if(_dataBuffer == nullptr)
+    THROW_EXCEPTION("InteropDataBuffer::setSpecial() - _dataBuffer is nullptr");
+  if(_closed)
+    return;  // Silently ignore if buffer was already closed
+  _dataBuffer->setSpecialBuffer(ptr, length);
+  // Update cached pointer
+  _cachedSpecialPtr = ptr;
+}
+
+void InteropDataBuffer::setPrimary(void* ptr, size_t length) {
+  if(_dataBuffer == nullptr)
+    THROW_EXCEPTION("InteropDataBuffer::setPrimary() - _dataBuffer is nullptr");
+  if(_closed)
+    return;  // Silently ignore if buffer was already closed
+  _dataBuffer->setPrimaryBuffer(ptr, length);
+  // Update cached pointer
+  _cachedPrimaryPtr = ptr;
+}
+
+void InteropDataBuffer::setDeviceId(int deviceId) {
+  if(_dataBuffer == nullptr || _closed)
+    return;
+  _dataBuffer->setDeviceId(deviceId);
+}
+int InteropDataBuffer::deviceId() const {
+  if(_dataBuffer == nullptr || _closed)
+    return 0;
+  return _dataBuffer->deviceId();
+}
+
+int InteropDataBuffer::useCount() const {
+  return 1;
+}
+
+void InteropDataBuffer::registerSpecialUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                           const std::vector<const InteropDataBuffer*>& readList) {
+  for (const auto& v : writeList) {
+    if (v == nullptr) continue;
+
+    v->getDataBuffer()->writeSpecial();
+  }
+}
+
+void InteropDataBuffer::prepareSpecialUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                          const std::vector<const InteropDataBuffer*>& readList,
+                                          bool synchronizeWritables) {
+  auto currentDeviceId = AffinityManager::currentDeviceId();
+  for (const auto& v : readList) {
+    if (v == nullptr || v->_closed) continue;
+
+    auto db = v->getDataBuffer();
+    if(db == nullptr) continue;
+
+    if (db->deviceId() != currentDeviceId) db->migrate();
+    db->syncToSpecial();
+  }
+
+  // we don't tick write list, only ensure the same device affinity
+  for (const auto& v : writeList) {
+    if (v == nullptr || v->_closed) continue;
+
+    auto db = v->getDataBuffer();
+    if(db == nullptr) continue;
+
+    // special case for legacy ops - views can be updated on host side, thus original array can be not updated
+    if (!db->isSpecialActual()) db->syncToSpecial();
+
+    if (db->deviceId() != currentDeviceId) db->migrate();
+  }
+}
+
+void InteropDataBuffer::registerPrimaryUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                           const std::vector<const InteropDataBuffer*>& readList) {
+
+}
+
+void InteropDataBuffer::preparePrimaryUse(const std::vector<const InteropDataBuffer*>& writeList,
+                                          const std::vector<const InteropDataBuffer*>& readList,
+                                          bool synchronizeWritables) {
+
+}
+
+void InteropDataBuffer::expand(size_t newlength) {
+  if(_dataBuffer == nullptr || _closed)
+    return;  // Cannot expand a closed or null buffer
+  _dataBuffer->expand(newlength * DataTypeUtils::sizeOf(_dataBuffer->getDataType()));
+}
+
+
+}  // namespace sd