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,137 @@
+/* ******************************************************************************
+ *
+ *
+ * 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
+//  @author Yurii Shyrma (iuriish@yahoo.com)
+
+#ifndef PAIRWISE_CU
+#define PAIRWISE_CU
+#include <loops/pairwise_instantiations.h>
+
+#include "../pairwise_transform.h"
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z, typename OpType>
+SD_KERNEL static void pairwiseSimpleShaped(void const* vx, sd::LongType const* xShapeInfo,
+                                           void const* vy, sd::LongType const* yShapeInfo,
+                                           void* vz, sd::LongType const* zShapeInfo,
+                                           void* vextraParams) {
+  auto x          = static_cast<X const*>(vx);
+  auto y          = static_cast<Y const*>(vy);
+  auto z          = static_cast<Z*>(vz);
+  auto extraParams= static_cast<Z*>(vextraParams);
+
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  __shared__ sd::LongType len;
+  __shared__ sd::LongType xRank;
+  __shared__ sd::LongType yRank;
+  __shared__ sd::LongType zRank;
+
+  __shared__ const sd::LongType* xShape;
+  __shared__ const sd::LongType* xStride;
+  __shared__ const sd::LongType* yShape;
+  __shared__ const sd::LongType* yStride;
+  __shared__ const sd::LongType* zShape;
+  __shared__ const sd::LongType* zStride;
+
+  if (threadIdx.x == 0) {
+    len     = shape::length(xShapeInfo);
+
+    xRank   = shape::rank(xShapeInfo);
+    xShape  = shape::shapeOf(xShapeInfo);
+    xStride = shape::stride(xShapeInfo);
+
+    yRank   = shape::rank(yShapeInfo);
+    yShape  = shape::shapeOf(yShapeInfo);
+    yStride = shape::stride(yShapeInfo);
+
+    zRank   = shape::rank(zShapeInfo);
+    zShape  = shape::shapeOf(zShapeInfo);
+    zStride = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  for (sd::LongType i = tid; i < len; i += (gridDim.x * blockDim.x)) {
+    sd::LongType xCoords[SD_MAX_RANK];
+    sd::LongType yCoords[SD_MAX_RANK];
+    sd::LongType zCoords[SD_MAX_RANK];
+    sd::LongType xOffset;
+    sd::LongType yOffset;
+    sd::LongType zOffset;
+
+    INDEX2COORDS(i, xRank, xShape, xCoords);
+    COORDS2INDEX(xRank, xStride, xCoords, xOffset);
+
+    INDEX2COORDS(i, yRank, yShape, yCoords);
+    COORDS2INDEX(yRank, yStride, yCoords, yOffset);
+
+    INDEX2COORDS(i, zRank, zShape, zCoords);
+    COORDS2INDEX(zRank, zStride, zCoords, zOffset);
+
+    z[zOffset] = OpType::op(x[xOffset], y[yOffset], extraParams);
+  }
+}
+
+namespace functions {
+namespace pairwise_transforms {
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+template <typename OpType>
+void SD_HOST PairWiseTransform<X, Y, Z>::intermediateShaped(dim3& launchDims,
+                                                            cudaStream_t* stream,
+                                                            void const* vx,
+                                                            sd::LongType const* xShapeInfo,
+                                                            void const* vy,
+                                                            sd::LongType const* yShapeInfo,
+                                                            void* vz,
+                                                            sd::LongType const* zShapeInfo,
+                                                            void* vextraParams) {
+  pairwiseSimpleShaped<X, Y, Z, OpType>
+  <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      vx, xShapeInfo,
+      vy, yShapeInfo,
+      vz, zShapeInfo,
+      vextraParams);
+  sd::DebugHelper::checkErrorCode(stream, "PairWiseTransform intermediateShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Y, typename Z>
+void SD_HOST PairWiseTransform<X, Y, Z>::executeCudaShaped(dim3& launchDims,
+                                                           cudaStream_t* stream,
+                                                           int opNum,
+                                                           void const* vx,
+                                                           sd::LongType const* xShapeInfo,
+                                                           void const* vy,
+                                                           sd::LongType const* yShapeInfo,
+                                                           void* vz,
+                                                           sd::LongType const* zShapeInfo,
+                                                           void* vextraParams) {
+  DISPATCH_BY_OPNUM_TTT(
+      intermediateShaped,
+      PARAMS(launchDims, stream, vx, xShapeInfo, vy, yShapeInfo, vz, zShapeInfo, vextraParams),
+      PAIRWISE_TRANSFORM_OPS);
+}
+
+}  // namespace pairwise_transforms
+}  // namespace functions
+
+#endif  // PAIRWISE_CU