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,215 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//
+
+#ifndef LIBND4J_HELPER_PTRMAP_H
+#define LIBND4J_HELPER_PTRMAP_H
+
+namespace sd {
+
+/**
+ * This class is a simple wrapper to represent batch arguments as single surface of parameters.
+ * So we pass batch parameters as single surface, and then we use this helper to extract arguments for each aggregates.
+ *
+ * Surface map format is simple:
+ * [0] we put numbers for num*Arguments
+ * [1] then we put indexing arguments, since their size is constant
+ * [2] here we put block of JVM IntArrays by value, batchLimit * maxIntArrays * maxArraySize;
+ * [3] then we put real arguments
+ * [4] then we put arguments pointers
+ * [5] then we put shape pointers
+ *
+ */
+template <typename T>
+class PointersHelper {
+ private:
+  int aggregates;
+  void *ptrGeneral;
+
+  // we enforce maximal batch size limit, to simplify
+#ifdef __CUDACC__
+  const int batchLimit = 8192;
+#else
+  const int batchLimit = 512;
+#endif
+
+  // we have 5 diff kinds of arguments: arguments, shapeArguments, intArrayArguments, indexArguments, realArguments
+  const int argTypes = 5;
+
+  int maxIntArrays;
+  int maxArraySize;
+
+  // right now we hardcode maximas, but we'll probably change that later
+  int maxIndexArguments;
+  int maxRealArguments;
+
+  // since that's pointers (which is 64-bit on 64bit systems), we limit number of maximum arguments to 1/2 of maxIndex
+  // arguments
+  int maxArguments;
+  int maxShapeArguments;
+
+  int sizeT;
+  int sizePtr;
+
+ public:
+  /**
+   * We accept single memory chunk and number of jobs stored.
+   *
+   * @param ptrToParams pointer to "surface"
+   * @param numAggregates actual number of aggregates being passed in
+   * @return
+   */
+  SD_HOST_DEVICE
+  PointersHelper(void *ptrToParams, int numAggregates, int maxArgs, int maxShapes, int maxIntArrays,
+                 int maxIntArraySize, int maxIdx, int maxReals) {
+    aggregates = numAggregates;
+    ptrGeneral = ptrToParams;
+
+    // ptrSize for hypothetical 32-bit compatibility
+    sizePtr = sizeof(ptrToParams);
+
+    // unfortunately we have to know sizeOf(T)
+    sizeT = sizeof(T);
+
+    this->maxIntArrays = maxIntArrays;
+    this->maxArraySize = maxIntArraySize;
+    this->maxIndexArguments = maxIdx;
+    this->maxArguments = maxArgs;
+    this->maxShapeArguments = maxShapes;
+    this->maxRealArguments = maxReals;
+  }
+
+  /**
+   * This method returns point
+   *
+   * @param aggregateIdx
+   * @return
+   */
+
+  SD_INLINE SD_HOST_DEVICE T **getArguments(int aggregateIdx) {
+    T **aPtr = (T **)getRealArguments(batchLimit);
+
+    return aPtr + (aggregateIdx * maxArguments);
+  }
+
+  /**
+   * This method returns number of array arguments for specified aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE int getNumArguments(int aggregateIdx) {
+    int *tPtr = (int *)ptrGeneral;
+    return tPtr[aggregateIdx * argTypes];
+  }
+
+  /**
+   * This method returns set of pointers to shape aruments for specified aggregates
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE sd::LongType **getShapeArguments(int aggregateIdx) {
+    sd::LongType **sPtr = (sd::LongType **)getArguments(batchLimit);
+
+    return sPtr + (aggregateIdx * maxShapeArguments);
+  }
+
+  /**
+   * This methor returns number of shape arguments for specified aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE int getNumShapeArguments(int aggregateIdx) {
+    int *tPtr = (int *)ptrGeneral;
+    return tPtr[aggregateIdx * argTypes + 1];
+  }
+
+  /**
+   * This method returns pointer to array of int/index arguments for specified aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE int *getIndexArguments(int aggregateIdx) {
+    // we skip first numeric num*arguments
+    int *ptr = ((int *)ptrGeneral) + (batchLimit * argTypes);
+
+    // and return index for requested aggregate
+    return ptr + (aggregateIdx * maxIndexArguments);
+  }
+
+  /**
+   * This method returns number of int/index arguments for specified aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE int getNumIndexArguments(int aggregateIdx) {
+    int *tPtr = (int *)ptrGeneral;
+    return tPtr[aggregateIdx * argTypes + 2];
+  }
+
+  /**
+   * This method returns pointer to array of jvm IntArray arguments
+   */
+  SD_INLINE SD_HOST_DEVICE int *getIntArrayArguments(int aggregateIdx, int argumentIdx) {
+    int *ptr = (int *)getIndexArguments(batchLimit);
+
+    return ptr + (aggregateIdx * maxIntArrays * maxArraySize) + (argumentIdx * maxArraySize);
+  }
+
+  /**
+   * This method returns number of jvm IntArray arguments
+   */
+  SD_INLINE SD_HOST_DEVICE int getNumIntArrayArguments(int aggregateIdx) {
+    int *tPtr = (int *)ptrGeneral;
+    return tPtr[aggregateIdx * argTypes + 4];
+  }
+
+  /**
+   * This method returns real arguments for specific aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE T *getRealArguments(int aggregateIdx) {
+    // we get pointer for last batchElement + 1, so that'll be pointer for 0 realArgument
+    T *ptr = (T *)getIntArrayArguments(batchLimit, 0);
+
+    return ptr + (aggregateIdx * maxRealArguments);
+  }
+
+  /**
+   * This methor returns number of real arguments for specified aggregate
+   *
+   * @param aggregateIdx
+   * @return
+   */
+  SD_INLINE SD_HOST_DEVICE int getNumRealArguments(int aggregateIdx) {
+    int *tPtr = (int *)ptrGeneral;
+    return tPtr[aggregateIdx * argTypes + 3];
+  }
+};
+}  // namespace sd
+
+#endif  // LIBND4J_HELPER_PTRMAP_H