chore: import upstream snapshot with attribution

libnd4j/include/helpers/cuda/ConstantHelper.cu

@@ -0,0 +1,190 @@
+/* ******************************************************************************
+ *
+ *
+ * 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/PrimaryPointerDeallocator.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <exceptions/cuda_exception.h>
+#include <execution/AffinityManager.h>
+#include <execution/LaunchContext.h>
+#include <helpers/ConstantHelper.h>
+#include <helpers/logger.h>
+#include <helpers/shape.h>
+#include <ops/specials.h>
+#include <system/selective_rendering.h>
+#define CONSTANT_LIMIT 49152
+
+__constant__ char deviceConstantMemory[CONSTANT_LIMIT];
+
+namespace sd {
+void * ConstantHelper::getConstantSpace() {
+  Pointer dConstAddr;
+  auto dZ = cudaGetSymbolAddress(reinterpret_cast<void **>(&dConstAddr), deviceConstantMemory);
+
+  if (dZ != 0) throw cuda_exception::build("cudaGetSymbolAddress(...) failed", dZ);
+
+  return dConstAddr;
+}
+
+int ConstantHelper::getCurrentDevice() { return AffinityManager::currentDeviceId(); }
+
+int ConstantHelper::getNumberOfDevices() { return AffinityManager::numberOfDevices(); }
+
+ConstantHelper::ConstantHelper() {
+  auto initialDevice = getCurrentDevice();
+
+  auto numDevices = getNumberOfDevices();
+  _devicePointers.resize(numDevices);
+  _deviceOffsets.resize(numDevices);
+  _cache.resize(numDevices);
+  _counters.resize(numDevices);
+
+  // filling all pointers
+  for (int e = 0; e < numDevices; e++) {
+    auto res = cudaSetDevice(e);
+    if (res != 0) throw cuda_exception::build("cudaSetDevice failed", res);
+    auto constant = getConstantSpace();
+
+    SD_MAP_IMPL<ConstantDescriptor, ConstantHolder *> devCache;
+
+    _devicePointers[e] = constant;
+    _deviceOffsets[e] = 0;
+    _cache[e] = devCache;
+    _counters[e] = 0L;
+  }
+
+  //
+  auto res = cudaSetDevice(initialDevice);
+  if (res != 0) throw cuda_exception::build("Final cudaSetDevice failed", res);
+}
+
+ConstantHelper::~ConstantHelper() {
+  for (const auto &v : _cache) {
+    for (const auto &c : v) {
+      delete c.second;
+    }
+  }
+}
+
+ConstantHelper &ConstantHelper::getInstance() {
+  static ConstantHelper instance;
+  return instance;
+}
+
+void *ConstantHelper::replicatePointer(void *src, size_t numBytes, memory::Workspace *workspace) {
+  std::lock_guard<std::mutex> lock(_mutex);
+
+  auto deviceId = getCurrentDevice();
+  Pointer constantPtr = nullptr;
+  LongType constantOffset = 0L;
+  if (_devicePointers[deviceId] == 0) {
+    auto constant = getConstantSpace();
+
+    // filling default ptr, which will be 0 probably
+    _devicePointers[deviceId] = constant;
+    _deviceOffsets[deviceId] = 0;
+    constantPtr = constant;
+  } else {
+    constantPtr = _devicePointers[deviceId];
+    constantOffset = _deviceOffsets[deviceId];
+  }
+
+  int8_t *ptr = nullptr;
+  ALLOCATE_SPECIAL(ptr, workspace, numBytes, int8_t);
+  auto res = cudaMemcpy(ptr, src, numBytes, cudaMemcpyHostToDevice);
+  if (res != 0) {
+    std::string errorMessage = "cudaMemcpy failed with error code " + std::to_string(res);
+    auto lastError = cudaGetLastError(); // get last error
+    if (lastError != cudaSuccess) {
+      errorMessage += "; last error: " + std::string(cudaGetErrorString(lastError));
+    }
+
+    THROW_EXCEPTION(errorMessage.c_str());
+
+  }
+
+  constantPtr = ptr;
+  return reinterpret_cast<int8_t *>(constantPtr) + constantOffset;
+}
+
+ConstantDataBuffer *ConstantHelper::constantBuffer(const ConstantDescriptor &descriptor, DataType dataType) {
+  const auto deviceId = getCurrentDevice();
+
+  // all cache modifications are synchronous
+  _mutexHolder.lock();
+
+  if (_cache[deviceId].count(descriptor) == 0) {
+    _cache[deviceId][descriptor] = new ConstantHolder();
+  }
+  auto holder = _cache[deviceId][descriptor];
+
+  // release cache lock
+  _mutexHolder.unlock();
+
+  ConstantDataBuffer *result;
+
+  // access to this holder instance is synchronous
+  std::lock_guard<std::mutex> lock(*holder->mutex());
+
+  if (holder->hasBuffer(dataType)) {
+    result = holder->getConstantDataBuffer(dataType);
+  } else {
+    auto numBytes = descriptor.length() * DataTypeUtils::sizeOf(dataType);
+    auto cbuff = std::make_shared<PointerWrapper>(new int8_t[numBytes], std::make_shared<PointerDeallocator>());
+    _counters[deviceId] += numBytes;
+
+    // create buffer with this dtype
+    if (descriptor.isFloat()) {
+      BUILD_DOUBLE_SELECTOR(
+          sd::DataType::DOUBLE, dataType, sd::SpecialTypeConverter::convertGeneric,
+          (nullptr, const_cast<double *>(descriptor.floatValues().data()), descriptor.length(), cbuff->pointer()),
+          (sd::DataType::DOUBLE, double), SD_COMMON_TYPES);
+    } else if (descriptor.isInteger()) {
+      auto int64DType = sd::DataType::INT64;
+      BUILD_DOUBLE_SELECTOR(sd::DataType::INT64, dataType, sd::SpecialTypeConverter::convertGeneric,
+                            (nullptr, const_cast<sd::LongType *>(descriptor.integerValues().data()),
+                                descriptor.length(), cbuff->pointer()),
+                            (sd::DataType::INT64, sd::LongType), SD_COMMON_TYPES);
+    }
+
+    // we don't have deallocator here.
+    // TODO: we probably want to make use deallocator here, if we're not using constant memory
+    auto dbuff = std::make_shared<PointerWrapper>(
+        replicatePointer(cbuff->pointer(), descriptor.length() * DataTypeUtils::sizeOf(dataType)));
+
+    ConstantDataBuffer *dataBuffer = new ConstantDataBuffer(cbuff, dbuff, descriptor.length(), dataType);
+
+    holder->addBuffer(*dataBuffer, dataType);
+    result = holder->getConstantDataBuffer(dataType);
+  }
+
+  return result;
+}
+
+LongType ConstantHelper::getCachedAmount(int deviceId) {
+  int numDevices = getNumberOfDevices();
+  if (deviceId > numDevices || deviceId < 0)
+    return 0L;
+  else
+    return _counters[deviceId];
+}
+}  // namespace sd

libnd4j/include/helpers/cuda/CudaShapeBufferCreator.cu

@@ -0,0 +1,64 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+#include <helpers/ConstantHelper.h>
+#include <helpers/cuda/CudaShapeBufferCreator.h>
+
+#include "array/CudaPointerDeallocator.h"
+#include "array/PrimaryPointerDeallocator.h"
+
+namespace sd {
+
+ConstantShapeBuffer* CudaShapeBufferCreator::create(const LongType* shapeInfo, int rank) {
+    const int shapeInfoLength = shape::shapeInfoLength(rank);
+    LongType* shapeCopy = new LongType[shapeInfoLength];
+    for(int i = 0; i < shapeInfoLength; i++) {
+        shapeCopy[i] = shapeInfo[i];
+    }
+
+    auto deallocator = std::shared_ptr<CudaPointerDeallocator>(
+        new CudaPointerDeallocator(),
+        [] (CudaPointerDeallocator* ptr) {
+          delete ptr;
+        }
+    );
+    
+    PointerWrapper hPtr(shapeCopy, deallocator);
+
+    // Create device pointer for CUDA
+    PointerWrapper dPtr(
+        ConstantHelper::getInstance().replicatePointer(shapeCopy,
+                                                   shapeInfoLength * sizeof(LongType)),
+        std::make_shared<CudaPointerDeallocator>());
+
+    if(dPtr.pointer() == nullptr) {
+        THROW_EXCEPTION("Failed to allocate device memory for shape buffer");
+    }
+    ConstantShapeBuffer *buffer = new ConstantShapeBuffer(hPtr, dPtr);
+    
+    return buffer;
+}
+
+CudaShapeBufferCreator& CudaShapeBufferCreator::getInstance() {
+    static CudaShapeBufferCreator instance;
+    return instance;
+}
+
+} // namespace sd

libnd4j/include/helpers/cuda/CudaShapeBufferCreator.h

@@ -0,0 +1,50 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+#ifndef LIBND4J_CUDASHAPEBUFFERCREATOR_H
+#define LIBND4J_CUDASHAPEBUFFERCREATOR_H
+
+#include <helpers/ShapeBufferCreator.h>
+#include <helpers/shape.h>
+#include <memory>
+
+namespace sd {
+
+/**
+ * CUDA implementation of the ShapeBufferCreator.
+ * Creates shape buffers with both host and device pointers.
+ */
+class CudaShapeBufferCreator : public ShapeBufferCreator {
+public:
+    /**
+     * Create a ConstantShapeBuffer for CUDA usage
+     */
+    ConstantShapeBuffer* create(const LongType* shapeInfo, int rank) override;
+    
+    // Singleton pattern implementation
+    static CudaShapeBufferCreator& getInstance();
+
+private:
+    CudaShapeBufferCreator() = default;
+};
+
+} // namespace sd
+
+#endif // LIBND4J_CUDASHAPEBUFFERCREATOR_H

libnd4j/include/helpers/cuda/MmulHelper.cu

@@ -0,0 +1,776 @@
+/*
+*  ******************************************************************************
+*  *
+*  *
+*  * 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)
+//
+#include <cublas_v2.h>
+#include <exceptions/cuda_exception.h>
+#include <helpers/PointersManager.h>
+#include <helpers/ShapeUtils.h>
+#include <ops/specials_cuda.h>
+
+#include <numeric>
+
+#include "../MmulHelper.h"
+#include "execution/cuda/LaunchDims.h"
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////////
+// MXK x KxN = MxN              -> actual sequence of axes doesn't matter
+template <typename T1, typename T2, typename T3>
+static SD_KERNEL void usualCudaGemm(const void* vA, const LongType* aShapeInfo, const void* vB,
+                                   const LongType* bShapeInfo, void* vC, const LongType* cShapeInfo,
+                                   const int aMaxis, const int aKaxis, const int bKaxis, const int bNaxis,
+                                   const int cMaxis, const int cNaxis, const double alpha, const double beta) {
+ // Cache shape information in shared memory
+ __shared__ LongType K;
+ __shared__ LongType cLen;
+ __shared__ LongType totalThreads;
+ __shared__ bool betaPresent;
+ __shared__ T3 alphaZ;
+ __shared__ T3 betaZ;
+ __shared__ const LongType* aShape;
+ __shared__ const LongType* bShape;
+ __shared__ const LongType* cShape;
+ __shared__ const LongType* aStride;
+ __shared__ const LongType* bStride;
+ __shared__ const LongType* cStride;
+ __shared__ LongType aRank;
+ __shared__ LongType bRank;
+ __shared__ LongType cRank;
+ __shared__ LongType* coords;
+
+ if (threadIdx.x == 0) {
+   extern __shared__ unsigned char shmem[];
+   coords = reinterpret_cast<LongType*>(shmem);
+
+   // Cache all shape information at start
+   aRank = shape::rank(aShapeInfo);
+   bRank = shape::rank(bShapeInfo);
+   cRank = shape::rank(cShapeInfo);
+
+   aShape = shape::shapeOf(aShapeInfo);
+   bShape = shape::shapeOf(bShapeInfo);
+   cShape = shape::shapeOf(cShapeInfo);
+
+   aStride = shape::stride(aShapeInfo);
+   bStride = shape::stride(bShapeInfo);
+   cStride = shape::stride(cShapeInfo);
+
+   cLen = shape::length(cShapeInfo);
+   K = aShape[aKaxis];
+   betaPresent = beta != 0;
+   totalThreads = gridDim.x * blockDim.x;
+   alphaZ = alpha;
+   betaZ = beta;
+ }
+ __syncthreads();
+
+ const T1* A = reinterpret_cast<const T1*>(vA);
+ const T2* B = reinterpret_cast<const T2*>(vB);
+ T3* C = reinterpret_cast<T3*>(vC);
+
+ auto aCoords = coords + threadIdx.x * 6;  // 6 = (aRank + bRank + cRank)
+ auto bCoords = aCoords + 2;
+ auto cCoords = bCoords + 2;
+
+ const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+ for (LongType i = tid; i < cLen; i += totalThreads) {
+   // evaluate C coordinates
+   INDEX2COORDS(i, cRank, cShape, cCoords);
+
+   // evaluate A coordinates
+   aCoords[aMaxis] = cCoords[cMaxis];
+   aCoords[aKaxis] = 0;
+
+   // evaluate B coordinates
+   bCoords[bKaxis] = 0;
+   bCoords[bNaxis] = cCoords[cNaxis];
+
+   LongType aOffset, bOffset, cOffset;
+   COORDS2INDEX(aRank, aStride, aCoords, aOffset);
+   COORDS2INDEX(bRank, bStride, bCoords, bOffset);
+
+   T3 val = A[aOffset] * B[bOffset];  // first iteration
+
+   for (LongType j = 1; j < K; ++j) {  // rest iterations
+     aOffset += aStride[aKaxis];
+     bOffset += bStride[bKaxis];
+     val = val + A[aOffset] * B[bOffset];
+   }
+
+   COORDS2INDEX(cRank, cStride, cCoords, cOffset);
+
+   if (betaPresent)
+     C[cOffset] = alphaZ * val + betaZ * C[cOffset];
+   else
+     C[cOffset] = alphaZ * val;
+ }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T1, typename T2, typename T3>
+SD_HOST static void usualGemm(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                             cudaStream_t* stream, const void* vA, const LongType* aShapeInfo, const void* vB,
+                             const LongType* bShapeInfo, void* vC, const LongType* cShapeInfo,
+                             const int aMaxis, const int aKaxis, const int bKaxis, const int bNaxis, const int cMaxis,
+                             const int cNaxis, const double alpha, const double beta) {
+ usualCudaGemm<T1, T2, T3><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
+     vA, aShapeInfo, vB, bShapeInfo, vC, cShapeInfo, aMaxis, aKaxis, bKaxis, bNaxis, cMaxis, cNaxis, alpha, beta);
+ DebugHelper::checkGlobalErrorCode("MMUL cuda gemv case failed(...) failed");
+}
+
+////////////////////////////////////////////////////////////////////////
+// MXN x N = M  -> actual sequence of {M,N} axes doesn't matter
+template <typename T1, typename T2, typename T3>
+static SD_KERNEL void usualCudaGemv(const void* vA, const LongType* aShapeInfo, const void* vX,
+                                   const LongType* xShapeInfo, void* vY, const LongType* yShapeInfo,
+                                   const int incx, const int incy, const int aMaxis, const double alpha,
+                                   const double beta) {
+
+ // Cache shape information in shared memory
+ __shared__ LongType M;
+ __shared__ LongType N;
+ __shared__ bool betaPresent;
+ __shared__ LongType totalThreads;
+ __shared__ LongType aNstride;
+ __shared__ LongType aMstride;
+ __shared__ T3 alphaZ;
+ __shared__ T3 betaZ;
+ __shared__ const LongType* aShape;
+ __shared__ const LongType* aStride;
+ __shared__ LongType aRank;
+
+ if (threadIdx.x == 0) {
+   N = shape::length(xShapeInfo);
+   M = shape::length(yShapeInfo);
+   aRank = shape::rank(aShapeInfo);
+
+   aShape = shape::shapeOf(aShapeInfo);
+   aStride = shape::stride(aShapeInfo);
+
+   aMstride = aStride[aMaxis];
+   aNstride = aStride[aMaxis == 0 ? 1 : 0];
+
+   totalThreads = gridDim.x * blockDim.x;
+   betaPresent = beta != 0;
+   alphaZ = alpha;
+   betaZ = beta;
+ }
+ __syncthreads();
+
+ const T1* A = reinterpret_cast<const T1*>(vA);
+ const T2* X = reinterpret_cast<const T2*>(vX);
+ T3* Y = reinterpret_cast<T3*>(vY);
+
+ const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+ for (LongType i = tid; i < M; i += totalThreads) {
+   // evaluate offsets
+   auto aOffset = i * aMstride;
+   auto xOffset = 0;
+
+   T3 val = A[aOffset] * X[xOffset];  // first iteration
+
+   for (LongType j = 1; j < N; ++j) {  // rest iterations
+     aOffset += aNstride;
+     xOffset += incx;
+     val = val + A[aOffset] * X[xOffset];
+   }
+
+   auto yOffset = i * incy;
+
+   if (betaPresent)
+     Y[yOffset] = alphaZ * val + betaZ * Y[yOffset];
+   else
+     Y[yOffset] = alphaZ * val;
+ }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T1, typename T2, typename T3>
+SD_HOST static void usualGemv(const int blocksPerGrid, const int threadsPerBlock, cudaStream_t* stream, const void* vA,
+                             const LongType* aShapeInfo, const void* vX, const LongType* xShapeInfo, void* vY,
+                             const LongType* yShapeInfo, const int incx, const int incy, const int aMaxis,
+                             const double alpha, const double beta) {
+ usualCudaGemv<T1, T2, T3><<<blocksPerGrid, threadsPerBlock, 512, *stream>>>(
+     vA, aShapeInfo, vX, xShapeInfo, vY, yShapeInfo, incx, incy, aMaxis, alpha, beta);
+ DebugHelper::checkGlobalErrorCode("MMUL cuda gemv case failed(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////////
+template <typename T1, typename T2, typename T3>
+static SD_KERNEL void usualCudaDot(const LongType length, const double alpha, const void* vX,
+                                  const LongType incx, const void* vY, const LongType incy, const double beta,
+                                  void* vZ) {
+ // Cache values in shared memory
+ __shared__ T3* pairwiseMul;
+ __shared__ T3 alphaZ;
+ __shared__ T3 betaZ;
+ __shared__ bool betaPresent;
+
+ if (threadIdx.x == 0) {
+   extern __shared__ unsigned char shmem[];
+   pairwiseMul = reinterpret_cast<T3*>(shmem);
+   alphaZ = alpha;
+   betaZ = beta;
+   betaPresent = beta != 0;
+ }
+ __syncthreads();
+
+ T1* X = reinterpret_cast<T1*>(const_cast<void*>(vX));
+ T2* Y = reinterpret_cast<T2*>(const_cast<void*>(vY));
+ T3* Z = reinterpret_cast<T3*>(vZ);
+
+ const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+ if (tid < length) {
+   pairwiseMul[tid] = X[tid * incx] * Y[tid * incy];
+ }
+
+ __syncthreads();
+
+ if (tid == 0) {
+   T3 sum = static_cast<T3>(0);
+   for (LongType i = 0; i < length; ++i) {
+     sum = sum + pairwiseMul[i];
+   }
+
+   if (betaPresent)
+     *Z = alphaZ * sum + betaZ * *Z;
+   else
+     *Z = alphaZ * sum;
+ }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T1, typename T2, typename T3>
+SD_HOST static void usualDot(const dim3& launchDims, cudaStream_t* stream,
+                            const LongType length, const double alpha, const void* vX, const LongType incx,
+                            const void* vY, const LongType incy, const double beta, void* vZ) {
+ usualCudaDot<T1, T2, T3><<<launchDims.x, launchDims.y, launchDims.z, *stream>>>(
+     length, alpha, vX, incx, vY, incy, beta, vZ);
+ DebugHelper::checkGlobalErrorCode("concat dot failed(...) failed");
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// [bS,M,K] x [bS,K,N] = [bS,M,N]
+// [bS,M,K] x    [K,N] = [bS,M,N]
+//    [M,K] x [bS,K,N] = [bS,M,N]
+// bS could stand for several axes
+template <typename T1, typename T2, typename T3>
+static SD_KERNEL void batchedCudaGemm(const void* vA, const LongType* aShapeInfo, const void* vB,
+                                     const LongType* bShapeInfo, void* vC, const LongType* cShapeInfo,
+                                     const LongType* aBatchDims, const LongType* bBatchDims,
+                                     const LongType* cBatchDims, const LongType aMaxis, const LongType aKaxis,
+                                     const LongType bKaxis, const LongType bNaxis, const LongType cMaxis,
+                                     const LongType cNaxis, const double alpha, const double beta) {
+
+ // Cache shape information in shared memory
+ __shared__ struct {
+   bool betaPresent;
+   LongType aRank, bRank, cRank, K;
+   LongType cLen, totalThreads;
+   T3 alphaZ, betaZ;
+   const LongType* aShape;
+   const LongType* bShape;
+   const LongType* cShape;
+   const LongType* aStride;
+   const LongType* bStride;
+   const LongType* cStride;
+   LongType* coords;
+ } shared;
+
+ if (threadIdx.x == 0) {
+   extern __shared__ unsigned char shmem[];
+   shared.coords = reinterpret_cast<LongType*>(shmem);
+   shared.cLen = shape::length(cShapeInfo);
+
+   shared.aRank = shape::rank(aShapeInfo);
+   shared.bRank = shape::rank(bShapeInfo);
+   shared.cRank = shape::rank(cShapeInfo);
+
+   shared.aShape = shape::shapeOf(aShapeInfo);
+   shared.bShape = shape::shapeOf(bShapeInfo);
+   shared.cShape = shape::shapeOf(cShapeInfo);
+
+   shared.aStride = shape::stride(aShapeInfo);
+   shared.bStride = shape::stride(bShapeInfo);
+   shared.cStride = shape::stride(cShapeInfo);
+
+   shared.K = shared.aShape[aKaxis];
+   shared.betaPresent = beta != 0;
+   shared.totalThreads = gridDim.x * blockDim.x;
+   shared.alphaZ = alpha;
+   shared.betaZ = beta;
+ }
+ __syncthreads();
+
+ const T1* A = reinterpret_cast<const T1*>(vA);
+ const T2* B = reinterpret_cast<const T2*>(vB);
+ T3* C = reinterpret_cast<T3*>(vC);
+
+ auto aCoords = shared.coords + threadIdx.x * (shared.aRank + shared.bRank + shared.cRank);
+ auto bCoords = aCoords + shared.aRank;
+ auto cCoords = bCoords + shared.bRank;
+
+ const auto tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+ for (LongType i = tid; i < shared.cLen; i += shared.totalThreads) {
+   // evaluate C coordinates
+   INDEX2COORDS(i, shared.cRank, shared.cShape, cCoords);
+
+   // calculate index of current batch
+   LongType batchInd;
+   if (cBatchDims != nullptr) {
+     COORDS2INDEX(shared.cRank - 2, shared.cStride, cCoords, batchInd);
+   }
+
+   // evaluate A coordinates
+   if (aBatchDims != nullptr) {
+     INDEX2COORDS(batchInd, shared.aRank - 2, shared.aShape, aCoords);
+   }
+   aCoords[aMaxis] = cCoords[cMaxis];
+   aCoords[aKaxis] = 0;
+
+   // evaluate B coordinates
+   if (bBatchDims != nullptr) {
+     INDEX2COORDS(batchInd, shared.bRank - 2, shared.bShape, bCoords);
+   }
+   bCoords[bKaxis] = 0;
+   bCoords[bNaxis] = cCoords[cNaxis];
+
+   LongType aOffset, bOffset, cOffset;
+   COORDS2INDEX(shared.aRank, shared.aStride, aCoords, aOffset);
+   COORDS2INDEX(shared.bRank, shared.bStride, bCoords, bOffset);
+
+   T3 val = A[aOffset] * B[bOffset];  // first iteration
+
+   for (LongType j = 1; j < shared.K; ++j) {  // rest iterations
+     aOffset += shared.aStride[aKaxis];
+     bOffset += shared.bStride[bKaxis];
+     val = val + A[aOffset] * B[bOffset];
+   }
+
+   COORDS2INDEX(shared.cRank, shared.cStride, cCoords, cOffset);
+
+   if (shared.betaPresent)
+     C[cOffset] = shared.alphaZ * val + shared.betaZ * C[cOffset];
+   else
+     C[cOffset] = shared.alphaZ * val;
+ }
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T1, typename T2, typename T3>
+SD_HOST static void batchedGemm(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem,
+                               cudaStream_t* stream, const void* vA, const LongType* aShapeInfo, const void* vB,
+                               const LongType* bShapeInfo, void* vC, const LongType* cShapeInfo,
+                               const LongType* aBatchDims, const LongType* bBatchDims, const LongType* cBatchDims,
+                               const LongType aMaxis, const LongType aKaxis, const LongType bKaxis,
+                               const LongType bNaxis, const LongType cMaxis, const LongType cNaxis, const double alpha,
+                               const double beta) {
+ batchedCudaGemm<T1, T2, T3><<<blocksPerGrid, threadsPerBlock, sharedMem, *stream>>>(
+     vA, aShapeInfo, vB, bShapeInfo, vC, cShapeInfo, aBatchDims, bBatchDims, cBatchDims, aMaxis, aKaxis, bKaxis,
+     bNaxis, cMaxis, cNaxis, alpha, beta);
+ DebugHelper::checkGlobalErrorCode("batch gemm failed(...) failed");
+}
+//////////////////////////////////////////////////////////////////////////////
+// MXK x KxN = MxN
+NDArray* MmulHelper::mmulMxM(NDArray* A, NDArray* B, NDArray* C, double alpha, double beta,
+                            const char outOrder) {
+ if (A->rankOf() != 2) THROW_EXCEPTION("MmulHelper::mmulMxM cuda: rank of A array is not equal 2 !");
+ if (B->rankOf() != 2) THROW_EXCEPTION("MmulHelper::mmulMxM cuda: rank of B array is not equal 2 !");
+
+ const auto M = A->sizeAt(0);
+ const auto K = A->sizeAt(1);
+ const auto N = B->sizeAt(1);
+
+ if (C != nullptr && C->rankOf() != 2)
+   THROW_EXCEPTION("MmulHelper::mmulMxM cuda: rank of C array is not equal 2 !");
+ if (B->sizeAt(0) != K) THROW_EXCEPTION("MmulHelper::mmulMxM cuda: B array has wrong number of rows !");
+ if (C != nullptr && C->sizeAt(0) != M)
+   THROW_EXCEPTION("MmulHelper::mmulMxM cuda: C array has wrong number of rows !");
+ if (C != nullptr && C->sizeAt(1) != N)
+   THROW_EXCEPTION("MmulHelper::mmulMxM cuda: C array has wrong number of columns !");
+
+ std::vector<LongType> cShape = {M, N};
+ if (C == nullptr)
+   C = new NDArray(outOrder, cShape, DataTypeUtils::pickPairwiseResultType(A->dataType(), B->dataType()),
+                   A->getContext());
+
+ if (C->isEmpty()) return C;
+
+ const int major = Environment::getInstance().capabilities()[AffinityManager::currentDeviceId()].first();
+
+ const auto aType = A->dataType();
+ const auto bType = B->dataType();
+ const auto cType = C->dataType();
+
+ const bool AB(aType == bType), AC(aType == cType), ABC(AB && AC);
+
+ const bool typeDouble = ABC && aType == DOUBLE;
+ const bool typeFloat = ABC && aType == FLOAT32;
+ const bool typeHalf = ABC && aType == HALF && major >= 6;
+ const bool typeIntFloat = AB && aType == INT8 && cType == FLOAT32 && major >= 6;
+ const bool typeHalfFloat = AB && aType == HALF && cType == FLOAT32 && major >= 6;
+
+ std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());
+
+ auto handle = reinterpret_cast<cublasHandle_t*>(A->getContext()->getCublasHandle());
+ auto stream = A->getContext()->getCudaStream();
+
+ auto status = cublasSetStream_v2(*handle, *stream);
+ if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", status);
+
+ if (!typeDouble && !typeFloat && !typeHalf && !typeIntFloat && !typeHalfFloat) {
+   dim3 dims = getMMulDims(C->lengthOf(),DataTypeUtils::sizeOf(cType));
+   NDArray::prepareSpecialUse({C}, {A, B});
+   BUILD_SINGLE_SELECTOR_THRICE(aType, usualGemm,
+                                (dims.y, dims.x, dims.z, stream, A->specialBuffer(),
+                                 A->specialShapeInfo(), B->specialBuffer(), B->specialShapeInfo(), C->specialBuffer(),
+                                 C->specialShapeInfo(), 0, 1, 0, 1, 0, 1, alpha, beta),
+                                SD_NUMERIC_TYPES)
+   NDArray::registerSpecialUse({C}, {A, B});
+
+   auto cudaResult = cudaStreamSynchronize(*stream);
+   if (cudaResult != 0) throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", cudaResult);
+
+ } else {
+   std::vector<NDArray*> toDelete;
+
+   NDArray *pA(const_cast<NDArray*>(A)), *pB(const_cast<NDArray*>(B)), *pC(const_cast<NDArray*>(C));
+
+   bool aMcont = M == 1 || A->strideAt(0) == 1;
+   bool aKcont = K == 1 || A->strideAt(1) == 1;
+   bool bKcont = K == 1 || B->strideAt(0) == 1;
+   bool bNcont = N == 1 || B->strideAt(1) == 1;
+   bool cMcont = M == 1 || C->strideAt(0) == 1;
+   bool cNcont = N == 1 || C->strideAt(1) == 1;
+
+   if (!aMcont && !aKcont) {
+     pA = A->dup('f');
+     toDelete.push_back(pA);
+     aMcont = true;
+   }
+   if (!bKcont && !bNcont) {
+     pB = B->dup('f');
+     toDelete.push_back(pB);
+     bKcont = true;
+   }
+   if (!cMcont) {
+     pC = C->dup('f');
+     toDelete.push_back(pC);
+     cMcont = true;
+   }
+
+   const bool transA = !aMcont;
+   const bool transB = !bKcont;
+
+   const int lda = (aMcont && aKcont) ? M : transA ? pA->strideAt(0) : pA->strideAt(1);
+   const int ldb = (bKcont && bNcont) ? K : transB ? pB->strideAt(0) : pB->strideAt(1);
+   const int ldc = (cMcont && cNcont) ? M : pC->strideAt(1);
+
+   const cublasOperation_t transAblas = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
+   const cublasOperation_t transBblas = transB ? CUBLAS_OP_T : CUBLAS_OP_N;
+
+   NDArray::prepareSpecialUse({pC}, {pA, pB});
+
+   // choose appropriate cuda gemm api depending on data types
+   if (typeDouble) {
+     status = cublasDgemm(*handle, transAblas, transBblas, M, N, K, &alpha, (double*)pA->specialBuffer(), lda,
+                          (double*)pB->specialBuffer(), ldb, &beta, (double*)pC->specialBuffer(), ldc);
+   } else if (typeFloat) {
+     float alphaF(alpha), betaF(beta);
+     status = cublasSgemm(*handle, transAblas, transBblas, M, N, K, &alphaF, (float*)pA->specialBuffer(), lda,
+                          (float*)pB->specialBuffer(), ldb, &betaF, (float*)pC->specialBuffer(), ldc);
+   } else if (typeHalf) {
+     float16 alphaH(alpha), betaH(beta);
+     status = cublasHgemm(*handle, transAblas, transBblas, M, N, K, &alphaH.data, (__half*)pA->specialBuffer(), lda,
+                          (__half*)pB->specialBuffer(), ldb, &betaH.data, (__half*)pC->specialBuffer(), ldc);
+   } else if (typeIntFloat) {
+     float alphaF(alpha), betaF(beta);
+     status = cublasSgemmEx(*handle, transAblas, transBblas, M, N, K, &alphaF, pA->specialBuffer(), CUDA_R_8I, lda,
+                            pB->specialBuffer(), CUDA_R_8I, ldb, &betaF, pC->specialBuffer(), CUDA_R_32F, ldc);
+   } else if (typeHalfFloat) {
+     float alphaF(alpha), betaF(beta);
+     status = cublasSgemmEx(*handle, transAblas, transBblas, M, N, K, &alphaF, pA->specialBuffer(), CUDA_R_16F, lda,
+                            pB->specialBuffer(), CUDA_R_16F, ldb, &betaF, pC->specialBuffer(), CUDA_R_32F, ldc);
+   }
+
+   if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", status);
+
+   NDArray::registerSpecialUse({pC}, {pA, pB});
+
+   auto cudaResult = cudaStreamSynchronize(*stream);
+   if (cudaResult != 0) throw cuda_exception::build("MmulHelper::mmulMxM cuda failed !", cudaResult);
+
+   if (C != pC) C->assign(pC);
+
+   for (int i = toDelete.size() - 1; i >= 0; --i) delete toDelete[i];
+ }
+
+ return C;
+}////////////////////////////////////////////////////////////////////////////
+// MXN x N = M
+NDArray* MmulHelper::mmulMxV(NDArray* A, NDArray* X, NDArray* Y, const double alpha, const double beta,
+                            const char outOrder) {
+ LongType xLenDim, yLenDim(0);
+
+ if (A->rankOf() != 2) THROW_EXCEPTION("MmulHelper::mmulMxV cuda: rank of A array is not equal 2 !");
+ if (!shape::isCommonVector(X->shapeInfo(), xLenDim))
+   THROW_EXCEPTION("MmulHelper::mmulMxV cuda: X array must be vector !");
+
+ const auto M = A->sizeAt(0);
+ const auto N = A->sizeAt(1);
+
+ if (Y != nullptr && !shape::isCommonVector(Y->shapeInfo(), yLenDim))
+   THROW_EXCEPTION("MmulHelper::mmulMxV cuda: Y array must be vector !");
+ if (X->lengthOf() != N) THROW_EXCEPTION("MmulHelper::mmulMxV cuda: X vector has wrong length !");
+ if (Y != nullptr && Y->lengthOf() != M)
+   THROW_EXCEPTION("MmulHelper::mmulMxV cuda: Y array has wrong length !");
+
+ std::vector<LongType> yShape = {M};
+ if (Y == nullptr)
+   Y = new NDArray(outOrder, yShape, DataTypeUtils::pickPairwiseResultType(A->dataType(), X->dataType()),
+                   A->getContext());
+
+ if (Y->isEmpty()) return Y;
+
+ const int incx = X->strideAt(xLenDim);
+ const int incy = Y->strideAt(yLenDim);
+
+ const auto aType = A->dataType();
+ const auto xType = X->dataType();
+ const auto yType = Y->dataType();
+
+ const bool AX(aType == xType), AY(aType == yType), AXY(AX && AY);
+
+ const bool typeDouble = AXY && aType == DOUBLE;
+ const bool typeFloat = AXY && aType == FLOAT32;
+
+ std::lock_guard<std::mutex> lock(*LaunchContext::deviceMutex());
+
+ auto handle = reinterpret_cast<cublasHandle_t*>(A->getContext()->getCublasHandle());
+ auto stream = A->getContext()->getCudaStream();
+
+ auto status = cublasSetStream_v2(*handle, *stream);
+ if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("MmulHelper::mmulMxV cuda failed !", status);
+
+ if (!typeDouble && !typeFloat) {
+   dim3 dims = getGemVDims(M);
+   NDArray::prepareSpecialUse({Y}, {A, X});
+
+   const int blocksPerGrid = dims.x;
+   const int threadsPerBlock = dims.y;
+   BUILD_SINGLE_SELECTOR_THRICE(
+       xType, usualGemv,
+       (blocksPerGrid,threadsPerBlock,stream, A->specialBuffer(), A->specialShapeInfo(), X->specialBuffer(),
+        X->specialShapeInfo(), Y->specialBuffer(), Y->specialShapeInfo(), incx, incy, 0, alpha, beta),
+       SD_NUMERIC_TYPES)
+   NDArray::registerSpecialUse({Y}, {A, X});
+
+   auto cudaResult = cudaStreamSynchronize(*stream);
+   if (cudaResult != 0) throw cuda_exception::build("MmulHelper::mmulMxV cuda failed !", cudaResult);
+
+ } else {
+   NDArray* pA(const_cast<NDArray*>(A));
+
+   bool aMcont = M == 1 || A->strideAt(0) == 1;
+   bool aNcont = N == 1 || A->strideAt(1) == 1;
+
+   if (!aMcont && !aNcont) {
+     pA = A->dup('f');  // dup() already returns NDArray*, no need for new
+     aMcont = true;
+   }
+
+   const bool transA = !aMcont;
+
+   const int lda = (aMcont && aNcont) ? M : transA ? pA->strideAt(0) : pA->strideAt(1);
+
+   const cublasOperation_t transAblas = transA ? CUBLAS_OP_T : CUBLAS_OP_N;
+
+   NDArray::prepareSpecialUse({Y}, {pA, X});
+
+   // choose appropriate cuda gemm api depending on data types
+   if (typeDouble) {
+     status = cublasDgemv(*handle, transAblas, transA ? N : M, transA ? M : N, &alpha, (double*)pA->specialBuffer(),
+                          lda, (double*)X->specialBuffer(), incx, &beta, (double*)Y->specialBuffer(), incy);
+   } else if (typeFloat) {
+     float alphaF(alpha), betaF(beta);
+     status = cublasSgemv(*handle, transAblas, transA ? N : M, transA ? M : N, &alphaF, (float*)pA->specialBuffer(),
+                          lda, (float*)X->specialBuffer(), incx, &betaF, (float*)Y->specialBuffer(), incy);
+   }
+
+   if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("MmulHelper::mmulMxV cuda failed !", status);
+
+   auto cudaResult = cudaStreamSynchronize(*stream);
+   if (cudaResult != 0) throw cuda_exception::build("MmulHelper::mmulMxV cuda failed !", cudaResult);
+
+   NDArray::registerSpecialUse({Y}, {pA, X});
+
+   if (pA != A) delete pA;
+ }
+
+ return Y;
+}////////////////////////////////////////////////////////////////////////////
+// (X * Y) = Z[0]
+NDArray* MmulHelper::dot(NDArray* X, NDArray* Y, NDArray* Z, const double alpha, const double beta) {
+ LongType xLenDim(0), yLenDim(0);
+
+ if (!shape::isCommonVector(X->shapeInfo(), xLenDim))
+   THROW_EXCEPTION("MmulHelper::dot cuda: X array must be vector !");
+ if (!shape::isCommonVector(Y->shapeInfo(), yLenDim))
+   THROW_EXCEPTION("MmulHelper::dot cuda: Y array must be vector !");
+ if (Z != nullptr && Z->lengthOf() > 1) {
+   THROW_EXCEPTION("MmulHelper::dot: Z array must be scalar !");
+ }
+
+ const auto length = X->lengthOf();
+
+ if (Y->lengthOf() != length)
+   THROW_EXCEPTION("MmulHelper::dot cuda: lengths of input vectors are different !");
+
+ if (Z == nullptr)
+   Z = new NDArray(DataTypeUtils::pickPairwiseResultType(X->dataType(), Y->dataType()), X->getContext());
+
+ const LongType incx = X->strideAt(xLenDim);
+ const LongType incy = Y->strideAt(yLenDim);
+
+ const auto xType = X->dataType();
+ const auto yType = Y->dataType();
+ const auto zType = Z->dataType();
+
+ if (!X->isActualOnDeviceSide()) X->syncToDevice();
+ if (!Y->isActualOnDeviceSide()) Y->syncToDevice();
+ if (!Z->isActualOnDeviceSide()) Z->syncToDevice();
+
+ cudaStream_t* stream = X->getContext()->getCudaStream();
+
+ dim3 dims = getMMulDims(length,DataTypeUtils::sizeOf(zType));
+
+ NDArray::prepareSpecialUse({Z}, {X, Y});
+
+
+ BUILD_SINGLE_SELECTOR_THRICE(xType, usualDot,
+                              (dims, stream, length, alpha, X->specialBuffer(), incx,
+                               Y->specialBuffer(), incy, beta, Z->specialBuffer()),
+                              SD_NUMERIC_TYPES);
+
+ auto cudaResult = cudaStreamSynchronize(*stream);
+ if (cudaResult != 0) throw cuda_exception::build("MmulHelper::dot cuda failed !", cudaResult);
+
+ NDArray::registerSpecialUse({Z}, {X, Y});
+
+ return Z;
+}
+///////////////////////////////////////////////////////////////////
+NDArray* MmulHelper::mmulNxN(NDArray* A, NDArray* B, NDArray* C, double alpha, double beta,
+                            const char outOrder) {
+ const LongType aRank = A->rankOf();
+ const LongType bRank = B->rankOf();
+
+ // input ranks validation
+ if (aRank > bRank && bRank != 2) {
+   THROW_EXCEPTION("MmulHelper::mmulNxN: rank of B array should be equal 2 !");
+ }
+ else if (bRank > aRank && aRank != 2) {
+   THROW_EXCEPTION("MmulHelper::mmulNxN: rank of A array should be equal 2 !");
+ }
+ else if (aRank == bRank) {
+   for (int i = 0; i < aRank - 2; ++i)
+     if (A->sizeAt(i) != B->sizeAt(i))
+       THROW_EXCEPTION(
+           "MmulHelper::mmulNxN: shapes of A and B arrays are not suitable for matrix multiplication !");
+ }
+
+ if (A->sizeAt(-1) != B->sizeAt(-2)) {
+   THROW_EXCEPTION("MmulHelper::mmulNxN: shapes of A and B arrays are not suitable for matrix multiplication !");
+ }
+ // validation of C array
+ auto* cExpectedShapePtr = aRank > bRank ? A->getShapeAsVector() : B->getShapeAsVector();
+ std::vector<LongType> cExpectedShape = *cExpectedShapePtr;
+ delete cExpectedShapePtr;
+ cExpectedShape[cExpectedShape.size() - 2] = A->sizeAt(-2);
+ cExpectedShape[cExpectedShape.size() - 1] = B->sizeAt(-1);
+
+ if (C != nullptr) {
+   if (!C->isSameShape(cExpectedShape))
+     THROW_EXCEPTION("MmulHelper::mmulNxN: shape of C array is not suitable for AxB matrix multiplication !");
+ } else
+   C = new NDArray(outOrder, cExpectedShape, DataTypeUtils::pickPairwiseResultType(A->dataType(), B->dataType()),
+                   A->getContext());
+
+ if (C->isEmpty()) return C;
+
+ const LongType cRank = C->rankOf();
+
+ const LongType aMaxis(aRank - 2), aKaxis(aRank - 1), bKaxis(bRank - 2), bNaxis(bRank - 1), cMaxis(cRank - 2),
+     cNaxis(cRank - 1);
+
+ const int threadsPerBlock = SD_MAX_NUM_THREADS / 8;
+ const int blocksPerGrid = (C->lengthOf() + threadsPerBlock - 1) / threadsPerBlock;
+ const int sharedMem = threadsPerBlock * sizeof(LongType) * (aRank + bRank + cRank) + 128;
+
+ PointersManager manager(A->getContext(), "MmulHelper::mmulNxN");
+
+ const LongType *aBatchDims(nullptr), *bBatchDims(nullptr), *cBatchDims(nullptr);
+
+ std::vector<LongType> aDimsVec = {aMaxis,aKaxis};
+ std::vector<LongType> *aDims = ShapeUtils::evalDimsToExclude(aRank, 2,aDimsVec.data());
+
+ std::vector<LongType> bDimsVec = {bKaxis, bNaxis};
+ std::vector<LongType> *bDims =  ShapeUtils::evalDimsToExclude(bRank,2, bDimsVec.data());
+
+
+ std::vector<LongType> cDimsVec = {cMaxis,2, cNaxis};
+ std::vector<LongType> *cDims = ShapeUtils::evalDimsToExclude(cRank, cDimsVec.size(),cDimsVec.data());
+ if (aRank > 2)
+   aBatchDims = reinterpret_cast<LongType*>(manager.replicatePointer(
+       aDims->data(), (aRank - 2) * sizeof(LongType)));
+ if (bRank > 2)
+   bBatchDims = reinterpret_cast<LongType*>(manager.replicatePointer(
+       bDims->data(), (bRank - 2) * sizeof(LongType)));
+ if (cRank > 2)
+   cBatchDims = reinterpret_cast<LongType*>(manager.replicatePointer(
+       cDims->data(), (cRank - 2) * sizeof(LongType)));
+
+ NDArray::prepareSpecialUse({C}, {A, B});
+ BUILD_SINGLE_SELECTOR_THRICE(
+     A->dataType(), batchedGemm,
+     (blocksPerGrid, threadsPerBlock, sharedMem, A->getContext()->getCudaStream(), A->specialBuffer(),
+      A->specialShapeInfo(), B->specialBuffer(), B->specialShapeInfo(), C->specialBuffer(), C->specialShapeInfo(),
+      aBatchDims, bBatchDims, cBatchDims, aMaxis, aKaxis, bKaxis, bNaxis, cMaxis, cNaxis, alpha, beta),
+     SD_NUMERIC_TYPES)
+ NDArray::registerSpecialUse({C}, {A, B});
+
+ manager.synchronize();
+
+ delete aDims;
+ delete bDims;
+ delete cDims;
+
+ return C;
+}
+
+} // namespace sd

libnd4j/include/helpers/cuda/PointersManager.cu

@@ -0,0 +1,125 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 Yurii Shyrma (iuriish@yahoo.com), created on 06.02.2019
+// @author raver119@gmail.com
+//
+#include <exceptions/cuda_exception.h>
+#include <helpers/PointersManager.h>
+#include <helpers/StringUtils.h>
+#include <helpers/logger.h>
+#include <memory/Workspace.h>
+
+#include "helpers/DebugHelper.h"
+
+namespace sd {
+
+//////////////////////////////////////////////////////////////////////////
+PointersManager::PointersManager(const LaunchContext* context, const std::string& funcName) {
+  _context = const_cast<LaunchContext*>(context);
+  _funcName = funcName;
+}
+//////////////////////////////////////////////////////////////////////////
+void* PointersManager::allocateDevMem(const size_t sizeInBytes) {
+  void* dst = nullptr;
+  if (_context->getWorkspace() == nullptr) {
+    cudaError_t cudaResult = cudaMalloc(reinterpret_cast<void**>(&dst), sizeInBytes);
+    if (cudaResult != 0)
+      throw cuda_exception::build(_funcName + ": cannot allocate global memory on device!", cudaResult);
+  } else {
+    dst = _context->getWorkspace()->allocateBytes(memory::MemoryType::DEVICE, sizeInBytes);
+  }
+  return dst;
+}
+
+//////////////////////////////////////////////////////////////////////////
+void* PointersManager::replicatePointer(const void* src, const size_t numberOfBytes) {
+  void* dst = allocateDevMem(numberOfBytes);
+  if (src) {
+    if (_context != nullptr)
+      cudaMemcpyAsync(dst, src, numberOfBytes, cudaMemcpyHostToDevice, *_context->getCudaStream());
+    else
+      cudaMemcpy(dst, src, numberOfBytes, cudaMemcpyHostToDevice);
+  }
+  _pOnGlobMem.emplace_back(dst);
+
+  return dst;
+}
+
+//////////////////////////////////////////////////////////////////////////
+void PointersManager::synchronize() const {
+  if (_context != nullptr) {
+    cudaError_t cudaResult = cudaStreamSynchronize(*_context->getCudaStream());
+    if (cudaResult != 0) throw cuda_exception::build(_funcName + ": cuda stream synchronization failed !", cudaResult);
+  } else {
+    sd_printf("<%s> syncStream isn't possible: no stream set!", _funcName.c_str());
+  }
+}
+
+//////////////////////////////////////////////////////////////////////////
+PointersManager::~PointersManager() {
+ // for (auto& p : _pOnGlobMem) cudaFree(p);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+static SD_KERNEL void printDevContentOnDev_(const void* pDev, const LongType len, const int tid) {
+  PointersManager::printDevContentOnDev<T>(pDev, len, tid);
+}
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+void PointersManager::printDevContentOnDevFromHost(const void* pDev, const LongType len, const int tid) {
+  printDevContentOnDev_<T><<<512, 512, 1024, *LaunchContext ::defaultContext()->getCudaStream()>>>(pDev, len, tid);
+  auto res = cudaStreamSynchronize(*LaunchContext ::defaultContext()->getCudaStream());
+  DebugHelper::checkGlobalErrorCode("concat general case failed(...) failed");
+
+}
+template void PointersManager::printDevContentOnDevFromHost<LongType>(const void* pDev, const LongType len,
+                                                                          const int tid);
+template void PointersManager::printDevContentOnDevFromHost<int>(const void* pDev, const LongType len,
+                                                                 const int tid);
+template void PointersManager::printDevContentOnDevFromHost<float>(const void* pDev, const LongType len,
+                                                                   const int tid);
+template void PointersManager::printDevContentOnDevFromHost<double>(const void* pDev, const LongType len,
+                                                                    const int tid);
+
+
+////////////////////////////////////////////////////////////////////////
+template <typename T>
+void PointersManager::printDevContentOnHost(const void* pDev, const LongType len) const {
+  printf("host print out\n");
+  void* pHost = operator new(sizeof(T) * len);
+
+  cudaMemcpyAsync(pHost, pDev, sizeof(T) * len, cudaMemcpyDeviceToHost, *_context->getCudaStream());
+  cudaError_t cudaResult = cudaStreamSynchronize(*_context->getCudaStream());
+  if (cudaResult != 0) THROW_EXCEPTION("PointersManager::printCudaHost: cudaStreamSynchronize failed!");
+
+  for (LongType i = 0; i < len; ++i) printf("%f, ", (double)reinterpret_cast<T*>(pHost)[i]);
+  printf("\n");
+
+  operator delete(pHost);
+}
+
+template void PointersManager::printDevContentOnHost<LongType>(const void* pDev, const LongType len) const;
+template void PointersManager::printDevContentOnHost<int>(const void* pDev, const LongType len) const;
+template void PointersManager::printDevContentOnHost<float>(const void* pDev, const LongType len) const;
+template void PointersManager::printDevContentOnHost<double>(const void* pDev, const LongType len) const;
+
+}  // namespace sd

libnd4j/include/helpers/cuda/TadCalculator.cu

@@ -0,0 +1,148 @@
+/*
+ *  ******************************************************************************
+ *  *
+ *  *
+ *  * 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
+ *  *****************************************************************************
+ */
+
+#include <helpers/ConstantHelper.h>
+#include <helpers/ConstantShapeHelper.h>
+#include <helpers/ShapeBufferCreatorHelper.h>
+#include <helpers/ShapeUtils.h>
+#include <helpers/cuda/CudaShapeBufferCreator.h>
+
+#include "array/CudaPointerDeallocator.h"
+#include "array/TadCalculator.h"
+
+namespace sd {
+
+TadCalculator::TadCalculator(LongType* originalShape)
+    : _originalShape(originalShape), _numTads(0) {}
+
+TadCalculator::~TadCalculator() {
+  if (_tadOffsets) {
+    delete _tadOffsets;
+    _tadOffsets = nullptr;
+  }
+  if (_tadShape) {
+    delete _tadShape;
+    _tadShape = nullptr;
+  }
+}
+
+
+void TadCalculator::createTadPack(const std::vector<LongType>& dimensions) {
+  if (!_originalShape) {
+    THROW_EXCEPTION("Original shape is null");
+  }
+
+  auto shapeInfo = ConstantShapeHelper::getInstance().createFromExisting(_originalShape);
+  const LongType rank = shape::rank(shapeInfo);
+
+  const std::vector<LongType>* dimsToExclude = ShapeUtils::evalDimsToExclude(rank, dimensions.size(), dimensions.data());
+  if (!dimsToExclude) {
+    THROW_EXCEPTION("Failed to evaluate dimensions to exclude");
+  }
+
+  bool hasSize1Dimension = false;
+  for (auto dim : dimensions) {
+    if (shape::shapeOf(shapeInfo)[dim] == 1) {
+      hasSize1Dimension = true;
+      break;
+    }
+  }
+
+  if (dimsToExclude->size() == 0 || dimsToExclude->size() == rank || hasSize1Dimension) {
+    const LongType totalElements = shape::length(shapeInfo);
+    
+    auto scalarShapeInfo = ConstantShapeHelper::getInstance().scalarShapeInfo(ArrayOptions::dataType(shapeInfo));
+    auto scalarShapeBuffer = CudaShapeBufferCreator::getInstance().create(scalarShapeInfo, 0);
+    
+    LongType* offsetsBuf = new LongType[totalElements];
+    
+    for (LongType i = 0; i < totalElements; ++i) {
+      offsetsBuf[i] = i;
+    }
+    
+    auto oPtr = std::make_shared<PointerWrapper>(offsetsBuf);
+    auto offDPtr = std::make_shared<PointerWrapper>(
+        ConstantHelper::getInstance().replicatePointer(oPtr->pointer(), totalElements * sizeof(LongType)),
+        std::make_shared<CudaPointerDeallocator>());
+    
+    _tadShape = scalarShapeBuffer;
+    _tadOffsets = new ConstantOffsetsBuffer(oPtr, offDPtr);
+    _numTads = totalElements;
+    
+    delete dimsToExclude;
+    return;
+  }
+
+  const LongType numOfSubArrs = ShapeUtils::getNumOfSubArrs(shapeInfo, *dimsToExclude);
+
+  if (numOfSubArrs > 0) {
+    const LongType subArrRank = (static_cast<size_t>(rank) == dimsToExclude->size() || false) ? rank : rank - dimsToExclude->size();
+
+    LongType* shapeInfoBuf = new LongType[shape::shapeInfoLength(subArrRank)];
+    LongType* offsetsBuf = new LongType[numOfSubArrs];
+
+    shape::calcSubArrsShapeInfoAndOffsets(
+        shapeInfo,
+        numOfSubArrs,
+        dimsToExclude->size(),
+        dimsToExclude->data(),
+        shapeInfoBuf,
+        offsetsBuf,
+        false);
+
+    ConstantShapeBuffer* shapesBuffer = CudaShapeBufferCreator::getInstance().create(shapeInfoBuf, subArrRank);
+
+    auto oPtr = std::make_shared<PointerWrapper>(offsetsBuf);
+    auto offDPtr = std::make_shared<PointerWrapper>(
+        ConstantHelper::getInstance().replicatePointer(oPtr->pointer(), numOfSubArrs * sizeof(LongType)),
+        std::make_shared<CudaPointerDeallocator>());
+    
+    _tadOffsets = new ConstantOffsetsBuffer(oPtr, offDPtr);
+    _tadShape = shapesBuffer;
+    _numTads = numOfSubArrs;
+    
+
+  } else {
+    const LongType subArrRank = rank;
+
+    LongType* shapeInfoBuf = new LongType[shape::shapeInfoLength(subArrRank)];
+
+    auto nonConstant = const_cast<LongType*>(shapeInfo);
+    shape::copyTo<LongType>(shape::shapeInfoLength(subArrRank), nonConstant, shapeInfoBuf);
+
+    ConstantShapeBuffer* shapesBuffer = CudaShapeBufferCreator::getInstance().create(shapeInfoBuf, subArrRank);
+
+    LongType* baseOffset = new LongType[1];
+    baseOffset[0] = 0;
+    auto oPtr = std::make_shared<PointerWrapper>(baseOffset);
+    auto offDPtr = std::make_shared<PointerWrapper>(
+        ConstantHelper::getInstance().replicatePointer(oPtr->pointer(), 1 * sizeof(LongType)),
+        std::make_shared<CudaPointerDeallocator>());
+    
+    _tadOffsets = new ConstantOffsetsBuffer(oPtr, offDPtr);
+    _tadShape = shapesBuffer;
+    _numTads = 1;
+
+  }
+
+  delete dimsToExclude;
+}
+
+} // namespace sd

libnd4j/include/helpers/cuda/cublasHelper.cu

@@ -0,0 +1,131 @@
+/* ******************************************************************************
+ *
+ *
+ * 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 <cublas_v2.h>
+#include <cusolverDn.h>
+#include <exceptions/cuda_exception.h>
+#include <execution/AffinityManager.h>
+#include <helpers/logger.h>
+
+#include "../cublasHelper.h"
+#include "config.h"
+
+#ifdef HAVE_CUDNN
+#include <cudnn.h>
+
+#endif
+
+namespace sd {
+std::mutex CublasHelper::_mutex;
+
+static void* handle_() {
+  auto _handle = new cublasHandle_t();
+  auto status = cublasCreate_v2(_handle);  // initialize CUBLAS context
+  if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("cuBLAS handle creation failed !", status);
+
+  return reinterpret_cast<void*>(_handle);
+}
+
+static void* solver_() {
+  auto cusolverH = new cusolverDnHandle_t();
+  auto status = cusolverDnCreate(cusolverH);
+  if (status != CUSOLVER_STATUS_SUCCESS) throw cuda_exception::build("cuSolver handle creation failed !", status);
+
+  return cusolverH;
+}
+
+static void* cudnn_() {
+#ifdef HAVE_CUDNN
+  auto cudnnH = new cudnnHandle_t();
+  auto status = cudnnCreate(cudnnH);
+  if (status != CUDNN_STATUS_SUCCESS) throw cuda_exception::build("cuDNN handle creation failed !", status);
+
+  return cudnnH;
+#endif
+  return nullptr;
+}
+
+static void destroyHandle_(void* handle) {
+  auto ch = reinterpret_cast<cublasHandle_t*>(handle);
+  auto status = cublasDestroy_v2(*ch);
+  if (status != CUBLAS_STATUS_SUCCESS) throw cuda_exception::build("cuBLAS handle destruction failed !", status);
+
+  delete ch;
+}
+
+CublasHelper::CublasHelper() {
+  auto numDevices = AffinityManager::numberOfDevices();
+  auto currentDevice = AffinityManager::currentDeviceId();
+  _cache.resize(numDevices);
+  _solvers.resize(numDevices);
+  _cudnn.resize(numDevices);
+  for (int e = 0; e < numDevices; e++) {
+    AffinityManager::setCurrentNativeDevice(e);
+
+    _cache[e] = handle_();
+    _solvers[e] = solver_();
+    _cudnn[e] = cudnn_();
+  }
+
+  // don't forget to restore back original device
+  AffinityManager::setCurrentNativeDevice(currentDevice);
+}
+
+CublasHelper::~CublasHelper() {
+  auto numDevices = AffinityManager::numberOfDevices();
+
+ // for (int e = 0; e < numDevices; e++) destroyHandle_(_cache[e]);
+}
+
+CublasHelper& CublasHelper::getInstance() {
+  static CublasHelper instance;
+  return instance;
+}
+
+void* CublasHelper::cudnn() {
+  auto deviceId = AffinityManager::currentDeviceId();
+  if (deviceId < 0 || deviceId > _cudnn.size())
+    throw cuda_exception::build("requested deviceId doesn't look valid", deviceId);
+
+  return _cudnn[deviceId];
+}
+
+void* CublasHelper::handle() {
+  auto deviceId = AffinityManager::currentDeviceId();
+  return handle(deviceId);
+}
+
+void* CublasHelper::solver() {
+  auto deviceId = AffinityManager::currentDeviceId();
+  if (deviceId < 0 || deviceId > _solvers.size())
+    throw cuda_exception::build("requested deviceId doesn't look valid", deviceId);
+
+  return _solvers[deviceId];
+}
+
+void* CublasHelper::handle(int deviceId) {
+  if (deviceId < 0 || deviceId > _cache.size())
+    throw cuda_exception::build("requested deviceId doesn't look valid", deviceId);
+
+  return _cache[deviceId];
+}
+}  // namespace sd