chore: import upstream snapshot with attribution

libnd4j/include/loops/cuda/transform/transform_any.cu

@@ -0,0 +1,192 @@
+/******************************************************************************
+*
+* 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 <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_any.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <execution/cuda/DeviceValidator.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+// The kernel that calls the transform CUDA method,
+// caching shape info in shared memory for offset computations.
+template <typename X, typename Z, typename OpType>
+__global__ void transformAnySimpleCached(
+    const void* x,
+    const sd::LongType* xShapeInfo,
+    sd::LongType xRank,
+    void* params,
+    void* z,
+    const sd::LongType* zShapeInfo,
+    sd::LongType zRank,
+    sd::LongType* allocationPointer,
+    void* reductionPointer,
+    const sd::LongType* tadShapeInfo,
+    const sd::LongType* tadOffsets)
+{
+  // Just delegate to transformCuda,
+  // which will do the shape caching logic for coords->offset conversions.
+  functions::transform::TransformAny<X, Z>::template transformCuda<OpType>(
+      x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace transform {
+
+////////////////////////////////////////////////////////////////////////////////
+// Implementation of the "executeTransformShaped" that calls the new cached kernel
+template <typename X, typename Y>
+SD_HOST void TransformAny<X, Y>::executeTransformShaped(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const int opNum,
+    const void* x,
+    const sd::LongType* xShape,
+    sd::LongType xRank,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShape,
+    sd::LongType zRank,
+    sd::LongType* allocationPointer,
+    void* reductionPointer,
+    const sd::LongType* tadShapeInfo,
+    const sd::LongType* tadOffsets)
+{
+  DISPATCH_BY_OPNUM_TT(
+      intermediateShaped,
+      PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
+             reductionPointer, tadShapeInfo, tadOffsets),
+      TRANSFORM_ANY_OPS);
+
+  sd::DebugHelper::checkErrorCode(stream, "transformAny executeTransformShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// The transformCuda method that uses shared memory for shape/stride caching,
+// then does coords->offset conversions.
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void TransformAny<X, Z>::transformCuda(
+    const void* vx,
+    const sd::LongType* xShapeInfo,
+    void* vparams,
+    void* vz,
+    const sd::LongType* zShapeInfo,
+    sd::LongType* allocationPointer,
+    void* vreductionPointer,
+    const sd::LongType* tadShapeInfo,
+    const sd::LongType* tadOffsets)
+{
+  // cast pointers
+  auto x      = reinterpret_cast<const X*>(vx);
+  auto z      = reinterpret_cast<Z*>(vz);
+  auto params = reinterpret_cast<X*>(vparams);
+
+  if (x == nullptr || z == nullptr) return;
+
+  // cache shape info in shared memory
+  __shared__ sd::LongType length;
+  __shared__ int xRank;
+  __shared__ const sd::LongType* xShapePtr;
+  __shared__ const sd::LongType* xStridePtr;
+
+  __shared__ int zRank;
+  __shared__ const sd::LongType* zShapePtr;
+  __shared__ const sd::LongType* zStridePtr;
+
+  if (threadIdx.x == 0) {
+    length      = shape::length(xShapeInfo);
+
+    xRank       = shape::rank(xShapeInfo);
+    xShapePtr   = shape::shapeOf(xShapeInfo);
+    xStridePtr  = shape::stride(xShapeInfo);
+
+    zRank       = shape::rank(zShapeInfo);
+    zShapePtr   = shape::shapeOf(zShapeInfo);
+    zStridePtr  = shape::stride(zShapeInfo);
+  }
+  __syncthreads();
+
+  // do the transform
+  const auto tid          = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto totalThreads = gridDim.x * blockDim.x;
+
+  for (sd::LongType i = tid; i < length; i += totalThreads) {
+    sd::LongType coordsX[SD_MAX_RANK];
+    sd::LongType coordsZ[SD_MAX_RANK];
+    sd::LongType offsetX;
+    sd::LongType offsetZ;
+
+    // convert i -> coords -> offset for x
+    INDEX2COORDS(i, xRank, xShapePtr, coordsX);
+    COORDS2INDEX(xRank, xStridePtr, coordsX, offsetX);
+
+    // convert i -> coords -> offset for z
+    INDEX2COORDS(i, zRank, zShapePtr, coordsZ);
+    COORDS2INDEX(zRank, zStridePtr, coordsZ, offsetZ);
+
+    z[offsetZ] = OpType::op(x[offsetX], params);
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST void TransformAny<X, Z>::intermediateShaped(
+    dim3 launchDims,
+    cudaStream_t* stream,
+    const void* x,
+    const sd::LongType* xShape,
+    sd::LongType xRank,
+    void* extraParams,
+    void* z,
+    const sd::LongType* zShape,
+    sd::LongType zRank,
+    sd::LongType* allocationPointer,
+    void* reductionPointer,
+    const sd::LongType* tadShapeInfo,
+    const sd::LongType* tadOffsets)
+{
+  // We call the new transformAnySimpleCached kernel
+  transformAnySimpleCached<X, Z, OpType>
+  <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+      x,
+      xShape,
+      xRank,
+      extraParams,
+      z,
+      zShape,
+      zRank,
+      allocationPointer,
+      reductionPointer,
+      tadShapeInfo,
+      tadOffsets);
+
+  sd::DebugHelper::checkErrorCode(stream, "transformAny(...) cached kernel failed");
+}
+
+BUILD_DOUBLE_TEMPLATE( class TransformAny, , SD_COMMON_TYPES, SD_COMMON_TYPES);
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cuda/transform/transform_bool.cu

@@ -0,0 +1,210 @@
+/******************************************************************************
+*
+* 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 <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_bool.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <execution/cuda/DeviceValidator.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+// Cached kernel that caches shape info in shared memory and uses cached variables
+template <typename X, typename Z, typename OpType>
+__global__ void transformBoolSimpleCached(
+   const void* x,
+   const sd::LongType* xShapeInfo,
+   sd::LongType xRank,
+   void* params,
+   void* z,
+   const sd::LongType* zShapeInfo,
+   sd::LongType zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Delegate the transform to the device function with cached shape info
+ functions::transform::TransformBool<X, Z>::template transformCuda<OpType>(
+     x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace transform {
+
+////////////////////////////////////////////////////////////////////////////////
+// Implementation of the "executeTransformShaped" that launches the cached kernel
+template <typename X, typename Y>
+SD_HOST void TransformBool<X, Y>::executeTransformShaped(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   const int opNum,
+   const void* x,
+   const sd::LongType* xShape,
+   long long int xRank,
+   void* extraParams,
+   void* z,
+   const sd::LongType* zShape,
+   long long int zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ DISPATCH_BY_OPNUM_TT(
+     intermediateShaped,
+     PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
+            reductionPointer, tadShapeInfo, tadOffsets),
+     TRANSFORM_BOOL_OPS);
+
+ sd::DebugHelper::checkErrorCode(stream, "transformBool executeTransformShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Device function that caches shape info and uses cached variables for computations
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void TransformBool<X, Z>::transformCuda(
+   const void* vx,
+   const sd::LongType* xShapeInfo,
+   void* vparams,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   sd::LongType* allocationPointer,
+   void* vreductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Cast pointers to appropriate types
+ auto x = reinterpret_cast<const X*>(vx);
+ auto z = reinterpret_cast<Z*>(vz);
+ auto params = reinterpret_cast<X*>(vparams);
+ auto reductionPointer = reinterpret_cast<Z*>(vreductionPointer);
+
+ // Check for special operations
+ if (OpType::requiresSpecial) {
+   OpType::execSpecialCuda(x,
+                           xShapeInfo,
+                           z,
+                           zShapeInfo,
+                           params,
+                           allocationPointer,
+                           reductionPointer,
+                           tadShapeInfo,
+                           tadOffsets);
+   return;
+ }
+
+ // Shared memory for caching shape information
+ __shared__ sd::LongType length;
+ __shared__ int xRankCached;
+ __shared__ const sd::LongType* xShapePtrCached;
+ __shared__ const sd::LongType* xStridePtrCached;
+
+ __shared__ int zRankCached;
+ __shared__ const sd::LongType* zShapePtrCached;
+ __shared__ const sd::LongType* zStridePtrCached;
+
+ // Thread 0 caches the shape information
+ if (threadIdx.x == 0) {
+   length = shape::length(xShapeInfo);
+
+   xRankCached = shape::rank(xShapeInfo);
+   xShapePtrCached = shape::shapeOf(xShapeInfo);
+   xStridePtrCached = shape::stride(xShapeInfo);
+
+   zRankCached = shape::rank(zShapeInfo);
+   zShapePtrCached = shape::shapeOf(zShapeInfo);
+   zStridePtrCached = shape::stride(zShapeInfo);
+ }
+ __syncthreads();
+
+ // Calculate thread ID and total threads
+ auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+ int totalThreads = gridDim.x * blockDim.x;
+
+ // Loop over all elements using cached shape info
+ for (sd::LongType i = tid; i < length; i += totalThreads) {
+   sd::LongType xCoords[SD_MAX_RANK];
+   sd::LongType zCoords[SD_MAX_RANK];
+   sd::LongType xOffset;
+   sd::LongType zOffset;
+
+   // Convert index to coordinates using cached shape info
+   INDEX2COORDS(i, xRankCached, xShapePtrCached, xCoords);
+   COORDS2INDEX(xRankCached, xStridePtrCached, xCoords, xOffset);
+
+   INDEX2COORDS(i, zRankCached, zShapePtrCached, zCoords);
+   COORDS2INDEX(zRankCached, zStridePtrCached, zCoords, zOffset);
+
+   // Apply the operation using cached offsets
+   z[zOffset] = OpType::op(x[xOffset], params);
+ }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Host function that launches the cached kernel
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST void TransformBool<X, Z>::intermediateShaped(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   const void* x,
+   const sd::LongType* xShape,
+   sd::LongType xRank,
+   void* extraParams,
+   void* z,
+   const sd::LongType* zShape,
+   sd::LongType zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Launch the cached kernel
+ transformBoolSimpleCached<X, Z, OpType>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShape,
+         xRank,
+         extraParams,
+         z,
+         zShape,
+         zRank,
+         allocationPointer,
+         reductionPointer,
+         tadShapeInfo,
+         tadOffsets
+     );
+
+ // Check for any errors during kernel execution
+ sd::DebugHelper::checkErrorCode(stream, "transformBool(...) cached kernel failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Macro to instantiate templates for TransformBool with common and bool types
+BUILD_DOUBLE_TEMPLATE( class TransformBool, , SD_COMMON_TYPES, SD_BOOL_TYPES);
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cuda/transform/transform_float.cu

@@ -0,0 +1,202 @@
+/******************************************************************************
+*
+*
+* 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 <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_float.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+#include <execution/cuda/DeviceValidator.h>
+
+using namespace simdOps;
+
+////////////////////////////////////////////////////////////////////////////////
+// The cached kernel that caches shape info in shared memory and uses cached variables
+template <typename X, typename Z, typename OpType>
+__global__ void transformFloatSimpleCached(
+   const void* x,
+   const sd::LongType* xShapeInfo,
+   sd::LongType xRank,
+   void* params,
+   void* z,
+   const sd::LongType* zShapeInfo,
+   sd::LongType zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Delegate the transform to the transformCuda method with cached shape info
+ functions::transform::TransformFloat<X, Z>::template transformCuda<OpType>(
+     x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace transform {
+
+////////////////////////////////////////////////////////////////////////////////
+// Implementation of the "executeTransformShaped" that launches the cached kernel
+template <typename X, typename Y>
+SD_HOST void TransformFloat<X, Y>::executeTransformShaped(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   int opNum,
+   const void* x,
+   const sd::LongType* xShape,
+   sd::LongType xRank,
+   void* extraParams,
+   void* z,
+   const sd::LongType* zShape,
+   sd::LongType zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ DISPATCH_BY_OPNUM_TT(
+     intermediateShaped,
+     PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
+            reductionPointer, tadShapeInfo, tadOffsets),
+     TRANSFORM_FLOAT_OPS);
+
+ sd::DebugHelper::checkErrorCode(stream, "transformFloat executeTransformShaped(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Device function that caches shape info and uses cached variables for computations
+template <typename X, typename Z>
+template <typename OpType>
+SD_DEVICE void TransformFloat<X, Z>::transformCuda(
+   const void* vx,
+   const sd::LongType* xShapeInfo,
+   void* vparams,
+   void* vz,
+   const sd::LongType* zShapeInfo,
+   sd::LongType* allocationPointer,
+   void* vreductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Cast pointers to appropriate types
+ auto x = reinterpret_cast<const X*>(vx);
+ auto z = reinterpret_cast<Z*>(vz);
+ auto params = reinterpret_cast<Z*>(vparams);
+ auto reductionPointer = reinterpret_cast<Z*>(vreductionPointer);
+
+ // Check for special operations
+ if (OpType::requiresSpecial) {
+   OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
+                           tadOffsets);
+   return;
+ }
+
+ // Shared memory for caching shape information
+ __shared__ sd::LongType length;
+ __shared__ int xRankCached;
+ __shared__ const sd::LongType* xShapePtrCached;
+ __shared__ const sd::LongType* xStridePtrCached;
+
+ __shared__ int zRankCached;
+ __shared__ const sd::LongType* zShapePtrCached;
+ __shared__ const sd::LongType* zStridePtrCached;
+
+ // Thread 0 caches the shape information
+ if (threadIdx.x == 0) {
+   length = shape::length(xShapeInfo);
+
+   xRankCached = shape::rank(xShapeInfo);
+   xShapePtrCached = shape::shapeOf(xShapeInfo);
+   xStridePtrCached = shape::stride(xShapeInfo);
+
+   zRankCached = shape::rank(zShapeInfo);
+   zShapePtrCached = shape::shapeOf(zShapeInfo);
+   zStridePtrCached = shape::stride(zShapeInfo);
+ }
+ __syncthreads();
+
+ // Calculate thread ID and total threads
+ auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+ int totalThreads = gridDim.x * blockDim.x;
+
+ // Loop over all elements using cached shape info
+ for (sd::LongType i = tid; i < length; i += totalThreads) {
+   sd::LongType xCoords[SD_MAX_RANK];
+   sd::LongType zCoords[SD_MAX_RANK];
+   sd::LongType xOffset;
+   sd::LongType zOffset;
+
+   // Convert index to coordinates using cached shape info
+   INDEX2COORDS(i, xRankCached, xShapePtrCached, xCoords);
+   COORDS2INDEX(xRankCached, xStridePtrCached, xCoords, xOffset);
+
+   INDEX2COORDS(i, zRankCached, zShapePtrCached, zCoords);
+   COORDS2INDEX(zRankCached, zStridePtrCached, zCoords, zOffset);
+
+   // Apply the operation using cached offsets
+   z[zOffset] = OpType::op(x[xOffset], params);
+ }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+// Host function that launches the cached kernel
+template <typename X, typename Z>
+template <typename OpType>
+SD_HOST void TransformFloat<X, Z>::intermediateShaped(
+   dim3 launchDims,
+   cudaStream_t* stream,
+   const void* x,
+   const sd::LongType* xShape,
+   sd::LongType xRank,
+   void* extraParams,
+   void* z,
+   const sd::LongType* zShape,
+   sd::LongType zRank,
+   sd::LongType* allocationPointer,
+   void* reductionPointer,
+   const sd::LongType* tadShapeInfo,
+   const sd::LongType* tadOffsets)
+{
+ // Launch the cached kernel
+ transformFloatSimpleCached<X, Z, OpType>
+     <<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+         x,
+         xShape,
+         xRank,
+         extraParams,
+         z,
+         zShape,
+         zRank,
+         allocationPointer,
+         reductionPointer,
+         tadShapeInfo,
+         tadOffsets
+     );
+
+ // Check for any errors during kernel execution
+ sd::DebugHelper::checkErrorCode(stream, "transformFloat(...) cached kernel failed");
+}
+
+BUILD_DOUBLE_TEMPLATE( class TransformFloat, , SD_COMMON_TYPES, SD_FLOAT_TYPES);
+
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cuda/transform/transform_same.cu

@@ -0,0 +1,136 @@
+/* ******************************************************************************
+*
+*
+* 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 <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_same.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+
+using namespace simdOps;
+
+template <typename X, typename OpType>
+SD_KERNEL void transformSameSimple(const void *x, const sd::LongType *xShapeInfo, long long int xRank, void *params, void *z,
+                                  const sd::LongType *zShapeInfo, long long int zRank,
+                                  sd::LongType *allocationPointer,
+                                  void *reductionPointer, const sd::LongType *tadShapeInfo,
+                                  const sd::LongType *tadOffsets) {
+ functions::transform::TransformSame<X>::template transformCuda<OpType>(
+     x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace transform {
+
+template <typename X>
+SD_HOST void TransformSame<X>::executeTransformShaped(dim3 launchDims, cudaStream_t *stream, const int opNum,
+                                                     const void *x, const sd::LongType *xShape, sd::LongType xRank,
+                                                     void *extraParams, void *z, const sd::LongType *zShape,
+                                                     sd::LongType zRank, sd::LongType *allocationPointer, void *reductionPointer,
+                                                     const sd::LongType *tadShapeInfo,
+                                                     const sd::LongType *tadOffsets) {
+ DISPATCH_BY_OPNUM_T(intermediateShaped,
+                     PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
+                            reductionPointer, tadShapeInfo, tadOffsets),
+                     TRANSFORM_SAME_OPS);
+
+ sd::DebugHelper::checkErrorCode(stream, "transformAny(...) failed");
+}
+
+template <typename X>
+template <typename OpType>
+SD_DEVICE void TransformSame<X>::transformCuda(const void *vx, const sd::LongType *xShapeInfo, void *vparams, void *vz,
+                                               const sd::LongType *zShapeInfo, sd::LongType *allocationPointer,
+                                               void *vreductionPointer, const sd::LongType *tadShapeInfo,
+                                               const sd::LongType *tadOffsets) {
+ auto x = static_cast<const X *>(vx);
+ auto z = static_cast<X *>(vz);
+ auto params = static_cast<X *>(vparams);
+ auto reductionPointer = static_cast<X *>(vreductionPointer);
+
+ if (OpType::requiresSpecial) {
+   OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
+                           tadOffsets);
+   return;
+ } else {
+   __shared__ sd::LongType length;
+
+   // Cache shape information for x buffer
+   __shared__ sd::LongType xRank;
+   __shared__ const sd::LongType* xShapePtr;
+   __shared__ const sd::LongType* xStridePtr;
+
+   // Cache shape information for z buffer
+   __shared__ sd::LongType zRank;
+   __shared__ const sd::LongType* zShapePtr;
+   __shared__ const sd::LongType* zStridePtr;
+
+   if (threadIdx.x == 0) {
+     length = shape::length(xShapeInfo);
+
+     // Cache x shape information
+     xRank = shape::rank(xShapeInfo);
+     xShapePtr = shape::shapeOf(xShapeInfo);
+     xStridePtr = shape::stride(xShapeInfo);
+
+     // Cache z shape information
+     zRank = shape::rank(zShapeInfo);
+     zShapePtr = shape::shapeOf(zShapeInfo);
+     zStridePtr = shape::stride(zShapeInfo);
+   }
+   __syncthreads();
+
+   auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+   int totalThreads = gridDim.x * blockDim.x;
+
+   for (sd::LongType i = tid; i < length; i += totalThreads) {
+     sd::LongType xCoords[SD_MAX_RANK];
+     sd::LongType zCoords[SD_MAX_RANK];
+     sd::LongType xOffset;
+     sd::LongType zOffset;
+
+     INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+     COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+     INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+     COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+     z[zOffset] = OpType::op(x[xOffset], params);
+   }
+ }
+};
+template <typename X>
+template <typename OpType>
+SD_HOST void TransformSame<X>::intermediateShaped(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                                 const sd::LongType *xShape, sd::LongType xRank, void *extraParams, void *z,
+                                                 const sd::LongType *zShape, sd::LongType zRank,
+                                                 sd::LongType *allocationPointer,
+                                                 void *reductionPointer, const sd::LongType *tadShapeInfo,
+                                                 const sd::LongType *tadOffsets) {
+ transformSameSimple<X, OpType><<<launchDims.x, launchDims.x, launchDims.z, *stream>>>(
+     x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+ sd::DebugHelper::checkErrorCode(stream, "transformSame(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( class TransformSame, , SD_COMMON_TYPES);
+}  // namespace transform
+}  // namespace functions

libnd4j/include/loops/cuda/transform/transform_strict.cu

@@ -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
+//
+#include <helpers/DebugHelper.h>
+#include <loops/legacy_ops.h>
+#include <loops/transform_strict.h>
+#include <system/Environment.h>
+#include <system/op_boilerplate.h>
+#include <types/types.h>
+
+
+using namespace simdOps;
+
+template <typename X, typename OpType>
+SD_KERNEL void transformStrictSimple(const void *x, const sd::LongType *xShapeInfo, long long int xRank, void *params, void *z,
+                                    const sd::LongType *zShapeInfo, long long int zRank,
+                                    sd::LongType *allocationPointer,
+                                    void *reductionPointer, const sd::LongType *tadShapeInfo,
+                                    const sd::LongType *tadOffsets) {
+ functions::transform::TransformStrict<X>::template transformCuda<OpType>(
+     x, xShapeInfo, params, z, zShapeInfo, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+}
+
+namespace functions {
+namespace transform {
+
+template <typename X>
+SD_HOST void TransformStrict<X>::executeTransformShaped(dim3 launchDims, cudaStream_t *stream, const int opNum,
+                                                       const void *x, const sd::LongType *xShape, sd::LongType xRank,
+                                                       void *extraParams, void *z, const sd::LongType *zShape,
+                                                       sd::LongType zRank, sd::LongType *allocationPointer, void *reductionPointer,
+                                                       const sd::LongType *tadShapeInfo,
+                                                       const sd::LongType *tadOffsets) {
+ DISPATCH_BY_OPNUM_T(intermediateShaped,
+                     PARAMS(launchDims, stream, x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer,
+                            reductionPointer, tadShapeInfo, tadOffsets),
+                     TRANSFORM_STRICT_OPS);
+
+ sd::DebugHelper::checkErrorCode(stream, "transformStrict(...) failed");
+}
+
+template <typename X>
+template <typename OpType>
+SD_DEVICE void TransformStrict<X>::transformCuda(const void *vx, const sd::LongType *xShapeInfo, void *vparams,
+                                                 void *vz, const sd::LongType *zShapeInfo,
+                                                 sd::LongType *allocationPointer,
+                                                 void *vreductionPointer, const sd::LongType *tadShapeInfo,
+                                                 const sd::LongType *tadOffsets) {
+ auto x = static_cast<const X *>(vx);
+ auto z = static_cast<X *>(vz);
+ auto params = static_cast<X *>(vparams);
+ auto reductionPointer = static_cast<X *>(vreductionPointer);
+
+ if (OpType::requiresSpecial) {
+   OpType::execSpecialCuda(x, xShapeInfo, z, zShapeInfo, params, allocationPointer, reductionPointer, tadShapeInfo,
+                           tadOffsets);
+   return;
+ } else {
+   __shared__ sd::LongType length;
+
+   // Cache shape information for x buffer
+   __shared__ sd::LongType xRank;
+   __shared__ const sd::LongType* xShapePtr;
+   __shared__ const sd::LongType* xStridePtr;
+
+   // Cache shape information for z buffer
+   __shared__ sd::LongType zRank;
+   __shared__ const sd::LongType* zShapePtr;
+   __shared__ const sd::LongType* zStridePtr;
+
+   if (threadIdx.x == 0) {
+     length = shape::length(xShapeInfo);
+
+     // Cache x shape information
+     xRank = shape::rank(xShapeInfo);
+     xShapePtr = shape::shapeOf(xShapeInfo);
+     xStridePtr = shape::stride(xShapeInfo);
+
+     // Cache z shape information
+     zRank = shape::rank(zShapeInfo);
+     zShapePtr = shape::shapeOf(zShapeInfo);
+     zStridePtr = shape::stride(zShapeInfo);
+   }
+   __syncthreads();
+
+   auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+   int totalThreads = gridDim.x * blockDim.x;
+
+   for (sd::LongType i = tid; i < length; i += totalThreads) {
+     sd::LongType xCoords[SD_MAX_RANK];
+     sd::LongType zCoords[SD_MAX_RANK];
+     sd::LongType xOffset;
+     sd::LongType zOffset;
+
+     INDEX2COORDS(i, xRank, xShapePtr, xCoords);
+     COORDS2INDEX(xRank, xStridePtr, xCoords, xOffset);
+     INDEX2COORDS(i, zRank, zShapePtr, zCoords);
+     COORDS2INDEX(zRank, zStridePtr, zCoords, zOffset);
+
+     z[zOffset] = OpType::op(x[xOffset], params);
+   }
+ }
+};
+template <typename X>
+template <typename OpType>
+SD_HOST void TransformStrict<X>::intermediateShaped(dim3 launchDims, cudaStream_t *stream, const void *x,
+                                                   const sd::LongType *xShape, sd::LongType xRank, void *extraParams, void *z,
+                                                   const sd::LongType *zShape, sd::LongType zRank,
+                                                   sd::LongType *allocationPointer,
+                                                   void *reductionPointer, const sd::LongType *tadShapeInfo,
+                                                   const sd::LongType *tadOffsets) {
+ transformStrictSimple<X, OpType><<<launchDims.x, launchDims.x, launchDims.z, *stream>>>(
+     x, xShape, xRank, extraParams, z, zShape, zRank, allocationPointer, reductionPointer, tadShapeInfo, tadOffsets);
+ sd::DebugHelper::checkErrorCode(stream, "transformStrict(...) failed");
+}
+
+BUILD_SINGLE_TEMPLATE( class TransformStrict, , SD_FLOAT_TYPES);
+}  // namespace transform
+}  // namespace functions