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,140 @@
+/* ******************************************************************************
+ *
+ * Copyright (c) 2024 Konduit K.K.
+ * 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.
+ *
+ * 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
+ ******************************************************************************/
+
+#include <array/TadDescriptor.h>
+#include <array/TadPack.h>
+#include <helpers/ConstantTadHelper.h>
+#include <helpers/ShapeUtils.h>
+
+namespace sd {
+
+ConstantTadHelper& ConstantTadHelper::getInstance() {
+  static ConstantTadHelper instance;
+  return instance;
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType dimension) {
+  return tadForDimensions(originalShape, &dimension, 1);
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, std::vector<LongType>* dimensions) {
+  if (dimensions == nullptr) {
+    THROW_EXCEPTION("Dimensions vector is null");
+  }
+  return tadForDimensions(originalShape, const_cast<LongType*>(dimensions->data()), dimensions->size());
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(TadDescriptor* descriptor) {
+  if (descriptor == nullptr) {
+    THROW_EXCEPTION("TadDescriptor is null");
+  }
+  return tadForDimensions(descriptor->originalShape(), descriptor->axis().data(),
+                          descriptor->axis().size());
+}
+
+std::shared_ptr<TadPack> ConstantTadHelper::tadForDimensions(LongType* originalShape, LongType* dimensions, LongType dimLength) {
+  if (originalShape == nullptr) {
+    THROW_EXCEPTION("Original shape is null");
+  }
+
+  if (dimensions == nullptr && dimLength > 0) {
+    THROW_EXCEPTION("Dimensions array is null but dimLength > 0");
+  }
+
+  // Check for empty array
+  if (shape::isEmptyConst(originalShape)) {
+    THROW_EXCEPTION("Cannot create TADs for empty array");
+  }
+
+  sd::LongType rank = shape::rank(originalShape);
+  if (rank < 0) {
+    THROW_EXCEPTION("Invalid shape rank");
+  }
+
+  // Check for zero-sized dimensions
+  for (LongType i = 0; i < rank; i++) {
+    if (shape::sizeAt(originalShape, i) == 0) {
+      THROW_EXCEPTION("Cannot create TADs for array with zero-sized dimensions");
+    }
+  }
+
+  // Handle zero dimension length case - treat entire array as single TAD
+  if (dimLength <= 0) {
+    // When no dimensions specified, create TAD along all dimensions
+    // This means the entire array is treated as a single TAD
+    std::vector<LongType> allDims;
+    for (LongType i = 0; i < rank; i++) {
+      allDims.push_back(i);
+    }
+
+    // Recursively call with all dimensions
+    return tadForDimensions(originalShape, allDims.data(), rank);
+  }
+
+  // Additional validation: check if dimensions are within valid range
+  for (LongType i = 0; i < dimLength; i++) {
+    LongType dim = dimensions[i];
+    if (dim < 0) dim += rank;  // Handle negative dimensions
+    if (dim < 0 || dim >= rank) {
+      THROW_EXCEPTION("Dimension index is out of bounds");
+    }
+  }
+
+  // Create non-temporary vector to satisfy the reference requirement
+  std::vector<LongType> dims(dimensions, dimensions + dimLength);
+
+  // The shared_ptr keeps the TadPack alive even if the cache tries to clear it
+  std::shared_ptr<TadPack> result = nullptr;
+  try {
+    result = _trie.getOrCreate(dims, originalShape);
+  } catch (const std::exception& e) {
+    THROW_EXCEPTION("Failed to create or retrieve TAD pack");
+  }
+
+  // DO NOT call checkAndCleanupCaches() here - would delete the pack we just created!
+  // Cleanup happens at NativeOps layer AFTER operations complete.
+  return result;
+}
+
+void ConstantTadHelper::clearCache() {
+  _trie.clear();
+}
+
+LongType ConstantTadHelper::getCachedEntries() const {
+  return _trie.getCachedEntries();
+}
+
+LongType ConstantTadHelper::getCachedBytes() const {
+  return _trie.getCachedBytes();
+}
+
+LongType ConstantTadHelper::getPeakCachedEntries() const {
+  return _trie.getPeakCachedEntries();
+}
+
+LongType ConstantTadHelper::getPeakCachedBytes() const {
+  return _trie.getPeakCachedBytes();
+}
+
+std::string ConstantTadHelper::toString(int maxDepth, int maxEntries) const {
+  return _trie.toString(maxDepth, maxEntries);
+}
+
+void ConstantTadHelper::getCachedPointers(std::unordered_set<void*>& out_pointers) const {
+  _trie.getCachedPointers(out_pointers);
+}
+
+} // namespace sd