diff --git a/libnd4j/include/ops/declarable/helpers/cpu/batchnorm.cpp b/libnd4j/include/ops/declarable/helpers/cpu/batchnorm.cpp index ada2c5d72..ad2e29a97 100644 --- a/libnd4j/include/ops/declarable/helpers/cpu/batchnorm.cpp +++ b/libnd4j/include/ops/declarable/helpers/cpu/batchnorm.cpp @@ -15,7 +15,7 @@ ******************************************************************************/ // -// @author Yurii Shyrma, created on 25.02.2018 +// @author Yurii Shyrma (iuriish@yahoo.com) // @@ -31,112 +31,160 @@ namespace helpers { ////////////////////////////////////////////////////////////////////////// template -static void batchnorm_(const NDArray* input, const NDArray* mean, const NDArray* variance, const NDArray* gamma, const NDArray* beta, NDArray* output, const std::vector& axes, const double epsilon) { +static void batchnorm_(const NDArray* input, const NDArray* mean, const NDArray* variance, const NDArray* gamma, const NDArray* beta, + NDArray* output, + const std::vector& axes, const double epsilon) { // formula: output = gamma * ((input - mean) / sqrt(variance + epsilon)) + beta - NDArray sigmaInvGam(mean); // do not copy mean's buffer, take only its shapeInfo - T eps = epsilon; + const T* x = input->bufferAsT(); + T* z = output->bufferAsT(); + const T* m = mean->bufferAsT(); + const T* v = variance->bufferAsT(); + const T* g = gamma == nullptr ? nullptr : gamma->bufferAsT(); + const T* b = beta == nullptr ? nullptr : beta->bufferAsT(); - if(gamma != nullptr) { - auto lambda = LAMBDA_TT(x, y, eps) {return x / nd4j::math::nd4j_sqrt(y + eps);}; - const_cast(gamma)->applyPairwiseLambda(*variance, lambda, sigmaInvGam); - } - else { - auto lambda = LAMBDA_T(x, eps) { return 1. / nd4j::math::nd4j_sqrt(x + eps); }; - const_cast(variance)->applyLambda(lambda, sigmaInvGam); - } + const bool xzSameOffset = shape::haveSameShapeAndStrides(input->getShapeInfo(), output->getShapeInfo()); - // auto sigmaInvGam = (*variance + epsilon).transform(transform::RSqrt); // sigmaInvGam = 1 / sqrt(variance + epsilon) - // if(gamma != nullptr) sigmaInvGam *= *gamma; - - const T* sigmaBuff = sigmaInvGam.bufferAsT(); - const T* meanBuff = mean->bufferAsT(); - const T* inBuff = input->bufferAsT(); - T* outBuff = output->bufferAsT(); + bool paramSameOffset = shape::haveSameShapeAndStrides(mean->getShapeInfo(), variance->getShapeInfo()); + if(paramSameOffset && gamma != nullptr) + paramSameOffset &= shape::haveSameShapeAndStrides(mean->getShapeInfo(), gamma->getShapeInfo()); + if(paramSameOffset && beta != nullptr) + paramSameOffset &= shape::haveSameShapeAndStrides(mean->getShapeInfo(), beta->getShapeInfo()); const Nd4jLong lenBig = input->lengthOf(); const Nd4jLong lenSmall = mean->lengthOf(); - const Nd4jLong* inShapeInfo = input->getShapeInfo(); - const Nd4jLong* meanShapeInfo = mean->getShapeInfo(); - uint inShapeInfoCast[MAX_RANK]; - uint meanShapeInfoCast[MAX_RANK]; - bool canCastIn = nd4j::DataTypeUtils::castShapeInfo(inShapeInfo, inShapeInfoCast); - bool canCastMean = nd4j::DataTypeUtils::castShapeInfo(meanShapeInfo, meanShapeInfoCast); - - const Nd4jLong step = lenBig / lenSmall; + const Nd4jLong steps = lenBig / lenSmall; std::vector dimsToExclude = ShapeUtils::evalDimsToExclude(input->rankOf(), axes); OmpLaunchHelper info(lenBig, lenSmall); - if(beta != nullptr) { - const T* betaBuff = beta->bufferAsT(); - auto func = PRAGMA_THREADS_DO { - const auto threadNum = thread_id; - Nd4jLong* inOffsets = new Nd4jLong[step]; - Nd4jLong* memBuff = new Nd4jLong[2 * inShapeInfo[0]]; + auto func = PRAGMA_THREADS_DO { - for (int j = 0; j < lenSmall; ++j) { + Nd4jLong* xOffsets = new Nd4jLong[steps]; + Nd4jLong* zOffsets = xzSameOffset ? xOffsets : new Nd4jLong[steps]; + Nd4jLong* auxBuff = new Nd4jLong[2 * input->rankOf()]; - const bool isOwner = j < info._numThreads ? threadNum == j : threadNum == j % info._numThreads; - if (!isOwner) continue; + for (int j = 0; j < lenSmall; ++j) { - const Nd4jLong start = j * step; - const Nd4jLong end = start + step; + const bool isOwner = (j < info._numThreads) ? thread_id == j : thread_id == (j % info._numThreads); - // calculate offset for mean, variance, gamma, beta (all of them have the same shape) - auto offsetSmall = shape::indexOffset(j, meanShapeInfo, meanShapeInfoCast, canCastMean); - // calculate offset for input and output (all of them have the same shape) - shape::outerArrayOffsets(inOffsets, j, inShapeInfo, meanShapeInfo, memBuff, dimsToExclude.data()); + if(!isOwner) + continue; - PRAGMA_OMP_SIMD - for (Nd4jLong i = 0; i < step; ++i) { - auto offsetBig = inOffsets[i]; - outBuff[offsetBig] = (inBuff[offsetBig] - meanBuff[offsetSmall]) * sigmaBuff[offsetSmall] + betaBuff[offsetSmall]; + const auto meanOffset = shape::getIndexOffset(j, mean->getShapeInfo()); + const auto varOffset = paramSameOffset ? meanOffset : shape::getIndexOffset(j, variance->getShapeInfo()); + + const auto meanVal = m[meanOffset]; + auto sigmaInvGam = static_cast(1) / nd4j::math::nd4j_sqrt(v[varOffset] + epsilon); + + if(g != nullptr) { + const auto gammaOffset = paramSameOffset ? meanOffset : shape::getIndexOffset(j, gamma->getShapeInfo()); + sigmaInvGam *= g[gammaOffset]; + } + + T betaVal = static_cast(0); + if(b != nullptr) { + const auto betaOffset = paramSameOffset ? meanOffset : shape::getIndexOffset(j, beta->getShapeInfo()); + betaVal = b[betaOffset]; + } + + // calculate offsets for input and output + shape::outerArrayOffsets(xOffsets, j, input->getShapeInfo(), mean->getShapeInfo(), auxBuff, dimsToExclude.data()); + if(!xzSameOffset) + shape::outerArrayOffsets(zOffsets, j, output->getShapeInfo(), mean->getShapeInfo(), auxBuff, dimsToExclude.data()); + + PRAGMA_OMP_SIMD + for (uint i = 0; i < steps; ++i) + z[zOffsets[i]] = (x[xOffsets[i]] - meanVal) * sigmaInvGam + betaVal; + } + + delete []auxBuff; + delete []xOffsets; + if(!xzSameOffset) + delete []zOffsets; + }; + + samediff::Threads::parallel_do(func, info._numThreads); +} + +////////////////////////////////////////////////////////////////////////// +template +static void batchnorm2_(const NDArray* input, const NDArray* mean, const NDArray* variance, const NDArray* gamma, const NDArray* beta, + NDArray* output, + const std::vector& axes, const double epsilon) { + + // formula: output = gamma * ((input - mean) / sqrt(variance + epsilon)) + beta + + const auto x = input->bufferAsT(); + auto z = output->bufferAsT(); + const auto m = mean->bufferAsT(); + const auto v = variance->bufferAsT(); + const auto g = gamma == nullptr ? nullptr : gamma->bufferAsT(); + const auto b = beta == nullptr ? nullptr : beta->bufferAsT(); + + // xRank == zRank, minRank = meanRank = varianceRank = gammaRank = betaRank + const uint xRank = input->rankOf(); + const uint minRank = mean->rankOf(); + const uint numAxes = axes.size(); + + const bool xzSameOffset = shape::haveSameShapeAndStrides(input->getShapeInfo(), output->getShapeInfo()); + + bool paramSameOffset = shape::haveSameShapeAndStrides(mean->getShapeInfo(), variance->getShapeInfo()); + if(paramSameOffset && gamma != nullptr) + paramSameOffset &= shape::haveSameShapeAndStrides(mean->getShapeInfo(), gamma->getShapeInfo()); + if(paramSameOffset && beta != nullptr) + paramSameOffset &= shape::haveSameShapeAndStrides(mean->getShapeInfo(), beta->getShapeInfo()); + + auto func = PRAGMA_THREADS_FOR { + + Nd4jLong coords[MAX_RANK]; + + for (auto i = start; i < stop; i += increment) { + + shape::index2coords(i, input->getShapeInfo(), coords); + + const auto xOffset = shape::getOffset(input->getShapeInfo(), coords); + const auto zOffset = xzSameOffset ? xOffset : shape::getOffset(output->getShapeInfo(), coords); + + if(minRank == xRank) { + for (uint i = 0, j = 0; i < xRank; ++i) { + if(j < numAxes && i != axes[j]) + coords[i] = 0; + else + ++j; } } - delete []inOffsets; - delete []memBuff; - }; + else // minRank = numAxes = 1 in this case + coords[0] = coords[axes[0]]; - samediff::Threads::parallel_do(func, info._numThreads); - } - else { - auto func = PRAGMA_THREADS_DO { - const auto threadNum = thread_id; - Nd4jLong* inOffsets = new Nd4jLong[step]; - Nd4jLong* memBuff = new Nd4jLong[2 * inShapeInfo[0]]; + const auto meanOffset = shape::getOffset(mean->getShapeInfo(), coords); + const auto varianceOffset = paramSameOffset ? meanOffset : shape::getOffset(variance->getShapeInfo(), coords); - for (int j = 0; j < lenSmall; ++j) { - const bool isOwner = j < info._numThreads ? threadNum == j : threadNum == j % info._numThreads; - if (!isOwner) continue; + T sigmaInvGam = 1. / nd4j::math::nd4j_sqrt(v[varianceOffset] + epsilon); - const Nd4jLong start = j * step; - const Nd4jLong end = start + step; - - // calculate offset for mean, variance, gamma, beta (all of them have the same shape) - auto offsetSmall = shape::indexOffset(j, meanShapeInfo, meanShapeInfoCast, canCastMean); - // calculate offset for input and output (all of them have the same shape) - shape::outerArrayOffsets(inOffsets, j, inShapeInfo, meanShapeInfo, memBuff, dimsToExclude.data()); - - PRAGMA_OMP_SIMD - for (Nd4jLong i = 0; i < step; ++i) { - auto offsetBig = inOffsets[i]; - outBuff[offsetBig] = (inBuff[offsetBig] - meanBuff[offsetSmall]) * sigmaBuff[offsetSmall]; - } + if(g != nullptr) { + const auto gammaOffset = paramSameOffset ? meanOffset : shape::getOffset(gamma->getShapeInfo(), coords); + sigmaInvGam *= g[gammaOffset]; } - delete []inOffsets; - delete []memBuff; - }; - samediff::Threads::parallel_do(func, info._numThreads); - } + z[zOffset] = (x[xOffset] - m[meanOffset]) * sigmaInvGam; + + if(b != nullptr) { + const auto betaOffset = paramSameOffset ? meanOffset : shape::getOffset(beta->getShapeInfo(), coords); + z[zOffset] += b[betaOffset]; + } + } + }; + + samediff::Threads::parallel_for(func, 0, input->lengthOf()); } ////////////////////////////////////////////////////////////////////////// void batchnorm(const NDArray* input, const NDArray* mean, const NDArray* variance, const NDArray* gamma, const NDArray* beta, NDArray* output, const std::vector& axes, const double epsilon) { + // batchnorm2_ is slower BUILD_SINGLE_SELECTOR(input->dataType(), batchnorm_, (input, mean, variance, gamma, beta, output, axes, epsilon), FLOAT_TYPES); } diff --git a/libnd4j/include/ops/declarable/helpers/cuda/batchnorm.cu b/libnd4j/include/ops/declarable/helpers/cuda/batchnorm.cu index d9188e3a8..eedbe1fdf 100644 --- a/libnd4j/include/ops/declarable/helpers/cuda/batchnorm.cu +++ b/libnd4j/include/ops/declarable/helpers/cuda/batchnorm.cu @@ -31,66 +31,66 @@ namespace helpers { ////////////////////////////////////////////////////////////////////////// -template -__global__ static void batchnormCuda(const void* vx, const Nd4jLong* xShapeInfo, - const void* vMean, const Nd4jLong* meanShapeInfo, - const void* vVariance, const Nd4jLong* varianceShapeInfo, - const void* vGamma, const Nd4jLong* gammaShapeInfo, - const void* vBeta, const Nd4jLong* betaShapeInfo, - void* vz, const Nd4jLong* zShapeInfo, - const Nd4jLong* xTadShapeInfo, const Nd4jLong* xTadOffsets, - const Nd4jLong* zTadShapeInfo, const Nd4jLong* zTadOffsets, - const T epsilon) { +// template +// __global__ static void batchnormCuda(const void* vx, const Nd4jLong* xShapeInfo, +// const void* vMean, const Nd4jLong* meanShapeInfo, +// const void* vVariance, const Nd4jLong* varianceShapeInfo, +// const void* vGamma, const Nd4jLong* gammaShapeInfo, +// const void* vBeta, const Nd4jLong* betaShapeInfo, +// void* vz, const Nd4jLong* zShapeInfo, +// const Nd4jLong* xTadShapeInfo, const Nd4jLong* xTadOffsets, +// const Nd4jLong* zTadShapeInfo, const Nd4jLong* zTadOffsets, +// const T epsilon) { - const auto x = reinterpret_cast(vx); - auto z = reinterpret_cast(vz); - const auto mean = reinterpret_cast(vMean); - const auto variance = reinterpret_cast(vVariance); - const auto gamma = reinterpret_cast(vGamma); - const auto beta = reinterpret_cast(vBeta); +// const auto x = reinterpret_cast(vx); +// auto z = reinterpret_cast(vz); +// const auto mean = reinterpret_cast(vMean); +// const auto variance = reinterpret_cast(vVariance); +// const auto gamma = reinterpret_cast(vGamma); +// const auto beta = reinterpret_cast(vBeta); - // maxRank = xRank = zRank, minRank = meanRank = varianceRank = gammaRank = betaRank - __shared__ Nd4jLong minLen, tadLen, totalThreads; +// // maxRank = xRank = zRank, minRank = meanRank = varianceRank = gammaRank = betaRank +// __shared__ Nd4jLong minLen, tadLen, totalThreads; - if (threadIdx.x == 0) { - totalThreads = gridDim.x * blockDim.x; +// if (threadIdx.x == 0) { +// totalThreads = gridDim.x * blockDim.x; - minLen = shape::length(meanShapeInfo); - tadLen = shape::length(xShapeInfo) / minLen; - } - __syncthreads(); +// minLen = shape::length(meanShapeInfo); +// tadLen = shape::length(xShapeInfo) / minLen; +// } +// __syncthreads(); - const auto tid = blockIdx.x * blockDim.x + threadIdx.x; +// const auto tid = blockIdx.x * blockDim.x + threadIdx.x; - for (uint i = tid; i < minLen; i += totalThreads) { +// for (uint i = tid; i < minLen; i += totalThreads) { - const auto meanOffset = shape::getIndexOffset(i, meanShapeInfo); - const auto varianceOffset = shape::getIndexOffset(i, varianceShapeInfo); +// const auto meanOffset = shape::getIndexOffset(i, meanShapeInfo); +// const auto varianceOffset = shape::getIndexOffset(i, varianceShapeInfo); - T sigmaInvGam = 1. / nd4j::math::nd4j_sqrt(variance[varianceOffset] + epsilon); +// T sigmaInvGam = 1. / nd4j::math::nd4j_sqrt(variance[varianceOffset] + epsilon); - if(gamma != nullptr) - sigmaInvGam *= gamma[shape::getIndexOffset(i, gammaShapeInfo)]; +// if(gamma != nullptr) +// sigmaInvGam *= gamma[shape::getIndexOffset(i, gammaShapeInfo)]; - auto betaOffset = 0; - if(beta != nullptr) - betaOffset = shape::getIndexOffset(i, betaShapeInfo); +// auto betaOffset = 0; +// if(beta != nullptr) +// betaOffset = shape::getIndexOffset(i, betaShapeInfo); - const auto xTad = x + xTadOffsets[i]; - auto zTad = z + zTadOffsets[i]; +// const auto xTad = x + xTadOffsets[i]; +// auto zTad = z + zTadOffsets[i]; - for (uint j = 0; j < tadLen; ++j) { +// for (uint j = 0; j < tadLen; ++j) { - const auto xTadOffset = shape::getIndexOffset(j, xTadShapeInfo); - const auto zTadOffset = shape::getIndexOffset(j, zTadShapeInfo); +// const auto xTadOffset = shape::getIndexOffset(j, xTadShapeInfo); +// const auto zTadOffset = shape::getIndexOffset(j, zTadShapeInfo); - zTad[zTadOffset] = (xTad[xTadOffset] - mean[meanOffset]) * sigmaInvGam; +// zTad[zTadOffset] = (xTad[xTadOffset] - mean[meanOffset]) * sigmaInvGam; - if(beta != nullptr) - zTad[zTadOffset] += beta[betaOffset]; - } - } -} +// if(beta != nullptr) +// zTad[zTadOffset] += beta[betaOffset]; +// } +// } +// } ////////////////////////////////////////////////////////////////////////// template @@ -110,13 +110,12 @@ __global__ static void batchnormCuda2(const void* vx, const Nd4jLong* xShapeInfo const auto gamma = reinterpret_cast(vGamma); const auto beta = reinterpret_cast(vBeta); - __shared__ int xRank, minRank; // xRank == zRank. minRank = meanRank = varianceRank = gammaRank = betaRank - __shared__ Nd4jLong xLen, totalThreads, *sharedMem; // xLen = zLen + __shared__ int xRank, minRank; // xRank == zRank, minRank = meanRank = varianceRank = gammaRank = betaRank + __shared__ Nd4jLong xLen, totalThreads; // xLen = zLen if (threadIdx.x == 0) { - extern __shared__ unsigned char shmem[]; - sharedMem = reinterpret_cast(shmem); + totalThreads = gridDim.x * blockDim.x; xLen = shape::length(xShapeInfo); @@ -125,7 +124,8 @@ __global__ static void batchnormCuda2(const void* vx, const Nd4jLong* xShapeInfo } __syncthreads(); - auto coords = sharedMem + threadIdx.x * xRank; + Nd4jLong coords[MAX_RANK]; + const auto tid = blockIdx.x * blockDim.x + threadIdx.x; for (uint i = tid; i < xLen; i += totalThreads) { @@ -166,24 +166,24 @@ __global__ static void batchnormCuda2(const void* vx, const Nd4jLong* xShapeInfo } /////////////////////////////////////////////////////////////////// -template -__host__ static void batchnormCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t *stream, - const void* vx, const Nd4jLong* xShapeInfo, - const void* vMean, const Nd4jLong* meanShapeInfo, - const void* vVariance, const Nd4jLong* varianceShapeInfo, - const void* vGamma, const Nd4jLong* gammaShapeInfo, - const void* vBeta, const Nd4jLong* betaShapeInfo, - void* vz, const Nd4jLong* zShapeInfo, - const Nd4jLong* xTadShapeInfo, const Nd4jLong* xTadOffsets, - const Nd4jLong* zTadShapeInfo, const Nd4jLong* zTadOffsets, - const double epsilon) { +// template +// __host__ static void batchnormCudaLauncher(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t *stream, +// const void* vx, const Nd4jLong* xShapeInfo, +// const void* vMean, const Nd4jLong* meanShapeInfo, +// const void* vVariance, const Nd4jLong* varianceShapeInfo, +// const void* vGamma, const Nd4jLong* gammaShapeInfo, +// const void* vBeta, const Nd4jLong* betaShapeInfo, +// void* vz, const Nd4jLong* zShapeInfo, +// const Nd4jLong* xTadShapeInfo, const Nd4jLong* xTadOffsets, +// const Nd4jLong* zTadShapeInfo, const Nd4jLong* zTadOffsets, +// const double epsilon) { - batchnormCuda<<>>(vx, xShapeInfo, vMean, meanShapeInfo, vVariance, varianceShapeInfo, vGamma, gammaShapeInfo, vBeta, betaShapeInfo, vz, zShapeInfo, xTadShapeInfo, xTadOffsets, zTadShapeInfo, zTadOffsets, static_cast(epsilon)); -} +// batchnormCuda<<>>(vx, xShapeInfo, vMean, meanShapeInfo, vVariance, varianceShapeInfo, vGamma, gammaShapeInfo, vBeta, betaShapeInfo, vz, zShapeInfo, xTadShapeInfo, xTadOffsets, zTadShapeInfo, zTadOffsets, static_cast(epsilon)); +// } /////////////////////////////////////////////////////////////////// template -__host__ static void batchnormCudaLauncher2(const int blocksPerGrid, const int threadsPerBlock, const int sharedMem, const cudaStream_t *stream, +__host__ static void batchnormCudaLauncher2(const int blocksPerGrid, const int threadsPerBlock, const cudaStream_t *stream, const void* vx, const Nd4jLong* xShapeInfo, const void* vMean, const Nd4jLong* meanShapeInfo, const void* vVariance, const Nd4jLong* varianceShapeInfo, @@ -193,42 +193,41 @@ __host__ static void batchnormCudaLauncher2(const int blocksPerGrid, const int t const int numDims, const int* dims, const double epsilon) { - batchnormCuda2<<>>(vx, xShapeInfo, vMean, meanShapeInfo, vVariance, varianceShapeInfo, vGamma, gammaShapeInfo, vBeta, betaShapeInfo, vz, zShapeInfo, numDims, dims, static_cast(epsilon)); + batchnormCuda2<<>>(vx, xShapeInfo, vMean, meanShapeInfo, vVariance, varianceShapeInfo, vGamma, gammaShapeInfo, vBeta, betaShapeInfo, vz, zShapeInfo, numDims, dims, static_cast(epsilon)); } ////////////////////////////////////////////////////////////////////////// void batchnorm(const NDArray* input, const NDArray* mean, const NDArray* variance, const NDArray* gamma, const NDArray* beta, NDArray* output, const std::vector& axes, const double epsilon) { - std::vector dimsToExclude = ShapeUtils::evalDimsToExclude(input->rankOf(), axes); + // std::vector dimsToExclude = ShapeUtils::evalDimsToExclude(input->rankOf(), axes); + + // auto packX = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(input->getShapeInfo(), dimsToExclude); + // auto packZ = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(output->shapeInfo(), dimsToExclude); + + // const int threadsPerBlock = MAX_NUM_THREADS / 2; + // const int blocksPerGrid = (mean->lengthOf() + threadsPerBlock - 1) / threadsPerBlock; + + // PointersManager manager(input->getContext(), "batchnorm"); + + // NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta}); + // BUILD_SINGLE_SELECTOR(input->dataType(), batchnormCudaLauncher, (blocksPerGrid, threadsPerBlock, input->getContext()->getCudaStream(), input->getSpecialBuffer(), input->getSpecialShapeInfo(), mean->getSpecialBuffer(), mean->getSpecialShapeInfo(), variance->getSpecialBuffer(), variance->getSpecialShapeInfo(), gamma ? gamma->getSpecialBuffer() : nullptr, gamma ? gamma->getSpecialShapeInfo() : nullptr, beta ? beta->getSpecialBuffer() : nullptr, beta ? beta->getSpecialShapeInfo() : nullptr, output->specialBuffer(), output->specialShapeInfo(), packX.platformShapeInfo(), packX.platformOffsets(), packZ.platformShapeInfo(), packZ.platformOffsets(), epsilon), FLOAT_TYPES); + // NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta}); + + // manager.synchronize(); - auto packX = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(input->getShapeInfo(), dimsToExclude); - auto packZ = nd4j::ConstantTadHelper::getInstance()->tadForDimensions(output->shapeInfo(), dimsToExclude); const int threadsPerBlock = MAX_NUM_THREADS / 2; - const int blocksPerGrid = (mean->lengthOf() + threadsPerBlock - 1) / threadsPerBlock; + const int blocksPerGrid = (input->lengthOf() + threadsPerBlock - 1) / threadsPerBlock; PointersManager manager(input->getContext(), "batchnorm"); + const int* dims = reinterpret_cast(manager.replicatePointer(axes.data(), axes.size() * sizeof(int))); + NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta}); - BUILD_SINGLE_SELECTOR(input->dataType(), batchnormCudaLauncher, (blocksPerGrid, threadsPerBlock, input->getContext()->getCudaStream(), input->getSpecialBuffer(), input->getSpecialShapeInfo(), mean->getSpecialBuffer(), mean->getSpecialShapeInfo(), variance->getSpecialBuffer(), variance->getSpecialShapeInfo(), gamma ? gamma->getSpecialBuffer() : nullptr, gamma ? gamma->getSpecialShapeInfo() : nullptr, beta ? beta->getSpecialBuffer() : nullptr, beta ? beta->getSpecialShapeInfo() : nullptr, output->specialBuffer(), output->specialShapeInfo(), packX.platformShapeInfo(), packX.platformOffsets(), packZ.platformShapeInfo(), packZ.platformOffsets(), epsilon), FLOAT_TYPES); + BUILD_SINGLE_SELECTOR(input->dataType(), batchnormCudaLauncher2, (blocksPerGrid, threadsPerBlock, input->getContext()->getCudaStream(), input->getSpecialBuffer(), input->getSpecialShapeInfo(), mean->getSpecialBuffer(), mean->getSpecialShapeInfo(), variance->getSpecialBuffer(), variance->getSpecialShapeInfo(), gamma ? gamma->getSpecialBuffer() : nullptr, gamma ? gamma->getSpecialShapeInfo() : nullptr, beta ? beta->getSpecialBuffer() : nullptr, beta ? beta->getSpecialShapeInfo() : nullptr, output->specialBuffer(), output->specialShapeInfo(), axes.size(), dims, epsilon), FLOAT_TYPES); NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta}); manager.synchronize(); - - - // const int threadsPerBlock = MAX_NUM_THREADS / 4; - // const int blocksPerGrid = (input->lengthOf() + threadsPerBlock - 1) / threadsPerBlock; - // const int sharedMem = sizeof(Nd4jLong) * threadsPerBlock * input->rankOf() + 128; - - // PointersManager manager(input->getContext(), "batchnorm"); - - // const int* dims = reinterpret_cast(manager.replicatePointer(axes.data(), axes.size() * sizeof(int))); - - // NDArray::prepareSpecialUse({output}, {input, mean, variance, gamma, beta}); - // BUILD_SINGLE_SELECTOR(input->dataType(), batchnormCudaLauncher2, (blocksPerGrid, threadsPerBlock, sharedMem, input->getContext()->getCudaStream(), input->getSpecialBuffer(), input->getSpecialShapeInfo(), mean->getSpecialBuffer(), mean->getSpecialShapeInfo(), variance->getSpecialBuffer(), variance->getSpecialShapeInfo(), gamma ? gamma->getSpecialBuffer() : nullptr, gamma ? gamma->getSpecialShapeInfo() : nullptr, beta ? beta->getSpecialBuffer() : nullptr, beta ? beta->getSpecialShapeInfo() : nullptr, output->specialBuffer(), output->specialShapeInfo(), axes.size(), dims, epsilon), FLOAT_TYPES); - // NDArray::registerSpecialUse({output}, {input, mean, variance, gamma, beta}); - - // manager.synchronize(); } diff --git a/libnd4j/tests_cpu/layers_tests/DeclarableOpsTests10.cpp b/libnd4j/tests_cpu/layers_tests/DeclarableOpsTests10.cpp index 66cc487e1..f9382f6c7 100644 --- a/libnd4j/tests_cpu/layers_tests/DeclarableOpsTests10.cpp +++ b/libnd4j/tests_cpu/layers_tests/DeclarableOpsTests10.cpp @@ -3431,6 +3431,35 @@ TEST_F(DeclarableOpsTests10, batchnorm_test6) { delete results; } +//////////////////////////////////////////////////////////////////// +TEST_F(DeclarableOpsTests10, batchnorm_test7) { + + NDArray input1('c', {3,3,15,15}, nd4j::DataType::FLOAT32); + NDArray input2('c', {3,15,15,3}, nd4j::DataType::FLOAT32); + input2.permutei({0,3,1,2}); + + NDArray mean ('c', {3}, {0, 0, 0}, nd4j::DataType::FLOAT32); + NDArray variance('c', {3}, {1, 1, 1}, nd4j::DataType::FLOAT32); + NDArray gamma ('c', {3}, {1, 1, 1}, nd4j::DataType::FLOAT32); + NDArray beta ('c', {3}, {0, 0, 0}, nd4j::DataType::FLOAT32); + + NDArray out1('c', {3,3,15,15}, nd4j::DataType::FLOAT32); + NDArray out2('c', {3,3,15,15}, nd4j::DataType::FLOAT32); + + input1.linspace(-1012, 1); + input2.assign(input1); + + nd4j::ops::batchnorm op; + + auto res1 = op.execute({&input1, &mean, &variance, &gamma, &beta}, {&out1}, {1e-5}, {1,1,1}, {}); + ASSERT_EQ(ND4J_STATUS_OK, res1); + + auto res2 = op.execute({&input2, &mean, &variance, &gamma, &beta}, {&out2}, {1e-5}, {1,1,1}, {}); + ASSERT_EQ(ND4J_STATUS_OK, res2); + + ASSERT_TRUE(out1.equalsTo(out2)); +} + /////////////////////////////////////////////////////////////////// TEST_F(DeclarableOpsTests10, bool_broadcast_test_1) { diff --git a/libnd4j/tests_cpu/layers_tests/PlaygroundTests.cpp b/libnd4j/tests_cpu/layers_tests/PlaygroundTests.cpp index a7bffb9a1..970c119ca 100644 --- a/libnd4j/tests_cpu/layers_tests/PlaygroundTests.cpp +++ b/libnd4j/tests_cpu/layers_tests/PlaygroundTests.cpp @@ -422,13 +422,50 @@ TEST_F(PlaygroundTests, my) { delete variableSpace; } -*/ + +#include TEST_F(PlaygroundTests, my) { - NDArray a('c',{2,3,4}, nd4j::DataType::DOUBLE); - a({0,0, 0,1, 0,1}).printShapeInfo(); - a({0,1, 0,0, 0,1}).printShapeInfo(); - a({0,0, 0,1, 0,1}).printShapeInfo(); + const int N = 10000; + const Nd4jLong dim0(128), dim1(128), dim2(128); + + NDArray input('c', {dim0,dim1,dim2}, nd4j::DataType::DOUBLE); + NDArray mean('c', {dim1}, nd4j::DataType::DOUBLE); + NDArray variance('c', {dim1}, nd4j::DataType::DOUBLE); + NDArray gamma('c', {dim1}, nd4j::DataType::DOUBLE); + NDArray beta ('c', {dim1}, nd4j::DataType::DOUBLE); + + NDArray output('c', {dim0,dim1,dim2}, nd4j::DataType::DOUBLE); + + input.linspace(-100, 0.1); + mean.linspace(-50, 0.15); + variance.linspace(-5, 0.2); + gamma = 1.5; + beta = -2.5; + + // warm up + ops::helpers::batchnorm(&input, &mean, &variance, &gamma, &beta, &output, {1}, 1e-5); + + auto timeStart = std::chrono::system_clock::now(); + for (int i = 0; i < N; ++i) + ops::helpers::batchnorm(&input, &mean, &variance, &gamma, &beta, &output, {1}, 1e-5); + + auto timeEnd = std::chrono::system_clock::now(); + auto time = std::chrono::duration_cast ((timeEnd - timeStart)/N).count(); + + printf("time: %li \n", time); + +} + + +*/ + + + + + + + + -} \ No newline at end of file